00_state_space_intuition.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fb9ed2ff-dcea-4a93-b574-0605640f7ddc",
   "metadata": {},
   "source": [
    "# Intuitions on State Space Model\n",
    "\n",
    "In this notebook, we will play around with the parameters of a state space model and generate various spike trains.\n",
    "These population activity patterns will be fun to look at. (hopefully)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63ec82d-8894-4fcf-aa88-1d83ed5c924d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c450d2fb-e3cd-459a-9b5f-934856891aed",
   "metadata": {},
   "source": [
    "## A simple 1-D latent process\n",
    "\n",
    "For illustration, we will use a sinusoid as the 1-D latent process.\n",
    "$$ x(t) = sin(2\\pi f\\cdot t) $$\n",
    "In this example, $x(t)$ represents the instantaneous state of the neural population of interest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d5c27a9-be0d-4ab0-be3c-7b798ad3d728",
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate a simple latent process\n",
    "nT = 1000\n",
    "T = 10\n",
    "frq = 0.3\n",
    "tr = np.linspace(0, T, nT)\n",
    "dt = tr[1] - tr[0]\n",
    "x = np.sin(2 * np.pi * frq * tr) # generate a sinusoid over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750cb1e3-18a5-4eb7-851c-dd479c6f674e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 3))\n",
    "plt.plot(tr, x); plt.title('1-D latent process'); plt.xlabel('time');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb978cc-72bc-4ede-8cde-cb5101d81177",
   "metadata": {},
   "source": [
    "## One Poisson neuron driven by the latent process\n",
    "\n",
    "We will generate spike trains from an inhomogeneous Poisson process with a time varying firing rate function $\\lambda(t)$.\n",
    "The spike count $y(t)$ in a small time bin of size $\\Delta$ is distributed as a Poisson distribution:\n",
    "$$ y(t) \\sim \\text{Poisson}(\\Delta\\lambda(t)) $$\n",
    "\n",
    "Importantly, the firing rate will be a function of $x(t)$, but not of past $x$ nor past $y$.\n",
    "$$ \\lambda(t) = g(x(t)) $$\n",
    "The only constraint for $g(\\cdot)$ is that the resulting firing rate has to be non-negative.\n",
    "A mathematically convenient function is the exponential function.\n",
    "\n",
    "$$ \\lambda(t) = \\exp(a x(t) + b) = \\exp(b)\\exp(a x(t)) $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e489539f-e458-49de-91a2-217f655d6ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 5\n",
    "b = -3\n",
    "lam = np.exp(a * x + b)\n",
    "y = np.random.poisson(lam*dt)\n",
    "\n",
    "plt.figure(figsize=(10, 2))\n",
    "plt.plot(tr, lam, label='firing rate');\n",
    "plt.eventplot(np.nonzero(y)[0]/nT*T, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ee64c02-2a4a-4594-9775-51634ace2794",
   "metadata": {},
   "source": [
    "## A population of Poisson neurons driven by a common 1-D latent process\n",
    "\n",
    "We can have more than one neuron that's driven by the same latent process.\n",
    "This way, we an have more observation dimensions than the latent state space dimension.\n",
    "Let's given them a random amount of \"drive\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57877799-aa62-41c3-86e7-05c90f71b35c",
   "metadata": {},
   "outputs": [],
   "source": [
    "nNeuron = 200\n",
    "C = 2 * np.random.randn(nNeuron)\n",
    "b = -2.0 + np.random.rand(nNeuron,1)\n",
    "lam = np.exp(np.outer(C, x) + b)\n",
    "y = np.random.poisson(lam*dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77fc44d-149c-4fbd-a65c-3a2aa4ea2100",
   "metadata": {},
   "source": [
    "We can make spike raster plot. But since we know the amount of drive (each value in $C$), we can sort the neurons accordingly as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c66547-3833-40c7-9375-e13265f15489",
   "metadata": {},
   "outputs": [],
   "source": [
    "cidx = np.argsort(C)\n",
    "\n",
    "raster = []\n",
    "rasterSorted = []\n",
    "for k in range(nNeuron):\n",
    "    raster.append(np.nonzero(y[k,:])[0]/nT*T)\n",
    "    rasterSorted.append(np.nonzero(y[cidx[k],:])[0]/nT*T)\n",
    "\n",
    "plt.subplots(1,2, figsize=(10, 4))\n",
    "plt.subplot(1,2,1)\n",
    "plt.eventplot(raster, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot'); plt.ylabel('neurons');\n",
    "plt.subplot(1,2,2)\n",
    "plt.eventplot(rasterSorted, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot (again)'); plt.ylabel('sorted neurons');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74846f2d-1c13-48a5-b278-736348c55aa9",
   "metadata": {},
   "source": [
    "## 2D latent space example\n",
    "\n",
    "Here we will build a 2D manifold with two independent processes.\n",
    "The first latent dimension will be same as above, but we will add $x_2(t)$ as a sawtooth function:\n",
    "$$ x_2(t) = t \\,\\, \\text{mod} \\, 1 $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f0a34da-dda3-4c8f-ad56-eb1cccabe19c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x2 = 1.5 * ((tr % 1) - 0.5)\n",
    "X = np.vstack([x, x2]) # (latent dim) x (time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c58c0bc-2d13-4cd8-aa4c-f2246e87ad7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplots(2,1,figsize=(10,4))\n",
    "plt.subplot(2,1,1);\n",
    "plt.plot(tr, x ); plt.ylabel('first latent dim'); plt.xlabel('time')\n",
    "plt.subplot(2,1,2);\n",
    "plt.plot(tr, x2); plt.ylabel('second latent dim'); plt.xlabel('time')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5f9ec7e-0f57-4925-b4d3-502d981264bc",
   "metadata": {},
   "source": [
    "Now that we have more than 1 latent dimension, we face a choice of how neurons relate to each of the latent dimes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "468d1ce3-bf62-4295-8d6c-e1ab3616023a",
   "metadata": {},
   "source": [
    "### Random projection observation\n",
    "\n",
    "Random projection assumes that each neuron is deriven by all the latent dimensions with a random amount.\n",
    "Under this assumption, the neural manifold is likely oblique to the axes, i.e., the neuron will be modulated by changes in any direction in the latent state space.\n",
    "Theoretical analysis of [Gao & Ganguli 2015] assumes random projections and showed that not many neurons need to be sampled (observed) to recover the manifold structure.\n",
    "In addition, the so-called *mixed-selectivity* appears as a result.\n",
    "\n",
    "- Gao, P., & Ganguli, S. (2015). On Simplicity and Complexity in the Brave New World of Large-Scale Neuroscience. Current Opinion in Neurobiology, 32, 148–155."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65628119-5f7f-4686-a342-0958f4386291",
   "metadata": {},
   "outputs": [],
   "source": [
    "dLatent = X.shape[0]\n",
    "C = 2 * np.random.randn(nNeuron, dLatent) # random projection\n",
    "lam = np.exp(C @ X + b)\n",
    "y = np.random.poisson(lam*dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d719f6ae-6e81-4f5b-af6e-86196b99940e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cidx1 = np.lexsort((C[:,0], C[:,1]), axis=0)\n",
    "cidx2 = np.lexsort((C[:,1], C[:,0]), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcfa3cb3-a58e-43d3-bea5-c288693a9e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "raster = []; rasterSorted1 = []; rasterSorted2 = []\n",
    "for k in range(nNeuron):\n",
    "    raster.append(np.nonzero(y[k,:])[0]/nT*T)\n",
    "    rasterSorted1.append(np.nonzero(y[cidx1[k],:])[0]/nT*T)\n",
    "    rasterSorted2.append(np.nonzero(y[cidx2[k],:])[0]/nT*T)\n",
    "\n",
    "plt.subplots(1,3, figsize=(10, 3))\n",
    "plt.subplot(1,3,1)\n",
    "plt.eventplot(raster, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot'); plt.ylabel('neurons');\n",
    "plt.subplot(1,3,2)\n",
    "plt.eventplot(rasterSorted1, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot (1)'); plt.ylabel('sorted neurons');\n",
    "plt.subplot(1,3,3)\n",
    "plt.eventplot(rasterSorted2, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot (2)'); plt.ylabel('sorted neurons');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "890bbe98-bb8c-404d-9b07-dff7d2a890cd",
   "metadata": {},
   "source": [
    "### Axis aligned observation\n",
    "\n",
    "Biologists have long loved neurons that are tuned specifically for a particular feature but not modulated by others.\n",
    "In our context, the neurons will be either driven by the first dimension or the second dimension of the latent process.\n",
    "Recent paper argues that this is optimal [Whittington et al. 2022].\n",
    "\n",
    " - Whittington, J. C. R., Dorrell, W., Ganguli, S., & Behrens, T. E. J. (2022). Disentangling with Biological Constraints: A Theory of Functional Cell Types. In arXiv [q-bio.NC]. arXiv. http://arxiv.org/abs/2210.01768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "534a8046-8d7c-41cf-925f-0c6456ad6c4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = 2.0 * np.random.randn(nNeuron, dLatent)\n",
    "b = -2.0 + np.random.rand(nNeuron,1)\n",
    "bidx = np.random.rand(nNeuron) < 0.5\n",
    "C[bidx, 0] = 0\n",
    "C[~bidx, 1] = 0\n",
    "b[bidx] += 1.5 # boost the firing rate a bit for the 2nd latent dim\n",
    "lam = np.exp(C @ X + b)\n",
    "y = np.random.poisson(lam*dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ac5362f-7904-49bf-a788-2d6cd019037a",
   "metadata": {},
   "outputs": [],
   "source": [
    "cidx = np.lexsort((C[:,1], C[:,0]), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a88bb2e4-b593-44d7-b3f6-12de43701f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "raster = []\n",
    "rasterSorted = []\n",
    "for k in range(nNeuron):\n",
    "    raster.append(np.nonzero(y[k,:])[0]/nT*T)\n",
    "    rasterSorted.append(np.nonzero(y[cidx[k],:])[0]/nT*T)\n",
    "\n",
    "plt.subplots(1,2, figsize=(10, 4))\n",
    "plt.subplot(1,2,1)\n",
    "plt.eventplot(raster, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot'); plt.ylabel('neurons');\n",
    "plt.subplot(1,2,2)\n",
    "plt.eventplot(rasterSorted, lw=0.5, color='k', label='spikes')\n",
    "plt.xlim(0, T); plt.xlabel('time'); plt.yticks([]); plt.title('raster plot (again)'); plt.ylabel('sorted neurons');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0bd59c8-85c3-458a-ba3b-24dce4e20287",
   "metadata": {},
   "source": [
    "## Some dynamical law governing the latent states\n",
    "\n",
    "So far, the latent states were given and not generated in a Markovian manner, that is, given the current state $x(t)$, the future states do not depend on $x(t-1)$ (or further past).\n",
    "In general, in discrete time, a dynamical law can be represented as a dynamical system:\n",
    "$$ \\frac{dx}{dt} = \\dot{x} = f(x(t)) $$\n",
    "where $f$ is a smooth function that represents the vector field."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79937bc",
   "metadata": {},
   "source": [
    "### van der Pol oscillator\n",
    "Van der pol oscillator is defined as 2D dimensional first order differential equations:\n",
    "    $$ \\dot{x} = y$$\n",
    "    $$ \\dot{y} = \\mu(1-x^2)y -x $$\n",
    "\n",
    "For our simulations we take a discrete time grid for convenience.\n",
    "We use an Euler integration of a Van der Pol oscillator with noisy transitions with $\\mu=1.5$, $\\tau_1=0.1$, $\\tau_2=0.1$, and $\\sigma=0.1$:\n",
    "\n",
    "$$ x_{t+1,1} = x_{t,1} + \\tau_1^{-1} \\Delta x_{t,2} + \\sigma \\epsilon$$\n",
    "$$ x_{t+1,2} = x_{t,2} + \\tau_2^{-1} \\Delta(\\mu (1-x_{t,1})^2 x_{t,2} - x_{t,1}) + \\sigma \\epsilon$$\n",
    "\n",
    "For the sake of brevity in the notebook, the generation code is provided in `code_pack/genereate_vdp_data.py` and we only load the saved data here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d21e424d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import h5py\n",
    "import scipy\n",
    "import scipy.ndimage\n",
    "from einops import rearrange\n",
    "\n",
    "import matplotlib\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",
    "\n",
    "# loading data from ./data/vdp_noisy.h5\n",
    "file_name = \"vanderpol/data/poisson_obs.h5\"\n",
    "\n",
    "# dynamics parameters\n",
    "data = h5py.File(file_name, 'r')\n",
    "system_parameters = {}\n",
    "system_parameters['mu'] = data['mu']\n",
    "system_parameters['tau_1'] = data['tau_1']\n",
    "system_parameters['tau_2'] = data['tau_2']\n",
    "system_parameters['sigma'] = data['sigma']\n",
    "system_parameters['scale'] = np.array(data['scale'])\n",
    "\n",
    "Y = np.array(data['Y'])\n",
    "X = np.array(data['X'])\n",
    "C = np.array(data['C'])\n",
    "b = np.array(data['bias'])\n",
    "\n",
    "n_trials = Y.shape[0]\n",
    "n_latents = X.shape[2]\n",
    "n_neurons = Y.shape[2]\n",
    "n_time_bins = Y.shape[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a279499d",
   "metadata": {},
   "source": [
    "### Visualizationing trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad83d305",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# plotting trajectories of the dataset\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",
    "\n",
    "# system_parameters_copy = copy.deepcopy(system_parameters)\n",
    "system_parameters['sigma'] = 0.0\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": "markdown",
   "id": "9ca8cf43",
   "metadata": {},
   "source": [
    "### Effect of tuning function\n",
    "\n",
    "TODO: explain different inverse link functions (soft plus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3ebf33",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "C_tilde = np.array(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(data['Y'])[0, :, :])\n",
    "events_softplus = raster_to_events(np.array(data['Y_softplus'])[0, :, :])\n",
    "events_axis_aligned = raster_to_events(np.array(data['Y_axis'])[0, :, idx].transpose())\n",
    "axs[0].eventplot(events, linewidths=0.5, color='k');\n",
    "axs[1].eventplot(events_softplus, linewidths=0.5, color='k');\n",
    "axs[2].eventplot(events_axis_aligned, linewidths=0.5, color='k');\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": "d8ab1d7f",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "Now we understand better the generative process of the model. But what we are interested is the opposite direction, that is, how do we infer the model parameters given just the observations (neural data)? This is the statistical inference problem of interest."
   ]
  }
 ],
 "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
}