notebooks/generate_recordings_with_allen_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate custom template library and recordings with models from the Allen Institute of Brain Science"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we show how to build a custom template library using cell models from the Allen Institute database.\n",
    "\n",
    "In particuar, we downloaded 3 cell models:\n",
    "\n",
    "- Cell ID [488695444](https://celltypes.brain-map.org/experiment/electrophysiology/488695444) (spiny - excitatory)\n",
    "- Cell ID [488680211](https://celltypes.brain-map.org/experiment/electrophysiology/488680211) (spiny - excitatory)\n",
    "- Cell ID [487667205](https://celltypes.brain-map.org/experiment/electrophysiology/487667205) (aspiny - inhibitory)\n",
    "\n",
    "To get the models, select `Select neuronal model` -> `Biophysical - perisomatic` (or `Biophysical - all active`, if available), download the zip file, and unzip it in a folder with the same name as the cell (e.g. neuronal_model_491623973) in the `allen_models` folder in the working directory. Note that the actual model id might be different than the cell id! In this notebook we assume that the created folder has the model ID (same ID as the downloaded zip)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import MEArec as mr\n",
    "import MEAutility as mu\n",
    "from pathlib import Path\n",
    "import neuron\n",
    "import LFPy\n",
    "import os\n",
    "import json\n",
    "import matplotlib.pylab as plt\n",
    "from pprint import pprint\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating custom templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to define a function to load the cell models in LFPy. The following function takes the cell folder as input, and it returns an `LFPy.Cell` object. \n",
    "In order to be used in MEArec to load custom models, the function needs to have the following arguments:\n",
    "\n",
    "- cell_model_folder: path to cell model folder (str)\n",
    "- dt: sampling period in s (float)\n",
    "- start_T: start time of NEURON simulation in s (default 0)\n",
    "- end_T: end time of NEURON simulation in s (default 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to load Allen cells in LFPy\n",
    "def return_allen_cell(cell_model_folder, dt=2**-5, start_T=0, end_T=1):    \n",
    "    cell_model_folder = Path(cell_model_folder)\n",
    "    cwd = os.getcwd()\n",
    "    os.chdir(cell_model_folder)\n",
    "    \n",
    "    # compile mechanisms\n",
    "    mod_folder = \"modfiles\"\n",
    "    os.chdir(mod_folder)\n",
    "    os.system('nrnivmodl')\n",
    "    os.chdir('..')\n",
    "    neuron.load_mechanisms(mod_folder)\n",
    "    params = json.load(open(\"fit_parameters.json\", 'r'))\n",
    "\n",
    "    celsius = params[\"conditions\"][0][\"celsius\"]\n",
    "    reversal_potentials = params[\"conditions\"][0][\"erev\"]\n",
    "    v_init = params[\"conditions\"][0][\"v_init\"]\n",
    "    active_mechs = params[\"genome\"]\n",
    "    neuron.h.celsius = celsius\n",
    "\n",
    "    cell_parameters = {\n",
    "        'morphology': 'reconstruction.swc',\n",
    "        'v_init': v_init,  # initial membrane potential\n",
    "        'passive': False,  # turn on NEURONs passive mechanism for all sections\n",
    "        'nsegs_method': 'lambda_f',  # spatial discretization method\n",
    "        'lambda_f': 200.,  # frequency where length constants are computed\n",
    "        'dt': dt,  # simulation time step size\n",
    "        'tstart': start_T,  # start time of simulation, recorders start at t=0\n",
    "        'tstop': end_T,  # stop simulation at 100 ms.\n",
    "    }\n",
    "\n",
    "    cell = LFPy.Cell(**cell_parameters)\n",
    "\n",
    "    for sec in neuron.h.allsec():\n",
    "        sec.insert(\"pas\")\n",
    "        sectype = sec.name().split(\"[\")[0]\n",
    "        for sec_dict in active_mechs:\n",
    "            if sec_dict[\"section\"] == sectype:\n",
    "                # print(sectype, sec_dict)\n",
    "                if not sec_dict[\"mechanism\"] == \"\":\n",
    "                    sec.insert(sec_dict[\"mechanism\"])\n",
    "                exec (\"sec.{} = {}\".format(sec_dict[\"name\"], sec_dict[\"value\"]))\n",
    "\n",
    "        for sec_dict in reversal_potentials:\n",
    "            if sec_dict[\"section\"] == sectype:\n",
    "                # print(sectype, sec_dict)\n",
    "                for key in sec_dict.keys():\n",
    "                    if not key == \"section\":\n",
    "                        exec (\"sec.{} = {}\".format(key, sec_dict[key]))\n",
    "    \n",
    "    os.chdir(cwd)\n",
    "\n",
    "    return cell"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also define a convenient function to plot the projections of the loaded cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cell_projections(cell):\n",
    "    fig = plt.figure()\n",
    "    ax_xy = fig.add_subplot(2,2,1)\n",
    "    ax_xz = fig.add_subplot(2,2,2)    \n",
    "    ax_yz = fig.add_subplot(2,2,3)    \n",
    "    \n",
    "    for i, (xs, xe, ys, ye, zs, ze) in enumerate(zip(cell.xstart, cell.xend, cell.ystart, cell.yend, cell.zstart, cell.zend)):\n",
    "        if i in cell.get_idx('soma'):\n",
    "            ax_xy.plot([xs, xe], [ys, ye], color='k', lw=5)\n",
    "            ax_xz.plot([xs, xe], [zs, ze], color='k', lw=5)\n",
    "            ax_yz.plot([ys, ye], [zs, ze], color='k', lw=5)\n",
    "        else:\n",
    "            ax_xy.plot([xs, xe], [ys, ye], color='k')\n",
    "            ax_xz.plot([xs, xe], [zs, ze], color='k')\n",
    "            ax_yz.plot([ys, ye], [zs, ze], color='k')\n",
    "        \n",
    "    ax_xy.axis('equal')\n",
    "    ax_xz.axis('equal')\n",
    "    ax_yz.axis('equal')\n",
    "    ax_xy.set_xlabel('x')\n",
    "    ax_xy.set_ylabel('y')\n",
    "    ax_xz.set_xlabel('x')\n",
    "    ax_xz.set_ylabel('z')\n",
    "    ax_yz.set_xlabel('y')\n",
    "    ax_yz.set_ylabel('z')\n",
    "\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's test that the cell model is loaded and run properly by our newly defined function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_folder = 'allen_models/neuronal_model_488462965/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell = return_allen_cell(cell_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot_cell_projections(cell)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! The cell is loaded properly. Let's now test that the simulation is actually working.\n",
    "In order to simulate a few spikes that will be used to simulate extracellular action potentials, we can use the `run_cell_model` function. By default, this function runs models from the Blue Brain Project repository, but we can use the `custom_return_cell_function` argument to load and simulate and arbitrary cell model. This function simulates the cell and can return the `LFPy.Cell` object, the somatic membrane potential, and the transmembrane currents for all compartments (when `save` is set to `False`.\n",
    "\n",
    "We first need to retrieve some parameters for template generation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_params = mr.get_default_templates_params()\n",
    "template_params['seed'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(template_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell, v, i = mr.run_cell_model(cell_folder, verbose=True, save=False, \n",
    "                               custom_return_cell_function=return_allen_cell, \n",
    "                               **template_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now plot the the somatic membrane potential and transmembrane current (the soma is compartment 0) for each spike:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax_v = fig.add_subplot(1,2,1)\n",
    "ax_i = fig.add_subplot(1,2,2)\n",
    "_ = ax_v.plot(v.T)\n",
    "_ = ax_i.plot(i[:, 0].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulating extracellular action potentials (EAPs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to define other parameters for the extracellular simulation. In this case, we will randomly rotate the cells in 3D and generate 10 templates (at random locations) for each cell model.\n",
    "Note that the `physrot` rotation is only implemented for BBP models, as it reproduces a specific physiological rotation depending on the cell type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_params['rot'] = '3drot'\n",
    "template_params['n'] = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can choose the probe that we want to use. In this case we'll use the default `Neuronexus-32` probe, with 32 electrodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(template_params['probe'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to list available probes we can use either the `MEAutility` package or the `MEArec` command line interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mu.return_mea_list())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mearec available-probes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `simulate_templates_one_cell` simulates and returns EAPs, locations of the soma, and rotations applied to each cell before computing the EAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eaps, locs, rots = mr.simulate_templates_one_cell(cell_folder, \n",
    "                                                  intra_save_folder='allen_sim', params=template_params,\n",
    "                                                  verbose=True, custom_return_cell_function=return_allen_cell)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `eaps` have a shape of (n_templates, n_electrodes, n_timepoints). The `locs` and `rots` have a shape of (n_templates, 3)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(eaps.shape)\n",
    "print(locs.shape)\n",
    "print(rots.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating EAPs for all cell models and assembling a template library"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now loop through all available cell models and build a template library. In order to do that, we also have to provide the cell type, that we can access from the `json` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_models = [p for p in Path('allen_models/').iterdir()]\n",
    "print(cell_models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we now the cell type of the different cells, let's build a dictionary to easily retrieve and save the cell type. At this point, we can choose how we want to characterize excitatory and inhibitory calls. Note that this information will need to be passed to the recording generation phase. For cells from the Allen Institute database we can use \"spiny\" for exctitatory and \"aspiny\" for inhibitory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_types = {'488462965': 'spiny', '489932682': 'spiny', '489932435': 'aspiny'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's initialize some variables that will contain our EAPs, locations, rotations, and cell_types:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "templates, template_locations, template_rotations, template_celltypes = [], [], [], []\n",
    "\n",
    "for cell in cell_models:\n",
    "    eaps, locs, rots = mr.simulate_templates_one_cell(cell_folder, intra_save_folder='allen_sim', \n",
    "                                                      params=template_params, verbose=True, \n",
    "                                                      custom_return_cell_function=return_allen_cell)\n",
    "    # find cell type\n",
    "    cell_type = None\n",
    "    for k, v in cell_types.items():\n",
    "        if k in str(cell):\n",
    "            cell_type = v\n",
    "            break\n",
    "    print(\"Cell\", cell, \"is\", cell_type)\n",
    "    \n",
    "    # if first cell, initialize the arrays\n",
    "    if len(templates) == 0:\n",
    "        templates = eaps\n",
    "        template_locations = locs\n",
    "        template_rotations = rots\n",
    "        template_celltypes = np.array([cell_type]*len(eaps))\n",
    "    else:\n",
    "        templates = np.vstack((templates, eaps))\n",
    "        template_locations = np.vstack((template_locations, locs))\n",
    "        template_rotations = np.vstack((template_rotations, rots))\n",
    "        template_celltypes = np.concatenate((template_celltypes, np.array([cell_type]*len(eaps))))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(templates.shape)\n",
    "print(template_locations.shape)\n",
    "print(template_rotations.shape)\n",
    "print(template_celltypes.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now build a `TemplateGenerator` object that can be stored as an `h5` file and used to simulate recordings. \n",
    "\n",
    "We first need to create two dictionaries, `temp_dict` and `info`, containing the templates and related information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_dict = {'templates': templates, \n",
    "             'locations': template_locations, \n",
    "             'rotations': template_rotations,\n",
    "             'celltypes': template_celltypes}\n",
    "info = {}\n",
    "info['params'] = template_params\n",
    "info['electrodes'] = mu.return_mea_info(template_params['probe'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can instantiate a `TemplateGenerator` object. We can alsu use the `MEArec` built-in plot functions to inspect the templates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempgen = mr.TemplateGenerator(temp_dict=temp_dict, info=info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.plot_templates(tempgen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can save the template library so that we can easily use it to build recordings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.save_template_generator(tempgen=tempgen, filename='allen/templates_allen.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating recordings\n",
    "\n",
    "We can now use the Allen template library to assemble recordings. In order to do so, we have to change the recording parameters `cell_types` to tell the simulator that \"spiny\" means excitatory and \"aspiny\" means inhibitory cells:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_params = mr.get_default_recordings_params()\n",
    "pprint(rec_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_params['cell_types'] = {'excitatory': ['spiny'], 'inhibitory': ['aspiny']}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now simulate a 30-s recording, with 5 excitatory cells and 2 inhibitory cells. Since we have only 30 templates in total (it's advised to simulate many more templates per cell to be able to generate multiple different recordings), we will reduce the minimum distance between cells to 5um and use a minimum amplitude of 30uV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_params['spiketrains']['duration'] = 30\n",
    "rec_params['spiketrains']['n_exc'] = 5\n",
    "rec_params['spiketrains']['n_inh'] = 2\n",
    "rec_params['templates']['min_dist'] = 5\n",
    "rec_params['templates']['min_amp'] = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen = mr.gen_recordings(params=rec_params, templates='allen/templates_allen.h5', verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now plot the generated spike trains, the selected templates, and the recordings!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ax_st = mr.plot_rasters(recgen.spiketrains)\n",
    "ax_temp = mr.plot_templates(recgen)\n",
    "ax_rec = mr.plot_recordings(recgen, start_time=0, end_time=5, overlay_templates=True, lw=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can save the generated recordings in `h5` format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.save_recording_generator(recgen=recgen, filename='allen/recordings_allen.h5')"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}