lfads_for_nlb.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <a href=\"https://githubtocolab.com/felixp8/lfads-nlb-tutorials/blob/main/lfads_for_nlb.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluating LFADS with the Neural Latents Benchmark '21\n",
    "\n",
    "Neural Latents Benchmark '21 (NLB'21) is a benchmark suite aimed at standardizing evaluation of latent variable models of neural spiking activity spanning a variety of tasks and brain areas. The primary objective of the challenge is to infer the firing rates of a set of held-out neurons given the spiking activity of held-in neurons, a procedure called co-smoothing.\n",
    "\n",
    "The benchmark suite features several datasets from experiments spanning a range of behaviors and brain regions, but they are all provided in the standard Neurodata Without Borders format and available on [DANDI](https://dandiarchive.org). The benchmark challenge itself is hosted on the platform [EvalAI](https://eval.ai), where model predictions can be submitted and automatically evaluated on private evaluation data.\n",
    "\n",
    "To facilitate participation in the competition, we provide the code package [`nlb_tools`](https://github.com/neurallatents/nlb_tools), which has functions for data preprocessing and submission preparation.\n",
    "\n",
    "This tutorial will walk through how to use NLB'21 to evaluate LFADS. Specifically, it will walk through:\n",
    "1. **Dataset download** - getting dataset files from DANDI on to your machine\n",
    "2. **Data loading and preprocessing** - using `nlb_tools` to extract the data we expect you to model\n",
    "3. **Modeling neural data** - using LFADS to perform inference on the extracted data\n",
    "4. **Submitting and evaluating model predictions** - packaging predictions for submission to EvalAI or local evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main.png?raw=true\" width=\"480\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 0: Setup\n",
    "\n",
    "First, we need to install and import the packages we need for this notebook. Note that, if running in Colab, you will need to restart the runtime after running the cell the first time and then re-run it for the proper package versions to be installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Import Dependencies\n",
    "\n",
    "# check if running in Google Colab\n",
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "except:\n",
    "    IN_COLAB = False\n",
    "\n",
    "import os\n",
    "\n",
    "# set up environment\n",
    "if IN_COLAB:\n",
    "    ! git clone https://github.com/felixp8/lfads-nlb-tutorials.git\n",
    "    os.chdir('lfads-nlb-tutorials')\n",
    "\n",
    "# nlb_tools\n",
    "if IN_COLAB:\n",
    "    !pip install git+https://github.com/neurallatents/nlb_tools.git \n",
    "else:\n",
    "    try:\n",
    "        from nlb_tools.nwb_interface import NWBDataset\n",
    "    except:\n",
    "        raise ImportError('Module `nlb_tools` not found')\n",
    "\n",
    "# DANDI CLI tool (optional, can use website instead)\n",
    "if IN_COLAB:\n",
    "    !pip install dandi==0.36.0\n",
    "else:\n",
    "    try:\n",
    "        import dandi\n",
    "    except:\n",
    "        raise ImportError('Module `dandi` not found. ' + \n",
    "            'However, it is optional for this tutorial, as you can download the data from the DANDI website directly'\n",
    "        )\n",
    "  \n",
    "if IN_COLAB:\n",
    "    # install necessary packages \n",
    "    ! git clone https://github.com/snel-repo/autolfads-tf2.git\n",
    "    os.chdir('autolfads-tf2')\n",
    "    ! pip install -e lfads-tf2\n",
    "    ! pip install -e tune-tf2\n",
    "    os.chdir('..')\n",
    "    # add lfads_tf2\n",
    "    os.environ['PYTHONPATH'] += \":/content/autolfads-tf2\"\n",
    "else:\n",
    "    try:\n",
    "        from lfads_tf2.utils import restrict_gpu_usage\n",
    "    except:\n",
    "        x = input('autolfads-tf2 is not installed. ' + \n",
    "            'We recommend setting up a conda environment and installing the package. ' + \n",
    "            'However, it can also be installed from the notebook. ' + \n",
    "            'Would you like to install it now? (y/n)'\n",
    "        )\n",
    "        if (x.strip() in ['Y', 'y']):\n",
    "            ! git clone https://github.com/snel-repo/autolfads-tf2.git\n",
    "            os.chdir('autolfads-tf2')\n",
    "            ! pip install -e .\n",
    "            os.chdir('..')\n",
    "            from lfads_tf2.utils import restrict_gpu_usage\n",
    "        else:\n",
    "            raise ImportError('Module `lfads_tf2` not found')\n",
    "\n",
    "import numpy as np\n",
    "import h5py\n",
    "import warnings\n",
    "warnings.simplefilter (action='ignore', category=FutureWarning)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Data Preparation\n",
    "\n",
    "NLB'21 features seven dataset sessions from four experiments, each posing a distinct challenge for latent variable models:\n",
    "* MC_Maze - large motor/premotor cortex dataset with rich, varied, but stereotyped behavior, over a hundred distinct conditions, thousands of trials, and over a hundred neurons per session. Dynamics have been shown to be well-approximated as autonomous, so models would not need to account for external inputs to perform well. We also include three additional sessions with artificially limited numbers of trials to gauge how well models scale to limited amounts of data.\n",
    "* MC_RTT - motor cortex dataset featuring self-paced and more naturalistic behavior. Inputs (new targets) occur irregularly and there are few repeated conditions, providing a valuable testing setting for BCI decoders.\n",
    "* Area2_Bump - sensory dataset displaying input-driven activity in response to both predictable and unpredictable sensory feedback.\n",
    "* DMFC_RSG - frontal cortex dataset during cognitive timing task. Cognitive areas do not have clearly delineated input or output functions/layers, and dynamics are likely a complex mix of input-driven and intrinsic. Additionally, the lack of clear moment-by-moment behavioral correlates makes supervised evaluation difficult."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/felixp8/lfads-nlb-tutorials/blob/main/images/nlb_data.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1. Dataset Download"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_download.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The datasets are available on the platform DANDI. They can be downloaded directly from the website or by using the DANDI CLI tool, as shown below. For this notebook, we will be using the MC_Maze_Large dataset, which is available from [here](https://dandiarchive.org/dandiset/000138). Links to the other datasets can be found on [our website](https://neurallatents.github.io/datasets)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!dandi download https://dandiarchive.org/dandiset/000138"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above line will download two files into the folder `./000138/sub-Jenkins/`. Next, we'll get the path of the downloaded files and list them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sub-Jenkins_ses-large_desc-test_ecephys.nwb',\n",
       " 'sub-Jenkins_ses-large_desc-train_behavior+ecephys.nwb']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "curr_path = os.getcwd()\n",
    "fpath = curr_path + '/000138/sub-Jenkins/'\n",
    "os.listdir(fpath) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The file with 'desc-train' in its name is for training, while the file with 'desc-test' in its name is for final model evaluation. As we take a look at the data, we will see the differences between these two files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2. Dataset Loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_dataload.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the NWB data into Python, we provide the `NWBDataset` class, which can load from the dataset files and perform simple preprocessing operations. To load a dataset, you instantiate an instance of NWBDataset and provide the path to the files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nlb_tools.nwb_interface import NWBDataset\n",
    "dataset = NWBDataset(fpath=fpath) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3. Data Format\n",
    "\n",
    "The loaded data are primarily stored in two pandas DataFrames: `NWBDataset.data` and `NWBDataset.trial_info`.\n",
    "\n",
    "#### `NWBDataset.data`\n",
    "`NWBDataset.data` contains the continuous recorded data, like spike counts and kinematics. Each row consists of measurements taken at a particular timestep. Most importantly, spiking data from held-in units are labeled `spikes` and data from held-out units is labeled `heldout_spikes`.\n",
    "\n",
    "In the training data, all fields are available:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr:last-of-type th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th>signal_type</th>\n",
       "      <th colspan=\"2\" halign=\"left\">cursor_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">eye_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">hand_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">hand_vel</th>\n",
       "      <th colspan=\"2\" halign=\"left\">heldout_spikes</th>\n",
       "      <th>...</th>\n",
       "      <th colspan=\"10\" halign=\"left\">spikes</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>channel</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>1031</th>\n",
       "      <th>1051</th>\n",
       "      <th>...</th>\n",
       "      <th>2741</th>\n",
       "      <th>2743</th>\n",
       "      <th>2761</th>\n",
       "      <th>2771</th>\n",
       "      <th>2781</th>\n",
       "      <th>2791</th>\n",
       "      <th>2801</th>\n",
       "      <th>2881</th>\n",
       "      <th>2941</th>\n",
       "      <th>2951</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>clock_time</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40</th>\n",
       "      <td>-5.200000</td>\n",
       "      <td>3.300000</td>\n",
       "      <td>0.6</td>\n",
       "      <td>1.1</td>\n",
       "      <td>-5.195095</td>\n",
       "      <td>-31.606258</td>\n",
       "      <td>-1.481366</td>\n",
       "      <td>0.261386</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.001000</th>\n",
       "      <td>-5.199120</td>\n",
       "      <td>3.299442</td>\n",
       "      <td>2.5</td>\n",
       "      <td>0.6</td>\n",
       "      <td>-5.196711</td>\n",
       "      <td>-31.605926</td>\n",
       "      <td>-1.727835</td>\n",
       "      <td>0.317635</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.002000</th>\n",
       "      <td>-5.198598</td>\n",
       "      <td>3.299110</td>\n",
       "      <td>2.5</td>\n",
       "      <td>0.4</td>\n",
       "      <td>-5.198551</td>\n",
       "      <td>-31.605623</td>\n",
       "      <td>-1.873343</td>\n",
       "      <td>0.342863</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.003000</th>\n",
       "      <td>-5.198598</td>\n",
       "      <td>3.299110</td>\n",
       "      <td>2.7</td>\n",
       "      <td>0.5</td>\n",
       "      <td>-5.200457</td>\n",
       "      <td>-31.605240</td>\n",
       "      <td>-2.053902</td>\n",
       "      <td>0.334682</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.004000</th>\n",
       "      <td>-5.199120</td>\n",
       "      <td>3.299442</td>\n",
       "      <td>2.8</td>\n",
       "      <td>0.8</td>\n",
       "      <td>-5.202659</td>\n",
       "      <td>-31.604953</td>\n",
       "      <td>-2.238392</td>\n",
       "      <td>0.280908</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.095000</th>\n",
       "      <td>-5.200000</td>\n",
       "      <td>3.300000</td>\n",
       "      <td>1.6</td>\n",
       "      <td>8.8</td>\n",
       "      <td>-5.239548</td>\n",
       "      <td>-31.618091</td>\n",
       "      <td>-1.010222</td>\n",
       "      <td>-1.636778</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.096000</th>\n",
       "      <td>-5.193734</td>\n",
       "      <td>3.308271</td>\n",
       "      <td>1.7</td>\n",
       "      <td>9.1</td>\n",
       "      <td>-5.240553</td>\n",
       "      <td>-31.619778</td>\n",
       "      <td>-0.964223</td>\n",
       "      <td>-1.878082</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.097000</th>\n",
       "      <td>-5.188833</td>\n",
       "      <td>3.311073</td>\n",
       "      <td>0.0</td>\n",
       "      <td>10.0</td>\n",
       "      <td>-5.241476</td>\n",
       "      <td>-31.621848</td>\n",
       "      <td>-0.942786</td>\n",
       "      <td>-2.159794</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.098000</th>\n",
       "      <td>-5.187298</td>\n",
       "      <td>3.309505</td>\n",
       "      <td>0.3</td>\n",
       "      <td>9.7</td>\n",
       "      <td>-5.242439</td>\n",
       "      <td>-31.624097</td>\n",
       "      <td>-1.023439</td>\n",
       "      <td>-2.277696</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:01:40.099000</th>\n",
       "      <td>-5.190729</td>\n",
       "      <td>3.305204</td>\n",
       "      <td>0.2</td>\n",
       "      <td>9.7</td>\n",
       "      <td>-5.243523</td>\n",
       "      <td>-31.626403</td>\n",
       "      <td>-0.917548</td>\n",
       "      <td>-2.321146</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 170 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "signal_type            cursor_pos           eye_pos        hand_pos  \\\n",
       "channel                         x         y       x     y         x   \n",
       "clock_time                                                            \n",
       "0 days 00:01:40         -5.200000  3.300000     0.6   1.1 -5.195095   \n",
       "0 days 00:01:40.001000  -5.199120  3.299442     2.5   0.6 -5.196711   \n",
       "0 days 00:01:40.002000  -5.198598  3.299110     2.5   0.4 -5.198551   \n",
       "0 days 00:01:40.003000  -5.198598  3.299110     2.7   0.5 -5.200457   \n",
       "0 days 00:01:40.004000  -5.199120  3.299442     2.8   0.8 -5.202659   \n",
       "...                           ...       ...     ...   ...       ...   \n",
       "0 days 00:01:40.095000  -5.200000  3.300000     1.6   8.8 -5.239548   \n",
       "0 days 00:01:40.096000  -5.193734  3.308271     1.7   9.1 -5.240553   \n",
       "0 days 00:01:40.097000  -5.188833  3.311073     0.0  10.0 -5.241476   \n",
       "0 days 00:01:40.098000  -5.187298  3.309505     0.3   9.7 -5.242439   \n",
       "0 days 00:01:40.099000  -5.190729  3.305204     0.2   9.7 -5.243523   \n",
       "\n",
       "signal_type                        hand_vel           heldout_spikes       \\\n",
       "channel                         y         x         y           1031 1051   \n",
       "clock_time                                                                  \n",
       "0 days 00:01:40        -31.606258 -1.481366  0.261386            0.0  0.0   \n",
       "0 days 00:01:40.001000 -31.605926 -1.727835  0.317635            0.0  0.0   \n",
       "0 days 00:01:40.002000 -31.605623 -1.873343  0.342863            0.0  0.0   \n",
       "0 days 00:01:40.003000 -31.605240 -2.053902  0.334682            0.0  0.0   \n",
       "0 days 00:01:40.004000 -31.604953 -2.238392  0.280908            0.0  0.0   \n",
       "...                           ...       ...       ...            ...  ...   \n",
       "0 days 00:01:40.095000 -31.618091 -1.010222 -1.636778            0.0  0.0   \n",
       "0 days 00:01:40.096000 -31.619778 -0.964223 -1.878082            0.0  0.0   \n",
       "0 days 00:01:40.097000 -31.621848 -0.942786 -2.159794            0.0  0.0   \n",
       "0 days 00:01:40.098000 -31.624097 -1.023439 -2.277696            0.0  0.0   \n",
       "0 days 00:01:40.099000 -31.626403 -0.917548 -2.321146            0.0  0.0   \n",
       "\n",
       "signal_type             ... spikes                                          \\\n",
       "channel                 ...   2741 2743 2761 2771 2781 2791 2801 2881 2941   \n",
       "clock_time              ...                                                  \n",
       "0 days 00:01:40         ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.001000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.002000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.003000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.004000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "...                     ...    ...  ...  ...  ...  ...  ...  ...  ...  ...   \n",
       "0 days 00:01:40.095000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.096000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.097000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:01:40.098000  ...    0.0  0.0  0.0  0.0  1.0  0.0  0.0  1.0  0.0   \n",
       "0 days 00:01:40.099000  ...    0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0   \n",
       "\n",
       "signal_type                  \n",
       "channel                2951  \n",
       "clock_time                   \n",
       "0 days 00:01:40         0.0  \n",
       "0 days 00:01:40.001000  0.0  \n",
       "0 days 00:01:40.002000  0.0  \n",
       "0 days 00:01:40.003000  0.0  \n",
       "0 days 00:01:40.004000  0.0  \n",
       "...                     ...  \n",
       "0 days 00:01:40.095000  0.0  \n",
       "0 days 00:01:40.096000  0.0  \n",
       "0 days 00:01:40.097000  0.0  \n",
       "0 days 00:01:40.098000  0.0  \n",
       "0 days 00:01:40.099000  0.0  \n",
       "\n",
       "[100 rows x 170 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.data.iloc[100000:100100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the test data, only held-in spikes are available, while other data is concealed with NaNs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe thead tr:last-of-type th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th>signal_type</th>\n",
       "      <th colspan=\"2\" halign=\"left\">cursor_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">eye_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">hand_pos</th>\n",
       "      <th colspan=\"2\" halign=\"left\">hand_vel</th>\n",
       "      <th colspan=\"2\" halign=\"left\">heldout_spikes</th>\n",
       "      <th>...</th>\n",
       "      <th colspan=\"10\" halign=\"left\">spikes</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>channel</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>1031</th>\n",
       "      <th>1051</th>\n",
       "      <th>...</th>\n",
       "      <th>2741</th>\n",
       "      <th>2743</th>\n",
       "      <th>2761</th>\n",
       "      <th>2771</th>\n",
       "      <th>2781</th>\n",
       "      <th>2791</th>\n",
       "      <th>2801</th>\n",
       "      <th>2881</th>\n",
       "      <th>2941</th>\n",
       "      <th>2951</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>clock_time</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.001000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.002000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.003000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.004000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.095000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.096000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.097000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.098000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0 days 00:00:01.099000</th>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 170 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "signal_type            cursor_pos     eye_pos     hand_pos     hand_vel      \\\n",
       "channel                         x   y       x   y        x   y        x   y   \n",
       "clock_time                                                                    \n",
       "0 days 00:00:01               NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.001000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.002000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.003000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.004000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "...                           ...  ..     ...  ..      ...  ..      ...  ..   \n",
       "0 days 00:00:01.095000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.096000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.097000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.098000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "0 days 00:00:01.099000        NaN NaN     NaN NaN      NaN NaN      NaN NaN   \n",
       "\n",
       "signal_type            heldout_spikes       ... spikes                      \\\n",
       "channel                          1031 1051  ...   2741 2743 2761 2771 2781   \n",
       "clock_time                                  ...                              \n",
       "0 days 00:00:01                   NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.001000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.002000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.003000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.004000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "...                               ...  ...  ...    ...  ...  ...  ...  ...   \n",
       "0 days 00:00:01.095000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.096000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.097000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.098000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "0 days 00:00:01.099000            NaN  NaN  ...    0.0  0.0  0.0  0.0  0.0   \n",
       "\n",
       "signal_type                                      \n",
       "channel                2791 2801 2881 2941 2951  \n",
       "clock_time                                       \n",
       "0 days 00:00:01         0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.001000  0.0  0.0  1.0  0.0  0.0  \n",
       "0 days 00:00:01.002000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.003000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.004000  0.0  0.0  0.0  0.0  0.0  \n",
       "...                     ...  ...  ...  ...  ...  \n",
       "0 days 00:00:01.095000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.096000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.097000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.098000  0.0  0.0  0.0  0.0  0.0  \n",
       "0 days 00:00:01.099000  0.0  0.0  0.0  0.0  0.0  \n",
       "\n",
       "[100 rows x 170 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.data.iloc[1000:1100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "#### `NWBDataset.trial_info`\n",
    "Each row of `trial_info` contains information about a particular experimental trial, such as when it begins and ends. As with the `NWBDataset.data`, almost all information is concealed in the test data. The field `split`, common to all of our provided datasets, indicates what data split a trial belongs to (explained in more detail in Section 2.4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>trial_id</th>\n",
       "      <th>start_time</th>\n",
       "      <th>end_time</th>\n",
       "      <th>move_onset_time</th>\n",
       "      <th>split</th>\n",
       "      <th>trial_type</th>\n",
       "      <th>trial_version</th>\n",
       "      <th>maze_id</th>\n",
       "      <th>success</th>\n",
       "      <th>target_on_time</th>\n",
       "      <th>go_cue_time</th>\n",
       "      <th>rt</th>\n",
       "      <th>delay</th>\n",
       "      <th>num_targets</th>\n",
       "      <th>target_pos</th>\n",
       "      <th>num_barriers</th>\n",
       "      <th>barrier_pos</th>\n",
       "      <th>active_target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0 days 00:00:00</td>\n",
       "      <td>0 days 00:00:00.700000</td>\n",
       "      <td>0 days 00:00:00.250000</td>\n",
       "      <td>test</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0 days 00:00:00.800000</td>\n",
       "      <td>0 days 00:00:01.500000</td>\n",
       "      <td>0 days 00:00:01.050000</td>\n",
       "      <td>test</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>0 days 00:00:01.600000</td>\n",
       "      <td>0 days 00:00:02.300000</td>\n",
       "      <td>0 days 00:00:01.850000</td>\n",
       "      <td>test</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>0 days 00:00:02.400000</td>\n",
       "      <td>0 days 00:00:03.100000</td>\n",
       "      <td>0 days 00:00:02.650000</td>\n",
       "      <td>test</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>0 days 00:00:03.200000</td>\n",
       "      <td>0 days 00:00:03.900000</td>\n",
       "      <td>0 days 00:00:03.450000</td>\n",
       "      <td>test</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>595</th>\n",
       "      <td>595</td>\n",
       "      <td>0 days 00:25:49.600000</td>\n",
       "      <td>0 days 00:25:52.636000</td>\n",
       "      <td>0 days 00:25:51.501000</td>\n",
       "      <td>train</td>\n",
       "      <td>8.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>38.0</td>\n",
       "      <td>True</td>\n",
       "      <td>0 days 00:25:50.405000</td>\n",
       "      <td>0 days 00:25:51.153000</td>\n",
       "      <td>348.0</td>\n",
       "      <td>748.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[[-105, 76]]</td>\n",
       "      <td>9.0</td>\n",
       "      <td>[[74, -102, 11, 53], [86, -44, 14, 11], [103, ...</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>596</th>\n",
       "      <td>596</td>\n",
       "      <td>0 days 00:25:52.700000</td>\n",
       "      <td>0 days 00:25:55.746000</td>\n",
       "      <td>0 days 00:25:54.595000</td>\n",
       "      <td>train</td>\n",
       "      <td>11.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>80.0</td>\n",
       "      <td>True</td>\n",
       "      <td>0 days 00:25:53.467000</td>\n",
       "      <td>0 days 00:25:54.116000</td>\n",
       "      <td>479.0</td>\n",
       "      <td>649.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>[[123, -81], [-130, -13], [123, 71]]</td>\n",
       "      <td>8.0</td>\n",
       "      <td>[[-65, -15, 14, 51], [-79, -55, 55, 6], [-103,...</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>597</th>\n",
       "      <td>597</td>\n",
       "      <td>0 days 00:25:55.800000</td>\n",
       "      <td>0 days 00:25:58.801000</td>\n",
       "      <td>0 days 00:25:57.701000</td>\n",
       "      <td>val</td>\n",
       "      <td>7.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>True</td>\n",
       "      <td>0 days 00:25:56.545000</td>\n",
       "      <td>0 days 00:25:57.410000</td>\n",
       "      <td>291.0</td>\n",
       "      <td>865.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>[[124, -79], [103, 83], [-105, 76]]</td>\n",
       "      <td>9.0</td>\n",
       "      <td>[[74, -102, 11, 53], [86, -44, 14, 11], [103, ...</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>598</th>\n",
       "      <td>598</td>\n",
       "      <td>0 days 00:25:58.900000</td>\n",
       "      <td>0 days 00:26:01.956000</td>\n",
       "      <td>0 days 00:26:00.777000</td>\n",
       "      <td>train</td>\n",
       "      <td>7.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>True</td>\n",
       "      <td>0 days 00:25:59.613000</td>\n",
       "      <td>0 days 00:26:00.479000</td>\n",
       "      <td>298.0</td>\n",
       "      <td>866.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>[[103, 83]]</td>\n",
       "      <td>9.0</td>\n",
       "      <td>[[74, -102, 11, 53], [86, -44, 14, 11], [103, ...</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>599</th>\n",
       "      <td>599</td>\n",
       "      <td>0 days 00:26:02</td>\n",
       "      <td>0 days 00:26:05.021000</td>\n",
       "      <td>0 days 00:26:03.837000</td>\n",
       "      <td>train</td>\n",
       "      <td>6.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>36.0</td>\n",
       "      <td>True</td>\n",
       "      <td>0 days 00:26:02.877000</td>\n",
       "      <td>0 days 00:26:03.426000</td>\n",
       "      <td>411.0</td>\n",
       "      <td>549.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>[[124, -79], [103, 83], [-105, 76]]</td>\n",
       "      <td>9.0</td>\n",
       "      <td>[[74, -102, 11, 53], [86, -44, 14, 11], [103, ...</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>600 rows × 18 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     trial_id             start_time               end_time  \\\n",
       "0           0        0 days 00:00:00 0 days 00:00:00.700000   \n",
       "1           1 0 days 00:00:00.800000 0 days 00:00:01.500000   \n",
       "2           2 0 days 00:00:01.600000 0 days 00:00:02.300000   \n",
       "3           3 0 days 00:00:02.400000 0 days 00:00:03.100000   \n",
       "4           4 0 days 00:00:03.200000 0 days 00:00:03.900000   \n",
       "..        ...                    ...                    ...   \n",
       "595       595 0 days 00:25:49.600000 0 days 00:25:52.636000   \n",
       "596       596 0 days 00:25:52.700000 0 days 00:25:55.746000   \n",
       "597       597 0 days 00:25:55.800000 0 days 00:25:58.801000   \n",
       "598       598 0 days 00:25:58.900000 0 days 00:26:01.956000   \n",
       "599       599        0 days 00:26:02 0 days 00:26:05.021000   \n",
       "\n",
       "           move_onset_time  split  trial_type  trial_version  maze_id success  \\\n",
       "0   0 days 00:00:00.250000   test         NaN            NaN      NaN     NaN   \n",
       "1   0 days 00:00:01.050000   test         NaN            NaN      NaN     NaN   \n",
       "2   0 days 00:00:01.850000   test         NaN            NaN      NaN     NaN   \n",
       "3   0 days 00:00:02.650000   test         NaN            NaN      NaN     NaN   \n",
       "4   0 days 00:00:03.450000   test         NaN            NaN      NaN     NaN   \n",
       "..                     ...    ...         ...            ...      ...     ...   \n",
       "595 0 days 00:25:51.501000  train         8.0            1.0     38.0    True   \n",
       "596 0 days 00:25:54.595000  train        11.0            2.0     80.0    True   \n",
       "597 0 days 00:25:57.701000    val         7.0            2.0     37.0    True   \n",
       "598 0 days 00:26:00.777000  train         7.0            1.0     37.0    True   \n",
       "599 0 days 00:26:03.837000  train         6.0            2.0     36.0    True   \n",
       "\n",
       "            target_on_time            go_cue_time     rt  delay  num_targets  \\\n",
       "0                      NaT                    NaT    NaN    NaN          NaN   \n",
       "1                      NaT                    NaT    NaN    NaN          NaN   \n",
       "2                      NaT                    NaT    NaN    NaN          NaN   \n",
       "3                      NaT                    NaT    NaN    NaN          NaN   \n",
       "4                      NaT                    NaT    NaN    NaN          NaN   \n",
       "..                     ...                    ...    ...    ...          ...   \n",
       "595 0 days 00:25:50.405000 0 days 00:25:51.153000  348.0  748.0          1.0   \n",
       "596 0 days 00:25:53.467000 0 days 00:25:54.116000  479.0  649.0          3.0   \n",
       "597 0 days 00:25:56.545000 0 days 00:25:57.410000  291.0  865.0          3.0   \n",
       "598 0 days 00:25:59.613000 0 days 00:26:00.479000  298.0  866.0          1.0   \n",
       "599 0 days 00:26:02.877000 0 days 00:26:03.426000  411.0  549.0          3.0   \n",
       "\n",
       "                               target_pos  num_barriers  \\\n",
       "0                                     NaN           NaN   \n",
       "1                                     NaN           NaN   \n",
       "2                                     NaN           NaN   \n",
       "3                                     NaN           NaN   \n",
       "4                                     NaN           NaN   \n",
       "..                                    ...           ...   \n",
       "595                          [[-105, 76]]           9.0   \n",
       "596  [[123, -81], [-130, -13], [123, 71]]           8.0   \n",
       "597   [[124, -79], [103, 83], [-105, 76]]           9.0   \n",
       "598                           [[103, 83]]           9.0   \n",
       "599   [[124, -79], [103, 83], [-105, 76]]           9.0   \n",
       "\n",
       "                                           barrier_pos  active_target  \n",
       "0                                                  NaN            NaN  \n",
       "1                                                  NaN            NaN  \n",
       "2                                                  NaN            NaN  \n",
       "3                                                  NaN            NaN  \n",
       "4                                                  NaN            NaN  \n",
       "..                                                 ...            ...  \n",
       "595  [[74, -102, 11, 53], [86, -44, 14, 11], [103, ...            0.0  \n",
       "596  [[-65, -15, 14, 51], [-79, -55, 55, 6], [-103,...            2.0  \n",
       "597  [[74, -102, 11, 53], [86, -44, 14, 11], [103, ...            1.0  \n",
       "598  [[74, -102, 11, 53], [86, -44, 14, 11], [103, ...            0.0  \n",
       "599  [[74, -102, 11, 53], [86, -44, 14, 11], [103, ...            0.0  \n",
       "\n",
       "[600 rows x 18 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.trial_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4. Data Splits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/data_splits.png?raw=true\" width=\"480\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset is divided into train, val, and test splits. The train and val splits are contained within the training data file, while the test split is entirely in the test data file. This means that all data fields are available in the train and val splits, but only held-in data is available in the test split. The NLB'21 challenge has two phases based on these splits:\n",
    "1. In the Validation Phase, models will be evaluated on their val split predictions. This phase is offered for sanity checking results and building familiarity with the EvalAI platform.\n",
    "2. In the Test Phase, models will be evaluated on their test split predictions. This is the phase that is displayed on the public leaderboard. In this phase, the val split does not strictly need to be used for model validation, despite its name.\n",
    "\n",
    "In this notebook, we will prepare a submission for the Validation Phase, though the code can be easily modified for the Test Phase."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5. Data Extraction\n",
    "\n",
    "#### 1.5.1. Resampling\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_resample.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The raw data is at 1 ms resolution, but the NLB'21 challenge expects submissions to be at 5 ms resolution, so we will resample the data before doing any other processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data shape: (1565021, 170)\n",
      "Bin width: 1 ms\n",
      "Resampled data shape: (313005, 170)\n",
      "Resampled bin width: 5 ms\n"
     ]
    }
   ],
   "source": [
    "print(f'Data shape: {dataset.data.shape}')\n",
    "print(f'Bin width: {dataset.bin_width} ms')\n",
    "dataset.resample(5)\n",
    "print(f'Resampled data shape: {dataset.data.shape}')\n",
    "print(f'Resampled bin width: {dataset.bin_width} ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5.2. Trial Alignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_alignment.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simplify data preparation, we abstract other formatting steps into the functions `make_train_input_tensors` and `make_eval_input_tensors`. These wrapper functions format the raw data into tensors to be used for model training and evaluation. The primary processing performed by the functions is trial alignment.\n",
    "\n",
    "Trial alignment involves choosing a particular trial event, such as a go cue, and taking a fixed window of data around each occurrence of that event. For all of the datasets in NLB'21, we have chosen trial alignments based on the experimental design and past analyses on the data. For the MC_Maze_Large dataset, the selected trial alignment is 250 ms before to 450 ms after movement onset. Our wrapper functions will apply this alignment when given the correct dataset name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nlb_tools.make_tensors import make_train_input_tensors, make_eval_input_tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_train_input_tensors` extracts the data available for model training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dict = make_train_input_tensors(dataset=dataset, \n",
    "                                      dataset_name='mc_maze_large', \n",
    "                                      trial_split='train', # trial_split=['train', 'val'], for Test phase\n",
    "                                      save_file=False, \n",
    "                                      include_forward_pred=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_eval_input_tensors` extracts the data used to evaluate the model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dict = make_eval_input_tensors(dataset=dataset,\n",
    "                                    dataset_name='mc_maze_large',\n",
    "                                    trial_split='val', # trial_split='test', for Test phase\n",
    "                                    save_file=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_naming.png?raw=true\" width=\"600\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both `make_train_input_tensors` and `make_eval_input_tensors` return dictionaries of tensors. \n",
    "\n",
    "The training dictionary contains:\n",
    "- train_spikes_heldin - spiking activity of held-in units on training trials \n",
    "- train_spikes_heldout - spiking activity of held-out units on training trials\n",
    "- train_spikes_heldin_forward - spiking activity of held-in units immediately after the trial period\n",
    "- train_spikes_heldout_forward - spiking activity of held-iout units immediately after the trial period\n",
    "\n",
    "The four different sets of data are visualized in the above figure. Each set of data is a 3D array with dimensions Trial x Time x Channel. \n",
    "\n",
    "<!---The tensor naming conventions are fairly straightforward. The tensors labeled 'heldin' contain spiking activity from held-in units. The tensors labeled 'heldout' contain spiking activity from held-out units. The tensors labeled 'forward' contain additional spiking activity occurring after each aligned trial window. All tensors have dimensions Batch x Time x Channel. --->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['train_spikes_heldin', 'train_spikes_heldout', 'train_spikes_heldin_forward', 'train_spikes_heldout_forward'])\n"
     ]
    }
   ],
   "source": [
    "print(train_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(375, 140, 122)\n",
      "(375, 140, 40)\n",
      "(375, 40, 122)\n",
      "(375, 40, 40)\n"
     ]
    }
   ],
   "source": [
    "print(train_dict['train_spikes_heldin'].shape)\n",
    "print(train_dict['train_spikes_heldout'].shape)\n",
    "print(train_dict['train_spikes_heldin_forward'].shape)\n",
    "print(train_dict['train_spikes_heldout_forward'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The shapes above indicate that there are 375 training trials, 140 time bins during the trial, 40 time bins after the trial, 122 held-in units, and 40 held-out units in this dataset.\n",
    "\n",
    "Next, we look at the data used for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['eval_spikes_heldin', 'eval_spikes_heldout'])\n"
     ]
    }
   ],
   "source": [
    "print(eval_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(125, 140, 122)\n"
     ]
    }
   ],
   "source": [
    "print(eval_dict['eval_spikes_heldin'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output tensors of `make_eval_input_tensors` follow the same dimension ordering and naming conventions as those in `train_dict`. The only tensor that will always be returned in `eval_dict` is `'eval_spikes_heldin'`, as that is the only data available in the test split.\n",
    "\n",
    "If you are using a language other than Python for your model, you will want to save these tensors as HDF5 files by changing the `save_file=False` lines in the above examples to `save_file=True`. The HDF5 files will have the same key-value pairs as the dicts and can be loaded into other programs like [MATLAB](https://www.mathworks.com/help/matlab/import_export/importing-hierarchical-data-format-hdf5-files.html) or [R](https://www.bioconductor.org/packages/devel/bioc/vignettes/rhdf5/inst/doc/rhdf5.html) scripts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Modeling\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_modeling.png?raw=true\" width=\"480\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will apply LFADS to the challenge. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1. Import LFADS packages\n",
    "\n",
    "We first import the necessary dependencies for LFADS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Import LFADS dependencies and necessary modules\n",
    "\n",
    "from lfads_tf2.utils import restrict_gpu_usage\n",
    "restrict_gpu_usage(gpu_ix=0)\n",
    "\n",
    "from lfads_tf2.models import LFADS\n",
    "from lfads_tf2.utils import load_data, merge_chops, load_posterior_averages\n",
    "from lfads_tf2.defaults import get_cfg_defaults\n",
    "from lfads_tf2.tuples import LoadableData\n",
    "\n",
    "import os.path as path\n",
    "import shutil\n",
    "\n",
    "import ray\n",
    "import yaml\n",
    "from lfads_tf2.utils import flatten\n",
    "from ray import tune\n",
    "from tune_tf2.models import create_trainable_class\n",
    "from tune_tf2.pbt.hps import HyperParam\n",
    "from tune_tf2.pbt.schedulers import MultiStrategyPBT\n",
    "from tune_tf2.pbt.trial_executor import SoftPauseExecutor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Input formatting\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/felixp8/lfads-nlb-tutorials/blob/main/images/lfads_nlb.png?raw=true\" width=\"800\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will next prepare our input and target data for training and evaluating LFADS. As seen in the figure above, we want the model to take held-in activity as input and predict firing rates for not only that held-in activity, but also held-out activity and future timesteps.\n",
    "\n",
    "Our training input data will have held-in and held-out data for in-trial and future timesteps, so that the model can learn the relationships between held-in and held-out channels and learn to forecast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "tlen = train_dict['train_spikes_heldin'].shape[1]\n",
    "num_heldin = train_dict['train_spikes_heldin'].shape[2]\n",
    "num_heldout = train_dict['train_spikes_heldout'].shape[2]\n",
    "fp_steps = train_dict['train_spikes_heldin_forward'].shape[1]\n",
    "train_spikes = np.hstack([\n",
    "    np.dstack([train_dict['train_spikes_heldin'], train_dict['train_spikes_heldout']]),\n",
    "    np.dstack([train_dict['train_spikes_heldin_forward'], train_dict['train_spikes_heldout_forward']]),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the input for evaluation, we will fill the array with zeros in place of the held-out channels and future timesteps, as we do not have access to the data. LFADS will need to predict that data without having seen it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trials = len(eval_dict['eval_spikes_heldin'])\n",
    "eval_spikes = np.hstack([\n",
    "    np.dstack([eval_dict['eval_spikes_heldin'], np.full((num_trials, tlen, num_heldout), 0.0)]),\n",
    "    np.full((num_trials, fp_steps, num_heldin + num_heldout), 0.0),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to save these files in the format that LFADS expects. As you may recall from [Running LFADS](), we will need to split the data into training and validation splits and save them in HDF5 files. \n",
    "\n",
    "We will have to create a \"train\" and \"val\" split for the evaluation data as well for compatibility with LFADS dataloading functions. However, because the model is not trained on these data, the splits are not meaningful and both sets of data will simply be passed through once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_ratio = 0.2 # 20% of trials for validation\n",
    "\n",
    "num_trials = len(train_spikes)\n",
    "valid_inds = np.arange(0, num_trials, int(round(1./valid_ratio)))\n",
    "train_inds = np.delete(np.arange(num_trials), valid_inds)\n",
    "\n",
    "with h5py.File('./data/mc_maze_large_train_lfads.h5', 'w') as h5file:\n",
    "    h5file.create_dataset('train_inds', data=train_inds)\n",
    "    h5file.create_dataset('valid_inds', data=valid_inds)\n",
    "    h5file.create_dataset('train_data', data=train_spikes[train_inds])\n",
    "    h5file.create_dataset('valid_data', data=train_spikes[valid_inds])\n",
    "\n",
    "num_trials = len(eval_spikes)\n",
    "valid_inds = np.arange(0, num_trials, int(1./valid_ratio))\n",
    "train_inds = np.delete(np.arange(num_trials), valid_inds)\n",
    "\n",
    "with h5py.File('./data/mc_maze_large_eval_lfads.h5', 'w') as h5file:\n",
    "    h5file.create_dataset('train_inds', data=train_inds)\n",
    "    h5file.create_dataset('valid_inds', data=valid_inds)\n",
    "    h5file.create_dataset('train_data', data=eval_spikes[train_inds])\n",
    "    h5file.create_dataset('valid_data', data=eval_spikes[valid_inds])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3. Model Training\n",
    "\n",
    "As in [Running LFADS](), you can either train a single LFADS model or train multiple to optimize hyperaparameters with AutoLFADS. We demonstrate both below. The YAML config files we use can be found in the `models/config/` directory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3.1. LFADS single model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build LFADS model\n",
    "# cfg_path = path.join('models/config/', 'lorenz.yaml')\n",
    "# model = LFADS(cfg_path=cfg_path)\n",
    "# train model for a bit\n",
    "# model.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3.2. AutoLFADS\n",
    "\n",
    "> **Note:** AutoLFADS is computationally expensive and it is not recommended to try running it in a single-GPU Google Colab environment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # the default configuration file for the LFADS model\n",
    "# CFG_PATH = path.join('models/config/', \"lorenz.yaml\")\n",
    "# # the directory to save PBT runs (usually '~/ray_results')\n",
    "# PBT_HOME = path.expanduser(\"./ray_results\")\n",
    "# # the name of this PBT run (run will be stored at {PBT_HOME}/{PBT_NAME})\n",
    "# RUN_NAME = \"lorenz_run\"  # the name of the PBT run\n",
    "# # the dataset to train the PBT model on\n",
    "# DATA_DIR = (\n",
    "#     \"data/\"\n",
    "# )\n",
    "# DATA_PREFIX = \"lfads\"\n",
    "\n",
    "# # the number of workers to use - make sure machine can handle all\n",
    "# NUM_WORKERS = 2\n",
    "# # the resources to allocate per model\n",
    "# RESOURCES_PER_TRIAL = {\"cpu\": 2, \"gpu\": 0.5}\n",
    "# # the hyperparameter space to search\n",
    "# HYPERPARAM_SPACE = {\n",
    "#     \"TRAIN.LR.INIT\": HyperParam(\n",
    "#         1e-5, 5e-3, explore_wt=0.3, enforce_limits=True, init=0.004\n",
    "#     ),\n",
    "#     \"MODEL.DROPOUT_RATE\": HyperParam(\n",
    "#         0.0, 0.6, explore_wt=0.3, enforce_limits=True, sample_fn=\"uniform\"\n",
    "#     ),\n",
    "#     \"MODEL.CD_RATE\": HyperParam(\n",
    "#         0.01, 0.7, explore_wt=0.3, enforce_limits=True, init=0.5, sample_fn=\"uniform\"\n",
    "#     ),\n",
    "#     \"TRAIN.L2.GEN_SCALE\": HyperParam(1e-4, 1e-0, explore_wt=0.8),\n",
    "#     \"TRAIN.L2.CON_SCALE\": HyperParam(1e-4, 1e-0, explore_wt=0.8),\n",
    "#     \"TRAIN.KL.CO_WEIGHT\": HyperParam(1e-6, 1e-4, explore_wt=0.8),\n",
    "#     \"TRAIN.KL.IC_WEIGHT\": HyperParam(1e-5, 1e-3, explore_wt=0.8),\n",
    "# }\n",
    "# PBT_METRIC = \"smth_val_nll_heldin\"\n",
    "# EPOCHS_PER_GENERATION = 25\n",
    "\n",
    "# # setup the data hyperparameters\n",
    "# dataset_info = {\"TRAIN.DATA.DIR\": DATA_DIR, \"TRAIN.DATA.PREFIX\": DATA_PREFIX}\n",
    "# # setup initialization of search hyperparameters\n",
    "# init_space = {name: tune.sample_from(hp.init) for name, hp in HYPERPARAM_SPACE.items()}\n",
    "# # load the configuration as a dictionary and update for this run\n",
    "# flat_cfg_dict = flatten(yaml.full_load(open(CFG_PATH)))\n",
    "# flat_cfg_dict.update(dataset_info)\n",
    "# flat_cfg_dict.update(init_space)\n",
    "\n",
    "# # Set the number of epochs per generation\n",
    "# tuneLFADS = create_trainable_class(EPOCHS_PER_GENERATION)\n",
    "# # connect to Ray cluster or start on single machine\n",
    "# ray.init(address=None)\n",
    "# # create the PBT scheduler\n",
    "# scheduler = MultiStrategyPBT(HYPERPARAM_SPACE, metric=PBT_METRIC)\n",
    "# # Create the trial executor\n",
    "# executor = SoftPauseExecutor(reuse_actors=True)\n",
    "# # Create the command-line display table\n",
    "# reporter = tune.CLIReporter(metric_columns=[\"epoch\", PBT_METRIC])\n",
    "# try:\n",
    "#     # run the tune job, excepting errors\n",
    "#     tune.run(\n",
    "#         tuneLFADS,\n",
    "#         name=RUN_NAME,\n",
    "#         local_dir=PBT_HOME,\n",
    "#         config=flat_cfg_dict,\n",
    "#         resources_per_trial=RESOURCES_PER_TRIAL,\n",
    "#         num_samples=NUM_WORKERS,\n",
    "#         sync_to_driver=\"# {source} {target}\",  # prevents rsync\n",
    "#         scheduler=scheduler,\n",
    "#         progress_reporter=reporter,\n",
    "#         trial_executor=executor,\n",
    "#         verbose=1,\n",
    "#         reuse_actors=True,\n",
    "#     )\n",
    "# except tune.error.TuneError:\n",
    "#     pass\n",
    "\n",
    "# # load the results dataframe for this run\n",
    "# pbt_dir = path.join(PBT_HOME, RUN_NAME)\n",
    "# df = tune.Analysis(pbt_dir).dataframe()\n",
    "# df = df[df.logdir.apply(lambda path: \"best_model\" not in path)]\n",
    "# # find the best model\n",
    "# best_model_logdir = df.loc[df[PBT_METRIC].idxmin()].logdir\n",
    "# best_model_src = path.join(best_model_logdir, \"model_dir\")\n",
    "# # copy the best model somewhere easy to find\n",
    "# best_model_dest = path.join(pbt_dir, \"best_model\")\n",
    "# shutil.copytree(best_model_src, best_model_dest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3.3. Loading a trained model\n",
    "\n",
    "To save time, we can load a pre-trained model instead. First, we'll download it from the GitHub repo and then we'll use it to load the trained LFADS model, as demonstrated in [Running LFADS](). The code we used to train this model can be found in our `nlb_tools` [repo](https://github.com/neurallatents/nlb_tools/blob/main/examples/baselines/autolfads/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_dir = 'models/lfads_mc_maze_large/'\n",
    "model = LFADS(model_dir=model_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4. Model Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_inference.png?raw=true\" width=\"480\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we'll generate our training and evaluation predictions by passing the data through the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_predictions, *_ = model.sample_and_average(save=False, merge_tv=True)\n",
    "\n",
    "loadpath = f'./data/mc_maze_large_eval_lfads.h5'\n",
    "h5file = h5py.File(loadpath, 'r')\n",
    "test_data = LoadableData(\n",
    "    train_data=h5file['train_data'][()].astype(np.float32),\n",
    "    valid_data=h5file['valid_data'][()].astype(np.float32),\n",
    "    train_ext_input=None,\n",
    "    valid_ext_input=None,\n",
    "    train_inds=h5file['train_inds'][()].astype(np.float32),\n",
    "    valid_inds=h5file['valid_inds'][()].astype(np.float32),\n",
    ")\n",
    "h5file.close()\n",
    "\n",
    "eval_predictions, *_ = model.sample_and_average(loadable_data=test_data, merge_tv=True, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Evaluation\n",
    "\n",
    "### 3.1. Online evaluation on EvalAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_submission.png?raw=true\" width=\"480\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1.1. File Preparation\n",
    "\n",
    "Now that we have predictions for the training and evaluation data, we can prepare a submission. The submission has a similar format to the returned data tensor dictionaries, but with an additional layer specifying the dataset.\n",
    "\n",
    "The dataset name and array names must be correct in order for the automated evaluation to work properly, as shown below. `'eval_rates_heldin_forward'` and `'eval_rates_heldout_forward'` are required only if you would like results on the optional forward prediction metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "tlen = train_dict['train_spikes_heldin'].shape[1]\n",
    "num_heldin = train_dict['train_spikes_heldin'].shape[2]\n",
    "\n",
    "submission = {\n",
    "    'mc_maze_large': {\n",
    "        'train_rates_heldin': training_predictions[:, :tlen, :num_heldin],\n",
    "        'train_rates_heldout': training_predictions[:, :tlen, num_heldin:],\n",
    "        'eval_rates_heldin': eval_predictions[:, :tlen, :num_heldin],\n",
    "        'eval_rates_heldout': eval_predictions[:, :tlen, num_heldin:],\n",
    "        'eval_rates_heldin_forward': eval_predictions[:, tlen:, :num_heldin],\n",
    "        'eval_rates_heldout_forward': eval_predictions[:, tlen:, num_heldin:]\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These dicts must be in an HDF5 format to be submitted to EvalAI. We have a function called `save_to_h5` to save these dictionaries to HDF5 files while preserving the structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from nlb_tools.make_tensors import save_to_h5\n",
    " \n",
    "# save_to_h5(submission, 'submission.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1.2. Submission Upload\n",
    "\n",
    "The files can be submitted through the EvalAI website or using their CLI tool. The CLI tool is recommended for large files (>300 MB), but there is no difference for smaller files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure EvalAI-CLI with your account credentials\n",
    "# !evalai set_token <auth_token>\n",
    "\n",
    "# Our challenge's id is 1256, and the phase ids are 2539 for Validation and 2540 for Test\n",
    "# So, to submit to the Validation phase of NLB'21:\n",
    "# !evalai challenge 1256 phase 2539 submit --file submission.h5\n",
    "\n",
    "# and if the file is large:\n",
    "# !evalai challenge 1256 phase 2539 submit --file submission.h5 --large"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See [this page](https://cli.eval.ai/) for more info on the EvalAI-CLI tool.\n",
    "\n",
    "Once your file is submitted, you can log in to EvalAI, go to our [challenge](https://eval.ai/web/challenges/challenge-page/1256/overview), and view the evaluation results in the 'My Submissions' tab. If your submission errored in evaluation, you can see the error output to assist in debugging."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. Local evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center> <img src=\"https://github.com/neurallatents/nlb_workshop/blob/main/nlb_technical/img/tutorial_diagram_main_evaluation.png?raw=true\" width=\"600\" /> </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the Test Phase, submissions must be uploaded to EvalAI for evaluation. For the Validation Phase, submissions can also be evaluated locally with provided data and functions.\n",
    "\n",
    "First, we prepare the data used for evaluation with `make_eval_target_tensors`. This function extracts all necessary evaluation data from the loaded dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nlb_tools.make_tensors import make_eval_target_tensors\n",
    "\n",
    "target_dict = make_eval_target_tensors(dataset=dataset, \n",
    "                                       dataset_name='mc_maze_large',\n",
    "                                       train_trial_split='train',\n",
    "                                       eval_trial_split='val',\n",
    "                                       include_psth=True,\n",
    "                                       save_file=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we can evaluate with the `evaluate` function. Every submission is scored on a number of metrics, each evaluating different aspects of the model output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'mc_maze_scaling_split': {'[500] co-bps': 0.3211858864626676,\n",
       "   '[500] vel R2': 0.856816843986107,\n",
       "   '[500] psth R2': 0.5743404432814287,\n",
       "   '[500] fp-bps': 0.18072337679652054}}]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from nlb_tools.evaluation import evaluate\n",
    "from pprint import pprint\n",
    "\n",
    "pprint(evaluate(target_dict, submission))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3. Evaluation metrics\n",
    "\n",
    "We evaluate submitted rate predictions on a number of metrics, as shown above. Each metric gauges a different aspect of model performance, providing a multi-faceted perspective on where models excel.\n",
    "\n",
    "#### 3.3.1. Co-smoothing\n",
    "\n",
    "Our primary metric is co-smoothing, as described above. To recap, co-smoothing evaluates a model's ability to predict the activity of a held-out set of neurons from just the activity of provided neurons. The rate predictions for held-out neurons are evaluated based on the negative log-likelihood of the observed spiking activity given predicted firing rates. The final score is standardized by the average firing rate of each neuron to compute a bits per spike value.\n",
    "\n",
    "The assumption behind co-smoothing is that the latent variables underlying held-out neuron activity can be inferred from the available neurons. It has been shown that latent representations are distributed across many neurons in a population, providing some support for this assumption.The advantage of co-smoothing is that it requires no additional data and makes very weak assumptions about the neural activity. It relies on already-available spiking activity, and it essentially only assumes that activity of neurons in a population are related. This makes it generally applicable across datasets that span experimental settings and brain areas. However, likelihood scores vary substantially across datasets due to differences in neuron count and firing rate statistics, making comparisons of performance across datasets difficult.\n",
    "\n",
    "#### 3.3.2. Decoding accuracy\n",
    "\n",
    "To evaluate decoding performance, we train a linear regression model from predicted firing rates to a behavioral variable, specifically arm velocity. The outputs of the regression model are evaluated with $R^2$, or variance explained, which takes a value between $(-\\infty, 1]$.\n",
    "\n",
    "Decoding measures how much information about a given external behavioral variable is contained in the rate predictions. This gives a very concrete and relevant measure of model performance, with practical applications for brain-computer interfaces. However, it is limiting in that it ignores other information that may be encoded in neural activity but which may not be as clearly related to externally measurable behavior. This is especially problematic in higher-order brain regions, like DMFC, where brain activity is believed to related to complex cognitive functions like problem solving or time interval estimation, which do not have moment-by-moment behavioral correlates.\n",
    "\n",
    "#### 3.3.3. Match to PSTH\n",
    "\n",
    "Traditional analyses of single-neuron firing rates relied on averaging smoothed spike counts across trials of the same condition, resulting in estimates of a neuron's true firing rate within a given condition. Our match to PSTH metric computes true PSTHs from spiking activity and compares it to average model firing rate predictions across trials of the same condition, evaluating the model predictions with $R^2$, or variance explained.\n",
    "\n",
    "Our implementation of match to PSTH gauges how well model predictions distinguish experimental conditions. However, by averaging across trials, the metric ignores single-trial variance. Models that accurately capture single-trial variance and models that fail to could perform equally well on this metric. As a result, it does a poor job distinguishing models past a particular level of performance, and it ignores an important advantage of more powerful single-trial modeling approaches.\n",
    "\n",
    "#### 3.3.4. Forward prediction\n",
    "\n",
    "Forward prediction tasks the models with predicting another 200 ms of data after each trial window. The predictions are evaluated by negative log-likelihood of spiking, just like co-smoothing evaluation. This evaluation includes both held-in and held-out neurons, as no activity for these future timesteps is provided as input to the model.\n",
    "\n",
    "Forward prediction assumes that future activity is predictable, which may not always hold true. In particular, datasets like MC_RTT, where new random inputs can occur unexpectedly at any point within and after a trial window, it should not be possible to predict neural activity after the end of the trial window. However, for datasets like MC_Maze, where dynamics during the reach are well-modeled as autonomous, future prediction is a reasonable measure for the quality of the model's inferred dynamical rules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary and Conclusions\n",
    "\n",
    "This tutorial gave an overview of the NLB'21 pipeline from loading raw data to submission and demonstrated how to make use of our provided code package, `nlb_tools`, to participate in NLB'21 with LFADS. It also described the principles motivating our evaluation strategy, which can be applied generally beyond the scope of just the benchmark.\n",
    "\n",
    "For additional helpful resources on NLB'21, we have a number of other tutorials and example scripts covering a variety of topics:\n",
    "* The notebooks in the [`nlb_tools` repo](https://github.com/neurallatents/nlb_tools) demonstrate application of classical methods like spike smoothing, GPFA, and SLDS to NLB'21.\n",
    "* Andrew Sedler's [nlb-lightning](https://github.com/arsedler9/nlb-lightning) package provides a convenient framework to develop and evaluate PyTorch Lightning models for NLB'21."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "02b1bc2cb6aba3c63676f81fd881bdec751fcef939c531164eae59d8a44feb6a"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}