ExplanatoryNotebook_bySteinmetz.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ExplanatoryNotebook_bySteinmetz.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Bionerdess/NMA2021/blob/main/ExplanatoryNotebook_bySteinmetz.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bEqdz1ZUMaj1"
      },
      "source": [
        "## Loading of Steinmetz data\n",
        "\n",
        "includes some visualizations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TLWjKq8bLDqm",
        "cellView": "form"
      },
      "source": [
        "#@title Data retrieval\n",
        "import os, requests\n",
        "\n",
        "fname = []\n",
        "for j in range(3):\n",
        "  fname.append('steinmetz_part%d.npz'%j)\n",
        "url = [\"https://osf.io/agvxh/download\"]\n",
        "url.append(\"https://osf.io/uv3mw/download\")\n",
        "url.append(\"https://osf.io/ehmw2/download\")\n",
        "\n",
        "for j in range(len(url)):\n",
        "  if not os.path.isfile(fname[j]):\n",
        "    try:\n",
        "      r = requests.get(url[j])\n",
        "    except requests.ConnectionError:\n",
        "      print(\"!!! Failed to download data !!!\")\n",
        "    else:\n",
        "      if r.status_code != requests.codes.ok:\n",
        "        print(\"!!! Failed to download data !!!\")\n",
        "      else:\n",
        "        with open(fname[j], \"wb\") as fid:\n",
        "          fid.write(r.content)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "raBVOEWgUK_B",
        "cellView": "form"
      },
      "source": [
        "#@title Import matplotlib and set defaults\n",
        "from matplotlib import rcParams \n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "rcParams['figure.figsize'] = [20, 4]\n",
        "rcParams['font.size'] =15\n",
        "rcParams['axes.spines.top'] = False\n",
        "rcParams['axes.spines.right'] = False\n",
        "rcParams['figure.autolayout'] = True"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sffzC_hyLgWZ"
      },
      "source": [
        "#@title Data loading\n",
        "import numpy as np\n",
        "\n",
        "alldat = np.array([])\n",
        "for j in range(len(fname)):\n",
        "  alldat = np.hstack((alldat, np.load('steinmetz_part%d.npz'%j, allow_pickle=True)['dat']))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5K7UT7dyj_6R"
      },
      "source": [
        "`alldat` contains 39 sessions from 10 mice, data from Steinmetz et al, 2019. Time bins for all measurements are 10ms, starting 500ms before stimulus onset. The mouse had to determine which side has the highest contrast. For each `dat = alldat[k]`, you have the fields below. For extra variables, check out the extra notebook and extra data files (lfp, waveforms and exact spike times, non-binned). \n",
        "\n",
        "* `dat['mouse_name']`: mouse name\n",
        "* `dat['date_exp']`: when a session was performed\n",
        "* `dat['spks']`: neurons by trials by time bins.    \n",
        "* `dat['brain_area']`: brain area for each neuron recorded. \n",
        "* `dat['ccf']`: Allen Institute brain atlas coordinates for each neuron. \n",
        "* `dat['ccf_axes']`: axes names for the Allen CCF. \n",
        "* `dat['contrast_right']`: contrast level for the right stimulus, which is always contralateral to the recorded brain areas.\n",
        "* `dat['contrast_left']`: contrast level for left stimulus. \n",
        "* `dat['gocue']`: when the go cue sound was played. \n",
        "* `dat['response_time']`: when the response was registered, which has to be after the go cue. The mouse can turn the wheel before the go cue (and nearly always does!), but the stimulus on the screen won't move before the go cue.  \n",
        "* `dat['response']`: which side the response was (`-1`, `0`, `1`). When the right-side stimulus had higher contrast, the correct choice was `-1`. `0` is a no go response. \n",
        "* `dat['feedback_time']`: when feedback was provided. \n",
        "* `dat['feedback_type']`: if the feedback was positive (`+1`, reward) or negative (`-1`, white noise burst).  \n",
        "* `dat['wheel']`: turning speed of the wheel that the mice uses to make a response, sampled at `10ms`. \n",
        "* `dat['pupil']`: pupil area  (noisy, because pupil is very small) + pupil horizontal and vertical position.\n",
        "* `dat['face']`: average face motion energy from a video camera. \n",
        "* `dat['licks']`: lick detections, 0 or 1.   \n",
        "* `dat['trough_to_peak']`: measures the width of the action potential waveform for each neuron. Widths `<=10` samples are \"putative fast spiking neurons\". \n",
        "* `dat['%X%_passive']`: same as above for `X` = {`spks`, `pupil`, `wheel`, `contrast_left`, `contrast_right`} but for  passive trials at the end of the recording when the mouse was no longer engaged and stopped making responses. \n",
        "* `dat['prev_reward']`: time of the feedback (reward/white noise) on the previous trial in relation to the current stimulus time. \n",
        "* `dat['reaction_time']`: ntrials by 2. First column: reaction time computed from the wheel movement as the first sample above `5` ticks/10ms bin. Second column: direction of the wheel movement (`0` = no move detected).  \n",
        "\n",
        "\n",
        "The original dataset is here: https://figshare.com/articles/dataset/Dataset_from_Steinmetz_et_al_2019/9598406"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JYhiUDTi1uAZ"
      },
      "source": [
        "# Make a plot of which brain areas are present in each dataset\n",
        "# note that region 4 (\"other ctx\" are neurons that were not able to be classified)\n",
        "# region 4 does not correspond to brain_group 4, which are all cortical neurons outside of visual cortex\n",
        "regions = [\"vis ctx\", \"thal\", \"hipp\", \"other ctx\", \"midbrain\", \"basal ganglia\", \"cortical subplate\", \"other\"]\n",
        "region_colors = ['blue', 'red', 'green', 'darkblue', 'violet', 'lightblue', 'orange', 'gray']\n",
        "brain_groups = [[\"VISa\", \"VISam\", \"VISl\", \"VISp\", \"VISpm\", \"VISrl\"], # visual cortex\n",
        "                [\"CL\", \"LD\", \"LGd\", \"LH\", \"LP\", \"MD\", \"MG\", \"PO\", \"POL\", \"PT\", \"RT\", \"SPF\", \"TH\", \"VAL\", \"VPL\", \"VPM\"], # thalamus\n",
        "                [\"CA\", \"CA1\", \"CA2\", \"CA3\", \"DG\", \"SUB\", \"POST\"], # hippocampal\n",
        "                [\"ACA\", \"AUD\", \"COA\", \"DP\", \"ILA\", \"MOp\", \"MOs\", \"OLF\", \"ORB\", \"ORBm\", \"PIR\", \"PL\", \"SSp\", \"SSs\", \"RSP\",\"TT\"], # non-visual cortex\n",
        "                [\"APN\", \"IC\", \"MB\", \"MRN\", \"NB\", \"PAG\", \"RN\", \"SCs\", \"SCm\", \"SCig\", \"SCsg\", \"ZI\"], # midbrain\n",
        "                [\"ACB\", \"CP\", \"GPe\", \"LS\", \"LSc\", \"LSr\", \"MS\", \"OT\", \"SNr\", \"SI\"], # basal ganglia \n",
        "                [\"BLA\", \"BMA\", \"EP\", \"EPd\", \"MEA\"] # cortical subplate\n",
        "                ]\n",
        "\n",
        "# Assign each area an index\n",
        "area_to_index = dict(root=0)\n",
        "counter = 1\n",
        "for group in brain_groups:\n",
        "    for area in group:\n",
        "        area_to_index[area] = counter\n",
        "        counter += 1\n",
        "\n",
        "# Figure out which areas are in each dataset\n",
        "areas_by_dataset = np.zeros((counter, len(alldat)), dtype=bool)\n",
        "for j, d in enumerate(alldat):\n",
        "    for area in np.unique(d['brain_area']):\n",
        "        i = area_to_index[area]\n",
        "        areas_by_dataset[i, j] = True\n",
        "\n",
        "# Show the binary matrix\n",
        "plt.figure(figsize=(8, 10))\n",
        "plt.imshow(areas_by_dataset, cmap=\"Greys\", aspect=\"auto\", interpolation=\"none\")\n",
        "\n",
        "# Label the axes\n",
        "plt.xlabel(\"dataset\")\n",
        "plt.ylabel(\"area\")\n",
        "\n",
        "# Add tick labels\n",
        "yticklabels = [\"root\"]\n",
        "for group in brain_groups:\n",
        "    yticklabels.extend(group)\n",
        "plt.yticks(np.arange(counter), yticklabels, fontsize=8)\n",
        "plt.xticks(np.arange(len(alldat)), fontsize=9)\n",
        "\n",
        "# Color the tick labels by region\n",
        "ytickobjs = plt.gca().get_yticklabels()\n",
        "ytickobjs[0].set_color(\"black\")\n",
        "counter = 1\n",
        "for group, color in zip(brain_groups, region_colors):\n",
        "    for area in group:\n",
        "        ytickobjs[counter].set_color(color)\n",
        "        counter += 1\n",
        "\n",
        "plt.title(\"Brain areas present in each dataset\")\n",
        "plt.grid(True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mmOarX5w16CR"
      },
      "source": [
        "#@title basic plots of population average\n",
        "\n",
        "# select just one of the recordings here. 11 is nice because it has some neurons in vis ctx. \n",
        "dat = alldat[11]\n",
        "print(dat.keys())\n",
        "\n",
        "dt = dat['bin_size'] # binning at 10 ms\n",
        "NT = dat['spks'].shape[-1]\n",
        "\n",
        "ax = plt.subplot(1,5,1)\n",
        "response = dat['response'] # right - nogo - left (-1, 0, 1)\n",
        "vis_right = dat['contrast_right'] # 0 - low - high\n",
        "vis_left = dat['contrast_left'] # 0 - low - high\n",
        "plt.plot(dt * np.arange(NT), 1/dt * dat['spks'][:,response>=0].mean(axis=(0,1))) # left responses\n",
        "plt.plot(dt * np.arange(NT), 1/dt * dat['spks'][:,response<0].mean(axis=(0,1))) # right responses\n",
        "plt.plot(dt * np.arange(NT), 1/dt * dat['spks'][:,vis_right>0].mean(axis=(0,1))) # stimulus on the right\n",
        "plt.plot(dt * np.arange(NT), 1/dt * dat['spks'][:,vis_right==0].mean(axis=(0,1))) # no stimulus on the right\n",
        "\n",
        "plt.legend(['left resp', 'right resp', 'right stim', 'no right stim'], fontsize=12)\n",
        "ax.set(xlabel  = 'time (sec)', ylabel = 'firing rate (Hz)');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1T40IuyU5D0B"
      },
      "source": [
        "nareas = 4 # only the top 4 regions are in this particular mouse\n",
        "NN = len(dat['brain_area']) # number of neurons\n",
        "barea = nareas * np.ones(NN, ) # last one is \"other\"\n",
        "for j in range(nareas):\n",
        "  barea[np.isin(dat['brain_area'], brain_groups[j])] = j # assign a number to each region"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lqMJRWZM2haT"
      },
      "source": [
        "#@title plots by brain region and visual conditions\n",
        "for j in range(nareas):\n",
        "  ax = plt.subplot(1,nareas,j+1)\n",
        "  \n",
        "  plt.plot(1/dt *  dat['spks'][barea==j][:,np.logical_and(vis_left==0, vis_right>0)].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt *  dat['spks'][barea==j][:,np.logical_and(vis_left>0 , vis_right==0)].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt *  dat['spks'][barea==j][:,np.logical_and(vis_left==0 , vis_right==0)].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt *  dat['spks'][barea==j][:,np.logical_and(vis_left>0, vis_right>0)].mean(axis=(0,1)))  \n",
        "  plt.text(.25, .92, 'n=%d'%np.sum(barea==j), transform=ax.transAxes)\n",
        " \n",
        "  if j==0:\n",
        "    plt.legend(['right only', 'left only', 'neither', 'both'], fontsize=12)\n",
        "  ax.set(xlabel = 'binned time', ylabel = 'mean firing rate (Hz)', title = regions[j])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vQA2tqAr-WSO"
      },
      "source": [
        "#@title plots by brain region and response type\n",
        "for j in range(nareas):\n",
        "  ax = plt.subplot(1,nareas,j+1)\n",
        "  plt.title(regions[j])\n",
        "  if np.sum(barea==j)==0:\n",
        "    continue\n",
        "  plt.plot(1/dt * dat['spks'][barea==j][:,response<0].mean(axis=(0,1)))  \n",
        "  plt.plot(1/dt * dat['spks'][barea==j][:,response>0].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt * dat['spks'][barea==j][:,response==0].mean(axis=(0,1)))\n",
        " \n",
        "  if j==0:\n",
        "    plt.legend(['resp = left', 'resp = right', 'resp = none'], fontsize=12)\n",
        "  ax.set(xlabel = 'time', ylabel = 'mean firing rate (Hz)')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vk69J9u-AV-2"
      },
      "source": [
        "#@title top PC directions from stimulus + response period, with projections of the entire duration\n",
        "from sklearn.decomposition import PCA \n",
        "\n",
        "droll = np.reshape(dat['spks'][:,:,51:130], (NN,-1)) # first 80 bins = 1.6 sec\n",
        "droll = droll - np.mean(droll, axis=1)[:, np.newaxis]\n",
        "model = PCA(n_components = 5).fit(droll.T)\n",
        "W = model.components_\n",
        "pc_10ms = W @ np.reshape(dat['spks'], (NN,-1))\n",
        "pc_10ms = np.reshape(pc_10ms, (5, -1, NT))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VmAKoW85hZz7"
      },
      "source": [
        "#@title The top PCs capture most variance across the brain. What do they care about? \n",
        "plt.figure(figsize= (20, 6))\n",
        "for j in range(len(pc_10ms)):\n",
        "  ax = plt.subplot(2,len(pc_10ms)+1,j+1)\n",
        "  pc1 = pc_10ms[j]\n",
        "\n",
        "  plt.plot(pc1[np.logical_and(vis_left==0, vis_right>0), :].mean(axis=0))  \n",
        "  plt.plot(pc1[np.logical_and(vis_left>0, vis_right==0), :].mean(axis=0))\n",
        "  plt.plot(pc1[np.logical_and(vis_left==0, vis_right==0), :].mean(axis=0))\n",
        "  plt.plot(pc1[np.logical_and(vis_left>0, vis_right>0), :].mean(axis=0))\n",
        "   \n",
        "  if j==0:\n",
        "    plt.legend(['right only', 'left only', 'neither', 'both'], fontsize=8)\n",
        "  ax.set(xlabel = 'binned time', ylabel = 'mean firing rate (Hz)')\n",
        "  plt.title('PC %d'%j)\n",
        "\n",
        "  ax = plt.subplot(2,len(pc_10ms)+1,len(pc_10ms)+1 + j+1)\n",
        "  \n",
        "  plt.plot(pc1[response>0, :].mean(axis=0))  \n",
        "  plt.plot(pc1[response<0, :].mean(axis=0))\n",
        "  plt.plot(pc1[response==0, :].mean(axis=0))\n",
        "\n",
        "  if j==0:\n",
        "    plt.legend(['resp = left', 'resp = right', 'resp = none'], fontsize=8)\n",
        "  ax.set(xlabel = 'binned time', ylabel = 'mean firing rate (Hz)')\n",
        "  plt.title('PC %d'%j)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6QZ-eFRwj3at"
      },
      "source": [
        "#@title now sort all trials by response latency and see if the PCs care about that.\n",
        "from scipy.stats import zscore\n",
        "\n",
        "isort = np.argsort(dat['response_time'].flatten())\n",
        "\n",
        "for j in range(len(pc_10ms)):\n",
        "  ax = plt.subplot(1,len(pc_10ms)+1,j+1)\n",
        "  pc1 = zscore(pc_10ms[j])\n",
        "  plt.imshow(pc1[isort, :], aspect='auto', vmax=2, vmin = -2, cmap = 'gray')\n",
        "  ax.set(xlabel = 'binned time', ylabel = 'trials sorted by latency')\n",
        "  plt.title('PC %d'%j)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C-90MWoY15kh"
      },
      "source": [
        "#@title correct vs incorrect trials\n",
        "# the following are the correct responses:\n",
        "# if vis_left > vis_right : response >0\n",
        "# if vis_left < vis_right : response <0\n",
        "# if vis_left = vis_right : response =0\n",
        "# trials below red line are incorrect\n",
        "is_correct = np.sign(response)==np.sign(vis_left-vis_right)\n",
        "\n",
        "# sort by correct, and then by response\n",
        "isort = np.argsort(-is_correct.astype('float32') + response/10) \n",
        "\n",
        "nwrong = np.sum(is_correct)\n",
        "for j in range(len(pc_10ms)):\n",
        "  ax = plt.subplot(1,len(pc_10ms)+1,j+1)\n",
        "  pc1 = zscore(pc_10ms[j])\n",
        "  plt.imshow(pc1[isort, :], aspect='auto', vmax=2, vmin = -2, cmap = 'gray')\n",
        "  ax.set(xlabel = 'binned time')\n",
        "  if j==0:\n",
        "    ax.set(ylabel = 'trials sorted by latency')  \n",
        "  plt.title('PC %d'%j)\n",
        "\n",
        "  plt.plot([0, NT], [nwrong, nwrong], 'r')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2HgiWh4Vb2Su"
      },
      "source": [
        "# plot the behavioral data (pupil area is noisy because it's very small)\n",
        "\n",
        "ax = plt.subplot(1,5,1)\n",
        "plt.plot(dat['pupil'][0, :].mean(0));\n",
        "ax.set(ylabel='pupil area', xlabel = 'binned time', title='Pupil dynamics')\n",
        "\n",
        "yl = [-10, 10]\n",
        "ax = plt.subplot(1,5,2)\n",
        "plt.plot(dat['wheel'][0, response>0].mean(0));\n",
        "ax.set(ylim=yl)\n",
        "ax.set(ylim=yl, ylabel='wheel turns', xlabel = 'binned time', title='Left choices');\n",
        "\n",
        "ax = plt.subplot(1,5,3)\n",
        "plt.plot(dat['wheel'][0, response<0].mean(0));\n",
        "ax.set(ylim=yl)\n",
        "ax.set(ylim=yl, ylabel='wheel turns', xlabel = 'binned time', title='Right choices');\n",
        "\n",
        "ax = plt.subplot(1,5,4)\n",
        "plt.plot(dat['wheel'][0, response==0].mean(0));\n",
        "ax.set(ylim=yl, ylabel='wheel turns', xlabel = 'binned time', title='No go choices');\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9JkyGhv0xvCQ"
      },
      "source": [
        "# plots by brain region and visual conditions for PASSIVE trials\n",
        "vis_left_p = dat['contrast_left_passive']\n",
        "vis_right_p = dat['contrast_right_passive']\n",
        "for j in range(nareas):\n",
        "  ax = plt.subplot(1,nareas,j+1)\n",
        "  plt.title(regions[j])\n",
        "  \n",
        "  plt.plot(1/dt *  dat['spks_passive'][barea==j][:,np.logical_and(vis_left_p==0, vis_right_p>0)].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt *  dat['spks_passive'][barea==j][:,np.logical_and(vis_left_p>0 , vis_right_p==0)].mean(axis=(0,1)))\n",
        "  #plt.plot(1/dt *  dat['spks_passive'][barea==j][:,np.logical_and(vis_left_p==0 , vis_right_p==0)].mean(axis=(0,1)))\n",
        "  plt.plot(1/dt *  dat['spks_passive'][barea==j][:,np.logical_and(vis_left_p>0, vis_right_p>0)].mean(axis=(0,1)))  \n",
        "  plt.text(.25, .92, 'n=%d'%np.sum(barea==j), transform=ax.transAxes)\n",
        " \n",
        "  if j==0:\n",
        "    plt.legend(['right only', 'left only', 'both'], fontsize=12)\n",
        "  ax.set(xlabel = 'binned time', ylabel = 'mean firing rate (Hz)')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WDb-nHAi3l3i"
      },
      "source": [
        "# for more variables check out the additional notebook (load_steinmetz_extra) which includes LFP, waveform shapes and exact spike times (non-binned)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}