03_variational_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5e672056-bf6f-4b2b-8c1a-7db0013555db",
   "metadata": {},
   "source": [
    "# Inferring both neural state trajectory and dynamics with VAEs\n",
    "\n",
    "We will look at two different variational methods to infer an interpretable structure from a streaming high-dimensional time series: sequential variational encoder (seqVAE) and variational joint filtering (VJF)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6444dc71-698d-4959-b5e9-0b28934e2af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from code_pack.generate_vdp_data import generate_van_der_pol, generate_noisy_van_der_pol\n",
    "from vjf.model import VJF\n",
    "from code_pack.plotting import plot_two_d_vector_field_from_data\n",
    "import math\n",
    "from einops import rearrange\n",
    "\n",
    "from vanderpol.config import get_cfg_defaults\n",
    "cfg = get_cfg_defaults()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b6a97a-a6b6-4a62-8edb-f377ec025cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup precision and random seeds\n",
    "torch.set_default_dtype(torch.double)  # using double precision\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5914fa1-ca6e-4045-8eeb-014fd2f6a968",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grid(n, lims):\n",
    "    xedges = np.linspace(*lims, n)\n",
    "    yedges = np.linspace(*lims, n)\n",
    "    X, Y = np.meshgrid(xedges, yedges)\n",
    "    grids = np.column_stack([X.reshape(-1), Y.reshape(-1)])\n",
    "    return X, Y, grids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e87e493-f242-4db9-8eaa-cff8f80a06a7",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8142194-ba57-4f56-aba9-d26a49af64ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = get_cfg_defaults()\n",
    "data = h5py.File('vanderpol/data/poisson_obs.h5')\n",
    "Y = torch.tensor(np.array(data['Y']), dtype=torch.float32)\n",
    "X = torch.tensor(np.array(data['X']), dtype=torch.float32)\n",
    "C = torch.tensor(np.array(data['C']), dtype=torch.float32)\n",
    "b = torch.tensor(np.array(data['bias']), dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f16d00a-7f82-46c3-b442-43382ba72d82",
   "metadata": {},
   "source": [
    "## Sequential Variational AutoEncoder\n",
    "\n",
    "We demonstrate that seqVAE can learn the latent dynamical system.\n",
    "For fast convergence, we fix the observation model parameters $C$ and $b$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e60eded1-1f6e-4bda-a13c-99c1a359bc7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, clear_output\n",
    "from vanderpol.SequentialVAE import NeuralVAE, SeqDataLoader\n",
    "\n",
    "n_epochs = 150\n",
    "batch_size = 25\n",
    "time_delta = 5e-3\n",
    "\n",
    "n_latents = X.shape[2]\n",
    "n_neurons = Y.shape[2]\n",
    "n_time_bins = Y.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95341dae-1134-47b0-be79-dab5f6ebfcde",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = NeuralVAE(cfg, time_delta, n_neurons, n_latents, n_time_bins)\n",
    "vae.manually_set_readout_params(C.clone(), b.clone())\n",
    "\n",
    "vae.decoder.C.bias.requires_grad_(False)\n",
    "vae.decoder.C.weight.requires_grad_(False)\n",
    "train_data_loader = SeqDataLoader((Y, X), batch_size)\n",
    "\n",
    "opt = torch.optim.Adam(vae.parameters(), lr=1e-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f2fd5f0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training the Sequence VAE\n",
    "We will now train the sequence VAE using synthetic Poisson observations with spiking rate modulated by Van der Pol dynamics.  The generative model is\n",
    "$$\n",
    "z_{t+1} = f(z_t) + w_t\\\\\n",
    "y_t \\mid z_t \\sim \\text{Poisson}(y_t \\mid \\Delta \\exp(C^\\top z_t + b))\n",
    "$$\n",
    "To perform one gradient step, the only thing we must do is evaluate the ELBO\n",
    "$$\n",
    "\\mathcal{L}(\\theta, \\phi, \\psi) = \\sum\\nolimits_t \\mathbb{E}_{q_{\\phi}(z_t)} \\log p_{\\psi}(y_t \\mid z_t) - \\mathbb{D}_{\\text{KL}} (q_{\\phi}(z_{1:T}) \\,||\\, p_{\\theta}(z_{1:T}))\n",
    "$$\n",
    "In convenient cases, we can calculate some of these expectations analytically.  However, for general non-linear and non-Gaussian dynamics we cannot.  Fortunately, our structured representation of the variational approximation makes it easy to sample trajectories from the current distribution so that we can use stochastic gradients instead.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f9008ab-5f73-4590-9793-f75f969d12e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "%matplotlib inline\n",
    "\n",
    "total_loss = []\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15, 4))\n",
    "axs_range = {'x_min': -1.5, 'x_max': 1.5,\n",
    "             'y_min': -1.5, 'y_max': 1.5}\n",
    "\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    batch_loss = []\n",
    "\n",
    "    for batch_idx, (y, x) in enumerate(train_data_loader):\n",
    "        loss, z, mu_t, log_var_t = vae(y, y, 1.0)\n",
    "        batch_loss.append(loss.item())\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(vae.parameters(), max_norm=1.0, norm_type=2)\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "\n",
    "    total_loss.append(np.mean(batch_loss))\n",
    "\n",
    "    if epoch % 1 == 0:\n",
    "        axs[0].cla()\n",
    "        axs[0].set_ylim(-1.75, 1.75)\n",
    "        axs[0].plot(z[:, 0, 0].detach().numpy())\n",
    "        axs[0].plot(x[0, :, 0].detach().numpy())\n",
    "        axs[0].set_xlabel('time'); axs[0].set_title('trial 0')\n",
    "\n",
    "        axs[1].set_xlim(0, epoch)\n",
    "        axs[1].cla()\n",
    "        axs[1].plot(total_loss)\n",
    "        axs[1].set_title('loss'); axs[1].set_xlabel('epoch'); axs[1].grid(True)\n",
    "\n",
    "        display(fig)\n",
    "        clear_output(wait=True)\n",
    "\n",
    "    if epoch % 5 == 0:\n",
    "        with torch.no_grad():\n",
    "            axs[2].cla()\n",
    "            dynamics_fn = torch.nn.Sequential(*[vae.decoder.p_mlp, vae.decoder.p_fc_mu])\n",
    "            plot_two_d_vector_field_from_data(dynamics_fn, axs[2], axs_range)\n",
    "            axs[2].set_title('phase portrait')\n",
    "\n",
    "dynamics_fn = torch.nn.Sequential(*[vae.decoder.p_mlp, vae.decoder.p_fc_mu])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6504146",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Evaluating K-step predictions\n",
    "We see how well the learned dynamics can do k-step predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "131d329e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# vdp parameters, we just compare the deterministic dynamics\n",
    "import scipy.stats\n",
    "\n",
    "system_parameters = {}\n",
    "system_parameters['mu'] = 1.5\n",
    "system_parameters['tau_1'] = 0.1\n",
    "system_parameters['tau_2'] = 0.1\n",
    "system_parameters['sigma'] = 0.0  # (no noise for comparison)\n",
    "system_parameters['scale'] = 1 / 0.4\n",
    "\n",
    "k_step_max = 5  # 3 or greater\n",
    "n_samples = 1000\n",
    "mse_per_sample = torch.zeros(n_samples, k_step_max, dtype=torch.float32)\n",
    "fig, axs = plt.subplots(1, k_step_max, figsize=((k_step_max-1)*4, 3))\n",
    "\n",
    "for k_steps in range(1, k_step_max + 1):\n",
    "\n",
    "    for n in range(n_samples):\n",
    "        x0 = torch.tensor([-1.5, -1.5], dtype=torch.float32) + 3 * torch.rand(2, dtype=torch.float32)\n",
    "\n",
    "        xs_pr = torch.zeros((k_steps, 2), dtype=torch.float32) # predicted xs passed through learned dynamics\n",
    "        xs_gt = generate_noisy_van_der_pol(x0, time_delta * np.arange(k_steps+1), system_parameters)[1:] # xs passed through ground truth dynamics\n",
    "\n",
    "        for i in range(k_steps):\n",
    "            if i == 0:\n",
    "                xs_pr[i] = dynamics_fn(x0)\n",
    "            else:\n",
    "                xs_pr[i] = dynamics_fn(xs_pr[i-1])\n",
    "\n",
    "        mse = torch.mean((torch.tensor(xs_gt) - xs_pr)**2)\n",
    "        mse_per_sample[n, k_steps-1] = mse\n",
    "\n",
    "        if k_steps > 1:\n",
    "            axs[k_steps-1].plot(xs_pr[:, 0].detach().numpy(), xs_pr[:, 1].detach().numpy())\n",
    "            axs[k_steps-1].set_title(f'k={k_steps}')\n",
    "\n",
    "mse_k_step = torch.mean(mse_per_sample, dim=0).detach()\n",
    "sem_k_step = scipy.stats.sem(mse_per_sample.detach(), axis=0)\n",
    "\n",
    "axs[0].set_title(f'k-step MSE')\n",
    "axs[0].plot(np.arange(1, k_step_max+1), mse_k_step, color='blue')\n",
    "axs[0].fill_between(np.arange(1, k_step_max+1),\n",
    "                    mse_k_step + 2 * sem_k_step,\n",
    "                    mse_k_step - 2 * sem_k_step,\n",
    "                    color='blue', linewidth=0, alpha=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfddb3e2",
   "metadata": {},
   "source": [
    "## Variational Joint Filtering\n",
    "A *real-time* method for variationally filtering latent trajectories and learning their underlying dynamics. An approximation to the filtering distribution is used at each time point to evaluate the time instantaneous ELBO; this makes it possible for us to infer latent states while *simultaneously* learning the dynamics\n",
    "\n",
    "See the paper [here](https://www.frontiersin.org/articles/10.3389/fncom.2020.00071/full).\n",
    "See the code [here](https://github.com/catniplab/vjf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93b5643a-6121-4639-876d-1556be4211c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = 5e-3  # time bin size\n",
    "n_trials = 1\n",
    "bin_size_ms = 5\n",
    "time_delta = bin_size_ms * 1e-3\n",
    "\n",
    "state0 = np.random.uniform(-1.0, 1.0, size=2)\n",
    "    \n",
    "system_parameters = {}\n",
    "system_parameters['mu'] = 1.5\n",
    "system_parameters['tau_1'] = 0.1\n",
    "system_parameters['tau_2'] = 0.1\n",
    "system_parameters['sigma'] = 0.1  # noise add into euler integration\n",
    "system_parameters['scale'] = 1 / 0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4bb61c8-8de0-4c02-83bb-0e3f5cf3c779",
   "metadata": {},
   "outputs": [],
   "source": [
    "#generate latent\n",
    "n_latents = 2\n",
    "n_neurons = 50\n",
    "n_time_bins = 2000\n",
    "t = delta * torch.arange(n_time_bins)\n",
    "X = generate_noisy_van_der_pol(state0, t, system_parameters)\n",
    "\n",
    "X = X @ np.array([[4,0],[0,4]]) \n",
    "\n",
    "#generate observations Guassian observations\n",
    "C = np.random.normal(size=(n_latents, n_neurons))  # loading matrix\n",
    "b = np.random.normal(size=n_neurons)               # bias\n",
    "\n",
    "# observation\n",
    "Y = X @ C + b\n",
    "Y = Y + np.random.normal(size=Y.shape) * 0.1  # add some noise\n",
    "\n",
    "n_epochs = 15\n",
    "batch_size = 25\n",
    "time_delta = 5e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6bd717-8584-4d67-8da9-68eef4a936db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup and fit VJF \n",
    "udim = 0\n",
    "n_rbf = 50  # number of radial basis functions for dynamical system\n",
    "hidden_sizes = [20]  # size of hidden layers of recognition model\n",
    "likelihood = 'gaussian'  # gaussian or poisson\n",
    "\n",
    "model_gauss = VJF.make_model(n_neurons, n_latents, udim=udim, n_rbf=n_rbf, hidden_sizes=hidden_sizes, likelihood=likelihood)\n",
    "m_gauss, logvar, _ = model_gauss.fit(Y, max_iter=150)  # fit and return list of state posterior tuples (mean, log variance)\n",
    "\n",
    "m_gauss = m_gauss.detach().numpy().squeeze()\n",
    "s_gauss = np.exp(0.5 * logvar.detach().numpy().squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e683f2-4952-441f-9816-83707cc00aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# regress to account for invariance\n",
    "X_hat = m_gauss #rearrange(m_vdp, 'batch time lat ->  (batch time) lat')\n",
    "S = np.linalg.pinv(X_hat) @ X.reshape(n_time_bins, n_latents)\n",
    "X_hat_tilde = X_hat @ S\n",
    "X_hat_tilde = X_hat_tilde.reshape(n_time_bins, n_latents)\n",
    "\n",
    "# Plot\n",
    "fig, axs = plt.subplots(2, 1, sharex='all')\n",
    "\n",
    "axs[0].plot(X[:, 0], color='black', alpha=0.75)\n",
    "axs[0].set_ylabel('dim0')\n",
    "axs[0].set_xlabel('time')\n",
    "axs[0].plot(X_hat_tilde[:, 0], color='red', alpha=0.75)\n",
    "axs[0].fill_between(np.arange(n_time_bins),\n",
    "                    X_hat_tilde[:, 0] + 2 * s_gauss[:, 0],\n",
    "                    X_hat_tilde[:, 0] - 2 * s_gauss[:, 0],\n",
    "                    color='red', alpha=0.1, linewidth=0)\n",
    "\n",
    "axs[1].plot(X[:, 1], color='black', alpha=0.75)\n",
    "axs[1].set_ylabel('dim1')\n",
    "axs[1].plot(X_hat_tilde[:, 1], color='red', alpha=0.75)\n",
    "axs[1].fill_between(np.arange(n_time_bins),\n",
    "                    X_hat_tilde[:, 1] + 2 * s_gauss[:, 1],\n",
    "                    X_hat_tilde[:, 1] - 2 * s_gauss[:, 1],\n",
    "                    color='red', alpha=0.1, linewidth=0)\n",
    "plt.legend([\"Data\", \"Fit\"], loc='upper right');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de006de8-07dd-4cab-b981-e52a38d2d9df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw the inferred velocity field\n",
    "\n",
    "ax = fig.add_subplot(223)\n",
    "r = np.mean(np.abs(m_gauss).max())  # determine the limits of plot\n",
    "\n",
    "Xm, Ym, XYm = grid(51, [-1.5*r, 1.5*r])\n",
    "Um, Vm = model_gauss.transition.velocity(torch.tensor(XYm)).detach().numpy().T  # get velocity\n",
    "Um = np.reshape(Um, Xm.shape)\n",
    "Vm = np.reshape(Vm, Ym.shape)\n",
    "plt.streamplot(Xm, Ym, Um, Vm)\n",
    "plt.plot(*m_gauss.T, color='C1', alpha=0.5, zorder=5)\n",
    "plt.title('Velocity field');"
   ]
  }
 ],
 "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}