notebooks/load_data_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "critical-coupon",
   "metadata": {},
   "source": [
    "### This tutorial shows how to get the data into dataloaders from deeplake and makes basic data overview.\n",
    "\n",
    "To see how to do it with dowloading data and working locally - check make_model tutorial or [last year repo](https://github.com/sinzlab/sensorium/blob/main/notebooks/dataset_tutorial/1_inspect_data.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "medium-rochester",
   "metadata": {},
   "outputs": [],
   "source": [
    "## TODO - uncomment for colab setup\n",
    "\n",
    "# !git clone https://github.com/ecker-lab/sensorium_2023\n",
    "# !pip install git+https://github.com/sinzlab/nnfabrik.git\n",
    "# !pip install git+https://github.com/sinzlab/neuralpredictors\n",
    "# !pip install pyarrow\n",
    "# !pip install wandb\n",
    "# !pip install GitPython\n",
    "# !pip install datajoint\n",
    "# !pip install deeplake\n",
    "# !pip install parquet\n",
    "# import sys\n",
    "# sys.path.append('/content/sensorium_2023')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "involved-apartment",
   "metadata": {},
   "outputs": [],
   "source": [
    "import deeplake\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib\n",
    "from IPython.display import HTML\n",
    "from matplotlib import animation\n",
    "import imageio\n",
    "import os\n",
    "import numpy as np\n",
    "import plotly.express as px\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "synthetic-spray",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import Compose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "atomic-vault",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "comfortable-ultimate",
   "metadata": {},
   "outputs": [],
   "source": [
    "mice = [\n",
    "    'dynamic29156-11-10-Video-8744edeac3b4d1ce16b680916b5267ce',\n",
    "    'dynamic29228-2-10-Video-8744edeac3b4d1ce16b680916b5267ce',\n",
    "    'dynamic29234-6-9-Video-8744edeac3b4d1ce16b680916b5267ce',\n",
    "    'dynamic29513-3-5-Video-8744edeac3b4d1ce16b680916b5267ce',\n",
    "    'dynamic29514-2-9-Video-8744edeac3b4d1ce16b680916b5267ce',\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spatial-leonard",
   "metadata": {},
   "source": [
    "## Load single dataset and visualise a video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "engaged-granny",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = deeplake.load(f'hub://sinzlab/Sensorium_2023_{mice[0]}_train')\n",
    "dataset.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "powerful-lesbian",
   "metadata": {},
   "outputs": [],
   "source": [
    "vid = dataset.videos[0].numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ordinary-permission",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "# an equal aspect (square), and turn the axes off\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.set_axis_off()\n",
    "\n",
    "# Images are generated and stored in a list to animate later\n",
    "images_vid = []\n",
    "images_as_array = []\n",
    "for image in vid[0,:,:,:]:\n",
    "    # Scatter plot each point using a dot of size 250 and color red\n",
    "    #images_vid.append([ax.imshow(image,cmap='gray')])\n",
    "    images_vid.append([ax.imshow(image, cmap='gray')])\n",
    "    images_as_array.append(image)\n",
    "# The animation can now be created using ArtistAnimation\n",
    "anim = animation.ArtistAnimation(fig,\n",
    "                                images_vid,\n",
    "                                interval= 1000 / 30,\n",
    "                                blit=True,\n",
    "                                repeat=False\n",
    "                                )\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "downtown-section",
   "metadata": {},
   "source": [
    "## Plot reponses and behaviour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bored-huntington",
   "metadata": {},
   "outputs": [],
   "source": [
    "responses = dataset.responses[0].numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "administrative-cathedral",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,  ax = plt.subplots(1,1, figsize=(6,6))\n",
    "n_frames = 200\n",
    "n_neurons = 200\n",
    "ax.imshow(responses[:n_frames, :n_neurons], cmap=\"gray\",vmin=0, vmax=5)\n",
    "\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "\n",
    "ax.set_xlabel(\"Natural Images\", fontsize=20, )\n",
    "ax.set_ylabel(\"Associated neuronal Response\", fontsize=20, );"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seventh-advance",
   "metadata": {},
   "outputs": [],
   "source": [
    "behaviors = dataset.behavior[0].numpy().T\n",
    "eye_positions = dataset.pupil_center[0].numpy().T\n",
    "all_behaviors = np.hstack([np.vstack(behaviors), np.vstack(eye_positions), ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "timely-lender",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all_behaviors.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "royal-brush",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(5,2, figsize=(14,12))\n",
    "\n",
    "behaviors = [\"Pupil size\", \n",
    "            \"Change of pupil size\", \n",
    "            \"Locomotion speed\",\n",
    "            \"horizontal eye position\",\n",
    "            \"vertical eye position\"]\n",
    "colormaps = [\"YlGn\", \"gray\", \"RdPu\", \"plasma\", \"plasma\"]\n",
    "\n",
    "\n",
    "for i in range(2):\n",
    "    for behavior in range(5):\n",
    "        if i == 0:\n",
    "            axs[behavior, i].imshow(np.tile(all_behaviors[:n_frames, behavior], (40,1)), cmap=colormaps[behavior])\n",
    "            axs[behavior, i].set_xlabel(\"Trial number\", fontsize=12)\n",
    "            axs[behavior, i].set_title(behaviors[behavior], fontsize=16)\n",
    "        else:\n",
    "            axs[behavior, i].plot(all_behaviors[:n_frames, behavior], 'k-')\n",
    "            axs[behavior, i].axis(\"off\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "constitutional-stability",
   "metadata": {},
   "source": [
    "## Get neurons 3D coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hydraulic-explanation",
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_motor_coordinates = dataset.info['cell_motor_coordinates']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "neither-blind",
   "metadata": {},
   "outputs": [],
   "source": [
    "def colorFader(c1,c2,mix=0): #fade (linear interpolate) from color c1 (at mix=0) to c2 (mix=1)\n",
    "    c1=np.array(mpl.colors.to_rgb(c1))\n",
    "    c2=np.array(mpl.colors.to_rgb(c2))\n",
    "    return mpl.colors.to_hex((1-mix)*c1 + mix*c2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "double-underground",
   "metadata": {},
   "outputs": [],
   "source": [
    "c1='#1f77b4' #blue\n",
    "c2='green' #green\n",
    "n=n_neurons\n",
    "\n",
    "colors = []\n",
    "for x in range(n):\n",
    "    colors.append(colorFader(c1, c2, x/n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "attractive-officer",
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_subset = cell_motor_coordinates[\n",
    "    np.random.choice(cell_motor_coordinates.shape[0], \n",
    "                     n_neurons, \n",
    "                     replace=False)\n",
    "]\n",
    "\n",
    "fig = px.scatter_3d(x=neuron_subset[:, 0], \n",
    "                    y=neuron_subset[:, 1], \n",
    "                    z=neuron_subset[:, 2],\n",
    "                   color=colors)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "described-bulgaria",
   "metadata": {},
   "source": [
    "## Add augmentations for dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fresh-colony",
   "metadata": {},
   "outputs": [],
   "source": [
    "def neuro_normalise(\n",
    "    data, dataset_statistic, inputs_mean=None, inputs_std=None, with_responses=True\n",
    "):\n",
    "    '''\n",
    "    adopted from neuropredictors NeuroNormalizer\n",
    "    https://github.com/sinzlab/neuralpredictors/blob/main/neuralpredictors/data/transforms.py#L309\n",
    "\n",
    "    - inputs are scaled by the training std of the stats_source and centered on the mean of the movie\n",
    "    - behavior is divided by the std if the std is greater than 1% of the mean std (to avoid division by 0)\n",
    "    - eye_position is z-scored\n",
    "    - reponses are divided by the per neuron std if the std is greater than\n",
    "            1% of the mean std (to avoid division by 0)\n",
    "\n",
    "    '''\n",
    "    data_len = data['videos'].shape[1]\n",
    "    inputs_mean = inputs_mean if inputs_mean else dataset_statistic['videos']['mean'][:, :, :data_len].transpose(2, 0, 1)\n",
    "    inputs_std  = inputs_std if inputs_std else dataset_statistic['videos']['std'][:, :, :data_len].transpose(2, 0, 1)\n",
    "    # -- inputs in deeplake are already stacked with behaviour in different channels, \n",
    "    # so we normalise only the first channel\n",
    "    data['videos'][0, :, :] = (data['videos'][0, :, :] - inputs_mean) / inputs_std\n",
    "    \n",
    "    if with_responses:\n",
    "    \n",
    "        s = np.array(dataset_statistic['responses']['std'])[:, :data_len]\n",
    "        threshold = 0.01 * s.mean()\n",
    "        idx = s > threshold\n",
    "        response_precision = np.ones_like(s) / threshold\n",
    "        response_precision[idx] = 1 / s[idx]\n",
    "        data['responses'] = data['responses'] * response_precision\n",
    "\n",
    "    eye_mean = np.array(dataset_statistic['pupil_center']['mean'])[:, :data_len]\n",
    "    eye_std = np.array(dataset_statistic['pupil_center']['std'])[:, :data_len]\n",
    "    data['pupil_center'] = (data['pupil_center'] - eye_mean) / eye_std\n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "serial-benchmark",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform(sample_in, dataset_statistic, frames=150, offset=-1, with_responses=True):\n",
    "    '''\n",
    "    As in our dataset videos are not of the same length, \n",
    "    to stack them in a batch we are going to subsample 'frames' frames randomly during loading data\n",
    "    and also normalise things (before subsampling)\n",
    "    \n",
    "    NOTE that the amount of frames should be bigger then 50 \n",
    "    otherwise, the assertion test would fail duringcomputing metric \n",
    "    as we delete first 50 frames of responses \n",
    "    (and the additional predictions if any)\n",
    "    '''\n",
    "    sample_in = neuro_normalise(sample_in, dataset_statistic, with_responses=with_responses)\n",
    "    \n",
    "    \n",
    "    t = sample_in['behavior'].shape[-1]\n",
    "    if offset < 0:\n",
    "        i = np.random.randint(0, t - frames)\n",
    "    \n",
    "    sample_in['videos'] = sample_in['videos'][:, i : i + frames, :, :]\n",
    "    sample_in['responses'] = sample_in['responses'][:, i : i + frames]\n",
    "    sample_in['behavior'] = sample_in['behavior'][:, i : i + frames]\n",
    "    sample_in['pupil_center'] = sample_in['pupil_center'][:, i : i + frames]\n",
    "    return sample_in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "peaceful-occupation",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataloader = dataset.pytorch(tensors=('videos', 'responses', \n",
    "                                      'behavior', 'pupil_center', 'id'), \n",
    "                        num_workers=0, \n",
    "                        batch_size=16, \n",
    "                        transform=partial(transform, dataset_statistic=dataset.info.statistics),\n",
    "                        shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "directed-night",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "batch = next(iter(dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "catholic-constraint",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch['videos'].shape, batch['responses'].shape, batch['behavior'].shape, batch['pupil_center'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "twelve-specialist",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "batch['id'][:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bulgarian-average",
   "metadata": {},
   "source": [
    "## Make the dataloaders according to the standard procedure we use for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "recognized-killing",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataloaders = {}\n",
    "\n",
    "rename_dict = {\n",
    "    'train' : 'train',\n",
    "    'val' : 'oracle',\n",
    "}\n",
    "\n",
    "for div in ['train', 'val']:\n",
    "    dataloaders[rename_dict[div]] = {}\n",
    "    for m in mice:\n",
    "        dataset = deeplake.load(f'hub://sinzlab/Sensorium_2023_{m}_{div}')\n",
    "        dataloaders[rename_dict[div]] = dataloader = dataset.pytorch(tensors=('videos', 'responses', \n",
    "                                      'behavior', 'pupil_center', 'id'), \n",
    "                        num_workers=0, \n",
    "                        batch_size=16, \n",
    "                        transform=partial(transform, \n",
    "                                          dataset_statistic=dataset.info.statistics, \n",
    "                                          with_responses=True),\n",
    "                        shuffle=True)\n",
    "        \n",
    "        \n",
    "for div in ['live_test_main', 'live_test_bonus', 'final_test_main', 'final_test_bonus']:\n",
    "    dataloaders[div] = {}\n",
    "    for m in mice:\n",
    "        dataset = deeplake.load(f'hub://sinzlab/Sensorium_2023_{m}_{div}')\n",
    "        dataloaders[div][m] = dataloader = dataset.pytorch(tensors=('videos', \n",
    "                                      'behavior', 'pupil_center', 'id'), \n",
    "                        num_workers=0, \n",
    "                        batch_size=16, \n",
    "                        transform=partial(neuro_normalise, \n",
    "                                          dataset_statistic=dataset.info.statistics,\n",
    "                                          with_responses=False),\n",
    "                        shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "turned-finnish",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataloaders"
   ]
  }
 ],
 "metadata": {
  "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}