{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# amLDS: example on synthetic data\n",
    "\n",
    "This notebook is a walk through the use of amLDS for estimating shared latent dynamics and stimulus decoding using synthetic \"neural\" data. <br> <br>\n",
    "1) we simulate a shared stimulus-specific latent neural responses according to the generative model <br>\n",
    "2) we generate the associated animal-specific high-dimensional measurements <br>\n",
    "3) we estimate the shared manifold, amLDS parameters, by EM using data from several animals <br>\n",
    "4) we estimate latents from unseen observations (test dataset) <br>\n",
    "5) we provide probabilistic estimate of the decoded stimulus identity via hierarchical inference  <br>\n",
    "\n",
    "Lastly, we estimate the latent dimensionality using the marginal likelihood and the reconstruction error for model comparison. Note that this should be the first step when using amLDS on data. <br>\n",
    "\n",
    "amLDS assumes that stimuli are encoded in the coordinated activity of populations of neurons, and this responses are shared across animals (Fig.A). The generative model si shown in Fig. B.\n",
    "\n",
    "![fig.png](funs/figs.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import relevant modules and amLDS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from funs.amLDS_engine import amLDS as amLDSm\n",
    "from funs.amLDS_utils  import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define ground truth amLDS parameters and hyperparameters\n",
    "\n",
    "The model hyperparameters are include the number of animals/subjects, stimuli and latent and observed dimensions. The assumption of amLDS generative model is that the latent dynamics are stimulus-specific and shared across animals. This forces the dimensionality of the latent manifold to be unique, however it allows for different obseration dimensions per animal. \n",
    "\n",
    "The model parameters can be divided in two: parameters driven the latent dynamics (**A**, **Q**, **b** and **Q0** ; which are stimulus specific but shared across animals), and parameters mapping the observed measurments (**C**, **R**; which are animal specific but shared across stimuli)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nStim   = 50; nSbj = 5\n",
    "sbjIDs  = ['sbj'+str(i)  for i in range(nSbj)]  \n",
    "stimIDs = ['stim'+str(i) for i in range(nStim)]\n",
    "\n",
    "zdim = 3                                     # same for all subjects and stimuli\n",
    "xdim = {key: 40 for key in sbjIDs}           # subject specific ; could be different for different subjects\n",
    "timepoints = 41                              # number of time steps ; it could vary on a trial by trial basis\n",
    "trialN = 50                                  # number of trials per stimulus condition and subject ; could vary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(10)\n",
    "# simulate latent parameters (stimulus specific) \n",
    "true_A_  = {key: A_init(zdim)  for key in stimIDs}\n",
    "true_Q_  = {key: Q_init(zdim)  for key in stimIDs}\n",
    "true_Q0_ = {key: Q_init(zdim)  for key in stimIDs}\n",
    "true_x0_ = {key: x0_init(zdim) for key in stimIDs}\n",
    "true_b_  = bs_init3D(stimIDs)\n",
    "\n",
    "# simulate observation parameters (subject specific)\n",
    "refC     = C_init(xdim[sbjIDs[0]], zdim)\n",
    "true_C_  = {key: C_inits(refC)     for key in sbjIDs}\n",
    "true_R_  = {key: R_init(xdim[key]) for key in sbjIDs}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample latent trajectories and observed measuraments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "idxx = 33 ; idxx_val = 1033                                   # set seeds for training and test sets\n",
    "amLDSsample = amLDSm(zdim=zdim, sbjID=sbjIDs, stimID=stimIDs) # initialize model class\n",
    "\n",
    "# sample from true latent parameters and structure observed data\n",
    "XX = {} ; ZZ = {} ; XX_val = {}\n",
    "\n",
    "for sbj in sbjIDs:\n",
    "    CC = {} ; cc = {} ; CC_val = {}\n",
    "    for stim in stimIDs:\n",
    "        samples = []\n",
    "        zamples = []\n",
    "        for i in range(trialN):\n",
    "            fooz, foox = amLDSsample.sample(A=true_A_[stim],   Q=true_Q_[stim], C=true_C_[sbj], R=true_R_[sbj], \n",
    "                                            x0=true_x0_[stim], b=true_b_[stim], timepoints= timepoints, seed=idxx)\n",
    "            samples.append( foox )\n",
    "            zamples.append( fooz )\n",
    "            idxx +=1\n",
    "        samples_val = []\n",
    "        for i in range(int(10)): # number of validation trials\n",
    "            samples_val.append(amLDSsample.sample(A=true_A_[stim], Q=true_Q_[stim],C=true_C_[sbj],R=true_R_[sbj], \n",
    "                                    x0=true_x0_[stim], b=true_b_[stim],timepoints= timepoints, seed=idxx_val)[1])\n",
    "            idxx_val +=1\n",
    "        CC[stim] = samples\n",
    "        cc[stim] = zamples\n",
    "        CC_val[stim] = samples_val\n",
    "    XX[sbj] = CC\n",
    "    ZZ[sbj] = cc\n",
    "    XX_val[sbj] = CC_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 68908 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_LatentsObserveds(ZZ, sbjIDs[-1], stimIDs, XX, sbjIDs[-1], stimIDs[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This are example latent space dynamics for the different stimuli (left; solid lines are averages over trials and dashed lines five example trials); and an example trial for a given animal (right)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## amLDS parameter learning and decoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n"
     ]
    }
   ],
   "source": [
    "amLDS = amLDSm(zdim=zdim, sbjID=sbjIDs, stimID=stimIDs, timepoints=timepoints)  # initialize model class\n",
    "amLDS.import_data(data=dict(XX))                                                # import data\n",
    "amLDS.import_validationData(data=dict(XX_val))                                  # import validation data\n",
    "amLDS.init_params(xdim=xdim[sbjIDs[0]], TL_FAinit=True, AQ_init=True)           # initialize parameters\n",
    "amLDS.LDS_EM_TranferLearning();                                                 # EM parameter learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "zz_amLDS = {}                                                                    # infere latents\n",
    "for sbj in sbjIDs:\n",
    "    zz_amLDS[sbj] = {stim: amLDS.getLatents(XX=None, sbj=sbj, stim=stim) for stim in stimIDs}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84964 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_LatentsLatents(ZZ, sbjIDs[-1], stimIDs, zz_amLDS)                           # latents visualization "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparison between simulated and recostructed trajectories. Note that there could be degeneracies on the latent space alignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoding stimulus class\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 58320 characters]\n",
      "text/plain": [
       "<Figure size 1224x288 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "amLDS.probability_decoding();                                                     # stimulus decoding\n",
    "amLDS.visualize_performanceCM(amLDS.sbj_ProbCM, annott=False)                     # CM visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63064 characters]\n",
      "text/plain": [
       "<Figure size 1440x576 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "amLDS.visualize_Prstim()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Decoder performance visualization as the confusion matrix (top) and as probability of stimulus for each individual trial (bottom).\n",
    "\n",
    "The decoder works accurately estimates even in a 50-label task. Change the model hyperparameters (e.g. increase the number of stimuli) to see the changes in performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter visualization\n",
    "\n",
    "In addition to recovering the latent dynamics, inspection of the learnt amLDS parameters could help undertand the underlying computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 269116 characters]\n",
      "text/plain": [
       "<Figure size 1440x360 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "amLDS.visualize_params('C') # change the input to 'A', 'Q0', 'Q', 'C', 'R' or 'b' to visualize other params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimate shared latent dimensionality\n",
    "\n",
    "When using amLDS on data, one first needs to estimate the dimensionality of the shared latent space. In the previous example the dimensionality is known, but in a general case it has to be estimated from data by model comparison. Here, we provide two wasys to test the dimensionality, negative marginal likelhood and reconstruction error. Both estimate the quality of the fit in withheld data as a function of latent dimensionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate new data for testing\n",
    "XX_test   = {}\n",
    "idxx_test = 12\n",
    "for sbj in sbjIDs:\n",
    "    CC_test = {}\n",
    "    for stim in stimIDs:\n",
    "        samples_test = []\n",
    "        for i in range(int(10)): # number of validation trials fro dimensionality estimation\n",
    "            samples_test.append(amLDSsample.sample(A=true_A_[stim],Q=true_Q_[stim],C=true_C_[sbj],R=true_R_[sbj], \n",
    "                                    x0=true_x0_[stim], b=true_b_[stim],timepoints= timepoints, seed=idxx_test)[1])\n",
    "            idxx_test +=1\n",
    "        CC_test[stim] = samples_test\n",
    "    XX_test[sbj] = CC_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Latent dimensions 1\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 1: 0.7885670582351696\n",
      "Latent dimensions 2\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 2: 0.4688128257224105\n",
      "Latent dimensions 3\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 3: 0.39606924564765833\n",
      "Latent dimensions 4\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 4: 0.41283024748572583\n",
      "Latent dimensions 5\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 5: 0.42970062099518813\n",
      "Latent dimensions 6\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 6: 0.44917506190322687\n",
      "Latent dimensions 7\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 7: 0.46882348486671876\n",
      "Latent dimensions 8\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 8: 0.486213645733472\n",
      "Latent dimensions 9\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 9: 0.5054265096266728\n",
      "Latent dimensions 10\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 10: 0.5319245449628485\n",
      "Latent dimensions 15\n",
      "Initializating C and R using population factor analysis\n",
      "Initializating b using latent estimates\n",
      "Initializating A and Q using latent estimates\n",
      "Learning amLDS parameters via EM\n",
      "Reconstruction error for latent space dimension 15: 0.6634089500600614\n"
     ]
    }
   ],
   "source": [
    "DIMS =  [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15] # dimensions to check\n",
    "LDS_err = []; LDS_LL = [];\n",
    "\n",
    "for dd in DIMS:\n",
    "    print('Latent dimensions '+str(dd))\n",
    "    amLDSd = amLDSm(zdim=dd, sbjID=sbjIDs, stimID=stimIDs, timepoints=timepoints) # initialize model class\n",
    "    amLDSd.import_data(data=dict(XX))                                             # import data\n",
    "    amLDSd.init_params(xdim=xdim[sbjIDs[0]], TL_FAinit=True, AQ_init=True)        # initialize parameters \n",
    "    amLDSd.LDS_EM_TranferLearning();                                              # EM parameter learning\n",
    "    LDS_err.append(amLDSd.reconstruction_error(XX_test))                          # estimate reconstruction error\n",
    "    print('Reconstruction error for latent space dimension '+str(dd)+': '+str(np.mean(LDS_err[-1])))\n",
    "    LDS_LL.append(amLDSd.marginal_loglikelihood(XX_test))                         # get posterior log-likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33220 characters]\n",
      "text/plain": [
       "<Figure size 216x252 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_estimatedLatentDimensionality(LDS_err, LDS_LL, DIMS, sbjIDs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dimensionality of the shared manifold is the dimension for which the reconstration error and negative log likelihood of the marginal are minimal."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
