{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# amLDS: application to mixtures and changes in stimulus intensity\n",
    "\n",
    "This notebook shows the ability of amLDS decoding to generalize in response to changes in the properties of the stimulus (Note: run the \"amLDS_example\" notebook to familiarize yourself with amLDS parameter learning and inference). In section A, we test the decoding performance when linearly interpolating between two stimulus conditions. In section B, we explore the generalization in performance when decreasing the concentration/intensity of the stimulus."
   ]
  },
  {
   "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\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "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": [
    "# Section A: mixtures\n",
    "\n",
    "In this section, we will only look at two stimuli and mixtures between them (feel free to tweak the different variables and test other settings). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define amLDS parameters and hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nStim   = 2; nSbj = 1\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 = 20                                  # 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, angles=20)\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}\n",
    "\n",
    "# force last two class conditions to have the same dynamic transitions\n",
    "true_A_[stimIDs[-2]]  = np.copy(true_A_[stimIDs[-1]])\n",
    "true_Q_[stimIDs[-2]]  = np.copy(true_Q_[stimIDs[-1]])\n",
    "true_Q0_[stimIDs[-2]] = np.copy(true_Q0_[stimIDs[-1]])\n",
    "true_x0_[stimIDs[-2]] = np.copy(true_x0_[stimIDs[-1]])"
   ]
  },
  {
   "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, 73468 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.init_params(xdim=xdim[sbjIDs[0]], TL_FAinit=True, AQ_init=True)           # initialize parameters\n",
    "amLDS.LDS_EM_TranferLearning();                                                 # EM parameter learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample mixtures and test decoding performance\n",
    "\n",
    "The following cells sample new responses for different mixture ratios of the two original stimuli and tests the decoder generalization performance to those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "mix_predict = []; zamples  = []; TEMP = []\n",
    "sbjT  = sbjIDs[-1]; stimA = stimIDs[-2]; stimB = stimIDs[-1]\n",
    "nTrmix = 20\n",
    "ratios = [0,.1,.2,.3,.4,.5,.6,.7,.8,.9,1]\n",
    "for i in ratios:\n",
    "    idxx  = 700\n",
    "    samples  = []\n",
    "    ratA = i; ratB = 1-i\n",
    "    \n",
    "    for jj in range(nTrmix):\n",
    "        fooz, foox = amLDSsample.sample(A = true_A_[stimA]*ratA +true_A_[stimB]*ratB,   \n",
    "                                        Q = true_Q_[stimA]*ratA +true_Q_[stimB]*ratB, \n",
    "                                        C = true_C_[sbjT], R=true_R_[sbjT], \n",
    "                                        x0= true_x0_[stimA]*ratA+true_x0_[stimB]*ratB, \n",
    "                                        b = true_b_[stimA]*ratA +true_b_[stimB]*ratB, \n",
    "                                        timepoints= timepoints, seed=idxx)\n",
    "        samples.append(foox)\n",
    "        idxx+=1\n",
    "        if i == .5:\n",
    "            zamples.append(fooz)\n",
    "\n",
    "    XXab = dict({sbjIDs[-1]: {stimA: samples}})\n",
    "    \n",
    "    tempPr = amLDS.probability_decoding(XXab, sbjIDs=[sbjIDs[-1]])[0][sbjIDs[-1]]\n",
    "    tempPr = tempPr[~np.isnan(tempPr).any(axis=1)]\n",
    "    TEMP.append(tempPr)\n",
    "    mix_predict.append(np.mean(np.argmax(tempPr, axis=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65980 characters]\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_mixtures(ZZ, stimA, stimB, sbjIDs[-1], zamples, nTrmix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Average (solid lines) and example latents (dash lines) for the two original stimuli (A and B) and the 50-50 mixture between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13384 characters]\n",
      "text/plain": [
       "<Figure size 360x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(facecolor='w', figsize=(5,4))\n",
    "plt.plot(ratios, np.array(mix_predict), '-', c='k', lw=3)\n",
    "plt.ylabel('Pr(A)'); plt.xlabel('A-B mix');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Average decoding performance for the different mixtures of A and B. In the case of pure A (on the left) the decoder perfectly identifies the responses as A, but as the stimulus morphs into B, the decoder smoothly transitions between the two stimulus labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, ax = plt.subplots(1, len(TEMP), facecolor='w', figsize=(20, 8))\n",
    "# for m in range(len(TEMP)):\n",
    "#     sns.heatmap(TEMP[m], cmap='binary', ax=ax[m], cbar=False);\n",
    "#     ax[m].set_title(str(ratios[m]*100)+'-'+str(100-ratios[m]*100))\n",
    "# ax[0].set_ylabel('trials')\n",
    "# ax[0].set_xlabel('stim');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Section B: concentration\n",
    "\n",
    "In this section, we will test the generazation of the amLDS decoder to changes in concentration. We will mimic changes in concentration by decreasing the magnitude of the latent dynamics (parameter **b**)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define amLDS parameters and hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "nStim   = 10; 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 = 10                                  # number of trials per stimulus condition and subject ; could vary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(21) #7\n",
    "# simulate true latent (stimulus specific) parameters\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_init(stimIDs)\n",
    "true_b_  = bs_init3D(stimIDs, angles=170)\n",
    "\n",
    "# simulate true observation (subject specific) parameters\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": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "trialSou = 30     # number of trials per stimulus condition and subject ; could vary\n",
    "# initialize model class\n",
    "amLDSsample = amLDSm(zdim=zdim, sbjID=sbjIDs, stimID=stimIDs)\n",
    "# sample from true latent parameters and structure observed data\n",
    "XX = {}\n",
    "idxx = 33\n",
    "for sbj in sbjIDs:\n",
    "    CC = {}\n",
    "    for stim in stimIDs:\n",
    "        samples = []\n",
    "        for i in range(trialSou):\n",
    "            samples.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)[1])\n",
    "            idxx +=1\n",
    "        CC[stim] = samples\n",
    "    XX[sbj] = CC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample latent trajectories at different concentrations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "trialVal = 20   # number of trials for testing\n",
    "# sample from true latent parameters and structure observed data\n",
    "idxx_val = 1033\n",
    "XX_valss = []\n",
    "paramsss = [[1.1, 1.0, 1.0] , [.9,.8,.8], [.7,.6,.6], [.5,.4,.4], [.3,.2,.2]]\n",
    "for par in paramsss:\n",
    "    conc_b  = bs_init3D(stimIDs, angles=170, mu1=par[0],mu2=par[1],mu3=par[2])\n",
    "    XX_val = {}\n",
    "    for sbj in sbjIDs:\n",
    "        CC_val = {}\n",
    "        for stim in stimIDs:\n",
    "            samples_val = []\n",
    "            for i in range(int(trialVal)):\n",
    "                samples_val.append(amLDSsample.sample(A=true_A_[stim], Q=true_Q_[stim], C=true_C_[sbj], \n",
    "                                                      R=true_R_[sbj],  x0=true_x0_[stim],b=conc_b[stim],\n",
    "                                                      timepoints=timepoints,seed=idxx_val)[1]) \n",
    "                idxx_val +=1\n",
    "            CC_val[stim] = samples_val\n",
    "        XX_val[sbj] = CC_val\n",
    "    XX_valss.append(XX_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54724 characters]\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "concZ = np.array([np.mean(XX_valss[i]['sbj0']['stim0'],axis=0) for i in range(len(XX_valss))])\n",
    "\n",
    "plt.figure(facecolor='w', figsize=(4,4))\n",
    "for i in range(concZ.shape[0]):\n",
    "    plt.plot(concZ[i, :, 0], concZ[i, :, 1], c='#1f77b4', alpha=1-(i/5),lw=3)\n",
    "plt.legend(range(5)); plt.xlabel('LD1'); plt.ylabel('LD2');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example average latent trajectories for an example stimulus at different concentrations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## amLDS parameter learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "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": [
    "TL_CN = amLDSm(zdim=zdim, sbjID=sbjIDs, stimID=stimIDs, timepoints=timepoints)# initialize model class\n",
    "TL_CN.import_data(data=XX)                                                    # import data\n",
    "TL_CN.init_params(xdim=xdim[sbjIDs[0]], TL_FAinit=True ,AQ_init=True)         # initialize parameters\n",
    "TL_CN.LDS_EM_TranferLearning();                                               # EM parameter learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## amLDS decoding across concentrations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoding stimulus class\n",
      "Decoding stimulus class\n",
      "Decoding stimulus class\n",
      "Decoding stimulus class\n",
      "Decoding stimulus class\n"
     ]
    }
   ],
   "source": [
    "AccCN = [TL_CN.probability_decoding(XX_test)[2] for XX_test in XX_valss]     # decode in test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21480 characters]\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_concentrationPerformance(paramsss, AccCN, sbjIDs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classification performance decreases gradually as one decreases concentration, while still allowing for generalization across a wide range of concentrations."
   ]
  }
 ],
 "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
}
