01_latent_neural_states.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "629a16c1-3b4d-4429-b108-ead0f56d4212",
   "metadata": {},
   "source": [
    "# Inferring Latent Neural States\n",
    "\n",
    "Let's analyze some neural data using popular dimensionality reduction methods.\n",
    "We will use the folloiwng methods with progressively better modeling assumptions.\n",
    "\n",
    "- PCA (Principal Components Analysis)\n",
    "  - Gaussian observation\n",
    "  - Independent identical gaussian noise per neuron  \n",
    "- GPFA (Gaussian Process Factor Analysis)\n",
    "  - Gaussian observation\n",
    "  - Unequal magnitude of noise per neuron\n",
    "  - Smoothness assumption on the latent trajectory\n",
    "- vLGP (varational latent Gaussian Process)\n",
    "  - Poisson observation\n",
    "  - Unequal magnitude of noise per neuron\n",
    "  - Smoothness assumption on the latent trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be30d4c6-7fb6-4275-b5d7-dbbb9529a703",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.ndimage\n",
    "import matplotlib.pyplot as plt\n",
    "from einops import rearrange\n",
    "from sklearn.linear_model import Ridge\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "\n",
    "from scipy.linalg import LinAlgWarning\n",
    "import warnings\n",
    "warnings.filterwarnings(action='ignore', category=LinAlgWarning, module='sklearn')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eb5c179-1577-48f1-852f-ecc5a7337c52",
   "metadata": {},
   "source": [
    "We will use two datasets and use the same analysis pipeline. The next two subsections load each dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "294ec105-85e4-46d6-858c-1825207c4dd0",
   "metadata": {},
   "source": [
    "## Load Monkey delayed-reaching task data\n",
    "\n",
    "We will visualize data collected from the motor cortex (M1) of a monkey during a delayed reaching task. This data has been preprocessed such that it is aligned to (-100, 450) ms as the monkey reaches toward a target location. \n",
    "\n",
    "#### Reference:\n",
    "Churchland, Mark; Kaufman, Matthew (2022) MC_Maze_Large: macaque primary motor and dorsal premotor cortex spiking activity during delayed reaching (Version 0.220113.0407) [Data set]. DANDI archive. https://doi.org/10.48324/dandi.000138/0.220113.0407"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a9865f6-11de-41cd-8af7-5d0e70060b4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "baseDir = 'mc_maze/data/'\n",
    "m5 = h5py.File(baseDir + 'monkey.hdf5', 'r')\n",
    "y = m5['spk-train'][()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "069854f9-d8b2-41c6-9c17-6b1a0008f568",
   "metadata": {},
   "outputs": [],
   "source": [
    "nTrial = m5['pos-train'].shape[0]\n",
    "nT = m5['pos-train'].shape[1]\n",
    "nNeuron = m5['spk-train'].shape[2]\n",
    "dt = 0.005  # 5 ms bin\n",
    "T = dt * nT\n",
    "n_latent_mc_maze = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24ee2558-2156-4002-a665-102e0f350acb",
   "metadata": {},
   "outputs": [],
   "source": [
    "kTrial = 100\n",
    "raster = []\n",
    "for kNeuron in range(nNeuron):\n",
    "    raster.append(np.nonzero(y[kTrial,:,kNeuron])[0]/nT*T)\n",
    "plt.eventplot(raster, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.title('raster plot'); plt.ylabel('neurons');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccf1d731-a06c-458e-bf88-dc5c65106882",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(100,120):\n",
    "    plt.plot(m5['pos-train'][i,:,0], m5['pos-train'][i,:,1])\n",
    "    \n",
    "plt.xlabel('X hand position'); plt.ylabel('Y hand position'); plt.grid(); plt.title('center out reaching trajectory')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23eb9e35-36ad-43f2-b38d-22657f1e29dd",
   "metadata": {},
   "source": [
    "## Load Van der Pol data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bd7f7a4-9b66-479b-b7df-8d991fc316ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading data from ./data/vdp_noisy.h5\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "from code_pack.plotting import plot_two_d_vector_field_from_data, raster_to_events\n",
    "from code_pack.generate_vdp_data import generate_van_der_pol, generate_noisy_van_der_pol\n",
    "file_name = \"vanderpol/data/poisson_obs.h5\"\n",
    "\n",
    "#load data\n",
    "vdp_data = h5py.File(file_name, 'r')\n",
    "\n",
    "# dynamics parameters\n",
    "system_parameters = {}\n",
    "system_parameters['mu'] = vdp_data['mu']\n",
    "system_parameters['tau_1'] = vdp_data['tau_1']\n",
    "system_parameters['tau_2'] = vdp_data['tau_2']\n",
    "system_parameters['sigma'] = vdp_data['sigma']\n",
    "system_parameters['scale'] = np.array(vdp_data['scale'])\n",
    "system_parameters['sigma'] = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4caa05a4-c351-46a6-bd1e-a4e59afd36ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# plotting trajectories of the dataset\n",
    "X = np.array(vdp_data['X'])\n",
    "fig, ax = plt.subplots(1, 1, figsize=(5,5))\n",
    "_ = ax.plot(X[0,:,0], X[0,:,1])\n",
    "ax.scatter(X[0, 0, 0], X[0, 0, 1], marker='o', color='red', zorder=10, s=100, label='start')\n",
    "ax.scatter(X[0, -1, 0], X[0, -1, 1], marker='x', color='red', zorder=10, s=100, label='end')\n",
    "dynamic_func = lambda inp : generate_noisy_van_der_pol(inp, np.array([0.0, 5e-3]), system_parameters)\n",
    "axs_range = {'x_min':-1.5, 'x_max':1.5, 'y_min':-1.5, 'y_max':1.5}\n",
    "plot_two_d_vector_field_from_data(dynamic_func, ax, axs_range)\n",
    "\n",
    "ax.legend()\n",
    "ax.set_title('sample trajectory (true state)');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a46ed3b4-d09f-4a58-8d0a-a9a51a2b6dde",
   "metadata": {},
   "outputs": [],
   "source": [
    "C_tilde = np.array(vdp_data['C_tilde'])\n",
    "idx = np.lexsort((C_tilde[:,0], C_tilde[:,1]), axis=0) # sort the loading\n",
    "\n",
    "# showing the spike raster generated from noisy Vdp\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15, 3), sharex=True, sharey=True)\n",
    "events = raster_to_events(np.array(vdp_data['Y'])[0,:,:])\n",
    "events_softplus = raster_to_events(np.array(vdp_data['Y_softplus'])[0,:,:])\n",
    "events_axis_aligned = raster_to_events(np.array(vdp_data['Y_axis'])[0,:,idx].transpose())\n",
    "axs[0].eventplot(events, linewidths=0.5, color='blue');\n",
    "axs[1].eventplot(events_softplus, linewidths=0.5, color='blue');\n",
    "axs[2].eventplot(events_axis_aligned, linewidths=0.5, color='blue');\n",
    "axs[0].set_title(f'$\\exp()$');\n",
    "axs[1].set_title(f'softplus$()$');\n",
    "axs[2].set_title(f'axis aligned');\n",
    "axs[0].set_xlabel(\"Time\");\n",
    "axs[0].set_ylabel(\"Neuron\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad53dbbe-e889-48b0-befc-b84dc0fc81c8",
   "metadata": {},
   "source": [
    "# Choose your dataset to analyze subsequently below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21fdd64-0400-4e83-a739-c56acd440b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Choose dataset\n",
    "use_data = 'vanderpol'\n",
    "\n",
    "if use_data=='vanderpol':\n",
    "    vdp_data = h5py.File(\"vanderpol/data/poisson_obs.h5\", 'r')\n",
    "    y = np.array(vdp_data['Y'])\n",
    "    dt = 5e-3  # time bin size\n",
    "    gaussian_filter_sigma = 12.0\n",
    "    nTrial = vdp_data['X'].shape[0]\n",
    "    nT = vdp_data['X'].shape[1]\n",
    "    nNeuron = vdp_data['Y'].shape[2]\n",
    "    n_latent = 2\n",
    "\n",
    "elif use_data=='monkey':\n",
    "    m5 = h5py.File(baseDir + 'monkey.hdf5', 'r')\n",
    "    y = m5['spk-train'][()]\n",
    "    dt = 0.005  # 5 ms bin\n",
    "    gaussian_filter_sigma = 0.050/dt\n",
    "    nTrial = m5['pos-train'].shape[0]\n",
    "    nT = m5['pos-train'].shape[1]\n",
    "    nNeuron = m5['spk-train'].shape[2]\n",
    "    n_latent = n_latent_mc_maze\n",
    "    \n",
    "else:\n",
    "    raise Exceptiontion(\"Typo?!\")\n",
    "\n",
    "T = dt * nT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73d41273-1524-4e08-9e47-b9e57e1fae1b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## PCA\n",
    "\n",
    "In order to perform PCA, we first concatenate the the trials such that the data is of the form (trial x time) x neurons. We then smooth the data with a gaussian kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca3744c-e0d6-44e7-a835-6727d947149b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# smoothing data with a gaussian kernel\n",
    "data_stacked = rearrange(y, 'trial time neurons -> (trial neurons) time')\n",
    "data_smooth = scipy.ndimage.gaussian_filter1d(input = data_stacked, sigma=gaussian_filter_sigma, axis=1)\n",
    "data_smooth = rearrange(data_smooth, '(trial neurons) time -> (trial time) neurons', trial=nTrial, neurons=nNeuron)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb1c1ae4-f0c8-4920-b51b-3b94f6981126",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_centered = data_smooth - np.mean(data_smooth, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd2f3a8-3e25-4eba-b97e-46c40fdf3e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "tidx = slice(nT,2*nT)\n",
    "fig, ax = plt.subplots(1, 1, figsize =(10, 5))\n",
    "tr = np.arange(0, T, dt)\n",
    "ax.plot(tr, data_centered[tidx, 0:10]);\n",
    "ax.set_xlabel(\"time\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8234b0a-2566-4d2b-bbd7-fc55ec03a4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PCA using SVD\n",
    "u, s, vh = np.linalg.svd(data_centered, full_matrices=False)\n",
    "u.shape, s.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83135f9c-8456-4a54-9b5d-e13a571b8d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "norm_sv = s**2/np.sum(s**2)\n",
    "top2sv = np.sum(norm_sv[:2])\n",
    "print(\"Total observations explained by the first two principal components: {0:.2f}%\".format(top2sv*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c39cb1-bbd1-4c01-8e46-7e03d83c9503",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(15, 3))\n",
    "\n",
    "axs[0].plot(norm_sv * 100, 'o-')\n",
    "axs[1].plot(norm_sv.cumsum() * 100, 'o-')\n",
    "axs[1].set_ylim([0, 100])\n",
    "axs[2].plot(20*np.log10(norm_sv), 'o-')\n",
    "\n",
    "[(axs[k].grid(), axs[k].set_title(f''), axs[k].set_xlabel(\"PC (ordered)\")) for k in range(3)]\n",
    "axs[0].set_ylabel(\"Variance explained per PC ($\\%$)\"); \n",
    "axs[1].set_ylabel(\"Cumulative variance explained ($\\%$)\");\n",
    "axs[2].set_ylabel(\"Variance explained (dB)\"); \n",
    "fig.suptitle(\"What's the dimensionality? Inspecting variance explained by each PC defined dim\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "035cf875-1ce2-412b-bc35-51392022dca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualizing top two PCs\n",
    "topLu = u[:, :n_latent]\n",
    "X_hat_PCA = rearrange(topLu, '(trial time) pcs -> trial time pcs', trial=nTrial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e720c871-6377-4141-817e-a8bddfab93c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in range(10):\n",
    "    plt.plot(X_hat_PCA[k,:,0],  X_hat_PCA[k,:,1])\n",
    "    plt.plot(X_hat_PCA[k,-1,0], X_hat_PCA[k,-1,1], 'o')\n",
    "    \n",
    "plt.xlabel('PC1'); plt.ylabel('PC2'); plt.grid(); plt.title('2D slice')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b38478b-5db5-4e5a-9790-7ab8f60faa94",
   "metadata": {},
   "source": [
    "## GPFA\n",
    "\n",
    "We are using the implementation included in the Elephant package:\n",
    "https://elephant.readthedocs.io/en/latest/reference/gpfa.html\n",
    "\n",
    " - Yu, B. M., Cunningham, J. P., Santhanam, G., Ryu, S. I., Shenoy, K. V., & Sahani, M. (2009). Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity. Journal of Neurophysiology, 102(1), 614–635."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cc9d640-1c4a-44f2-99e8-0e15bee6ddd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from elephant.gpfa import GPFA\n",
    "import neo\n",
    "import quantities as pq\n",
    "warnings.filterwarnings(action='ignore', category=UserWarning, module='elephant')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1b8430f-9c18-42ea-8985-55559ac9b577",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---- Convert to neo.SpikeTrains ---- #\n",
    "def array_to_spiketrains(array, bin_size):\n",
    "    \"\"\"Convert B x T x N spiking array to list of list of SpikeTrains\"\"\"\n",
    "    stList = []\n",
    "\n",
    "    for trial in range(array.shape[0]):\n",
    "        trialList = []\n",
    "        for channel in range(array.shape[2]):\n",
    "            times = np.nonzero(array[trial, :, channel])[0]\n",
    "            counts = array[trial, times, channel].astype(int)\n",
    "            times = np.repeat(times, counts)\n",
    "            st = neo.SpikeTrain(times*bin_size, t_stop=array.shape[1]*bin_size)\n",
    "            trialList.append(st)\n",
    "        stList.append(trialList)\n",
    "    return stList\n",
    "\n",
    "Y_st_train = array_to_spiketrains(y, dt*pq.s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b2a7c3-7005-461c-b574-a26909505dd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ---- Run GPFA ---- #\n",
    "gpfa = GPFA(bin_size=(dt * pq.s), x_dim=n_latent)\n",
    "gpfa_val_result = gpfa.fit_transform(Y_st_train, returned_data=['latent_variable', 'VsmGP'])\n",
    "length_scales = gpfa.params_estimated['gamma']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50f882d1-4bb2-4b3a-9d38-158c0186f8c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hat_GPFA = rearrange(np.stack(gpfa_val_result['latent_variable'], 0), 'trials lat time -> trials time lat')\n",
    "P_hat_GPFA = rearrange(np.stack(gpfa_val_result['VsmGP'], 0)[:, np.arange(X_hat_GPFA.shape[1]), np.arange(X_hat_GPFA.shape[1])], 'trials time lat -> trials time lat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "378f94ff-92cb-4101-99cf-09e8653f4781",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in range(10):\n",
    "    plt.plot(X_hat_GPFA[k,:,0], X_hat_GPFA[k,:,1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc83ff87-7a68-4ba8-b5c2-7a14962db01b",
   "metadata": {},
   "source": [
    "## vLGP\n",
    "\n",
    " - Zhao, Y., & Park, I. M. (2017). Variational Latent Gaussian Process for Recovering Single-Trial Dynamics from Population Spike Trains. Neural Computation, 29(5), 1293–1316.\n",
    " - Nam, H. (2015). Poisson Extension of Gaussian Process Factor Analysis for Modelling Spiking Neural Populations (J. Macke (ed.)). Eberhard-Karls-Universität Tübingen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "536afea6-c90a-42de-861f-eff8fb37d885",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vlgpax.kernel import RBF\n",
    "from vlgpax import Session, vi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e352d8e-36c2-4f44-b224-21ab19fcb761",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = Session(dt)\n",
    "\n",
    "# Session is the top level container of data. Two arguments, binsize and unit of time, are required at construction.\n",
    "for i, yy in enumerate(y):\n",
    "    session.add_trial(i + 1, y = yy)  # Add trials to the session.\n",
    "\n",
    "# Build the model\n",
    "kernel = RBF(scale = 1., lengthscale = 25 * dt)  # RBF kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13a998a-1f0c-4aed-9d41-e2a16a0e769b",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_seed = 20221011\n",
    "np.random.seed(random_seed)\n",
    "session, params = vi.fit(session, n_factors=n_latent, kernel=kernel, seed=random_seed, max_iter=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0564d5b4-b277-4724-b409-293e507fd4f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_hat_VLGP = rearrange(session.z, '(trials time) lat -> trials time lat', time=nT)\n",
    "P_hat_VLGP = rearrange(session.v, '(trials time) lat -> trials time lat', time=nT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac9bca0-3010-491d-8c06-dde289aa8422",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplots(1,3,figsize=(12,4))\n",
    "\n",
    "plt.subplot(1,3,1)\n",
    "for k in range(10):\n",
    "    plt.plot(X_hat_PCA[k,:,0],  X_hat_PCA[k,:,1])\n",
    "    plt.plot(X_hat_PCA[k,-1,0], X_hat_PCA[k,-1,1], 'o')\n",
    "plt.xticks([]); plt.yticks([]); plt.gca().axis('equal')\n",
    "plt.title('PCA')\n",
    "    \n",
    "plt.subplot(1,3,2)\n",
    "for k in range(10):\n",
    "    plt.plot(X_hat_GPFA[k,:,0],  X_hat_GPFA[k,:,1])\n",
    "    plt.plot(X_hat_GPFA[k,-1,0], X_hat_GPFA[k,-1,1], 'o')\n",
    "plt.xticks([]); plt.yticks([]); plt.gca().axis('equal')\n",
    "plt.title(\"GPFA\");\n",
    "        \n",
    "plt.subplot(1,3,3)\n",
    "for k in range(10):\n",
    "    plt.plot(X_hat_VLGP[k,:,0],  X_hat_VLGP[k,:,1])\n",
    "    plt.plot(X_hat_VLGP[k,-1,0], X_hat_VLGP[k,-1,1], 'o')\n",
    "plt.xticks([]); plt.yticks([]); plt.gca().axis('equal')\n",
    "plt.title(\"vLGP\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1308911",
   "metadata": {},
   "source": [
    "### Evaluating LVMs\n",
    "One simple metric to evaluate the capability of these LVMs is the expected log-likelihood given by\n",
    "\n",
    "$$\n",
    "\\ell(m_{1:T}, P_{1:T}) = \\sum\\nolimits_t \\mathbb{E}_{q(z_t; \\,m_t, P_t)} \\log p(y_t \\mid z_t)\n",
    "$$\n",
    "\n",
    "Recall that for PCA, the `posterior' degrades into a point estimate, so that we should supply $P_t = 0$ when evaluating PCA\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d471fb",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import code_pack.utils as utils\n",
    "\n",
    "readout_iter = 500\n",
    "C_GPFA = utils.estimate_readout_matrix(y, X_hat_GPFA, None, dt, readout_iter)\n",
    "ell_GPFA = utils.expected_ll_poisson(y, X_hat_GPFA, P_hat_GPFA, C_GPFA, dt)\n",
    "\n",
    "C_PCA = utils.estimate_readout_matrix(y, X_hat_PCA, None, dt, readout_iter)\n",
    "ell_PCA = utils.expected_ll_poisson(y, X_hat_PCA, np.zeros_like(P_hat_GPFA), C_PCA, dt)\n",
    "\n",
    "C_vLGP = utils.estimate_readout_matrix(y, np.asarray(X_hat_VLGP), None, dt, readout_iter)\n",
    "ell_vLGP = utils.expected_ll_poisson(y, np.asarray(X_hat_VLGP), np.asarray(P_hat_VLGP), C_vLGP, dt)\n",
    "\n",
    "print(f'PCA ell: {ell_PCA} bits/spk/neuron')\n",
    "print(f'GPFA ell: {ell_GPFA} bits/spk/neuron')\n",
    "print(f'VLGP ell: {ell_vLGP} bits/spk/neuron')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e98a0f35-6181-4950-acab-cc803b983b95",
   "metadata": {},
   "source": [
    "Run the following code only for the monkey data (MC_MAZE).\n",
    "\n",
    "We can compare the decoding performance for recovering the velocity of the arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d05c88f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if(use_data=='monkey'):\n",
    "    velocity_train = np.asarray(m5['vel-train'])\n",
    "\n",
    "    rates_PCA = dt * np.exp(C_PCA(torch.tensor(X_hat_PCA, dtype=torch.float32)).detach().numpy())\n",
    "    print(velocity_train.shape, rates_PCA.shape, X_hat_PCA.shape)\n",
    "    gscv_PCA = GridSearchCV(Ridge(), {'alpha': np.logspace(-7, -1, 100)})\n",
    "    gscv_PCA.fit(rates_PCA.reshape(-1, nNeuron), velocity_train.reshape(-1, 2))\n",
    "\n",
    "    rates_GPFA = dt * np.exp(C_GPFA(torch.tensor(X_hat_GPFA, dtype=torch.float32)).detach().numpy())\n",
    "    gscv_GPFA = GridSearchCV(Ridge(), {'alpha': np.logspace(-7, -1, 100)})\n",
    "    gscv_GPFA.fit(rates_GPFA.reshape(-1, nNeuron), velocity_train.reshape(-1, 2));\n",
    "\n",
    "    rates_vLGP = dt * np.exp(C_vLGP(torch.tensor(np.asarray(X_hat_VLGP), dtype=torch.float32)).detach().numpy())\n",
    "    gscv_vLGP = GridSearchCV(Ridge(), {'alpha': np.logspace(-7, -1, 100)})\n",
    "    gscv_vLGP.fit(rates_vLGP.reshape(-1, nNeuron), velocity_train.reshape(-1, 2));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43b15b91",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if(use_data=='monkey'):\n",
    "    predicted_velocity_pca = np.zeros((nTrial, nT, 2))\n",
    "    predicted_velocity_gpfa = np.zeros((nTrial, nT, 2))\n",
    "    predicted_velocity_vlgp = np.zeros((nTrial, nT, 2))\n",
    "\n",
    "    for n in range(nTrial):\n",
    "        predicted_velocity_pca[n] = gscv_PCA.predict(dt * np.exp(C_PCA(torch.tensor(X_hat_PCA[n], dtype=torch.float32)).detach().numpy()))\n",
    "        predicted_velocity_gpfa[n] = gscv_GPFA.predict(dt * np.exp(C_GPFA(torch.tensor(X_hat_GPFA[n], dtype=torch.float32)).detach().numpy()))\n",
    "        predicted_velocity_vlgp[n] = gscv_vLGP.predict(dt * np.exp(C_vLGP(torch.tensor(np.asarray(X_hat_VLGP)[n], dtype=torch.float32)).detach().numpy()))\n",
    "\n",
    "    predicted_trajectory_pca = np.cumsum(predicted_velocity_pca, axis=1) * dt\n",
    "    predicted_trajectory_gpfa = np.cumsum(predicted_velocity_gpfa, axis=1) * dt\n",
    "    predicted_trajectory_vlgp = np.cumsum(predicted_velocity_vlgp, axis=1) * dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "939eb5c6",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if(use_data=='monkey'):\n",
    "    fig, axs = plt.subplots(1, 4, figsize=(15, 3))\n",
    "\n",
    "    for i in range(nTrial - 1):\n",
    "        line_color = np.asarray(m5['colors-train'])[i]\n",
    "        axs[0].plot(np.asarray(m5['pos-train'])[i,:,0],np.asarray(m5['pos-train'])[i,:,1], color=line_color)\n",
    "        axs[1].plot(predicted_trajectory_pca[i, :, 0], predicted_trajectory_pca[i, :, 1], color=line_color)\n",
    "        axs[2].plot(predicted_trajectory_gpfa[i, :, 0], predicted_trajectory_gpfa[i, :, 1], color=line_color)\n",
    "        axs[3].plot(predicted_trajectory_vlgp[i, :, 0], predicted_trajectory_vlgp[i, :, 1], color=line_color)\n",
    "\n",
    "    axs[0].set_title(f'True trajectories')\n",
    "    axs[1].set_title(f'PCA R2: {gscv_PCA.best_score_:.3f}')\n",
    "    axs[2].set_title(f'GPFA R2: {gscv_GPFA.best_score_:.3f}')\n",
    "    axs[3].set_title(f'vLGP R2: {gscv_vLGP.best_score_:.3f}')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}