paper/svca.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 neuropop import dimensionality\n",
    "\n",
    "# path to paper code\n",
    "sys.path.insert(0, '/github/rastermap/paper')\n",
    "from loaders import load_fish_data, load_visual_data, load_alexnet_data\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### compute SVCA for large datasets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scovs = []\n",
    "\n",
    "dat = np.load(os.path.join(root, \"data/\", \"spont_data.npz\"))\n",
    "spks = dat[\"spks\"]\n",
    "spks = zscore(spks, axis=1)\n",
    "scov, varcov = dimensionality.SVCA(spks)\n",
    "scovs.append(scov)\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",
    "scov, varcov = dimensionality.SVCA(spks)\n",
    "scovs.append(scov)\n",
    "\n",
    "### folder with \"subject_8\" folder\n",
    "data_path = os.path.join(root, \"data\")\n",
    "spks, F, xyz, stims, swimming, eyepos = load_fish_data(data_path, subject=8)\n",
    "spks = zscore(spks, axis=1)\n",
    "scov, varcov = dimensionality.SVCA(spks)\n",
    "scovs.append(scov)\n",
    "\n",
    "filename = os.path.join(root, \"data/\", \"TX61_3x.npz\")\n",
    "stim_filename = os.path.join(root, \"data/\", \"text5k_3x.mat\")\n",
    "out = load_visual_data(filename, stim_filename)\n",
    "spks = out[0]\n",
    "spks = zscore(spks, axis=1)\n",
    "scov, varcov = dimensionality.SVCA(spks)\n",
    "scovs.append(scov)\n",
    "\n",
    "env_ids = [\"PongNoFrameskip-v4\", \"SpaceInvadersNoFrameskip-v4\", \n",
    "           \"EnduroNoFrameskip-v4\", \"SeaquestNoFrameskip-v4\"]\n",
    "\n",
    "for env_id in env_ids:\n",
    "    dat = np.load(os.path.join(root, \"simulations/\", f\"qrdqn_{env_id}.npz\"))\n",
    "    spks = dat[\"spks\"]\n",
    "\n",
    "    x_std = spks.std(axis=1)\n",
    "    igood = x_std > 1e-3\n",
    "    print(igood.mean())\n",
    "\n",
    "    S = zscore(spks[igood], axis=1)\n",
    "    scov, varcov = dimensionality.SVCA(S)\n",
    "    scovs.append(scov)\n",
    "\n",
    "filename = os.path.join(root, \"data\", \"ann_fvs_Grayscale(224)_TX61_3X.npz\")\n",
    "sresp, ilayer, ipos, iconv, nmax = load_alexnet_data(filename)\n",
    "for l in range(5):\n",
    "    S = zscore(sresp[ilayer==l].copy(), axis=1)\n",
    "    scov, varcov = dimensionality.SVCA(S)\n",
    "    scovs.append(scov)\n",
    "\n",
    "scovs_all = np.nan * np.zeros((13, 1024))\n",
    "for k, scov in enumerate(scovs):\n",
    "    scovs_all[k, :len(scov)] = scov\n",
    "\n",
    "np.save(os.path.join(root, \"results\", \"scovs.npy\"), scovs_all)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### make figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fig_utils import *\n",
    "\n",
    "scovs_all = np.load(os.path.join(root, \"results\", \"scovs.npy\"))\n",
    "\n",
    "fig = plt.figure(figsize=(14, 7.5))\n",
    "grid = plt.GridSpec(3,5, figure=fig, left=0.06, right=0.98, top=0.96, bottom=0.08, \n",
    "                    wspace = 0.8, hspace = 0.8)\n",
    "transl = mtransforms.ScaledTranslation(-13 / 72, 20 / 72, fig.dpi_scale_trans)\n",
    "il = 0\n",
    "\n",
    "titles = [\"spontaneous activity, mouse\", \"virtual reality, mouse\", \"zebrafish wholebrain activity\", \"flashed images, mouse\",\n",
    "    \"QRDQN - Pong\", \"QRDQN - SpaceInvaders\", \"QRDQN - Enduro\", \"QRDQN - Seaquest\",\n",
    "    \"Alexnet - layer 1\", \"Alexnet - layer 2\", \"Alexnet - layer 3\", \"Alexnet - layer 4\", \"Alexnet - layer 5\"]\n",
    "\n",
    "for k, scov in enumerate(scovs_all):\n",
    "    if k<8:\n",
    "        if k==2 or k==3:\n",
    "            ki = 3 if k==2 else 2\n",
    "        else:\n",
    "            ki = k\n",
    "        ax = plt.subplot(grid[ki//4, ki%4])\n",
    "    else:\n",
    "        ax = plt.subplot(grid[2, k-8])\n",
    "    ss = scov.copy()\n",
    "    ss /= ss[0]\n",
    "    #ss *= len(ss)\n",
    "    alpha, ypred = dimensionality.get_powerlaw(ss, np.arange(11, 500))\n",
    "    #print(alpha)\n",
    "    ax.loglog(np.arange(1, len(ss)+1), ss, color=[0,0.5,1])\n",
    "    ax.plot(np.arange(len(ypred))+1, ypred, color=\"k\", lw=1)\n",
    "    ax.text(0.5, 0.7, rf\"$\\alpha$={alpha:.2f}\", transform=ax.transAxes)\n",
    "    ax.set_title(titles[k], fontsize=\"medium\")#, fontsize=\"small\")\n",
    "    ax.set_ylim([1/1024, 1.5])\n",
    "    ax.set_xlim([1, 1024])\n",
    "    ax.set_xticks([1, 10, 100, 1000])\n",
    "    ax.set_xlabel(\"PC #\")\n",
    "    ax.set_ylabel(\"shared covariance\")\n",
    "    ax.xaxis.get_minor_locator().set_params(numticks=99, subs=np.arange(0.1, 1, 0.1))\n",
    "\n",
    "fig.savefig(os.path.join(root, \"figures\", \"suppfig_powerlaws.pdf\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}