GG3_project.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ad6a89a2",
   "metadata": {
    "tags": []
   },
   "source": [
    "# IIA project GG3: Neural Data Analysis\n",
    "\n",
    "Easter 2024<br>\n",
    "Project Leader: Yashar Ahmadian (ya311)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad546cd7",
   "metadata": {},
   "source": [
    "## Important dates\n",
    "\n",
    "Project start: __Thursday May 9 2024 9:30am GMT+1 (UK summer time)__ \n",
    "\n",
    "Interim report deadline: 🔥__Friday 17 May 2024, 4pm__🔥 (electronic submission via Moodle)\n",
    "<br>\n",
    "(Interim report should contain report on tasks under the header \"**Week 1**\".\n",
    "\n",
    "Presentations: __Monday 27 May 2024, 11am-12:30pm in LT5__\n",
    "\n",
    "Final project report deadline: __Friday 7 June 2024,  4pm__ (electronic submission via Moodle)\n",
    "\n",
    "\n",
    "## Project notes\n",
    "\n",
    "- You should spend about 20 hours a week on the project, basically half of your time.\n",
    "- Project is to be carried out in **Google Colab** or on your own computer. You can download this notebook and use it with a normal Jupyter server, or duplicate it here in your **Colab** account. If you do the latter, you can share and show your work easily. The computational resources on **Colab** are limited, so you may find it more convenient to run the programs on your own computer, especially in the later parts of the project when computations will be heavier. When you need to ask a question about a specific piece of code, you can still use the **Colab** to share a notebook. \n",
    "- Weekly sessions will be held on Mondays 11:00-13:00, and Thursdays 9:00-11:00 and 14:00-16:00 all in **LR5** (EXCEPT for Monday 5th June which will be in LR10). \n",
    "- Attendance is compulsory for the first/introductory session and all Monday sessions. Thursday sessions will be optional (although this is still subject to change), but attending them is a good way to get answers to questions, some help with coding. It also provides space for teamwork with your teammates. \n",
    "- You are strongly encouraged to seek verbal feedback after your interim report - there will be a special session for this on Monday May 20 and Thursday May 23. \n",
    "- Project carries 80 marks overall:\n",
    "  - 20 marks for interim report (individual)\n",
    "  - 20 marks for presentation (group based)\n",
    "  - 40 marks for final report (individual)<br>\n",
    "  \n",
    "  \n",
    "### Project reports\n",
    "  - Should be clearly broken down by _Tasks_ (see below), any notes you wish to make in how you or your group structured and carried out the tasks, and most importantly your __results__ in the form of completely labelled graphs, and __accompanying conclusions__ you draw from your results. \n",
    "  - **Page limits:** interim report no more than 5 pages, and final report about 10-15 pages (no more than 15) -- this is excluding appendices such as attached code, but _including_ figures. The final report can be an extension of the interim report, but make sure you take into account the feedback you receive for your interim report.  \n",
    "  - When deciding what to include in your report, how to organise it and what to emphasize, please prioritise communicating understanding over formalities - I would like give you marks for doing the right thing and showing that you did it and understand it. If I have to wade through pages of undigested data and graphs shown just because it was there, I will feel less generous. The length requirements are only guidelines. \n",
    "  - Take a look at [this page](http://teaching.eng.cam.ac.uk/node/444/#hdr-9) and [this](http://teaching.eng.cam.ac.uk/node/340) for further guidance and recommendations for writing reports.\n",
    "  - __All code__ that you used during to project must be attached as an appendix to your reports. If you modified one of the provided `.py` file (and you used that modified version for that report), include it. \n",
    "  - A jupyter or **Colab** notebook are acceptable as a report, as long as it is \"clean\" (its main section includes text and figures) and reads like a report, and (importantly) can be converted to a PDF, so you can upload it to the Moodle submission protal. \n",
    "  - Incude [cover sheets](http://teaching.eng.cam.ac.uk/node/4171) provided by the Teaching Office\n",
    "  \n",
    "### Presentations:\n",
    "- Each group will jointly prepare and present a **12 minute (strict!)** presentation, broken up into three **4-minute parts** each delivered by one of the team members.<br>\n",
    "There will be 3 minutes of question time after each talk, and so overall the session should take about an hour and a half. \n",
    "\n",
    "- Since many tasks are not intrinsically divided between team members, the part presented by a student need not be something they solely contributed to. \n",
    "- The presentations will be delivered at the beginning of Week 3, **Monday 27 May** during the regular session time in **LR5**.\n",
    "- I will give guidelines and recommendations for making good presentations in due course. \n",
    "\n",
    "## Timeline\n",
    "\n",
    "See the Approximate Timeline section below.\n",
    "\n",
    "## Survey\n",
    "\n",
    "The **online survey** should be completed at the end of the project period. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f4b8e2-1bbf-4ef0-8e61-356e6533a848",
   "metadata": {},
   "source": [
    "# Neuroscience Background"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa168e21-72a4-43ca-9193-b06b23daf9ec",
   "metadata": {},
   "source": [
    "The background is also provided as a Jupyter notebook [accessible here](https://github.com/ahmadianlab/gg3_nda/blob/main/Background.ipynb). \n",
    "\n",
    "**Note:** The main point of the Background handout is to introduce some terminology (which appear all in boldface),<br>\n",
    "and mathematical notation that will be used in the next section, \"What is the right model of LIP?\"<br> \n",
    "Deep understanding of this Background  section is not required for carrying out the project. But, apart from the<br>\n",
    "above reason, you are encourged to read it to understand the scientific motivations and significance of this problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ca6c12-5c30-4cc2-bbb1-29fb05547262",
   "metadata": {
    "tags": []
   },
   "source": [
    "# What is the right model of LIP?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d0a6cd-7d16-43f4-9549-cb99481b2d01",
   "metadata": {},
   "source": [
    "### Alternative hypothesis: stepping model\n",
    "\n",
    "As we saw in the Background section, classic studies suggested that LIP neurons which exhibit ramping activity in<br>\n",
    "their trial-averaged PSTH's are involved in evidence accumulation. However, the story became more complicated,<br>\n",
    "when in 2015, [Latimer et al.](https://www.science.org/doi/10.1126/science.aaa4056) provided evidence that most LIP neurons are better modelled<br>\n",
    " as neurons with a \"stepping firing rate\". In this alternative model\n",
    " the rate does not continuously ramp up or down<br>  (albeit via a random walk) as in a drif-diffusion model.\n",
    "Rather, the rate is piece-wise constant:<br> it starts relatively low, but at some time point it jumps (\"steps\") up discontinuously to <br>\n",
    "a higher firing rate level. The jump point is random and varies from trial to trial, according to some distribution. <br>\n",
    "\n",
    "#### -------------------------------   Figure 1   -------------------------------\n",
    "<img src=\"figs/latimer-step-ramp.png\" width=600 />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ff6abb3-eeb2-4630-8dfb-5ae866a8b81b",
   "metadata": {},
   "source": [
    "We will refer to these two competing hypotheses or models as the **ramping** and **stepping models**, respectively<br>\n",
    "(other common synonyms for the ramping model are \"the drif-diffusion model\", mentioned above, and \"the diffusion-to-bound model\";<br>\n",
    "we will also use **jump model** as synonymous with the stepping model.) \n",
    "\n",
    "In this project we aim to develop tools that allow us to reject or accept one of these hypotheses<br>\n",
    "based on observed spike trains. Understanding which of the two is a more accurate description of LIP activity <br>\n",
    "is scientifically significant. The ramping hypothesis suggests that LIP cortex is responsible<br>\n",
    "for accumulating evidence to inform and make decisions. On the other hand, the binary nature of the stepping model<br>\n",
    "suggests that LIP is downstream of the evidence accumulating area, and may simply reflect, in its activity, the decision already made<br>\n",
    "in an upstream area."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3d78609-a056-4ed8-a703-7baf368e61e7",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Two generative models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "807d3cc0-3549-4c0f-a568-cb20b12ef24b",
   "metadata": {
    "tags": []
   },
   "source": [
    "\n",
    "So far, our two \"models\" have mostly remained conceptual and qualitative. At this high, conceptual level <br>\n",
    "I will therefore refer to them as hypotheses instead: the ramping hypothesis vs. stepping hypothesis. <br>\n",
    "However, in order to use the powerful tools of probability theory and machine learning, we need to <br>\n",
    "turn these conceptual hypotheses into well-defined mathematical models.\n",
    "\n",
    "### The common, abstract model structure\n",
    "The ramping and stepping models to be described here and simulated in the project,  are examples of<br>\n",
    "**probabilistic generative models**. Each model has a set of **parameters** (such as the drift rate of the ramping model),<br>\n",
    "and it stochastically generates data, in our case spike trains. The systematic behaviour of these spike trains depends<br>\n",
    " on the various model parameters. Mathematically, this stochastic relationship between the parameters and data<br>\n",
    " is given by a conditional probability distribution\n",
    "\n",
    "$P(\\mathrm{data}| \\Theta, M)$\n",
    "\n",
    "where $\\Theta$ denotes the set of parameters and $M$ denotes the model (in our case $M$ = ramping, or $M =$ stepping).<br>\n",
    "This conditional probability, when viewed as a function of $\\Theta$, is called the model's **likelihood function**.<br>\n",
    "By (observed) \"data\" we mean a set of spike-trains recorded (in our case simulated) over many trials:\n",
    "\n",
    "$\n",
    "\\text{data} \\equiv \\{(n_t)_{t=1}^T\\}.\n",
    "$\n",
    "\n",
    "$n_t$ will sometimes be referred to as **observed variables**.\n",
    "\n",
    "**Latent variables:** As generative models, the two models can also be simulated to generate spike trains. In order to do this, <br>\n",
    "the two models first generate a firing rate function or time-series, $r_t$. The spike count, $n_t$, in a given<br>\n",
    "time bin is then stochstically generated based solely on $r_t$. The rate sequence $r_t$ is itself a stochastic<br>\n",
    "process, and depends on a set of *latent variables*. Latent variables are random variables that are so called<br>\n",
    "because they are not directly observed by us (data-)scientists and engineers, but need to be inferred from observed<br>\n",
    "data (the spike trains). In the simple version of the stepping model, with which we will start, there is only <br>\n",
    "a single latent variable: the stepping time. The ramping model, on the other hand, generates a whole sequence of <br>\n",
    "latent variables in each trial: these are the values of the ramping stochastic process, which is closely tied to the<br>\n",
    "firing rate.\n",
    "\n",
    "Since the latent variables (unlike the model parameters) vary from trial to trial, in each trial they need to be<br>\n",
    "inferred from a single spike-train. By contrast, parameters which control the systematic behaviour of the model will<br>\n",
    "be inferred from the entire dataset, i.e. the collection of spike trains in all trials.\n",
    "\n",
    "**Discrete vs continuous time:** Both models are implemented in discrete time. Thus the varible $t$ above is an integer<br>\n",
    "(index for the) time-step. We will denote the (fixed) total number of time steps in a trial by $T$. Real trials<br>\n",
    "last on the order of 1 second, and we would want our time steps or time bins to be around 1 to 10 milliseconds. <br>\n",
    "So, correspondingly, $T$ will be rather large, we will experiment with $T=$ 100 to 1000. For various purposes, <br>\n",
    "we will need to convert from discrete to continuous time in seconds. For that purpose we will fix the trial duration at<br>\n",
    "1 second and thus interpret each time-step to have duration $1/T$ seconds; we will denote this by $dt$ here and in the code<br>\n",
    "(thus $dt = 1/T$ seconds).\n",
    "\n",
    "We will now describe the probabilistic structure of the two models in some detail.\n",
    "\n",
    "### Stepping model\n",
    "\n",
    "**Latent variables:** This is the simpler one of the two. The only latent variable of this model is the step time or **jump time**. <br> \n",
    "I will denote the step time in trial $j$ by $\\tau_j$. Since we work in discrete time, $\\tau_j$ is a (non-negative) integer.<br>\n",
    "In `models.py` the corresponding variable is called `jump` or (when containing the value of multiple trials) `jumps`.<br>\n",
    "In each trial, the step time is sampled from some probability distribtion:\n",
    "\n",
    "$\\tau \\sim P(\\tau)$\n",
    "\n",
    "In the provided code this distribution is a so-called **negative binomial distribution** (see [this](https://en.wikipedia.org/wiki/Negative_binomial_distribution)) with two parameters: $m$ and $r$.<br>\n",
    "$m$ sets the average step time, and $r$ ... that's left for you to figure out.<br>\n",
    "\n",
    "As we said above, in each trial, the firing rate sequence of this model is piece-wise constant. If we denote the jump time of trial<br>\n",
    "$j$ by $\\tau_j$, then for $t < \\tau_j$, $r_t = R_0$ and for $t \\geq \\tau_j$, $r_t = R_h > R_0$, where the two constants $R_0$ and $R_h$ are <br>\n",
    "part of the model parameters. We will refer to them as pre- and post-step firing rates. \n",
    "\n",
    "Finally, given the rate sequence, $r_t$, the spike counts in different timesteps are generated indpendently from a [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution):\n",
    "\n",
    "$n_t \\sim \\mathrm{Poiss}(r_t dt)$\n",
    "\n",
    "Note  that since we measure rates in Hz, $dt$ has to be in seconds, in order to get the right dimensionless parameter (mean spike count) of <br>\n",
    "of the Poisson distribution. \n",
    "\n",
    "**Fit parameters:** $m, r, \\text{ and } x_0$.\n",
    "\n",
    "These are the parameters which you aim to infer from spike train datasets.<br>\n",
    "$x_0$ is equivalent to $R_0$ and is given by $R_0/R_h$; think of it as the noramalised pre-step rate<br>\n",
    "(we use $x_0$ instead of $R_0$, to match the similar parameter in the ramping model). By definition $0< x_0 < 1$.\n",
    "\n",
    "**\"Fixed\" parameters:** We will take $R_h$ as known/fixed, and will not infer it from data. For many project tasks we will fix it at $R_h = 50$ Hz.<br>\n",
    "Though in early tasks you will explore its effects by varying it.\n",
    "\n",
    "### Ramping model\n",
    "\n",
    "**Latent variables:** This model, which approxmiates the continuous time drift-diffusion model, has a whole sequence of latent variables<br>\n",
    "which we will denote by $x_t$. This variable is the so-called **decision variable**. The update equations for $x_t$ are discretised versions<br>\n",
    "of the equation in Figure 3:\n",
    "\n",
    "$x_{t+1} = x_t + \\beta dt + \\sigma \\sqrt{dt} \\epsilon_t \\qquad\\qquad\\qquad$            Eq. (1)\n",
    "\n",
    "where \n",
    "\n",
    "$\\epsilon_t \\overset{iid}{\\sim} \\mathcal{N}(0,1) \\qquad\\qquad\\qquad$            Eq. (2)\n",
    "\n",
    "(i.e. $\\epsilon_t$ is sampled independently in each time step from the standard normal distribution, $\\mathcal{N}(0,1)$, in other words, it has<br>\n",
    "a Gaussian distribution with mean 0 and variance 1.) The initial condition is set via\n",
    "\n",
    "$\n",
    "x_1 = x_0 + \\sigma \\sqrt{dt} \\epsilon_0\n",
    "$\n",
    "\n",
    "where $\\epsilon_0$ is again standard normal, and $x_0$ is a model parameter (and not a latent variable, <br>\n",
    "as it is the same across all trials). (Note that, due to python indices starting from 0, the equation above will (implicitly)<br>\n",
    "appear as `x[0] = x0 + sigma * np.random.randn()` in the code).\n",
    "\n",
    "The firing rate in this model is a rectified and scaled version of $x_t$:\n",
    "\n",
    "$r_t = R_h [x_t]_+ = R_h \\max(0, x_t)$\n",
    "\n",
    "It is not hard to see that the sequential variables $x_t$ form a Markov chain (this has to do with the fact that $\\epsilon_t$<br> in different\n",
    "trials are independent), and therefore the ramping model is an example of a **hidden Markov model (HMM)**.<br>\n",
    "In fact, if you have taken 3F8, you will realise that $x_t$ is *almost* an AR(1) Gaussian Process. I said almost an AR(1) Gaussian process, because<br>\n",
    "\n",
    "*when $x_t$ reaches 1, it will get stuck there for the rest of the trial*. (Equivalently, after this point, the firing rate, $r_t$, stays at its<br>\n",
    "maximal level $R_h$.)\n",
    "\n",
    "This reflects the interpretation of $x_t$ as a decision variable, which upon reaching a pre-set bound or threshold, triggers<br>\n",
    "the decision; in our case the bound is 1. \n",
    "\n",
    "Similar to the stepping model, given the rate sequence, $r_t$, the spike counts in different timesteps are generated indpendently from a Poisson distribution\n",
    "\n",
    "$n_t \\sim \\mathrm{Poiss}(r_t dt)$.\n",
    "\n",
    "**Fit parameters:** $\\beta, \\sigma, \\text{ and } x_0$.\n",
    "\n",
    "$\\beta$ and $\\sigma$ control the systematic drift vs. stochasticity of the ramping variable $x_t$.<br>\n",
    "Similar to the stepping model, $x_0$ sets the initial rate, $r_0$, via $r_0 = R_h x_0$. And again $0< x_0 < 1$.\n",
    "\n",
    "**\"Fixed\" parameters:**  $R_h$, maximal rate, to be treated as in the stepping time.\n",
    "\n",
    "\n",
    "### Ignored stimulus dependence\n",
    "\n",
    "In the full version of the ramping model, the magnitude and sign of $\\beta$ depends on the coherence and the direction of motion of the RDM stimulus<br>\n",
    "in that trial. However, for simplicity, in this project we assume $\\beta$ is fixed in all trials and assume it is positive.<br>\n",
    "\n",
    "Similarly, in the full version of the stepping model, the post-jump rate can take two possible values $R_h > R_0$, as described above, <br>\n",
    "or $R_l < R_0$. We can call $R_h$ and $R_l$ the up or down rates, and call their normalized values of 1 and $R_l/R_h < 1$ up and down states.<br>\n",
    "The probability with with the model transition up or down after the jump time can again depend on the coherence and direction of motion of the <br>\n",
    "RDM stimulus in a trial. But again, to simplify the model, we ignore this fact. In fact, for most of the project we work with a stepping model<br>\n",
    "without a down state. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbe9c313-952c-4beb-8c39-0100c23d5768",
   "metadata": {},
   "source": [
    "# Approximate timeline\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57be6ffc-eaf1-4d02-97cd-cd8a8d8d1f79",
   "metadata": {
    "tags": []
   },
   "source": [
    "Our task in this project is to slowly build up techniques to ultimately reject or accept one or the other hypothesis based on \n",
    "recorded (or, in our case, simulated) spike trains. We will do this by progressively moving from lower to higher levels of \n",
    "probabilistic inference. \n",
    "\n",
    "\n",
    "- (week 1) explore the behaviour of the two models based on simulator code provided to you in `models.py`.<br>\n",
    "And take preliminary steps towards developing a discrete-state HMM approximation to them, which allows us to apply  powerful inference tools. \n",
    "\n",
    "- (week 2) develop tools to carry out **single-trial inference**  of the models' latent variables from observed  \n",
    "on single spike trains, taking advantage of their HMM formulation.\n",
    "\n",
    "- (weeks 2-3) Assuming model $M$ is the true model underlying data, use Bayesian or maximum-likelihood inference to infer or estimate<br>\n",
    "model parameters, $\\Theta$, based on observed data, that is, many trials of simulated spike trains. \n",
    "\n",
    "- (weeks 3-4) Use Bayesian inference to select/reject one or the other hypothesis/model, given a dataset of spike trains.<br> \n",
    "\n",
    "Depending on feedback and pace of progress, in week 4 we will also investigate the consequences of model mismatch. Since \"all models are wrong (but some are useful)\",<br> \n",
    "what can we say about the possibility of reaching wrong conclusions regarding our alternative conceptual hypotheses, due to <br>\n",
    "some arbitrary choices we had to make in translating those conceptual models to concrete mathematical models?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c33dd59-3fc8-4a7f-b6bf-7fdf540300d5",
   "metadata": {},
   "source": [
    "# Running on Google Colab\n",
    "As I said in the Intro lecture, the project Jupyter notebook is accessible via [Google Colab](https://colab.research.google.com/github/ahmadianlab/gg3_nda/blob/main/GG3_project.ipynb#scrollTo=ad546cd7). You can choose to run your notebook on either of those, or you can download it from there (or from the [GitHub](https://github.com/ahmadianlab/gg3_nda/blob/main/GG3_project.ipynb)) and run things on your own machine, which may be faster. \n",
    "\n",
    "If you use Colab, you will have to \"Save a copy in Drive\" in order to be able to save your changes (if you don't have Google Drive, then either sign up for it, or you need to download and work on your laptop).\n",
    "\n",
    "Finally, the `.py` modules (including Week 1's `models.py`) are accessible and can be downloaded from the [GitHub](https://github.com/ahmadianlab/gg3_nda/) page. <br>\n",
    "Alternatively, on Colab, you will have to run the following cell to import them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1aa7c4a2-98f7-4f6d-8841-b74dc713750c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mode = \"colab\" # change this to \"local\" if you are on your own computer\n",
    "\n",
    "if mode == \"local\":\n",
    "    import models\n",
    "elif mode == \"colab\":\n",
    "    import requests\n",
    "    url = 'https://github.com/ahmadianlab/gg3_nda/blob/main/models.py?raw=true'\n",
    "    r = requests.get(url)\n",
    "    with open('models.py', 'w') as f:\n",
    "        f.write(r.text)\n",
    "    import models\n",
    "else:\n",
    "    raise Exception(\"mode must be either local or colab\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe0d9133-5bd6-4f34-a680-0268aecef4d6",
   "metadata": {},
   "source": [
    "# Week 1\n",
    "### model simulation and behaviour\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df778a9b-3b9c-4330-b5fa-53a9945c2051",
   "metadata": {},
   "source": [
    "### Task 1.1\n",
    "\n",
    "Study the code in `models.py`, specifically the implementations of the two models in the `StepModel` and `RampModel` classes. <br>\n",
    "The main part to study (and relate to the mathematical discussion above) is their `simulate` method/function. You create <br>\n",
    "an object instance of each model by providing the model parameters (both \"fit\" and \"fixed\" parameters, as named above)<br>\n",
    "to the class constructors: e.g. `ramp = RampModel(beta=..., sigma=..., ...)`. <br>\n",
    "(Ignore the other input arguments in the class constructor `__init__` for now, and leave them at their default values.)<br>\n",
    "Once a model object is created you can use its `simulate` method to get an array of spike trains over multiple trials. <br>\n",
    "(For usage see the docstring (or run help via `ramp.simulate?`.) `simulate` will also return the generated latent variables, <br>\n",
    "and, optionally, the firing rates in different trials. \n",
    "\n",
    "Visualise the simulated spike trains by writing code to make so-called \"spike raster\" plots. See the bottom row of Figure 5<br>\n",
    "above for example spike raster: different rows represent the spike trains in different trials, and spikes are shown by dots. <br>\n",
    "(you can put a dot for every nonzero $n_t$, even if the nonzero value is more than 1; this is unlikely if you keep `Rh` below<br>\n",
    "50 Hz and use a `T` of at least 100 (recommended). At this stage it should not be time-consuming to use higher `T`'s as well,<br>\n",
    "e.g. `T = 1000` (corresponding to 1 millisecond time-steps). If you are simulating hundreds of trials, you don't want to include<br>\n",
    "all of them in the raster. Use your common sense to decide how many trials to include in the raster; this is a visualisation tool used to get<br>\n",
    "an idea of how spike trains behave qualitatively by seeing a good number of example. \n",
    "\n",
    "Vary the parameters of each model and generate spike rasters in different regions of the parameter space, trying to find<br>\n",
    "qualitatively different behavior. The default values of the parameters give you a first guess or the right order of magnitude for the <br>\n",
    "value of different parameters. (For `m` and `r` of the step model, note that they should scale with the `T` that you will be using for the simulation;<br>\n",
    "in particular, for more interesting/relevant results, you would want to start with values of `m` around `T / 2` so that the steps happen on average in the middle of the trial.)\n",
    "\n",
    "What systematic patterns can you detect? \n",
    "\n",
    "In your spike raster plots mark the jump times in different trials (by some marker superposed over the spike). <br>\n",
    "Also make histograms of jump times. What is the effect of the `r` parameter on the behaiour of the stepping model?<br>\n",
    "\n",
    "Similarly make plots of the trajectories of $x_t$ or $r_t$ (of the ramp model) in several trials, in a single plot.<br>\n",
    "You can extract the time when $x_t$ of the ramping model hits its upper bound of 1 (equivalently $r_t$ reaches $R_h$), and histogram that as well.<br> \n",
    "How do `beta` and `sigma` affect this histogram or the behaviour of the $x_t$ trajectories?\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e20ba518-9a6b-4eac-87f1-b3d2c7477733",
   "metadata": {},
   "source": [
    "### Task 1.2\n",
    "\n",
    "PSTH (peri-stimulus time histogram) is an important data analysis tool used in neuroscience.<br> \n",
    "This  is a statistical estimate of the trial-averaged firing rate as a function of time, based on recordings of spike trains in<br>\n",
    "multiple experimental trials. It is obtained by binning/histogramming spikes (e.g. using `np.histogram`) in different time bins<br>\n",
    "and averaging the resulting spike counts over many trials (you can also divide by `dt` to turn mean spike counts into fiting rate, in units of Hz).<br>\n",
    "\n",
    "Write code to construct and plot PSTH's in different regions of each model's parameter space. Note how the PSTH <br>\n",
    "fluctuates randomly from dataset to dataset. It is better to do some sort of (temporal) smoothing in order to reduce these<br>\n",
    "fluctuations and the jaggedness of the PSTH. You can use either a sliding window (e.g. a boxcar window/functin) averaging, or simply <br>\n",
    "use time bins that are larger than the oridinal time steps (e.g. 50 milliseconds -- or 5 timesteps if you are using a `dt` of 10 ms,<br> \n",
    "corresponding to `T = 100`). The smooth ramping firing rate curves in Figure 4 of the [Background](https://github.com/ahmadianlab/gg3_nda/blob/main/Background.ipynb)\n",
    "are examples of smoothed PSTH's.\n",
    "\n",
    "Even with the smoothing there will be fluctuations in the PSTH from dataset to dataset. How does the strength of these fluctuations depend on (or scale with)<br>\n",
    "the number of trials (in each dataset)? Try to be quantitative about this. <br>\n",
    "For the rest of this task use a high number of trials (e.g. 5000) to minimise these fluctuations. (But note <br>\n",
    "that in real experiments the number of trials rarely exceeds a few hundred -- so for later tasks we will bring the number down.)\n",
    "\n",
    "\n",
    "Finally, try to find parameter regimes that make the PSTH of the stepping model very close to that of the ramp model. (First make sure<br>\n",
    "the ramp model's PSTH look qualitatively like the classic ramping PSTH's in the LIP experiments.) In which parameter regions<br>\n",
    "does this fail drastically, and in which regimes are the two PSTH's nearly indistinguishable?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb66917-0d16-4cb3-8946-63243c4918aa",
   "metadata": {},
   "source": [
    "### Task 1.3\n",
    "\n",
    "The PSTH is an example of a so-called first-order statistic, in that it is the average of spike counts, $n_t$, which is the first moment of their marginal distribution.\n",
    "\n",
    "You can also evaluate higher order statistics, such as the variance of $n_t$ (across trials).\n",
    "(Instead of smoothing, for evaluating the variance use larger time bins -- e.g. 50 or 100 milliseconds).\n",
    "However, instead of directly plotting and exploring the variance, we will plot a more useful quantity: the **Fano factor**, which is the ratio of the variance of $n_t$ to its mean (obviously both evaluated in the same time bin). This quantity is 1 for the Poisson distribution (the default choice for the emission distribution of both models).<br>\n",
    "Evalute and plot the Fano Factor as a function of time, and again investigate how it changes in different parameter regimes, and importantly<br>\n",
    "whether and how it behaves differently in the two models. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdf1821f-ef7e-436e-a443-18e45c38883f",
   "metadata": {},
   "source": [
    "### Task 1.4\n",
    "\n",
    "(This is a more open-ended and less guided task compared to the previous ones. Use brainstorming in the group and come up with creative ideas to address this, working together.)\n",
    "\n",
    "In this task you will explore an informal or relatively ad-hoc version of what we intend to do eventually using the systematic approach of Bayesian inference. <br>\n",
    "The aim is to find an intelligent but ad-hoc (in the sense that it is not Bayesian and does not rely on the two models' likelihood function, but only relies on \n",
    "observed statistics) method for telling the two models apart, i.e. deciding which model generated a dataset, ***only*** based on the dataset that the method receives.\n",
    "\n",
    "Relying on the 1st and 2nd order statistics that you have explored (or other 2nd or higher order statistics you may come up it),\n",
    "construct a criterion (or multiple alternative criteria which you would then compare) for deciding between the two models. \n",
    "Try to come up with criteria that are relatively robust, i.e. are not handcrafted to tell the two models apart only for fixed choices (or very limited ranges) of parameters.\n",
    "Further, to make this challenging, try to find/construct a method/criterion that would work decently (without expecting magic) when the dataset was generated by model parameters in respective regimes in which the two models are intuitively least distinguishable based on their generated spike trains (e.g. based on the spike train statistics you explored in previous tasks). \n",
    "\n",
    "You have to test your criterion by running it on several datasets, once generated by the ramp model, and in another round, generted by the step model. \n",
    "And then quantify what percent of datasets in each case where classified/decided correctly. Use a number of trials not more than 400 for each dataset that you will run your criterion on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7af18d4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Week 2\n",
    "### Latent variable inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7794027e-ee22-44ab-9da2-d784dddedffc",
   "metadata": {},
   "source": [
    "In Week 2, your task is to formulate the two models as (time-homogeneous) **Hidden Markov Models (HMM)** with *discrete* states. To recap, <br>\n",
    "a HMM is described by latent variables that are (1) temporal/sequential and (2) form a **Markov chain (MC)**.<br> \n",
    "If we denote these variable at time step $t$ by $x_t$, the Markovian property means that conditioned on $x_t$, all future states<br>\n",
    "$(x_{t+1}, x_{t+2}, \\ldots)$ are independent of states at times before $t$.<br>\n",
    "Colloquially, conditioned on the present, the future is independent of the past.<br> \n",
    "\n",
    "In an HMM, the Markovian state variables are not directly observerd. Instead at each time, we observe an observed variable, $n_t$, which<br>\n",
    "only depends of the Markov state, $x_t$, at the same time step. For us this dependence is given by the Poisson distribution describing the <br>\n",
    "spike emissions, conditioned on the rates, with the latter being determined by $x_t$ (in other words, the rate at time $t$ is a<br>\n",
    "deterministic instantaneous functions of $x_t$).<br> Figure 2 shows the graphical model for the HMM, with such Poisson observation (aka emission) distributions.  \n",
    "\n",
    "Here, we aim to design Markov chains with discrete states to approximate the behavior of the latent variables of the step and ramp models.<br>\n",
    "In the case of the step model, the discretization is actually exact, as the model really just has two levels of rates (although we will see that the correspondence between Markov states and possible firing rate levels is rather complicated). <br>\n",
    "It is for the ramp model  that a discretization will be an approximation (since in the original formulation of this model, the state variables, $x_t$, are continuous variables). \n",
    "\n",
    "The reason for formulating the models as discrete state HMMs is that in this case we can use the powerful and efficient \n",
    "[forward-backward algorithm](https://en.wikipedia.org/wiki/Forward%E2%80%93backward_algorithm) to calculate<br>\n",
    "(1) the Bayesian *posterior* estimate (the posterior mean) of the state variables, $x_t$, given the observations, $n_t$, and<br>\n",
    "(2) calculate the model likelihood function,  $P(n_{1:t}|\\Theta, \\mathcal{M})$, i.e. the probability of the observed spike train conditioned on the model parameters $\\Theta$ (for each model $\\mathcal{M} = $ ramp or step).\n",
    "\n",
    "\n",
    "<img src=\"figs/HMM_graph.png\" width=600/>\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168eedf8-5a3d-41e6-87f4-e06d1e91523f",
   "metadata": {},
   "source": [
    "### Task 2.1: \n",
    "**Forming an HMM approximation to the ramp model**\n",
    "\n",
    "$\\newcommand{\\T}{\\mathcal{T}}$\n",
    "\n",
    "The state variable of the (original) ramp model is continuous and, because the update rule $x_{t+1} = x_t + \\beta dt + \\sigma\\sqrt{dt}\\epsilon_t$<br>\n",
    "involves the Gaussian variables $\\epsilon_t$, their transition probabilities $P(x_{t+1}| x_{t})$ are Gaussian (with mean and variance possibly depending on $x_t$, as well as the model parameters). Your first sub-task is to work out this distribution.\n",
    "\n",
    "First, we will now assume $x_t$ does not go below 0. Thus if $x_t$ is currently zero, and the proposed change according to normal update rule of $x_t$ is negative, $x_t$ will remain at 0; only if the proposed change is positive it is actually implemented. (Note that this is a redefinition of the model's latent variables, it does not affect the behaviour of rates and thus spikes; and thus it does really change the model.)\n",
    "\n",
    "Next, we will approximate $x_t$ by assuming it takes values on a regular grid of $K$ points going from 0 to 1, inclusive of both ends. Let's use $K=50$ or $K=100$.<br>\n",
    "These $K$ points form the discrete states of the discrete-state HMM that approximates the original ramp model. We will denote these states by their index $s$ going from 0 to $K-1$.\n",
    "The corresponding value of $x_t$ is then given by \n",
    "\n",
    "$x_t = \\frac{s_t}{K-1} \\qquad \\text{where} \\quad s_t \\in \\{0, \\ldots, K - 1\\}$.\n",
    "\n",
    "By evaluating the relevant probabilities, based on the Gaussian distribution you derived above, form the transition matrix for the Markov chain, defined as \n",
    "\n",
    "$\n",
    "\\T_{s,s'} = P(s_{t+1}= s'| s_{t}= s)\n",
    "$\n",
    "\n",
    "Note that the state *transitioned to* (i.e. the one at $t+1$) is the column index of the matrix. Thus the rows of $\\T$ have to sum up to 1.<br> \n",
    "You may need to (and it certainly would not hurt to) **enforce this constraint by hand** after constructing the (intial) matrix using the evaluated Gaussian probabilities. <br>\n",
    "\n",
    "Also note transitions out of the last state $s = K-1$ need special consideration, as according to the original model, once the variable $x$ reaches 1, it stays there. \n",
    "\n",
    "Next, you will need to form the initial state distribution, $\\pi$, as an array of $K$ values (summing to 1!) giving the probabilities of different possibilities of $s_0$. This should approximate the equation $x[0] = x_0 + \\sigma\\sqrt{dt} \\epsilon_0$.\n",
    "\n",
    "Once the transition matrix $\\T$ and initial state distribution $\\pi$ are formed, we can simulate the chain. In order to <br> do this, you will first draw $s_0$ from the initial state distribution, then successively sample from the <br>\n",
    "appropriate distribution according to the transition matrix $\\T$ (and depending on the current state $s_t$). To sample <br>\n",
    "the discrete (integer) $s$ from a distribution over its $K$ possibilities, you can use `np.random.choice`.\n",
    "\n",
    "For different choices of the model parameters, $\\beta$, $\\sigma$ and $x_0$, simulate several trials of this chain and (after appropriate rescaling) plot the trajectories $x_t$. Based on the trajectory $x_t$, calculate the firing rate trajectory $r_t$. Compare these rate trajectories with corresponding simulated rate trajectories of the original (continuous state) ramp model, to make sure your implementation is accurate *enough*. \n",
    "\n",
    "Note that for small enough values of $\\sigma$ the Markov chain approximation will produce trajectories that get stuck at the inital state. Why is this? For the case $\\beta=0$, estimate, solely in terms of $K$ and $T$ (or $dt$), the *order of magnitude* (or \"scaling\") of the value of $\\sigma$  (up to a constant of proportionality) below which trajectories tend to get stuck. In the rest of the project use values of $\\sigma$ above this value, unless indicated otherwise (e.g. when use of specific ranges for parameters are instructed). \n",
    "\n",
    "**Note:** Note also that if you use values of $\\sigma$ that are too small, and depending on your code for constructing the transition matrix, your code for generating $\\T$ or $\\pi$ may run into numerical truncation issues resulting in `NaN` values. Something that could help is implementing things first in terms of log-probabilities, using the numerically stable function `scipy.special.logsumexp` in the normalization step (when you normalize rows of $\\T$ or the vector $\\pi$), and only in the end exponentiating to obtain the actual $\\T$ or $\\pi$. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9552450-cd09-451f-9951-0882b1dd4f56",
   "metadata": {},
   "source": [
    "### Task 2.2: \n",
    "**Forming an HMM approximation to the step model**\n",
    "\n",
    "Intuitively, given its discrete (binary) rate levels, the step model can actually be exactly formulated as a discrete state HMM.<br>\n",
    "\n",
    "First, implement a *time-homogeneous* Markov chain representing the step state with two states. How would you choose the transition probabilities?<br> Hint: think of the \n",
    "parameter $p$ of the Negative Binomial distribution (see the Wikipedia article), which in terms of $m$ and $r$ is given by $r/(m+r)$.\n",
    "\n",
    "Simulate the Markov chain for several trials and plot the corresponding $x_t$ trajectories. \n",
    "\n",
    "Also evaluate the jump times (time-steps) and make histograms of these jump times. How do the histograms appear? Do they resemble any of the histograms of jump times<br> (i.e. histograms \n",
    "for different values of $m$ and $r$) that  you made in Week 1? \n",
    "\n",
    "What do you think is wrong with the 2-state Markov chain approximation to the step model? **§**<br>\n",
    "Read about the [Negative Binomial distribution](https://en.wikipedia.org/wiki/Negative_binomial_distribution) and its \"meaning\" (*for the case of positive integer $r$*), to get clues for constucting an exact Markov Chain formulation of the step model (Hint: use $r + 1$ states!).**§§**\n",
    "\n",
    "Again, simulate several trials of this new chain, plot state trajectories, and form histograms of jump times for different values of $r$. How do these compare with the histograms you obtained in Week 1 for jump times of the original step model?\n",
    "\n",
    "\n",
    "**§:** An exact fomrulation of the step model is possible as a *time-inhomogenous* 2-state Markov chain (MC). You can experiment with that of course (this is optional). However, we will be using a time-homogenous HMM, and so in the above sub-task I am asking you to construct a time-homogenous Markov chain. If you did implement a time-inhomogeneous MC formulation, feel free to write about it in your final report. But make sure you do investigate the time-homogeneous version (which is in general not a correct formulation of the original model) as well, and answer the questions for that (too).\n",
    "\n",
    "**§§:** The suggested MC implementation of the step model will differ from the Week 1 step model, in which step time had the distribution $\\mathrm{NB}(\\tau | m, r)$ by a time shift by $r$ time-steps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f821da5-d58b-4fb0-a830-9e9e519db23e",
   "metadata": {},
   "source": [
    "### Task 2.3: \n",
    "**Inference of hidden states**\n",
    "\n",
    "Henceforth, for the rest of the activities this week (as well as most of the activities of the next two weeks), we will be solely working with the discrete-state HMM versions of the two models that you have implemented (instead of using the `models.py` simulators). We will also limit the values of the parameter $r$ to positive integers.\n",
    "\n",
    "The `hmm_expected_states` function (see its doc/help) in the `inference.py` (run the next code cell to import this on Colab **§**) module implements the forward-backward algorithm (FBA) to calculate the posterior probabilities $P(s_t | n_{1:T})$, as well as the log-likelihood $\\ln P(n_{1:T})$.**§§**<br> \n",
    "One of the inputs to `hmm_expected_states` is the array, `ll`, of the logs of the conditional observation probabilities $ll[t, s] = \\log P(n_t|s_t = s)$.**§§§** Use the function `inference.poisson_logpdf` from the provided new module `inference.py`, to construct this array based on the observed spike counts of one or several trials (see the help of this function).\n",
    "\n",
    "\n",
    "\n",
    "Use this function to obtain the posterior probabilities, $P(s_t | n_{1:T})$, for your finite-state HMM implementations of both models.**§§§§**\n",
    "\n",
    "- Write code to calcualte the posterior expectation of $x_t$ (i.e. $\\mathbb{E}[x_t | n_{1:T}]$), for the ramp model, based on the posterior probabilities $P(s_t | n_{1:T})$. Generate several trial spike trains using the discrete-state HMM ramp model (with the corresponding $x_t$ trajectories retained), and for each trial infer and plot $\\mathbb{E}[x_t | n_{1:T}]$, together with the ground-truth simulated $x_t$. Repeat this in different regions of model parameter space (including low and high values of $x_0$ and $R_h$). In what parameter regimes is the inference more accurate, and vice versa? Provide intuitive/qualitative explanations for your observations. \n",
    "\n",
    "- Repeat the above for the step model, with the following modifications. The aspect of the step model's hidden states that we really care about is whether or not the ''neuron\" has jumped to the upper rate level. Calculate the probability of being in the upper rate level based on the posterior state probabilities $P(s_t | n_{1:T})$, and again make plots of it for various simulated spike-train trials. Visualise the true jump time on these plots. You can take the time point when the posterior probability of being in the upper rate level exceeds 0.5 as the estimated/inferred jump time, and mark that on the plots as well.<br>\n",
    "For your report, try to combine different trials in one plot or figure, in a compact but nice way.<br> Based on these plots, comment (include both accounts of your observations and your qualitative explanations for them) on the accuracy  of the inference in different regions of the model parameter space. \n",
    "\n",
    "- The `hmm_expected_states` function has an optional boolean input `filter` which is false by default. When true, the function calculates $P(s_t | n_{1:t})$ instead of $P(s_t | n_{1:T})$: i.e. the posterior conditioned only on observations up to and including the \"current\" time-step $t$. This corresponds to the so-called filtering problem (as in the Kalman filter; the default case, `filter = False`, corresponds to \"smoothing\"). Filtering is appropriate for applications where inference has to be performed online, in which case, to infer $x_t$ we do not have the luxury of having access to future observations -- without a time-machine, that is! In our case (as engineers studying the computational mechanisms in area LIP), we do of course have access to the entire spike-train; hence the default option. Nevertheless, carry out a theoretical study of the differences in inference accuracy, for both models, using smoothing vs. filtering. What qualitative differences do you observe between the inference accuracy using filtering vs. smoothing? \n",
    "\n",
    "In all of the above sub-tasks,  quantiatify the latent-state inference accuracy by evaluating the average error (both over trials and over time, if the latter makes sense) of the posterior estimates of the $x_t$ trajectories (for the ramp model) or the jump times (for the step model). You can then make heatmap or contour plots of these errors as a function of two parameters -- and different plots for different choices of parameter pairs. (The `tricontourf` function of `matplotlib` is very useful for this purpose.)\n",
    "\n",
    "______________________________________________________________________________________________________________________________________\n",
    "**§:** The module `inference.py` makes use of the [Numba package](http://numba.pydata.org/) to speed up computations by so-called just-in-time (JIT) compilation.  Numba can be imported in Colab, but to use it locally, you will need to install it, following these [instructions](https://numba.readthedocs.io/en/stable/user/installing.html).\n",
    "Also note that JIT makes a function run slowly the first time you use it. So the first time you run a JITed function (e.g. `hmm_expected_states`) it's better to run it on a \"light\" test case (e.g. a single spike train, rather than several spike trains, or on a short one), and only then run it for the true use case.  \n",
    "\n",
    "\n",
    "**§§:** This (log) probability depends implicitly on model parameters, hence the name (log) likelihood; it is called `normalizer` in the code for reasons having to do with the fact that it normalizes the message products involved in the FBA.\n",
    "\n",
    "**§§§:** Note that at inference time, the $n_t$ are known and fixed, but the MC states are unknown; thus, as part of the inference procedure, we need to evaluate the observation (log-)probabilities for *all* possible states, $s$, at *every* times $t$. For more details see my [notes on the FBA](https://github.com/ahmadianlab/gg3_nda/blob/main/fwdbwd.ipynb).\n",
    "\n",
    "\n",
    "**§§§§:** Note that the observed spike trains used here (and later in Weeks 3 & 4) should be generated using the discrete-state HMM version of the two models. For this, first generate the latent variables as in tasks 2.1 and 2.2, and calculate the firing rates, $r_t$, based on them. Then, given the rates, sample the spike counts, $n_t$, from the corresponding Poisson distribution (you can look up how this was done in `models.py`).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff6ebefd-e7d3-48c4-bc96-eae663f8c4ca",
   "metadata": {},
   "source": [
    "# Week 3\n",
    "\n",
    "This week we will proceed to the inference of model parameters (aka learning or fitting of model parameters), and preliminary model selection.\n",
    "\n",
    "For Bayesian inference of parameters, we would like to compute the posterior distribution over model parameters, $\\Theta$, given a *dataset* of observed spike trains:\n",
    "\n",
    "$P(\\Theta\\, |\\, \\mathrm{data}\\!=\\!\\{(n_t)_{t=1}^T\\}, \\mathcal{M})$\n",
    "\n",
    "where the observed data are the set of spike trains in multiple trials and $\\mathcal{M}$ ($ =$ *ramp* or *step*) denotes the model. According to the Bayes rule<br>\n",
    "\n",
    "\n",
    "$P(\\Theta| \\mathrm{data}, \\mathcal{M}) \\propto P(\\mathrm{data}|\\Theta, \\mathcal{M}) P(\\Theta| \\mathcal{M}), \\qquad\\qquad$             **(3.1)**\n",
    "\n",
    "where $P(\\Theta| \\mathcal{M})$ is the prior distribution over parameters and $P(\\mathrm{data}|\\Theta, \\mathcal{M})$ is the model's likelihood function.<br>\n",
    "Calculating the posterior distribution is hard  in general, since the  expression on the right hand side of Eq. (3.1)<br>\n",
    "is not normalized (below we will refer to this expression  as the **unnormalized posterior**), and calculating <br>\n",
    "the normalization constant requires integrating or summing over the possibly high-dimensional space of parameters.  <br>\n",
    "The same is true if we want to calculate the expectation of some function of the parameters (in particular their posterior expectation $\\mathbb{E}[\\Theta| \\text{data}, \\mathcal{M}]$):\n",
    "again we have to integrate over the parameters.<br>\n",
    "In general, the computational cost of this integration grows exponentially in the number of parameters \n",
    "(or the dimension of parameter space). Thus Bayesian methods almost inevitably rely on one or another type of approximation\n",
    " for calculating the posterior.\n",
    "\n",
    "In our project, the models have only 3 unknown parameters to be inferred, and the above problem is not prohibitive. \n",
    "However, if, e.g., we approximate the integral by a sum over a discretized 3D grid of points, such that the grid partitions\n",
    " each parameter's range of values (as determined by the prior) into, say, only $M = 10$ sub-intervals, we will have to sum \n",
    "over $M^3$ points, which requires $M^3 = 1000$ evaluations of the likelihood. <br>\n",
    "In our case, at each point of the grid, the likelihood or its logarithm, can be calculated using the forward pass\n",
    "of the forward-backward algorithm (those curious about the details: see the section Model Log-likelihood of the [notes on the FBA](https://github.com/ahmadianlab/gg3_nda/blob/main/fwdbwd.ipynb)).\n",
    "Note that this requires order $T$ matrix multiplications, each involving order $K^2$ arithmetic operations for $K$ Markov states;\n",
    "thus evaluation of the log-likelikehood for each value of parameters has cost $O(K^2 T)$.\n",
    "Thus computing the likelihood is relatively expensive (even if on our laptops it can be evaluated in a fraction of a second),\n",
    "and evaluating the posterior on such a grid has computational cost $O(M^d K^2 T)$, where for us $d=3$, though, in some of the tasks below, $d=2$.                                                                                                                \n",
    "\n",
    "Use the function `hmm_normalizer` of `inference.py` to calculate the model **log-** likelihood using the forward pass.\n",
    "\n",
    "**Prior distributions:**\n",
    "Unless instructed otherwise, use a uniform prior distribution on the parameters, on specified ranges. For $\\sigma$ of the ramp model<br>\n",
    "it will be more sensible to use a uniform prior on its logarithm. In fact, to make things easier you should work with $\\ln \\sigma$, when<br>\n",
    "evaluating the unnormalized posterior, and in the grid method, construct the (regular) grid for $\\ln \\sigma$, instead of $\\sigma$ <br>\n",
    "(but obviously exponentiate $\\ln \\sigma$ when passing to your likelihood function, and when calculating posterior estimates of the parameters<br>\n",
    "calculate the posterior mean or error of the exponential of the $\\ln \\sigma$ values). Sensible ranges are as follows:\n",
    "\n",
    "$\\beta \\in [0, 4]$\n",
    "\n",
    "$\\ln\\sigma \\in [\\ln(0.04), \\ln(4)]$\n",
    "\n",
    "$r \\in \\{1,2, \\ldots, 10\\}$ (note that you can still further \"discretize\" or coarse-grain this range of integers into a fewer number of \"bins\".)\n",
    "\n",
    "$m \\in [0,T]$\n",
    "\n",
    "$x_0 \\in [0, 0.5]$\n",
    "\n",
    "**Hyperparameters:** Unless otherwise stated, use $T = 100$ and for the ramp model: $K = 100$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e775ef-0e59-42de-b66d-ebc2d6376053",
   "metadata": {},
   "source": [
    "#### Note on the step model's likelihood evaluation:\n",
    "Before we proceed to the tasks, a note on the $K = r + 1$ hidden-state HMM implementation of the step model (you can ignore this note<br>\n",
    "if you have opted to use the time-inhomogeneous Markov chain implementation instead). As you may have noted,<br>\n",
    "the jump times in the $(r+1)$-state implementation will have a delay of $r$ steps compared to the draws form the $\\mathrm{NB}(m, r)$ distribution<br>\n",
    "as used in Week 1. For simulating a trial with $T$ time steps, this issue can be dealt with by a time-shift,<br>\n",
    "at the expense of simulating the Markov chain for $r$ extra steps: simulate $T+r$ steps but only consider the last $T$ steps,<br>\n",
    "and only generate spikes in those final $T$ steps.<br>\n",
    "This week, however, we would like to evaluate the likelihoods of the two models, on the *same* set of spike trains <br>\n",
    "on the interval $t = 1:T$, using the forward pass of FBA. In this case, the necessary compensation for the extra, observation-less<br>\n",
    "$r$ steps can be made by the following modification to the input argument `ll` that you will pass to the function `hmm_normalizer`:  <br>\n",
    "\n",
    "- evaluate this array as normal, as the $T \\times K$ matrix of observed probabilities of $n_{1:T}$ in all possible hidden states.<br>\n",
    "- then prepend $r$ rows with all elements equal to 0 to the top of this matrix to make it a $(r+T)\\times K$ matrix instead.\n",
    "\n",
    "(If curious, you can verify that this is the correct way of generalizing the forward pass to  cases with \"missing observations\" by inspecting<br>\n",
    "the [details of the forward pass](https://github.com/ahmadianlab/gg3_nda/blob/main/fwdbwd.ipynb).) Note that this compenstation\n",
    "is exactly equivalent to running the forward pass on the original $T$ steps (that have spike observations), but using $\\mathcal{T}^{r} \\pi$ --where $\\mathcal{T}$\n",
    "is transition matrix--<br> as the initial Markov chain probability distribution instead of $\\pi$. (This second method is actually computationally preferable,\n",
    "since --for a given values of $r$ or $m$-- you calculate $\\mathcal{T}^{r} \\pi$ once, before using it in the FWA to be run on all spike trains.)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a548c05c-4eae-4ed9-a3a3-46f78217328f",
   "metadata": {},
   "source": [
    "### Task 3.1\n",
    "**Grid approximation**\n",
    "\n",
    "We will start this task by fixing the value of $x_0$ to 0.2 (for both models); so we will not infer it, but  assume it known.<br>\n",
    "Construct a regular 2D grid on each model's parameter space (on the ranges given above, and using $\\ln \\sigma$ instead of $\\sigma$) with $M$ grid points along each coordinate axis,<br>\n",
    "thinking of each grid point as <u>*the center*</u> of its \"assigned\" sub-interval.<br> \n",
    "Start by relatively high values of $M$, e.g. $M=30$.<br>\n",
    "\n",
    "(For one calculation of the posterior over all grid points, going to higher $M$'s is also feasible, but since you want to experiment and explore with different parameters <br>\n",
    " and hyperparameters, you should compromise on using higher values of $M$. But ultimately use your own judgement, in particular by using $M < 30$ if necessary. <br>\n",
    "Similarly, in the sub-tasks below, make reasonable choices of the number of different cases of true parameters, etc., that you explore.<br> \n",
    "It is probably wise to start small-scale, in a first pass through all sub-tasks, and after you have a rough overview of computational times, inference behaviour, etc,<br> \n",
    "do  more thorough final runs going over a larger number of possibilities and/or using higher $M$ values, which would take longer to compute.)\n",
    "\n",
    "For each model simulate a dataset of $N_{\\rm trial}$ spike trains, and for each dataset, evaluate *that* model's log-likelihood on all grid points<br>\n",
    "(by summing the log-likelihood of all trials), and form the *normalized* posterior by using Bayes' rule Eq. (3.1).\n",
    "\n",
    "**3.1.1:** Visualize the (approximate) posterior on the grid (e.g. using `plt.imshow`) with a  super-imposed point representing the true parameters. <br>\n",
    "Repeat for different true parameter values (within the ranges of the prior distribution) and for different number of trials (over a range from 1 to 400). <br>\n",
    "Document your observations of systematic dependence of the posterior on the number of trials. Similarly, note any systematic changes in behaviour<br>\n",
    "for different values of the true parameters. \n",
    "\n",
    "**3.1.2:** In addition to visualisation, evaluate the posterior expectations of the parameters as their estimates, as well as their posterior standard deviations (as measures of their<br>\n",
    "posterior uncertainties or  \"error bars\" for their estimates).  Given the true parameter value, you can also evaluate the actual estimation error of different parameters, <br>\n",
    "for  dataset of different sizes (different number of trials) or choices of true parameters (averaging the error over a few, even just 2-3, datasets when $N_{trial}$ is low). <br>\n",
    "Make appropriate plots. \n",
    "\n",
    "Also quantitatively compare the estimation errors with the posterior uncertainties of different parameters. \n",
    "What do you expect for this relationship, and what do you observe?\n",
    "\n",
    "**3.1.3:** Repeat the above (using judgment to scale down $M$ or the number of explored cases) for the case where $x_0$ is unknown as well, and has to be inferred."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b453f72-920f-4fdb-93e0-383c7b4b14a8",
   "metadata": {},
   "source": [
    "### Task 3.2\n",
    "**Model selection (grid based)**\n",
    "\n",
    "Repeat the sub-tasks in Task 3.1, for the two \"cross-cases\" where you compute the posterior and posterior expectations of one models's parameters,<br>\n",
    "given spike-trains generated by the other model. How do things change in each case? Document any interesting observations of systematic behaviour.\n",
    "\n",
    "One (or, arguably, the) Bayesian method for model selection is the computation of so-called Bayes factors, i.e. the model posteriors (for $\\mathcal{M} = $ *ramp* or *step*):\n",
    "\n",
    "$P(\\mathcal{M}| \\mathrm{data}) \\propto P(\\mathrm{data}| \\mathcal{M}) P(\\mathcal{M})$\n",
    "\n",
    "where $P(\\mathcal{M})$ denotes the prior probability of each model (and $\\mathrm{data}=\\{(n_t)_{t=1}^T\\}$), which we will take to be 0.5 throughout. In this case, the posterior<br>\n",
    "model probabilities are proportional to their so-called **marginal likelihoods** $P(\\mathrm{data}| \\mathcal{M})$, and model selection can be entirely based on the **marginal likelihood ratio (MLR)**,<br>\n",
    "also known as the **Bayes factor**: <br>\n",
    "\n",
    "$R = \\frac{P(\\mathrm{data}| \\mathcal{M}=ramp)}{P(\\mathrm{data}| \\mathcal{M}=step)}\\qquad\\qquad$     **(3.2)**\n",
    "\n",
    "or its logarithm, with the ramp model chosen if $R>1$ or $\\ln R > 0$, and vice versa.\n",
    "\n",
    "$P(\\mathrm{data}| \\mathcal{M})$ is called marginal likelihood because its calculation involves integrating out (or marginalizing) model parameters:\n",
    "\n",
    "$P(\\mathrm{data}| \\mathcal{M}) = \\int P(\\mathrm{data}|\\Theta, \\mathcal{M}) P(\\Theta| \\mathcal{M})\\, d\\Theta.\\qquad\\qquad$   **(3.3)**\n",
    "\n",
    "Comparison with Eq. (3.1) shows that this is *exactly* the normalizing constant (the denominator of Bayes rule) for the unnormalized posterior over model parameters. <br>\n",
    "(The denominator of the Bayes rule 3.1, is also known as \"evidence\", as it quantifies the strength of evidence in favor of the model.)<br>\n",
    "So you have already evaluated the marignal likelihoods in Task 3.1 when you normalized the posteriors (up to approximating the integral with a sum over grid points).<br>\n",
    "(**However,** you need to make sure that you have included the appropriate (approximate) $d\\Theta$'s in the formula. This is the volume of each square in your grid, and for <br>\n",
    "integer $r$ the number of integers within the coarse-grid, if using a coarse-grid. Also note that this depends on your choice of normalization of the prior; ultimately what <br>\n",
    "matters is that if you replace $P(\\mathrm{data}|\\Theta, \\mathcal{M})$ in Eq. (3.3) with 1, the integral or grid-sum should evaluate to 1.)\n",
    "\n",
    "\n",
    "Do model selection, based on the MLR, for datasets of different sizes (number of trials) and quantify the two different error rates. <br>\n",
    "\n",
    "- **3.2.1:** First do this by sampling the (\"true\") parameters, used for generating a dataset from each model, from that model's prior parameter distribution.\n",
    "\n",
    "- **3.2.2:** Repeat this with \"mismatched\" parameter priors: without changing the parameter-sampling distribution (i.e. keep sampling parameters as you were doing in 3.2.1), use different priors for parameter inference.<br>\n",
    "Specifically, use truncated Gaussian prior distributions (uncorrelated/independent across different parameters) that are centered at different locations in the parameter space: for parameters $m$ and $\\sigma$ center the Gaussian in the middle of the ranges given above, but for $r$ and $\\beta$ use distributions centered at the two extremes (for a total of four possibilities). Try different values for the corresponding prior standard deviations (SD). Speficially, for all parameters choose their SD to be the *same* fraction of the length of their total range, but try 2-3 different values for this (shared) fraction; choose sensible/informative fractions.**§** <br>\n",
    "What happens to the error rates? Describe how this depends on the dataset size (try 50, 100, 500, 1000 trials), as well as on the degree of deviation between the data-generation distribution and the priors.<br>\n",
    "\n",
    "\n",
    "**§:** Note that we are still using the discritized parameter grid (used to approximate the marginal likelihood integral by a discrete sum), so by Gaussian distribution we really mean truncated Gaussians. More precisely let the prior probability of each choice of $\\Theta$ in the grid be given by the expression for the Gaussian probability density, but then normalize this prior  to sum to 1 over the parameter grid. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d11ef8b8-33aa-4eae-8b27-82020c6a274f",
   "metadata": {},
   "source": [
    "# Week 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17a7edef-5b6d-4298-a90d-bc7b6597ab14",
   "metadata": {},
   "source": [
    "In this task we will explore the effect of **model mismatch** on model selection, and the related concept of **brittleness** (a term coined in Chandrasekaran et al. 2018).<br>\n",
    "As may be intuitive, mismatch of either model with the true data generative process would lead to additional systematic errors in model selection (as a means to hypothesis testing).\n",
    "Brittleness refers to a special case of this: in order to quantitatively evaluate empirical evidence for a conceptual hypothesis (such as whether a LIP neuron \"continuously\n",
    "accumulates evidence by ramping up its activity\", or whether \"its firing rate abruptly jumps during a trial at a random time\") we need to concoct a precise mathematical model to encapsulate\n",
    "the alternative conceptual hypotheses. In order to carry out probabilistic/Bayesian inference we further need these models to be probabilistic models, with a likelihood function\n",
    "and a prior distribution over parameters. In the process of designing such mathematical models, in general, and probabilistic models, in particular, various choices need to be made\n",
    "for different aspects of the model which are independent of the conceptual hypotheses and are not constrained by them.\n",
    "\n",
    "For example, in our case, a model that is identical to our step model, but has a jump time distribution that is not negative binomial, but, say, a uniform distribution, is still consistent with the<br>\n",
    "conceptual hypothesis of a discontinous jump in the firing rate, as opposed to a continuous ramp.<br>  \n",
    "As another example, our two competing hypotheses differ only in their predictions for the behaviour of latent firing rates, and are silent on the precise mechanism or statistical form of <br>\n",
    "spiking based on those firing rates. However, in order to build a probabilistic model of our observed data (spike counts) we needed to make a modeling choice and decided to model spike counts as<br> Poisson distributed (conditional on rates).<br>\n",
    "\n",
    "But what if reality deviates from these assumptions? What effect would such deviations have on our inferences about the true mechanism underlying the data?<br>\n",
    "Finally, what can we do to reduce the brittleness of hypothesis testing? This task is intended to provide insight about these questions. \n",
    "\n",
    "For all the tasks below, use $M = 10$, and unless otherwise instructed, $N_{\\rm trials} = 25$, to save on computational time. For number of datasets (to evaluate error rates)<br>\n",
    "use 10 datasets at a minimum, and no need to exceed 100. (In general, be strategic about choice of the number of datasets. If you see, based on small-scale explorations, that a trend that<br>\n",
    "is the subject of a sub-task or question is visible for a smaller number of datasets, you don't need to use more datasets.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e511d802-50e8-4762-90ac-41293a3df57b",
   "metadata": {},
   "source": [
    "### Task 4.1\n",
    "**Non-Poissonian spiking**\n",
    "\n",
    "Here you will study the effect of (unaccounted) non-Poissonian spiking on the model selection. The function `gamma_isi_point_process` in `models.py` allows you<br>\n",
    "to simulate non-Poissonian spikes within a certain family for which inter-spike interval (ISI)  distributions deviate from the exponential distribution<br>\n",
    "which is the ISI distribution under Poisson spiking (this corresponds to the special case of `shape = 1` for the parameter of `gamma_isi_point_process`).<br>\n",
    "(You are not directly interested in ISI's and their distribution for the tasks.) \n",
    "\n",
    "Generalise your code implementing the HMM-based ramp and step models to emit spikes according to the new distribution family (inspection of the `emit` methods of the ramp and step model<br>\n",
    "classes in `models.py`  shows you how to use `gamma_isi_point_process`).\n",
    "\n",
    "**4.1.1** Investigate the behavior of the two types of model-selection error rate (the mis-classification rate for each case of true model) with increasing values of the<br>\n",
    "`shape` parameter of `gamma_isi_point_process` (do not use values of `shape` less than 1, and no need to go above 5). Try this for a few different choices of true model parameters<br>\n",
    "(making sure they are chosen to produce ramp-like PSTH's). In this sub-task use $N_{\\rm trials} = 25$.\n",
    "\n",
    "How does deviation from Poisson  (with `shape` > 1) affect the two error rates? Does it bias the model selection to one or the other model systematically? If so, towards which model?  \n",
    "\n",
    "**4.1.2** The model-mismatch due to non-Poisson spiking is an example of mismatch of a model's likelihood with the true data-generating model. A mismatch of model prior with the <br>\n",
    "true distribution from which true model parameters are sampled (in the real world, this could correspond to the variation of parameters across real LIP neurons, for example)<br>\n",
    "can also affect model selection. You have already investigated this in 3.2.2. Investigate the difference in the effect of number of trials on model misclassification biases, <br>\n",
    "for the two cases of prior mismatch and unaccounted non-Poissonian spiking (choose a sensible degree of prior mismatch based on your explorations in 3.2.2).<br>\n",
    "If there is a systematic difference, try to provide an explanation. <br>\n",
    "(You can use $N_{\\rm trials} = 10, 20, 50, 200$ --and higher if possible-- and you can fix the true parameters, including `shape`, to reasonable values; reasonable according to your previous findings.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71139b34-403e-481d-877e-8bc16061987b",
   "metadata": {},
   "source": [
    "### Task 4.2\n",
    "\n",
    "**Understanding the effects of non-Poissonian spiking**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10be085e-b90c-45a1-b6ab-806ed944aba7",
   "metadata": {},
   "source": [
    "This task only involves simulation (spike train generation) of the HMM models (or the models from `models.py`, up to you) with `gamma_isi_point_process` used for spike emission.<br>\n",
    "By remaking the sort of plots you made in Week 1 (e.g. raster plots and plots of PSTH and Fano Factor vs. time) for the non-Poissonian model and <br>\n",
    "observing its behavior, answer the following questions:\n",
    "\n",
    "- How does `shape` affect spiking quality and the resulting Fano Factor behaviour?\n",
    "- Based on your observations, try to explain the effects  that you observed in **4.1.1** of  non-Poisson spiking on model selection errors. \n"
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}