NBLAST

# Import libs and initialise API objects
from vfb_connect.cross_server_tools import VfbConnect
import navis.interfaces.neuprint as neu

import pandas as pd
import navis

navis.set_pbars(jupyter=False)

vc = VfbConnect()
client = neu.Client('https://neuprint.janelia.org', dataset='hemibrain:v1.1')

import requests
from tqdm import tqdm

def get_vfb_neurons(vfb_ids, template='JRC2018Unisex'):
    """Load neurons from VFB as navis.TreeNeurons."""
    vfb_ids = list(navis.utils.make_iterable(vfb_ids))
    inds = vc.neo_query_wrapper.get_anatomical_individual_TermInfo(short_forms=vfb_ids)
    nl = []
    # Note: queries should be parallelized in the future 
    # -> for pymaid I use request-futures which works really well
    for i in tqdm(inds, desc='Loading VFB neurons', leave=False):        
        if not ('has_image' in i['term']['core']['types']):
            continue
        label = i['term']['core']['label']
        image_matches = [x['image'] for x in i['channel_image']]
        if not image_matches:
            continue
        for imv in image_matches:
            if imv['template_anatomy']['label'] == template:
                r = requests.get(imv['image_folder'] + '/volume.swc')
                ### Slightly dodgy warning - could mask network errors
                if not r.ok:
                    warnings.warn("No '%s' file found for '%s'." % (image_type, label))
                    continue

                # `id` should ideally be unique but that's not enforced
                n = navis.TreeNeuron(r.text, name=label, id=i['term']['core']['short_form'])   

                # This registers attributes that you want to show in the summary
                # alternatively we could have a generic attribute like so:
                # n.template = template
                n._register_attr("template", template, temporary=False)

                # I assume all available templates are in microns?
                # @Robbie -do we have this metadata on templates?  If not, we should add ASAP.
                n.units = '1 micron'

                nl.append(n)
    return navis.NeuronList(nl)

# Search for similar neurons using the VFB ID of an ellipsoid body neuron in FAFB
# We got that ID from a search on the VFB website
similar_to_EPG6L1 = vc.get_similar_neurons('VFB_001012ay')
similar_to_EPG6L1 

# Read those neurons from VFB 
query = get_vfb_neurons('VFB_001012ay')
matches = get_vfb_neurons(similar_to_EPG6L1.id.values)

matches

# Define colors such that query is black 
import seaborn as sns 
colors = ['k'] + sns.color_palette('muted', len(matches))

navis.plot3d([query, matches], color=colors)

# Find instances of EB-PB-gall neurons
eb_pb_ids = pd.DataFrame.from_records(vc.get_instances("adult ellipsoid body-protocerebral bridge-gall neuron"))
eb_pb_ids.head()

# Get the neurons from VFB
eb_pb = get_vfb_neurons(eb_pb_ids.id.values)
eb_pb

eb_dps = navis.make_dotprops(eb_pb, resample=1, k=5)

# Note this NeuronList /contains dotprops now
eb_dps

sk = eb_pb[0].copy()
dp = eb_dps[0].copy()

# Slightly offset the dotprops by a micron 
dp.points[:, 0] += 1

navis.plot3d([sk, dp], color=['k', 'r'])

scores = navis.nblast_allbyall(eb_dps)
scores.head()

# Symmetrize scores
scores_mean = (scores + scores.T) / 2

scores_mean.head()

dist_mean = (scores_mean - 1) * -1
dist_mean.head()

from scipy.spatial.distance import squareform
from scipy.cluster.hierarchy import linkage, dendrogram, fcluster

# This turns the square distance matrix into 1-d vector form
dist_sq = squareform(dist_mean)

# This does the actual hierarchical clustering
Z = linkage(dist_sq, method='ward', optimal_ordering=True)

# Plot a quick dendrogram
import matplotlib.pyplot as plt 
import seaborn as sns 

# Generate the canvas
fig, ax = plt.subplots(figsize=(12, 5))

# Make some more meaningful labels 
vfb2name = dict(zip(eb_pb.id, eb_pb.name))
labels = dist_mean.index.map(lambda x: vfb2name[x].split(' - ')[0])

# Plot the dendrogram
dn = dendrogram(Z, ax=ax, labels=labels, color_threshold=.65, above_threshold_color='lightgrey')

# Make xticks biggger
ax.set_xticklabels(ax.get_xticklabels(), size=11)

# Clean axes
sns.despine(bottom=True)

# Make cluster matching the color threshold in the dendrogram
clusters = fcluster(Z, t=.65, criterion='distance')
clusters

# We have to downsample the neurons for visualization 
# because DeepNote enforces a limit on the size of plots
eb_pb_ds = eb_pb.downsample(10)

# Make a color per cluster 
palette = sns.color_palette('tab20', max(clusters))
cmap = dict(zip(range(1, max(clusters) + 1), palette))

colors = [cmap[i] for i in clusters]

navis.plot3d(eb_pb_ds, color=colors)