medulla/Medulla7column_Neuprint_to_NeuroArch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading NeuroArch Database with Medulla 7 Column Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial provides code to load NeuroArch database with Medulla 7 Column Dataset. Requirement before running the notebook:\n",
    "- Installed [NeuroArch](https://github.com/fruitflybrain/neuroarch), [OrientDB Community Version](https://www.orientdb.org/download) version 3.1.x, and [pyorient](https://github.com/fruitflybrain/pyorient). The [NeuroNLP Docker image](https://hub.docker.com/r/fruitflybrain/neuronlp) and [FlyBrainLab Docker image](https://hub.docker.com/r/fruitflybrain/fbl) all have a copy of the software requirement ready.\n",
    "- Download the [Neuprint database dump for the Medulla 7 Column dataset](https://github.com/connectome-neuprint/neuPrint/raw/master/fib25_neo4j_inputs.zip).\n",
    "- Download the neuron skeletons from [ConnectomeHackathon2015 repository](https://github.com/janelia-flyem/ConnectomeHackathon2015) and rename the `skeletons` folder to `swc` and move it under the same directory as this notebook.\n",
    "- Download the [GSE116969 transcriptome cell expression data](https://www.ncbi.nlm.nih.gov/geo/download/?acc=GSE116969&format=file&file=GSE116969%5FdataTable7b%2Egenes%5Fx%5Fcells%5Fp%5Fexpression%2Emodeled%5Fgenes%2Etxt%2Egz) and uncompress it to the same folder.\n",
    "- Have 3GB free disk space (for Neuprint dump and NeuroArch database).\n",
    "\n",
    "A backup of the database created by this notebook can be downloaded [here](https://drive.google.com/file/d/1lX_nGmqfQ7YouO4TTGw7Q3fRaAZn1T1S/view?usp=sharing). To restore it in OrientDB, run\n",
    "```\n",
    "/path/to/orientdb/bin/console.sh \"create database plocal:../databases/medulla admin admin; restore database /path/to/medulla7column_fib25_na_v2.0.1_backup.zip\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import subprocess\n",
    "import csv\n",
    "from collections import Counter\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import h5py\n",
    "\n",
    "import neuroarch.models as models\n",
    "import neuroarch.na as na"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from requests import HTTPError\n",
    "from neuprint import Client\n",
    "c = Client('neuprint-examples.janelia.org', dataset='medulla7column', token = '')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('fib25/Neuprint_Neurons_fib25.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm(df['bodyId:long']):\n",
    "    try:\n",
    "        s = c.fetch_skeleton(i, format='swc')\n",
    "    except HTTPError:\n",
    "        continue\n",
    "    with open('swc_new/{}.swc'.format(i), 'w') as f:\n",
    "        f.write(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Neuron and Synapse Attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(chunk):\n",
    "    neurons = []\n",
    "    for i, row in chunk.iterrows():\n",
    "        li = [row['bodyId:long'], row['pre:int'], row['post:int'], row['status:string'],\\\n",
    "              row['statusLabel:string'], int(row['cropped:boolean']) if not np.isnan(row['cropped:boolean']) else row['cropped:boolean'], row['instance:string'], \\\n",
    "              row['type:string']]\n",
    "        neurons.append(li)\n",
    "    return neurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunksize = 100000\n",
    "\n",
    "with open('neurons.csv', 'w') as f:\n",
    "    writer = csv.writer(f)\n",
    "    writer.writerow(['bodyID','pre','post','status','statusLabel','cropped','instance','type'])\n",
    "    df = pd.read_csv('fib25/Neuprint_Neurons_fib25.csv')\n",
    "    neurons = process(df)\n",
    "    writer.writerows(neurons)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neurons = pd.read_csv('neurons.csv')\n",
    "traced_neuron_id = neurons['bodyID'].to_numpy()\n",
    "\n",
    "chunksize = 1000000\n",
    "pre_syn = np.empty((int(1e7),3), np.int64)\n",
    "post_syn = np.empty((int(1e7),3), np.int64)\n",
    "\n",
    "pre_count = 0\n",
    "post_count = 0\n",
    "count = 0\n",
    "for chunk in pd.read_csv('fib25/Neuprint_SynapseSet_to_Synapses_fib25.csv', chunksize=chunksize):\n",
    "    ids = chunk[':START_ID']\n",
    "    pre_site = np.array([[n, int(i.split('_')[0]), int(i.split('_')[1])] \\\n",
    "                         for n,i in enumerate(ids) if i.split('_')[2] == 'pre'])\n",
    "    post_site = np.array([[n, int(i.split('_')[0]), int(i.split('_')[1])] \\\n",
    "                          for n,i in enumerate(ids) if i.split('_')[2] == 'post'])\n",
    "    pre_site_known = pre_site[np.logical_and(\n",
    "                              np.isin(pre_site[:,1], traced_neuron_id),\n",
    "                              np.isin(pre_site[:,2], traced_neuron_id)),0]\n",
    "    post_site_known = post_site[np.logical_and(\n",
    "                                np.isin(post_site[:,1], traced_neuron_id),\n",
    "                                np.isin(post_site[:,2], traced_neuron_id)),0]\n",
    "    retrieved_pre_site = chunk.iloc[pre_site_known]\n",
    "    pre_site = np.array([[row[':END_ID(Syn-ID)'], int(row[':START_ID'].split('_')[0]), int(row[':START_ID'].split('_')[1])] \\\n",
    "                         for i, row in retrieved_pre_site.iterrows()])\n",
    "    retrieved_post_site = chunk.iloc[post_site_known]\n",
    "    post_site = np.array([[row[':END_ID(Syn-ID)'], int(row[':START_ID'].split('_')[0]), int(row[':START_ID'].split('_')[1])] \\\n",
    "                         for i, row in retrieved_post_site.iterrows()])\n",
    "    if pre_site.size:\n",
    "        pre_syn[pre_count:pre_count+pre_site.shape[0], :] = pre_site\n",
    "        pre_count += pre_site.shape[0]\n",
    "    if post_site.size:\n",
    "        post_syn[post_count:post_count+post_site.shape[0], :] = post_site\n",
    "        post_count += post_site.shape[0]\n",
    "    count += chunksize\n",
    "    print(count, pre_count, post_count)\n",
    "\n",
    "pre_syn = pre_syn[:pre_count,:]\n",
    "post_syn = post_syn[:post_count,:]\n",
    "\n",
    "ind = np.argsort(pre_syn[:,0])\n",
    "pre_syn_sorted = pre_syn[ind, :]\n",
    "ind = np.argsort(post_syn[:,0])\n",
    "post_syn_sorted = post_syn[ind, :]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract synapse (pre-site) to synapse (post-site) connection\n",
    "# use only the post synaptic site to get all the synapses because one presynaptic site can have multiple postsynaptic sites\n",
    "post_syn_index = post_syn_sorted[:,0].copy()\n",
    "\n",
    "df = pd.read_csv('fib25/Neuprint_Synapse_Connections_fib25.csv')\n",
    "post_ids = df[':END_ID(Syn-ID)']\n",
    "used = np.where(post_ids.isin(post_syn_index).to_numpy())[0]\n",
    "connections = df.iloc[used].to_numpy()\n",
    "ind = np.argsort(connections[:,1])\n",
    "connections = connections[ind, :]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract synapse details\n",
    "# with h5py.File('syn_pre_post_sorted_by_synapse_id.h5', 'r') as f:\n",
    "#     pre_syn_sorted = f['pre_syn_sorted'][:]\n",
    "#     post_syn_sorted = f['post_syn_sorted'][:]\n",
    "chunksize = 100000\n",
    "\n",
    "pre_syn_index = list(set(pre_syn_sorted[:,0].copy()))\n",
    "pre_syn_index.extend(list(post_syn_sorted[:,0].copy()))\n",
    "syn_index = np.array(sorted(pre_syn_index))\n",
    "del pre_syn_index#, pre_syn_sorted, post_syn_sorted\n",
    "\n",
    "synapse_array = np.empty((len(syn_index), 6), np.int64)\n",
    "\n",
    "synapse_count = 0\n",
    "count = 0\n",
    "\n",
    "for chunk in pd.read_csv('fib25/Neuprint_Synapses_fib25.csv', chunksize=chunksize):\n",
    "    ids = chunk[':ID(Syn-ID)']\n",
    "    \n",
    "    start_id = ids.iloc[0]\n",
    "    stop_id = ids.iloc[-1]\n",
    "    pre_start = np.searchsorted(syn_index, start_id, side='left')\n",
    "    pre_end = np.searchsorted(syn_index, stop_id, side='right')\n",
    "    if pre_start >= len(syn_index):\n",
    "        pre_index = []\n",
    "    else:\n",
    "        if pre_end >= len(syn_index):\n",
    "            pre_index = syn_index[pre_start:pre_end] #same as syn_index[pre_start:]\n",
    "        else:\n",
    "            pre_index = syn_index[pre_start:pre_end]\n",
    "    pre_used_synapse = chunk.loc[ids.isin(pre_index)]\n",
    "    li = np.empty((pre_index.size, 6), np.int64)\n",
    "    i = 0\n",
    "    for _, row in pre_used_synapse.iterrows():\n",
    "        location = eval(row['location:point{srid:9157}'].replace('x', \"'x'\").replace('y', \"'y'\").replace('z', \"'z'\"))\n",
    "        li[i,:] = [row[':ID(Syn-ID)'], # synpase id\n",
    "                     0 if row['type:string'] == 'pre' else 1, #synapse type\n",
    "                     int(row['confidence:float']*1000000), #confidence\n",
    "                     location['x'], location['y'], location['z']]\n",
    "        i += 1\n",
    "    synapse_array[synapse_count:synapse_count+pre_index.shape[0],:] = li\n",
    "    synapse_count += pre_index.shape[0]\n",
    "    count += chunksize\n",
    "    print(count, len(pre_used_synapse))\n",
    "synapse_array = synapse_array[:synapse_count,:]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reorder synapses\n",
    "\n",
    "synapse_connections = connections\n",
    "    \n",
    "ids = synapse_array[:,0]\n",
    "syn_id_dict = {j: i for i, j in enumerate(ids)}\n",
    "# ids = pre_syn_sorted[:,0]\n",
    "# pre_syn_id_dict = {j: i for i, j in enumerate(ids)} # map syn id to pre_syn_sorted\n",
    "ids = post_syn_sorted[:,0]\n",
    "post_syn_id_dict = {j: i for i, j in enumerate(ids)} # map syn id to post_syn_sorted\n",
    "\n",
    "synapse_dict = {}\n",
    "wrong_synapse = 0\n",
    "for i, pair in tqdm(enumerate(synapse_connections)):\n",
    "    pre_syn_id = pair[0]\n",
    "    post_syn_id = pair[1]\n",
    "    post_id = post_syn_id_dict[post_syn_id]\n",
    "    post_info = synapse_array[syn_id_dict[post_syn_id]]\n",
    "    post_neuron_id, pre_neuron_id = post_syn_sorted[post_id, 1:]\n",
    "\n",
    "    #if len(np.where((pre_syn_sorted == (pre_syn_id, pre_neuron_id, post_neuron_id)).all(axis=1))[0]) != 1:\n",
    "    #    print(pre_syn_id, post_syn_id)\n",
    "    # pre_id = pre_syn_id_dict[pre_syn_id]\n",
    "    pre_info = synapse_array[syn_id_dict[pre_syn_id]]\n",
    "\n",
    "    if pre_neuron_id not in synapse_dict:\n",
    "        synapse_dict[pre_neuron_id] = {}\n",
    "    pre_dict = synapse_dict[pre_neuron_id]\n",
    "    if post_neuron_id not in synapse_dict[pre_neuron_id]:\n",
    "        pre_dict[post_neuron_id] =  {'pre_synapse_ids': [],\n",
    "                                     'post_synapse_ids': [],\n",
    "                                     'pre_confidence': [],\n",
    "                                     'post_confidence': [],\n",
    "                                     'pre_x': [],\n",
    "                                     'pre_y': [],\n",
    "                                     'pre_z': [],\n",
    "                                     'post_x': [],\n",
    "                                     'post_y': [],\n",
    "                                     'post_z': [],\n",
    "                                     }\n",
    "    info_dict = pre_dict[post_neuron_id]\n",
    "    info_dict['pre_synapse_ids'].append(pre_syn_id)\n",
    "    info_dict['post_synapse_ids'].append(post_syn_id)\n",
    "    info_dict['pre_confidence'].append(pre_info[2])\n",
    "    info_dict['post_confidence'].append(post_info[2])\n",
    "    info_dict['pre_x'].append(pre_info[3])\n",
    "    info_dict['pre_y'].append(pre_info[4])\n",
    "    info_dict['pre_z'].append(pre_info[5])\n",
    "    info_dict['post_x'].append(post_info[3])\n",
    "    info_dict['post_y'].append(post_info[4])\n",
    "    info_dict['post_z'].append(post_info[5])\n",
    "\n",
    "with open('synapses.csv', 'w') as f:\n",
    "    writer = csv.writer(f)\n",
    "    writer.writerow(['pre_id','post_id','N','pre_confidence','post_confidence',\\\n",
    "                     'pre_x','pre_y','pre_z','post_x','post_y','post_z'])\n",
    "    for pre, k in tqdm(synapse_dict.items()):\n",
    "        for post, v in k.items():\n",
    "            writer.writerow([pre, post, len(v['pre_x']), str(v['pre_confidence']), \\\n",
    "                             str(v['post_confidence']), str(v['pre_x']), str(v['pre_y']), str(v['pre_z']), \\\n",
    "                             str(v['post_x']), str(v['post_y']), str(v['post_z'])])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data to NeuroArch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "medulla = na.NeuroArch('medulla', mode = 'o', version = \"2.0.1\",\n",
    "                       maintainer_name = \"\", maintainer_email = \"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "species = medulla.add_Species('Drosophila melanogaster', stage = 'adult',\n",
    "                                sex = 'female',\n",
    "                                synonyms = ['fruit fly', 'common fruit fly', 'vinegar fly'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "version = 'fib25'\n",
    "datasource = medulla.add_DataSource('Medulla7column', version = version,\n",
    "                                      url = 'https://www.janelia.org/project-team/flyem/research/previous-connectomes-analyzed/seven-column-connectome-fib-sem',\n",
    "                                      description = 'data obtained from https://github.com/connectome-neuprint/neuPrint/blob/922a107df827a2fedd671438595603c4d15eafa7/fib25_neo4j_inputs.zip; neuron skeleton from https://github.com/janelia-flyem/ConnectomeHackathon2015',\n",
    "                                      species = species)\n",
    "medulla.default_DataSource = datasource\n",
    "transcriptome_datasource = medulla.add_DataSource('GSE116969', version = '1.0',\n",
    "                                                  url = 'https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=GSE116969',\n",
    "                                                  description = 'Fred P Davis, Aljoscha Nern, Serge Picard, Michael B Reiser, Gerald M Rubin, Sean R Eddy, Gilbert L Henry, A genetic, genomic, and computational resource for exploring neural circuit function. eLife 2020;9:e50901. DOI: 10.7554/eLife.50901',\n",
    "                                                  species = species)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "medulla.add_Neuropil('MED(L)', synonyms = ['left medulla'])\n",
    "\n",
    "for i in range(1, 11):\n",
    "    medulla.add_Subregion('MED-M{}(L)'.format(i),\\\n",
    "                          synonyms = ['left medulla M{} stratum'.format(i), 'left medulla stratum M{}'.format(i)],\n",
    "                          neuropil = 'MED(L)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt_df = pd.read_csv('GSE116969_dataTable7b.genes_x_cells_p_expression.modeled_genes.txt', sep = '\\t', index_col = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_list = pd.read_csv('neurons.csv')\n",
    "new_swc_dir = 'swc_new'\n",
    "swc_dir = 'swc'\n",
    "uname_dict = {}\n",
    "columns = {}\n",
    "new_morph = []\n",
    "old_morph = []\n",
    "to_combine = []\n",
    "\n",
    "added_neurons = []\n",
    "added_fragments = []\n",
    "\n",
    "\n",
    "for i, row in tqdm(neuron_list.iterrows()):\n",
    "    bodyID = row['bodyID']\n",
    "    cell_type = row['type']\n",
    "    name = row['instance']\n",
    "    segment = False\n",
    "    \n",
    "    if not isinstance(name, str):\n",
    "        if isinstance(cell_type, str):\n",
    "            name = '{}-{}'.format(cell_type, bodyID)\n",
    "        else:\n",
    "            segment = True\n",
    "    else:\n",
    "        if ' home' in name:\n",
    "            name = name.replace(' home', '-home')\n",
    "        if not isinstance(cell_type, str):\n",
    "            if name.split('-')[0] == 'tan':\n",
    "                cell_type = 'tangential'\n",
    "                if name not in uname_dict:\n",
    "                    uname_dict[name] = 0\n",
    "                uname_dict[name] += 1\n",
    "                name = '{}-{}'.format(name, uname_dict[name])\n",
    "            elif name.split('-')[0].startswith('out'):\n",
    "                cell_type = 'output'\n",
    "            elif name in ['Tm23/24', 'Tm23/24-F', 'Dm8-out', 'TmY16?']:\n",
    "                cell_type = name\n",
    "                if name not in uname_dict:\n",
    "                    uname_dict[name] = 0\n",
    "                uname_dict[name] += 1\n",
    "                name = '{}-{}'.format(name, uname_dict[name])\n",
    "            elif name.split('-')[0].startswith('glia'):\n",
    "                cell_type = 'glia'\n",
    "                if 'glia' not in uname_dict:\n",
    "                    uname_dict['glia'] = 0\n",
    "                uname_dict['glia'] += 1\n",
    "                name = '{}-{}'.format('glia', uname_dict['glia'])\n",
    "            elif name.startswith('sdk'):\n",
    "                cell_type = 'sdk'\n",
    "                if 'sdk' not in uname_dict:\n",
    "                    uname_dict['sdk'] = 0\n",
    "                uname_dict['sdk'] += 1\n",
    "                name = '{}-{}'.format('sdk', uname_dict['sdk'])\n",
    "            else:\n",
    "                cell_type = name.split('-')[0]\n",
    "                if '-' not in name or 'like' in name:\n",
    "                    if name not in uname_dict:\n",
    "                        uname_dict[name] = 0\n",
    "                    uname_dict[name] += 1\n",
    "                    name = '{}-{}'.format(name, uname_dict[name])\n",
    "        else:\n",
    "            if name in ['Tm23/24', 'Tm23/24-F', 'Dm8-out', 'TmY16?']:\n",
    "                if name not in uname_dict:\n",
    "                    uname_dict[name] = 0\n",
    "                uname_dict[name] += 1\n",
    "                name = '{}-{}'.format(name, uname_dict[name])\n",
    "            elif cell_type == 'cell':\n",
    "                if name not in uname_dict:\n",
    "                    uname_dict[name] = 0\n",
    "                uname_dict[name] += 1\n",
    "                name = '{}-{}'.format(name, uname_dict[name])\n",
    "            elif '-' not in name or 'like' in name:\n",
    "                if name not in uname_dict:\n",
    "                    uname_dict[name] = 0\n",
    "                uname_dict[name] += 1\n",
    "                name = '{}-{}'.format(name, uname_dict[name])\n",
    "    if not segment:\n",
    "        info = {}\n",
    "        column = name.split('-')[-1]\n",
    "        if column in ['ant', 'post']:\n",
    "            column = name.split('-')[-2]\n",
    "        circuit = None\n",
    "        if (len(column) == 1 and column.isalpha()) or column == 'home':\n",
    "            circuit = columns.get(column, None)\n",
    "            if circuit is None:\n",
    "                circuit = medulla.add_Circuit('Column {}'.format(column), 'Column', neuropil = 'MED(L)')\n",
    "                columns[column] = circuit\n",
    "    \n",
    "        neurotransmitter = []\n",
    "        if cell_type in nt_df.columns or cell_type == 'R8':\n",
    "            if cell_type in ['R{}'.format(i) for i in range(1,7)]:\n",
    "                gene_expression_type = 'R1_6'\n",
    "            elif cell_type == 'R8':\n",
    "                gene_expression_type = 'R8_Rh5'\n",
    "            else:\n",
    "                gene_expression_type = cell_type\n",
    "\n",
    "            if nt_df[gene_expression_type]['Hdc'] > 0.9:\n",
    "                neurotransmitter.append('histamine')\n",
    "            if nt_df[gene_expression_type]['Gad1'] > 0.9:\n",
    "                neurotransmitter.append('GABA')\n",
    "            if nt_df[gene_expression_type]['VAChT'] > 0.9:\n",
    "                neurotransmitter.append('acetylcholine')\n",
    "            if nt_df[gene_expression_type]['VGlut'] > 0.9:\n",
    "                neurotransmitter.append('glutamate')\n",
    "            if nt_df[gene_expression_type]['ple'] > 0.9 \\\n",
    "                    and nt_df[gene_expression_type]['Ddc'] > 0.9 \\\n",
    "                    and nt_df[gene_expression_type]['Vmat'] > 0.9 \\\n",
    "                    and nt_df[gene_expression_type]['DAT'] > 0.9:\n",
    "                neurotransmitter.append('dopamine')\n",
    "\n",
    "    if os.path.exists('{}/{}.swc'.format(new_swc_dir, bodyID)):\n",
    "        morphology = {'type': 'swc', 'filename': '{}/{}.swc'.format(new_swc_dir, bodyID), 'scale': 0.001*10}\n",
    "        new_morph.append(bodyID)\n",
    "    elif os.path.exists('{}/{}.swc'.format(swc_dir, bodyID)):\n",
    "        morphology = {'type': 'swc', 'filename': '{}/{}.swc'.format(swc_dir, bodyID), 'scale': 0.001*10}\n",
    "        old_morph.append(bodyID)\n",
    "    else:\n",
    "        morphology = None\n",
    "        if segment: # no name, not traced, no morph\n",
    "            to_combine.append(bodyID)\n",
    "            continue\n",
    "        else:\n",
    "            segment = True\n",
    "            if name.lower().startswith('out-'):\n",
    "                to_combine.append(bodyID)\n",
    "                continue\n",
    "    \n",
    "    if isinstance(row['statusLabel'], str):\n",
    "        info['Tracing Status'] = row['statusLabel']\n",
    "    else:\n",
    "        info['Tracing Status'] = 'Untraced'\n",
    "    \n",
    "    if not segment:\n",
    "        medulla.add_Neuron(name, # uname\n",
    "                           cell_type, # name\n",
    "                           referenceId = str(bodyID), #referenceId\n",
    "                           info = info if len(info) else None,\n",
    "                           morphology = morphology,\n",
    "                           neurotransmitters = neurotransmitter if len(neurotransmitter) else None,\n",
    "                           neurotransmitters_datasources = [transcriptome_datasource]*len(neurotransmitter) if len(neurotransmitter) else None,\n",
    "                           circuit = circuit)\n",
    "        added_neurons.append(bodyID)\n",
    "    else:\n",
    "        if cell_type == 'unknown' or not isinstance(cell_type, str):\n",
    "            cell_type = 'segment'\n",
    "            name = 'segment_{}'.format(bodyID)\n",
    "        else:\n",
    "            name = 'segment_{}'.format(name)\n",
    "            if str(bodyID) not in name:\n",
    "                name = 'segment_{}'.format(bodyID)\n",
    "        medulla.add_NeuronFragment(name,\n",
    "                                   cell_type,\n",
    "                                   referenceId = str(bodyID),\n",
    "                                   info = info if len(info) else None,\n",
    "                                   morphology = morphology)\n",
    "        added_fragments.append(bodyID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neurons = medulla.sql_query('select from NeuronAndFragment').node_objs\n",
    "# set the cache so there is no need for database access.\n",
    "for neuron in neurons:\n",
    "    medulla.set('NeuronAndFragment', neuron.uname, neuron, medulla.default_DataSource)\n",
    "neuron_ref_to_obj = {int(neuron.referenceId): neuron for neuron in neurons}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synapse_df = pd.read_csv('synapses.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = {}\n",
    "morph_dict = {}\n",
    "\n",
    "for i, row in tqdm(synapse_df.iterrows()):\n",
    "    pre = int(row['pre_id'])\n",
    "    post = int(row['post_id'])\n",
    "    if pre not in neuron_ref_to_obj:\n",
    "        pre = -1\n",
    "    if post not in neuron_ref_to_obj:\n",
    "        post = -1\n",
    "    if pre == -1 and post == -1:\n",
    "        continue\n",
    "\n",
    "    pre_conf = np.array(eval(row['pre_confidence']))/1e6\n",
    "    post_conf = np.array(eval(row['post_confidence']))/1e6\n",
    "    NHP = np.sum(np.logical_and(post_conf>=0.7, pre_conf>=0.7))\n",
    "    \n",
    "    if pre == -1:\n",
    "        if post not in tmp:\n",
    "            tmp[post] = {}\n",
    "        if 'pre' not in tmp[post]:\n",
    "            tmp[post]['pre'] = {'pre_x': [], 'pre_y': [], 'pre_z': [], 'post_x': [], 'post_y': [], 'post_z': [],\n",
    "                                'pre_confidence': [], 'post_confidence': [],\n",
    "                                'N': 0, 'NHP': 0}\n",
    "        tmp[post]['pre']['pre_x'].append(np.array(eval(row['pre_x']))*0.01)\n",
    "        tmp[post]['pre']['pre_y'].append(np.array(eval(row['pre_y']))*0.01)\n",
    "        tmp[post]['pre']['pre_z'].append(np.array(eval(row['pre_z']))*0.01)\n",
    "        tmp[post]['pre']['post_x'].append(np.array(eval(row['post_x']))*0.01)\n",
    "        tmp[post]['pre']['post_y'].append(np.array(eval(row['post_y']))*0.01)\n",
    "        tmp[post]['pre']['post_z'].append(np.array(eval(row['post_z']))*0.01)\n",
    "        tmp[post]['pre']['pre_confidence'].append(pre_conf)\n",
    "        tmp[post]['pre']['post_confidence'].append(post_conf)\n",
    "        tmp[post]['pre']['N'] += row['N']\n",
    "        tmp[post]['pre']['NHP'] += NHP\n",
    "\n",
    "\n",
    "    elif post == -1:\n",
    "        if pre not in tmp:\n",
    "            tmp[pre] = {}\n",
    "        if 'post' not in tmp[pre]:\n",
    "            tmp[pre]['post'] = {'pre_x': [], 'pre_y': [], 'pre_z': [], 'post_x': [], 'post_y': [], 'post_z': [],\n",
    "                                'pre_confidence': [], 'post_confidence': [],\n",
    "                                'N': 0, 'NHP': 0}\n",
    "        tmp[pre]['post']['pre_x'].append(np.array(eval(row['pre_x']))*0.01)\n",
    "        tmp[pre]['post']['pre_y'].append(np.array(eval(row['pre_y']))*0.01)\n",
    "        tmp[pre]['post']['pre_z'].append(np.array(eval(row['pre_z']))*0.01)\n",
    "        tmp[pre]['post']['post_x'].append(np.array(eval(row['post_x']))*0.01)\n",
    "        tmp[pre]['post']['post_y'].append(np.array(eval(row['post_y']))*0.01)\n",
    "        tmp[pre]['post']['post_z'].append(np.array(eval(row['post_z']))*0.01)\n",
    "        tmp[pre]['post']['pre_confidence'].append(pre_conf)\n",
    "        tmp[pre]['post']['post_confidence'].append(post_conf)\n",
    "        tmp[pre]['post']['N'] += row['N']\n",
    "        tmp[pre]['post']['NHP'] += NHP\n",
    "\n",
    "    else:\n",
    "        pre_neuron = neuron_ref_to_obj[pre]\n",
    "        post_neuron = neuron_ref_to_obj[post]\n",
    "        content = {'type': 'swc'}\n",
    "        content['x'] = (np.array(eval(row['pre_x'])+eval(row['post_x']))/1000.*10).tolist()\n",
    "        content['y'] = (np.array(eval(row['pre_y'])+eval(row['post_y']))/1000.*10).tolist()\n",
    "        content['z'] = (np.array(eval(row['pre_z'])+eval(row['post_z']))/1000.*10).tolist()\n",
    "        content['r'] = [0]*len(content['x'])\n",
    "        content['parent'] = [-1]*(len(content['x'])//2) + [i+1 for i in range(len(content['x'])//2)]\n",
    "        content['identifier'] = [7]*(len(content['x'])//2) + [8]*(len(content['x'])//2)\n",
    "        content['sample'] = [i+1 for i in range(len(content['x']))]\n",
    "        content['confidence'] = pre_conf.tolist() + post_conf.tolist()\n",
    "\n",
    "        synapse = medulla.add_Synapse(pre_neuron, post_neuron, N = row['N'], NHP = NHP)\n",
    "        morph_dict[synapse._id] = {'morphology': content}\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "medulla.flush_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "version = '2.0'\n",
    "species = medulla.sql_query('select from Species').node_objs[0]\n",
    "notional_datasource = medulla.add_DataSource('notional', version = version,\n",
    "                                         species = species)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, v in tqdm(tmp.items()):\n",
    "    if 'pre' in v:\n",
    "        post_neuron = neuron_ref_to_obj[n]\n",
    "        cell_type = 'combined_untraced_segments'\n",
    "        name = 'segments_presynaptic_to_{}'.format(post_neuron.uname)\n",
    "        pre_neuron = medulla.add_NeuronFragment(\n",
    "                                     name,\n",
    "                                     cell_type,\n",
    "                                     data_source = notional_datasource)\n",
    "        \n",
    "        content = {'type': 'swc'}\n",
    "        content['x'] = np.concatenate(v['pre']['pre_x'] + v['pre']['post_x']).tolist()\n",
    "        content['y'] = np.concatenate(v['pre']['pre_y'] + v['pre']['post_y']).tolist()\n",
    "        content['z'] = np.concatenate(v['pre']['pre_z'] + v['pre']['post_z']).tolist()\n",
    "        content['r'] = [0]*len(content['x'])\n",
    "        content['parent'] = [-1]*(len(content['x'])//2) + [i+1 for i in range(len(content['x'])//2)]\n",
    "        content['identifier'] = [7]*(len(content['x'])//2) + [8]*(len(content['x'])//2)\n",
    "        content['sample'] = [i+1 for i in range(len(content['x']))]\n",
    "        content['confidence'] = np.concatenate(v['pre']['pre_confidence'] + v['pre']['post_confidence']).tolist()\n",
    "        synapse = medulla.add_Synapse(pre_neuron, post_neuron, N = v['pre']['N'], NHP = v['pre']['NHP'])\n",
    "        morph_dict[synapse._id] = {'morphology': content}\n",
    "    if 'post' in v:\n",
    "        pre_neuron = neuron_ref_to_obj[n]\n",
    "        cell_type = 'combined_untraced_segments'\n",
    "        name = 'segments_postsynaptic_to_{}'.format(pre_neuron.uname)\n",
    "        post_neuron = medulla.add_NeuronFragment(\n",
    "                                     name,\n",
    "                                     cell_type,\n",
    "                                     data_source = notional_datasource)\n",
    "        content = {'type': 'swc'}\n",
    "        content['x'] = np.concatenate(v['post']['pre_x'] + v['post']['post_x']).tolist()\n",
    "        content['y'] = np.concatenate(v['post']['pre_y'] + v['post']['post_y']).tolist()\n",
    "        content['z'] = np.concatenate(v['post']['pre_z'] + v['post']['post_z']).tolist()\n",
    "        content['r'] = [0]*len(content['x'])\n",
    "        content['parent'] = [-1]*(len(content['x'])//2) + [i+1 for i in range(len(content['x'])//2)]\n",
    "        content['identifier'] = [7]*(len(content['x'])//2) + [8]*(len(content['x'])//2)\n",
    "        content['sample'] = [i+1 for i in range(len(content['x']))]\n",
    "        content['confidence'] = np.concatenate(v['post']['pre_confidence'] + v['post']['post_confidence']).tolist()\n",
    "        synapse = medulla.add_Synapse(pre_neuron, post_neuron, N = v['post']['N'], NHP = v['post']['NHP'])\n",
    "        morph_dict[synapse._id] = {'morphology': content}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "medulla.flush_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for rid, data in tqdm(morph_dict.items()):\n",
    "    if 'morphology' in data:\n",
    "        medulla.add_morphology(rid, data['morphology'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure out Arborization Data from Loaded Synapse Positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def strata_arborization(z):\n",
    "    sep = np.array([20, 28, 35, 38, 42, 44, 55, 58, 70])\n",
    "    return Counter(np.digitize(z, sep))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neurons = medulla.sql_query('select from NeuronAndFragment').node_objs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_dict = {}\n",
    "\n",
    "for neuron in neurons:\n",
    "    neuron_dict[neuron.uname] = {'axons': Counter(), 'dendrites': Counter(), 'obj': neuron} \n",
    "    \n",
    "for neuron in tqdm(neurons):\n",
    "    outgoing_synapses = medulla.sql_query(\"\"\"select expand(out('SendsTo')) from {}\"\"\".format(neuron._id)).node_objs\n",
    "    for synapse in outgoing_synapses:\n",
    "        morphology = [n for n in synapse.out('HasData') if isinstance(n, models.MorphologyData)][0]\n",
    "        arborization = []\n",
    "        arborization.append({'type': 'neuropil',\n",
    "                             'synapses': {'MED(L)': len(morphology.x)}})\n",
    "        s = strata_arborization(morphology.z[:(len(morphology.z)//2)])\n",
    "        arborization.append({'type': 'subregion',\n",
    "                             'synapses': {'MED-M{}(L)'.format(k+1): v for k, v in s.items()}})\n",
    "        neuron_dict[neuron.uname]['axons'] += s\n",
    "        try:\n",
    "            neuron_dict[synapse.out('SendsTo')[0].uname]['dendrites'] += s\n",
    "        except KeyError:\n",
    "            pass\n",
    "        medulla.add_synapse_arborization(synapse, arborization)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, neuron in tqdm(enumerate(neurons)):\n",
    "    if not isinstance(neuron, models.Neuron):\n",
    "        continue\n",
    "    arborization = []\n",
    "    arborization.append({'type': 'neuropil',\n",
    "                         'dendrites': {'MED(L)': sum(neuron_dict[neuron.uname]['dendrites'].values())},\n",
    "                         'axons': {'MED(L)': sum(neuron_dict[neuron.uname]['axons'].values())}})\n",
    "    arborization.append({'type': 'subregion',\n",
    "                         'dendrites': {'MED-M{}(L)'.format(k+1): v for k, v in neuron_dict[neuron.uname]['dendrites'].items()},\n",
    "                         'axons': {'MED-M{}(L)'.format(k+1): v for k, v in neuron_dict[neuron.uname]['axons'].items()}})\n",
    "    medulla.add_neuron_arborization(neuron, arborization)\n",
    "    \n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}