paper/fig3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "import numpy as np\n",
    "import torch\n",
    "from scipy.stats import zscore\n",
    "from neuropop import nn_prediction\n",
    "from rastermap import Rastermap, utils\n",
    "\n",
    "### use cuda version of torch if available\n",
    "# (otherwise use >>> device = torch.device('cpu'))\n",
    "device = torch.device('cuda')\n",
    "\n",
    "# path to paper code\n",
    "sys.path.insert(0, '/github/rastermap/paper')\n",
    "import fig3\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 spont 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/\", \"spont_data.npz\"))\n",
    "spks, U, sv, V = dat[\"spks\"], dat[\"U\"], dat[\"sv\"], dat[\"V\"]\n",
    "xpos, ypos = dat[\"xpos\"], dat[\"ypos\"]\n",
    "tcam, tneural = dat[\"tcam\"], dat[\"tneural\"]\n",
    "run, beh, beh_names = dat[\"run\"], dat[\"beh\"], dat[\"beh_names\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### predict neural activity from behavior and run rastermap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "torch.cuda.manual_seed(0)\n",
    "\n",
    "### fit linear and non-linear model from behavior to neural activity\n",
    "Vfit = V.copy() * sv\n",
    "for i in range(2):\n",
    "    if i==1:\n",
    "        pred_model = nn_prediction.PredictionNetwork(n_in=beh.shape[-1], n_kp=22, identity=False, \n",
    "                                                n_filt=10, n_latents=0,\n",
    "                                n_out=Vfit.shape[-1], n_core_layers=1,\n",
    "                                relu_wavelets=False, relu_latents=False)\n",
    "    else:\n",
    "        pred_model = nn_prediction.PredictionNetwork(n_in=beh.shape[-1], n_kp=22, n_out=Vfit.shape[-1], )\n",
    "    pred_model.to(device)\n",
    "\n",
    "    y_pred_all, ve_all, itest = pred_model.train_model(beh, Vfit, tcam, tneural, delay=-1,\n",
    "                                                        learning_rate=1e-3, n_iter=400,\n",
    "                                                    device=device, verbose=True)\n",
    "    if i==1:\n",
    "        y_pred_nn = y_pred_all.copy()\n",
    "\n",
    "### run rastermap\n",
    "model = Rastermap(n_clusters=100, \n",
    "                    n_PCs=128, \n",
    "                    locality=0.0,\n",
    "                    time_lag_window=5,\n",
    "                    ).fit(spks)\n",
    "cc_nodes = model.cc.copy()\n",
    "isort = model.isort\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### bin full data into superneurons\n",
    "nbin = 50\n",
    "sn = zscore(utils.bin1d(spks[isort], nbin, axis=0), axis=1)\n",
    "np.random.seed(0)\n",
    "sn_rand = zscore(utils.bin1d(spks[np.random.permutation(spks.shape[0])], nbin, axis=0), axis=1)\n",
    "\n",
    "# sort in time\n",
    "model2 = Rastermap(n_clusters=100, locality=0.,\n",
    "                    n_PCs=128).fit(sn.T)\n",
    "isort2 = model2.isort\n",
    "\n",
    "### bin test data and prediction into superneurons\n",
    "sn_test = utils.bin1d(spks[isort][:,itest.flatten()], nbin, axis=0)\n",
    "sn_pred_test = utils.bin1d(U[isort] @ y_pred_nn.T, nbin, axis=0)\n",
    "sn_pred_test -= sn_test.mean(axis=1, keepdims=True)\n",
    "sn_pred_test /= sn_test.std(axis=1, keepdims=True)\n",
    "sn_test = zscore(sn_test, axis=1)\n",
    "cc_pred = (sn_test * zscore(sn_pred_test, axis=1)).mean(axis=1)\n",
    "# sort and bin PCs for maxstim estimation\n",
    "U_sn = utils.bin1d(U[isort], nbin)\n",
    "\n",
    "### maxstim estimation for superneurons (receptive fields)\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "torch.cuda.manual_seed(0)\n",
    "ms_model = nn_prediction.MaxStimModel(pred_model)\n",
    "ms_model.requires_grad = False\n",
    "u = torch.from_numpy(U_sn).to(device)\n",
    "u.requires_grad = False\n",
    "xr = ms_model.train_batch(u, n_iter=200, learning_rate=1e-2)\n",
    "rfs = xr.detach().cpu().numpy()\n",
    "\n",
    "\n",
    "# save results\n",
    "np.savez(os.path.join(root, \"results\", \"spont_proc.npz\"), sn_test=sn_test, \n",
    "            sn_pred_test=sn_pred_test, itest=itest, \n",
    "            sn=sn, sn_rand=sn_rand, isort2=isort2,\n",
    "            rfs=rfs, isort=isort, cc_nodes=cc_nodes,\n",
    "            xpos=xpos, ypos=ypos,\n",
    "            tcam=tcam, tneural=tneural, \n",
    "            run=run, beh=beh, beh_names=beh_names)\n"
   ]
  },
  {
   "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",
    "fig3.fig3(root, save_figure=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### make supp fig "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3.suppfig_random(root, save_figure=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### supp behavior sorting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuropop.utils import resample_data\n",
    "\n",
    "# resample behavior for correlation sorting\n",
    "beh = resample_data(beh, tcam, tneural)\n",
    "\n",
    "whisk_speed = (np.diff(beh[:,1:3], axis=0)**2).sum(axis=-1)**0.5\n",
    "nose_speed = (np.diff(beh[:,3:5], axis=0)**2).sum(axis=-1)**0.5\n",
    "vars = np.stack((run[:-1], whisk_speed, nose_speed, beh[:-1,0]), axis=-1)\n",
    "isorts_beh = []\n",
    "corrs_beh = []\n",
    "sn_beh = []\n",
    "itrain = np.ones(spks.shape[1]-1, \"bool\")\n",
    "itrain[itest.flatten()] = False\n",
    "for i in range(vars.shape[1]):\n",
    "    corr = (spks[:,:-1][:,itrain] * zscore(vars[itrain,i])).mean(axis=1)\n",
    "    isort = corr.argsort()\n",
    "    isorts_beh.append(isort)\n",
    "    corrs_beh.append(corr)\n",
    "    sn_beh.append(zscore(utils.bin1d(spks[isort][:,itest.flatten()], nbin, axis=0), axis=1))\n",
    "corrs_beh = np.array(corrs_beh)\n",
    "isorts_beh = np.array(isorts_beh)\n",
    "sn_beh = np.array(sn_beh)\n",
    "\n",
    "\n",
    "np.savez(os.path.join(root, \"results/spont_corrs_beh.npz\"), \n",
    "         vars=vars,\n",
    "         sn_beh=sn_beh,\n",
    "         corrs_beh=corrs_beh,\n",
    "         isorts_beh=isorts_beh)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make figure\n",
    "fig3.suppfig_beh(root, save_figure=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### supp locality changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imp \n",
    "from metrics import distance_matrix\n",
    "from tqdm import tqdm\n",
    "\n",
    "localities = [0, 0.25, 0.5, 0.75, 1.0]\n",
    "ccs_spont, isorts_spont, BBts, scores_spont = [], [], [], []\n",
    "\n",
    "for k, loc in tqdm(enumerate(localities)):\n",
    "    model = Rastermap(n_clusters=100, n_PCs=128, locality=loc,\n",
    "                        time_lag_window=5, verbose=False).fit(spks)\n",
    "    cc_nodes = model.cc.copy()\n",
    "    BBt = model.BBt\n",
    "    ccs_spont.append(cc_nodes)\n",
    "    isorts_spont.append(model.isort)\n",
    "    BBts.append(BBt.copy())\n",
    "\n",
    "    inds = np.triu_indices(BBt.shape[0], k=1)\n",
    "    cc = model.cc[inds[0], inds[1]].copy()\n",
    "    bb = BBt[inds[0], inds[1]]\n",
    "    score_global = (zscore(bb) * zscore(cc)).mean()\n",
    "\n",
    "    cc = model.cc.copy()\n",
    "    cc -= np.diag(np.diag(cc))\n",
    "    score_local = 0\n",
    "    ntot = 0\n",
    "    lscores = np.nan * np.zeros((len(cc), 2))\n",
    "    for j in range(len(cc)):\n",
    "        if j < len(cc)-1:\n",
    "            csort = cc[j].argsort()[::-1]\n",
    "            lscores[j, 0] = np.nonzero(csort==(j+1))[0][0]\n",
    "            ntot += 1\n",
    "        if j > 0:\n",
    "            csort = cc[:,j].argsort()[::-1]\n",
    "            lscores[j, 1] = np.nonzero(csort==(j-1))[0][0]\n",
    "            ntot += 1\n",
    "        if lscores[j,0]==0 or lscores[j,1]==0:\n",
    "            score_local += 1\n",
    "            if lscores[j,0]==1 or lscores[j,1]==1:\n",
    "                score_local += 1\n",
    "    score_local /= ntot \n",
    "\n",
    "    scores_spont.append(np.array([score_global, score_local]))\n",
    "    print(scores_spont[-1])\n",
    "    \n",
    "\n",
    "dat = np.load(os.path.join(root, \"data/\", \"corridor_neur.npz\"))\n",
    "xpos, ypos, spks = dat[\"xpos\"], dat[\"ypos\"], dat[\"spks\"]\n",
    "spks = zscore(spks, axis=1)\n",
    "\n",
    "isorts_vr, ccs_vr, BBts, scores_vr = [], [], [], []\n",
    "for loc in tqdm(localities):\n",
    "    model = Rastermap(n_clusters=100, n_PCs=200, \n",
    "                        time_lag_window=10, \n",
    "                        locality=loc, bin_size=100).fit(spks)\n",
    "    isort = model.isort \n",
    "    cc_nodes = model.cc\n",
    "    BBt = model.BBt\n",
    "    isorts_vr.append(isort)\n",
    "    ccs_vr.append(cc_nodes)\n",
    "    BBts.append(BBt)\n",
    "\n",
    "    inds = np.triu_indices(BBt.shape[0], k=1)\n",
    "    cc = model.cc[inds[0], inds[1]].copy()\n",
    "    bb = BBt[inds[0], inds[1]]\n",
    "    score_global = (zscore(bb) * zscore(cc)).mean()\n",
    "\n",
    "    cc = model.cc.copy()\n",
    "    cc -= np.diag(np.diag(cc))\n",
    "    score_local = 0\n",
    "    ntot = 0\n",
    "    lscores = np.nan * np.zeros((len(cc), 2))\n",
    "    for j in range(len(cc)):\n",
    "        if j < len(cc)-1:\n",
    "            csort = cc[j].argsort()[::-1]\n",
    "            lscores[j, 0] = np.nonzero(csort==(j+1))[0][0]\n",
    "            ntot += 1\n",
    "        if j > 0:\n",
    "            csort = cc[:,j].argsort()[::-1]\n",
    "            lscores[j, 1] = np.nonzero(csort==(j-1))[0][0]\n",
    "            ntot += 1\n",
    "        if lscores[j,0]==0 or lscores[j,1]==0:\n",
    "            score_local += 1\n",
    "            if lscores[j,0]==1 or lscores[j,1]==1:\n",
    "                score_local += 1\n",
    "    score_local /= ntot \n",
    "    \n",
    "    scores_vr.append(np.array([score_global, score_local]))\n",
    "    print(scores_vr[-1])\n",
    "\n",
    "print(BBts[0][0,1], BBts[-1][0,1])\n",
    "\n",
    "np.savez(os.path.join(root, \"results/asym_vr_spont.npz\"), \n",
    "         ccs_vr=np.array(ccs_vr), \n",
    "         ccs_spont=np.array(ccs_spont),\n",
    "         isorts_vr=np.array(isorts_vr), \n",
    "         isorts_spont=np.array(isorts_spont), \n",
    "         localities=localities, \n",
    "         BBts=np.array(BBts),\n",
    "         scores_vr=np.array(scores_vr),\n",
    "         scores_spont=np.array(scores_spont))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make figure\n",
    "fig3.suppfig_locality(root, save_figure=True)"
   ]
  }
 ],
 "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
}