paper/fig2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import numpy as np\n",
    "from scipy.stats import zscore\n",
    "from rastermap import Rastermap, utils\n",
    "\n",
    "# path to paper code\n",
    "sys.path.insert(0, '/github/rastermap/paper')\n",
    "from loaders import tuning_curves_VR\n",
    "import fig2\n",
    "\n",
    "# path to directory with data etc\n",
    "### *** CHANGE THIS TO WHEREEVER YOU ARE DOWNLOADING THE DATA ***\n",
    "root = \"/media/carsen/ssd2/rastermap_paper/\"\n",
    "# (in this folder we have a \"data\" folder and a \"results\" folder)\n",
    "os.makedirs(os.path.join(root, \"data\"), exist_ok=True)\n",
    "os.makedirs(os.path.join(root, \"results\"), exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load virtual reality task data\n",
    "\n",
    "(this data will be available upon publication of the paper)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat = np.load(os.path.join(root, \"data/\", \"corridor_neur.npz\"))\n",
    "corridor = np.load(os.path.join(root, \"data/\", \"corridor_behavior.npz\"))\n",
    "\n",
    "xpos, ypos, spks = dat[\"xpos\"], dat[\"ypos\"], dat[\"spks\"]\n",
    "spks = zscore(spks, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### run rastermap and compute tuning curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Rastermap(n_clusters=100, n_PCs=200, \n",
    "                    time_lag_window=10, locality=0.75).fit(spks)\n",
    "isort = model.isort \n",
    "cc_nodes = model.cc\n",
    "bin_size = 100\n",
    "sn = zscore(utils.bin1d(spks[isort], bin_size, axis=0), axis=1)\n",
    "corridor_tuning = tuning_curves_VR(sn, corridor[\"VRpos\"], corridor[\"corridor_starts\"])\n",
    "\n",
    "# sort in time\n",
    "model2 = Rastermap(n_clusters=100, n_splits=0, locality=0.,\n",
    "                             n_PCs=200).fit(sn.T)\n",
    "isort2 = model2.isort\n",
    "\n",
    "np.savez(os.path.join(root, \"results\", \"corridor_proc.npz\"),\n",
    "         sn=sn, xpos=xpos, ypos=ypos, isort=isort, isort2=isort2,\n",
    "        cc_nodes=cc_nodes, corridor_tuning=corridor_tuning)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### make figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# root path has folder \"results\" with saved results\n",
    "# will save figures to \"figures\" folder\n",
    "os.makedirs(os.path.join(root, \"figures/\"), exist_ok=True)\n",
    "fig2.fig2(root)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### supplementary analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import metrics\n",
    "\n",
    "ys = [metrics.run_TSNE(model.Usv), \n",
    "      metrics.run_UMAP(model.Usv)]\n",
    "\n",
    "snys = []\n",
    "ctunings = []\n",
    "for k in range(2):\n",
    "    isorty = ys[k][:,0].argsort()\n",
    "    sny = zscore(utils.bin1d(spks[isorty], 100, axis=0))\n",
    "    ctuning = tuning_curves_VR(sny, corridor[\"VRpos\"], corridor[\"corridor_starts\"])\n",
    "    snys.append(sny)\n",
    "    ctunings.append(ctuning)\n",
    "\n",
    "np.savez(os.path.join(root, \"results\", \"corridor_supp.npz\"),\n",
    "         snys=snys, ctunings=ctunings, \n",
    "         corridor_starts=corridor[\"corridor_starts\"], \n",
    "         corridor_widths=corridor[\"corridor_widths\"], \n",
    "         reward_inds=corridor[\"reward_inds\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = np.load(os.path.join(root, \"results\", \"corridor_supp.npz\"))\n",
    "fig = fig3._suppfig_vr_algs(**d)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('rastermap')",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "998540cc2fc2836a46e99cd3ca3c37c375205941b23fd1eb4b203c48f2be758f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}