docs/source/notebooks/ODE_API_introduction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import pymc3 as pm\n",
    "from pymc3.ode import DifferentialEquation\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.integrate import odeint\n",
    "import arviz as az\n",
    "import theano\n",
    "\n",
    "plt.style.use('seaborn-darkgrid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GSoC 2019: Introduction of pymc3.ode API\n",
    "by [Demetri Pananos](https://dpananos.github.io/posts/2019/08/blog-post-21/)\n",
    "\n",
    "Ordinary differential equations (ODEs) are a convenient mathematical framework for modelling the temporal dynamics of a system in disciplines from engineering to ecology. Though most analyses focus on bifurcations and stability of fixed points, parameter and uncertainty estimates are more salient in systems of practical interest, such as population pharmacokinetics and pharmacodynamics.\n",
    "\n",
    "\n",
    "Both parameter estimation and uncertainty propagation are handled elegantly by the Bayesian framework.  In this notebook, I showcase how PyMC3 can be used to do inference for differential equations using the `ode` submodule.  \n",
    "\n",
    "\n",
    "# Catching Up On Differential Equations\n",
    "\n",
    "A differential equation is an equation relating an unknown function's derivative to itself.  We usually write differentual equations as \n",
    "\n",
    "$$ \\mathbf{y}' = f(\\mathbf{y},t,\\mathbf{p}) \\quad \\mathbf{y}(t_0) = \\mathbf{y}_0 $$\n",
    "\n",
    "Here, $\\mathbf{y}$ is the unknown function, $t$ is time, and $\\mathbf{p}$ is a vector of parameters.  The function $f$ can be either scalar or vector valued.\n",
    "\n",
    "Only a small subset of differential equations have an analytical solution.  For most differential equations of applied interest, numerical methods must be used to obtain approximate solutions.\n",
    "\n",
    "\n",
    "# Doing Bayesian Inference With Differential Equations\n",
    "\n",
    "PyMC3 uses Hamiltonian Monte Carlo (HMC) to obtain samples from the posterior distribution.  HMC requires derivatives of the ODE's solution with respect to the parameters $p$.  The `ode` submodual automatically computes appropriate derivatives so you don't have to.  All you have to do is \n",
    "\n",
    "* Write the differential equation as a python function\n",
    "* Write the model in PyMC3\n",
    "* Hit the Inference Button $^{\\text{TM}}$\n",
    "\n",
    "Let's see how this is done in practice with a small example.\n",
    "\n",
    "# A Differential Equation For Freefall\n",
    "\n",
    "An object of mass $m$ is brought to some height and allowed to fall freely until it reaches the ground. A differential equation describing the object's speed over time is \n",
    "\n",
    "$$ y' = mg - \\gamma y $$\n",
    "\n",
    "The force the object experiences in the downwards direction is $mg$, while the force the object experiences in the opposite direction (due to air resistance) is proportional to how fast the object is presently moving. Let's assume the object starts from rest (that is, that the object's inital velocity is 0).  This may or may not be the case.  To showcase how to do inference on intial conditions, I will first assume the object starts from rest, and then relax that assumption later.\n",
    "\n",
    "Data on this object's speed as a function of time is shown below.  The data may be noisy because of our measurement tools, or because the object is an irregular shape, thus leading to times during freefall when the object is more/less aerodynamic.  Let's use this data to estimate the proportionality constant for air restistance.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For reproducibility\n",
    "np.random.seed(20394)\n",
    "\n",
    "def freefall(y, t, p):    \n",
    "    return 2.0*p[1] - p[0]*y[0]\n",
    "\n",
    "# Times for observation\n",
    "times = np.arange(0,10,0.5)\n",
    "gamma,g, y0, sigma = 0.4, 9.8, -2, 2\n",
    "y = odeint(freefall, t=times, y0=y0, args=tuple([[gamma,g]]))\n",
    "yobs = np.random.normal(y,2)\n",
    "\n",
    "fig, ax = plt.subplots(dpi=120)\n",
    "plt.plot(times,yobs, label='observed speed', linestyle='dashed', marker='o', color='red')\n",
    "plt.plot(times,y, label='True speed', color='k', alpha=0.5)\n",
    "plt.legend()\n",
    "plt.xlabel('Time (Seconds)')\n",
    "plt.ylabel(r'$y(t)$');\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To specify and ordinary differential equation with pyMC3, use the `DifferentialEquation` class.  This class takes as arguments:\n",
    "\n",
    "* `func`: A function specifying the differential equation (i.e. $f(\\mathbf{y},t,\\mathbf{p})$).\n",
    "* `times`: An array of times at which data was observed.\n",
    "* `n_states`: The dimension of $f(\\mathbf{y},t,\\mathbf{p})$.\n",
    "* `n_theta`: The dimension of $\\mathbf{p}$.\n",
    "* `t0`: Optional time to which the initial condition belongs.\n",
    "\n",
    "The argument `func` needs to be written as if `y` and `p` are vectors.  So even when your model has one state and/or one parameter, you should explicitly write `y[0]` and/or `p[0]`.\n",
    "\n",
    "Once the model is specified, we can use it in our pyMC3 model by passing paramerters and inital conditions.  `DifferentialEquation` returns a flattened solution, so you will need to reshape it to the same shape as your observed data in the model.\n",
    "\n",
    "Shown below is a model to estimate $\\gamma$ in the ODE above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ode_model = DifferentialEquation(\n",
    "    func=freefall,\n",
    "    times=times,\n",
    "    n_states=1, n_theta=2,\n",
    "    t0=0\n",
    ")\n",
    "\n",
    "with pm.Model() as model:\n",
    "    # Specify prior distributions for soem of our model parameters\n",
    "    sigma = pm.HalfCauchy('sigma',1)    \n",
    "    gamma = pm.Lognormal('gamma',0,1)\n",
    "    \n",
    "    # If we know one of the parameter values, we can simply pass the value.\n",
    "    ode_solution = ode_model(y0=[0], theta=[gamma, 9.8])\n",
    "    # The ode_solution has a shape of (n_times, n_states)\n",
    "    \n",
    "    Y = pm.Normal('Y', mu=ode_solution, sd=sigma, observed=yobs)\n",
    "    \n",
    "    prior = pm.sample_prior_predictive()\n",
    "    trace = pm.sample(2000, tune=1000, cores=1)\n",
    "    posterior_predictive = pm.sample_posterior_predictive(trace)\n",
    "    \n",
    "    data = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=posterior_predictive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(data);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our estimates of the proportionality constant and noise in the system are incredibly close to their actual values!\n",
    "\n",
    "We can even estimate the acceleration due to gravity by specifying a prior for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model2:    \n",
    "    sigma = pm.HalfCauchy('sigma',1)\n",
    "    gamma = pm.Lognormal('gamma',0,1)\n",
    "    # A prior on the acceleration due to gravity\n",
    "    g = pm.Lognormal('g',pm.math.log(10),2)\n",
    "    \n",
    "    # Notice now I have passed g to the odeparams argument\n",
    "    ode_solution = ode_model(y0=[0], theta=[gamma, g])\n",
    "    \n",
    "    Y = pm.Normal('Y', mu=ode_solution, sd=sigma, observed=yobs)\n",
    "\n",
    "    prior = pm.sample_prior_predictive()\n",
    "    trace = pm.sample(2000, tune=1000, target_accept=0.9, cores=1)\n",
    "    posterior_predictive = pm.sample_posterior_predictive(trace)\n",
    "    \n",
    "    data = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=posterior_predictive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(data);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The uncertainty in the acceleration due to gravity has increased our uncertainty in the prportionality constant.\n",
    "\n",
    "Finally, we can do inference on the initial condition.  If this object was brought to it's initial height by an airplane, then turbulent air might have made the airplane move up or down, thereby changing the inital velocity of the object. \n",
    "\n",
    "Doing inference on the inital condition is as easy as specifying a prior for the inital condition, and then passing the inital condition to `ode_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as model3:    \n",
    "    sigma = pm.HalfCauchy('sigma',1)\n",
    "    gamma = pm.Lognormal('gamma',0,1)\n",
    "    g = pm.Lognormal('g',pm.math.log(10),2)\n",
    "    # Initial condition prior.  We think it is at rest, but will allow for perturbations in initial velocity.\n",
    "    y0 = pm.Normal('y0', 0, 2)\n",
    "    \n",
    "    ode_solution = ode_model(y0=[y0], theta=[gamma, g])\n",
    "    \n",
    "    Y = pm.Normal('Y', mu=ode_solution, sd=sigma, observed=yobs)\n",
    "    \n",
    "    prior = pm.sample_prior_predictive()\n",
    "    trace = pm.sample(2000, tune=1000, target_accept=0.9, cores=1)\n",
    "    posterior_predictive = pm.sample_posterior_predictive(trace)\n",
    "    \n",
    "    data = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=posterior_predictive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(data, figsize=(13,3));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that by explicitly modelling the initial condition, we obtain a much better estimate of the acceleration due to gravity than if we had insisted that the object started at rest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Non-linear Differential Equations\n",
    "\n",
    "The example of an object in free fall might not be the most appropriate since that differential equation can be solved exactly. Thus, `DifferentialEquation` is not needed to solve that particular problem.  There are, however, many examples of differential equations which cannot be solved exactly.  Inference for these models is where `DifferentialEquation` truly shines.\n",
    "\n",
    "Consider the SIR model of infection.  This model describes the temporal dynamics of a disease spreading through a homogenously mixed closed population.  Members of the population are placed into one of three cateories: Susceptible, Infective, or Recovered.  The differential equations are...\n",
    "\n",
    "\n",
    "$$ \\dfrac{dS}{dt} = - \\beta SI \\quad S(0) = S_0 $$\n",
    "$$ \\dfrac{dI}{dt} = \\beta SI - \\lambda I \\quad I(0) = I_0 $$\n",
    "$$ \\dfrac{dR}{dt} = \\lambda I \\quad R(0) = R_0 $$\n",
    "\n",
    "With the constraint that $S(t) + I(t) + R(t) = 1 \\, \\forall t$. Here, $\\beta$ is the  rate  of infection per susceptible and per infective, and $\\lambda$ is the rate of recovery.\n",
    "\n",
    "If we knew $S(t)$ and $I(t)$, then we could determine $R(t)$, so we can peel off the differential equation for $R(t)$ and work only with the first two.  \n",
    "\n",
    "\n",
    "In the SIR model, it is straight-forward to see that $\\beta, \\gamma$ and $\\beta/2, \\gamma/2$ will produce the same qualitative dynamics but on much different time scales.  To study the *quality* of the dynamics, regardless of time scale, applied mathematicians will *non-dimensionalize* differential equations.  Non-dimensionalization is the process of introducing scaleless variables into the differential equation to understand the system's dynamics under families of equivalent paramterizations.\n",
    "\n",
    "To non-dimensionalize this system, let's scale time by $1/\\lambda$ (we do this because people stay infected for an average of $1/\\lambda$ units of time.  It is a straight forward argument to show this.  For more, see [1]).  Let $t = \\tau/\\lambda$, where $\\tau$ is a unitless variable.  Then...\n",
    "\n",
    "\n",
    "$$ \\dfrac{dS}{d\\tau} = \\dfrac{dt}{d\\tau} \\dfrac{dS}{dt} = \\dfrac{1}{\\lambda}\\dfrac{dS}{dt} = -\\dfrac{\\beta}{\\lambda}SI$$\n",
    "\n",
    "and \n",
    "\n",
    "$$ \\dfrac{dI}{d\\tau} = \\dfrac{dt}{d\\tau} \\dfrac{dI}{dt} = \\dfrac{1}{\\lambda}\\dfrac{dI}{dt} = \\dfrac{\\beta}{\\lambda}SI - I$$\n",
    "\n",
    "The quantity $\\beta/\\lambda$ has a very special name.  We call it *The R-Nought* ($\\mathcal{R}_0$).  It's interpretation is that if we were to drop a single infected person into a population of suceptible individuals, we would expect $\\mathcal{R}_0$ new infections.  If $\\mathcal{R}_0>1$, then an epidemic will take place.  If $\\mathcal{R}_0\\leq1$ then there will be no epidemic (note, we can show this more rigoursly by studying eigenvalues of the system's Jacobain.  For more, see [2]).\n",
    "\n",
    "This non-dimensionalization is important because it gives us information about the parameters.  If we see an epidemic has occured, then we know that $\\mathcal{R}_0>1$ which means $\\beta> \\lambda$. Furthermore, it might be hard to place a prior on $\\beta$ because of beta's interpretation.  But since $1/\\lambda$ has a simple interpretation, and since $\\mathcal{R}_0>1$, we can obtain $\\beta$ by computing $\\mathcal{R}_0\\lambda$. \n",
    "\n",
    "Side note: I'm going to choose a likelihood which certainly violates these constraints, just for exposition on how to use `DifferentialEquation`.  In reality, a likelihood which respects these constraints should be chosen.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SIR(y, t, p):\n",
    "    ds = -p[0]*y[0]*y[1]\n",
    "    di = p[0]*y[0]*y[1] - p[1]*y[1]    \n",
    "    return [ds, di]\n",
    "\n",
    "times = np.arange(0,5,0.25)\n",
    "\n",
    "beta,gamma = 4,1.0\n",
    "# Create true curves\n",
    "y = odeint(SIR, t=times, y0=[0.99, 0.01], args=((beta,gamma),), rtol=1e-8)\n",
    "# Observational model.  Lognormal likelihood isn't appropriate, but we'll do it anyway\n",
    "yobs = np.random.lognormal(mean=np.log(y[1::]), sigma=[0.2, 0.3])\n",
    "\n",
    "\n",
    "plt.plot(times[1::],yobs, marker='o', linestyle='none')\n",
    "plt.plot(times, y[:,0], color='C0', alpha=0.5, label=f'$S(t)$')\n",
    "plt.plot(times, y[:,1], color ='C1', alpha=0.5, label=f'$I(t)$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sir_model = DifferentialEquation(\n",
    "    func=SIR, \n",
    "    times=np.arange(0.25, 5, 0.25), \n",
    "    n_states=2,\n",
    "    n_theta=2,\n",
    "    t0=0,\n",
    ")\n",
    "\n",
    "with pm.Model() as model4:    \n",
    "    sigma = pm.HalfCauchy('sigma', 1, shape=2)\n",
    "    \n",
    "    # R0 is bounded below by 1 because we see an epidemic has occured\n",
    "    R0 = pm.Bound(pm.Normal, lower=1)('R0', 2,3)\n",
    "    lam = pm.Lognormal('lambda',pm.math.log(2),2)\n",
    "    beta = pm.Deterministic('beta', lam*R0)\n",
    "    \n",
    "    sir_curves = sir_model(y0=[0.99, 0.01], theta=[beta, lam])\n",
    "    \n",
    "    Y = pm.Lognormal('Y', mu=pm.math.log(sir_curves), sd=sigma, observed=yobs)\n",
    "\n",
    "    prior = pm.sample_prior_predictive()\n",
    "    trace = pm.sample(2000,tune=1000, target_accept=0.9, cores=1)\n",
    "    posterior_predictive = pm.sample_posterior_predictive(trace)\n",
    "    \n",
    "    data = az.from_pymc3(trace=trace, prior = prior, posterior_predictive = posterior_predictive)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(data,round_to=2, credible_interval=0.95);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen from the posterior plots, $\\beta$ is well estimated by leveraging knoweldege about the non-dimensional parameter $\\mathcal{R}_0$ and $\\lambda$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusions & Final Thoughts\n",
    "\n",
    "ODEs are a really good model for continuous temporal evolution.  With the addition of `DifferentialEquation` to PyMC3, we can now use bayesian methods to estimate the parameters of ODEs.\n",
    "\n",
    "`DifferentialEquation` is not as fast as compared to Stan's `integrate_ode_bdf`.  However, the ease of use of `DifferentialEquation` will allow practioners to get up and running much quicker with Bayesian estimation for ODEs than Stan (which has a steep learning curve).  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "\n",
    "1. Earn, D. J., et al. Mathematical epidemiology. Berlin: Springer, 2008.\n",
    "2. Britton, Nicholas F. Essential mathematical biology. Springer Science & Business Media, 2012.\n"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}