notebooks/generate_templates_and_recordings_drift.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating templates and recordings with non-rigid drifts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to generate templates and recordings from scratch using MEArec.\n",
    "\n",
    "The last part of the notebook showcases some plotting routines available through the package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import MEArec as mr\n",
    "import MEAutility as mu\n",
    "import yaml\n",
    "from pprint import pprint\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load default configuration files\n",
    "\n",
    "First, let's load the default configuration of MEArec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "default_info, mearec_home = mr.get_default_config()\n",
    "pprint(default_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating and saving templates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define cell_models folder\n",
    "cell_folder = default_info['cell_models_folder']\n",
    "template_params = mr.get_default_templates_params()\n",
    "pprint(template_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's change a few parameters and generate templates. We need to generate templates with drift. In order to use the advanced drift features, we have to make sure that all templates drift roughly in the same direction and by the same distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_params[\"drifting\"] = True\n",
    "template_params[\"drift_steps\"] = 30\n",
    "# this ensures that all cells drift on the same z trajectory, with a small xy variation\n",
    "template_params[\"drift_xlim\"] = [-5, 5]\n",
    "template_params[\"drift_ylim\"] = [-5, 5]\n",
    "template_params[\"drift_zlim\"] = [100, 100]\n",
    "template_params[\"max_drift\"] = 200\n",
    "\n",
    "# let's generate 50 templates per cell models (total 650 templates)\n",
    "template_params[\"n\"] = 20\n",
    "template_params['probe'] = 'Neuronexus-32'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "templates_path = Path(\"data/test_drift_templates.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not templates_path.is_file():\n",
    "    # the templates are not saved, but the intracellular simulations are saved in 'templates_folder'\n",
    "    tempgen = mr.gen_templates(cell_models_folder=cell_folder, params=template_params, n_jobs=13, verbose=1)\n",
    "    # this will take a few minutes...\n",
    "else:\n",
    "    tempgen = mr.load_templates(templates_path, return_h5_objects=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `tempgen` variable is a `TemplateGenerator` object. It contains the `templates`, `locations`, `rotations`, and `celltypes` of the generated templates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Templates shape', tempgen.templates.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot locations\n",
    "probe = mu.return_mea(info=tempgen.info[\"electrodes\"])\n",
    "\n",
    "ax_probe = mu.plot_probe(probe)\n",
    "for loc in tempgen.locations[::5]:\n",
    "    ax_probe.plot([loc[0, 1], loc[-1, 1]], [loc[0, 2], loc[-1, 2]], alpha=0.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now save the `TemplateGenerator` object in h5 format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not templates_path.is_file():\n",
    "    # save templates in h5 format\n",
    "    mr.save_template_generator(tempgen, filename=templates_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating and saving drift recordings with 1 drift signal\n",
    "\n",
    "Once the templates have been generated, we can use them to generate recordings. Let's fisrt load and take a look at the default parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recordings_params = mr.get_default_recordings_params()\n",
    "pprint(recordings_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The default drift parameters can be used to generate a **single drift signal**. We'll see at the end of the notebook how to generate more complex drifting modes.\n",
    "\n",
    "Similarly to the templates generation, we can change th eparameters that we pass to the `gen_recordings` function.\n",
    "In this case we will keep the default parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10 min\n",
    "recordings_params[\"spiketrains\"][\"duration\"] = 600\n",
    "\n",
    "# 100 Excitatory, 20 inhibitory (the main difference is morphology and avg firing rates)\n",
    "recordings_params[\"spiketrains\"][\"n_exc\"] = 15\n",
    "recordings_params[\"spiketrains\"][\"n_inh\"] = 5\n",
    "\n",
    "# set template selection params\n",
    "recordings_params[\"templates\"][\"min_amp\"] = 30\n",
    "recordings_params[\"templates\"][\"min_dist\"] = 20 # um \n",
    "\n",
    "# other settings\n",
    "recordings_params[\"recordings\"][\"filter\"] = True\n",
    "\n",
    "# noise level and model\n",
    "recordings_params[\"recordings\"][\"noise_level\"] = 10\n",
    "recordings_params[\"recordings\"][\"noise_mode\"] = \"distance-correlated\"\n",
    "\n",
    "# set chunk duration (IMPORTANT for RAM usage and parallelization)\n",
    "recordings_params[\"recordings\"][\"chunk_duration\"] = 10\n",
    "\n",
    "# drifting option\n",
    "recordings_params[\"recordings\"][\"drifting\"] = True\n",
    "recordings_params[\"recordings\"][\"slow_drift_velocity\"] = 30\n",
    "recordings_params[\"recordings\"][\"slow_drift_amplitude\"] = 30\n",
    "recordings_params[\"recordings\"][\"t_start_drift\"] = 100\n",
    "recordings_params[\"recordings\"][\"t_end_drift\"] = 500\n",
    "\n",
    "recordings_params[\"recordings\"][\"drift_mode_probe\"] = 'rigid'\n",
    "\n",
    "\n",
    "recordings_params[\"recordings\"][\"filter\"] = False\n",
    "\n",
    "\n",
    "# (optional) set seeds for reproducibility \n",
    "# (e.g. if you want to maintain underlying activity, but change e.g. noise level)\n",
    "recordings_params['seeds']['spiketrains'] = None\n",
    "recordings_params['seeds']['templates'] = None\n",
    "recordings_params['seeds']['convolution'] = None\n",
    "recordings_params['seeds']['noise'] = None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "recgen_rigid = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                 params=recordings_params, verbose=True, \n",
    "                                 n_jobs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_rigid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recordings_params[\"recordings\"][\"drift_mode_probe\"] = 'non-rigid'\n",
    "recordings_params[\"recordings\"][\"non_rigid_gradient_mode\"] = 'linear'\n",
    "recordings_params[\"recordings\"][\"non_rigid_linear_direction\"] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen_non_rigid = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                     params=recordings_params, verbose=True,\n",
    "                                     n_jobs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(recgen_rigid.drift_list[0]['drift_vector_idxs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_non_rigid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recordings_params[\"recordings\"][\"drift_mode_probe\"] = 'non-rigid'\n",
    "recordings_params[\"recordings\"][\"non_rigid_gradient_mode\"] = 'linear'\n",
    "recordings_params[\"recordings\"][\"non_rigid_linear_direction\"] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen_non_rigid_rev = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                         params=recordings_params, verbose=True,\n",
    "                                         n_jobs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_non_rigid_rev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recordings_params[\"recordings\"][\"drift_mode_probe\"] = 'non-rigid'\n",
    "recordings_params[\"recordings\"][\"non_rigid_gradient_mode\"] = 'step'\n",
    "recordings_params[\"recordings\"][\"non_rigid_step_factors\"] = (1, 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen_non_rigid_step = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                          params=recordings_params, verbose=True,\n",
    "                                          n_jobs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_non_rigid_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save recordings in h5 format\n",
    "mr.save_recording_generator(recgen_rigid, filename='data/test_drift_recordings_rigid.h5')\n",
    "mr.save_recording_generator(recgen_non_rigid, filename='data/test_drift_recordings_nonrigid.h5')\n",
    "mr.save_recording_generator(recgen_non_rigid, filename='data/test_drift_recordings_nonrigid_rev.h5')\n",
    "mr.save_recording_generator(recgen_non_rigid, filename='data/test_drift_recordings_nonrigid_step.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen_rigid_loaded = mr.load_recordings('data/test_drift_recordings_rigid.h5')\n",
    "recgen_non_rigid_loaded = mr.load_recordings('data/test_drift_recordings_nonrigid.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use multiple drift signals\n",
    "\n",
    "In order to construct more complicated drifts, we can use a combination of multiple drift signals.\n",
    "In this example, we show how to combine 3 drift signals:\n",
    "1. a slow rigid drift with velocity 10um/min\n",
    "2. a fast rigi drift with events every 60s\n",
    "3. an additional non-rigid fast sine oscillation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_slow = mr.get_default_drift_dict()\n",
    "drift_dict_fast = mr.get_default_drift_dict()\n",
    "drift_dict_slow_nonrigid = mr.get_default_drift_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_slow[\"drift_mode_probe\"] = \"rigid\"\n",
    "drift_dict_slow[\"drift_fs\"] = 5\n",
    "drift_dict_slow[\"slow_drift_velocity\"] = 10\n",
    "drift_dict_slow[\"slow_drift_amplitude\"] = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_fast[\"drift_mode_speed\"] = \"fast\"\n",
    "drift_dict_fast[\"fast_drift_period\"] = 60\n",
    "drift_dict_fast[\"fast_drift_max_jump\"] = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_slow_nonrigid[\"drift_mode_probe\"] = \"non-rigid\"\n",
    "drift_dict_slow_nonrigid[\"drift_mode_speed\"] = \"slow\"\n",
    "drift_dict_slow_nonrigid[\"slow_drift_waveform\"] = \"sine\"\n",
    "drift_dict_slow_nonrigid[\"slow_drift_velocity\"] = 80\n",
    "drift_dict_slow_nonrigid[\"slow_drift_amplitude\"] = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dicts = [drift_dict_slow, drift_dict_fast, drift_dict_slow_nonrigid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "recgen_mixed = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                 params=recordings_params, verbose=True,\n",
    "                                 n_jobs=1, drift_dicts=drift_dicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_mixed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.save_recording_generator(recgen_mixed, filename='data/test_drift_recordings_mixed.h5')\n",
    "recgen_mixed_loaded = mr.load_recordings('data/test_drift_recordings_mixed.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recgen_mixed.template_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use an external drift signal\n",
    "\n",
    "In this section we combine a slow drift with an externally-defined fast drift with an exponential decay shape.\n",
    "The fast drift is only transient and is applied to neurons only in a section of the probes. To do that, we use `template_ids` to pre-select templates and we define `external_drift_vectors` based on their locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_external = mr.get_default_drift_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we define a time vector from 100 tp 150 seconds, with a 0.1s period\n",
    "drift_times = np.arange(100, 150, 0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau = 5\n",
    "amplitude = 30\n",
    "drift_vector_um = np.zeros_like(drift_times)\n",
    "# let's leave the first 50 samples as zeros\n",
    "start_index = 50\n",
    "drift_vector_um[start_index:] = amplitude * np.exp(-(drift_times[:-start_index] - drift_times[0]) / tau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(drift_times, drift_vector_um)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we compute drift vectors. We will apply this drift only to cells between 0-50 um in depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_ids = recgen_mixed.template_ids\n",
    "template_locations = tempgen.locations[template_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_locations.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_steps = template_locations.shape[1]\n",
    "# start position is half the drift range\n",
    "drift_direction = 2 # depth\n",
    "drift_factors = np.zeros(len(template_ids))\n",
    "mask = np.logical_and(template_locations[:, drift_steps // 2, drift_direction] <= 50,\n",
    "                      template_locations[:, drift_steps // 2, drift_direction] >= 0) \n",
    "drift_factors[mask] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dict_external[\"external_drift_times\"] = drift_times\n",
    "drift_dict_external[\"external_drift_vector_um\"] = drift_vector_um\n",
    "drift_dict_external[\"external_drift_factors\"] = drift_factors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_dicts = [drift_dict_slow, drift_dict_external]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "recgen_mixed_ext = mr.gen_recordings(templates='data/test_drift_templates.h5', \n",
    "                                     params=recordings_params, verbose=True,\n",
    "                                     n_jobs=10, drift_dicts=drift_dicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mr.plot_cell_drifts(recgen_mixed_ext)"
   ]
  }
 ],
 "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
}