Continuous time Markov processes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Continuous-time Markov processes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A **Markov process** is a stochastic model for a system, in which the evolution of the system is determined only by the _current_ state of the system, not by any state in the past _(except perhaps through derivatives of the dynamical quantities, where those derivatives exist - often they do not)_. Such a system may be said to possess the **Markov property** ([Wikipedia](https://en.wikipedia.org/wiki/Markov_property)), or to be 'memoryless'.\n",
    "\n",
    "The most famous and elementary types of Markov processes are Markov chains ([Wikipedia](https://en.wikipedia.org/wiki/Markov_chain)), which have a discrete set of states (which may be infinite) and evolution between the states takes place in discrete, uniform units of time. However, for some systems a **continuous-time Markov process** may be more appropriate. Under continuous-time evolution, we speak not of 'transition probabilities' but rather of 'transition rates'. This is most important when we care about the dependence of the dynamical properties (_e.g._ the population of a species) on time, and when the rate of change depends on the quantity (_e.g._ higher population implies shorter periods between birth or death events).\n",
    "\n",
    "In the examples below, the dynamical quantities are still assumed to take discrete values; it is only the time between state evolutions that is assumed to be continuous. This makes these CTMPs different from Brownian processes, where the dynamical quantities (_e.g._ position of a particle) are also continuous. \n",
    "\n",
    "At each state, the transition rates to other possible states are given by specifying a transition function, and typically these transitions are assumed to be Poisson processes - with the first event to 'trigger' determining the next state that is visited. The time to transition is therefore distributed according to an exponential distribution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import random\n",
    "import warnings\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from numpy.random import default_rng\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets\n",
    "\n",
    "logging.getLogger().setLevel(logging.WARN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One way to numerically simulate a CTMP is to use an _event-driven simulation_. At a given state, the transition rates are calculated (as specified by the transition function); and random exponential variables are sampled, one for each of the transition rates. These r.v.s are the times until each of the possible transitions triggers. The 'shortest time wins' and the state evolves to the corresponding state. \n",
    "\n",
    "To avoid numerical overflow or underflow, we terminate the solution if transition rates become too large (and therefore evolution times too small)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evolution(fun, tspan, state0, *, args=(), t_eval=None, maxrate=1e6):\n",
    "    \"\"\"Event-driven simulation of a continuous-time Markov process.\"\"\"\n",
    "    # TODO specify max rate/min step, to interrupt in case of very high rates\n",
    "    tmin, tmax = tspan\n",
    "    \n",
    "    state = state0\n",
    "    t = tmin\n",
    "    ts = [t]\n",
    "    history = [state0]\n",
    "    \n",
    "    while t < tmax:\n",
    "        transitions = fun(t, state, *args)\n",
    "        transitions = {k: transitions[k] for k in transitions if transitions[k] > 0}\n",
    "        if not transitions:\n",
    "            break\n",
    "            \n",
    "        if maxrate:\n",
    "            transitions = {k: transitions[k] for k in transitions if transitions[k] < maxrate}\n",
    "            if not transitions:\n",
    "                warnings.warn(\n",
    "                    f\"Maximum rate {maxrate} exceeded by all possible transitions at time {t}.\", \n",
    "                    RuntimeWarning\n",
    "                )\n",
    "                break\n",
    "                \n",
    "        logging.info(transitions)\n",
    "\n",
    "        \n",
    "        times_to_move = {k: np.random.exponential(1/transitions[k]) for k in transitions}\n",
    "        logging.info(times_to_move)       \n",
    "\n",
    "        new_state = min(times_to_move, key=lambda k: times_to_move[k])\n",
    "        time_elapsed = times_to_move[new_state]\n",
    "        t += time_elapsed\n",
    "        state = new_state\n",
    "        ts.append(t)\n",
    "        history.append(state)\n",
    "            \n",
    "    ts = np.array(ts)\n",
    "    history = np.array(history)\n",
    "    \n",
    "    if t_eval is None:\n",
    "        return ts, history\n",
    "    \n",
    "    interp = interp1d(ts, history, kind='previous', axis=0)\n",
    "    return np.array(t_eval), np.array([interp(t) for t in t_eval])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stochastic Lotka-Volterra\n",
    "\n",
    "The [Lotka-Volterra equations](https://en.wikipedia.org/wiki/Lotka%E2%80%93Volterra_equations) are a family of simple models for the populations of a predator-prey system. The original model is a deterministic dynamical system, consisting of a pair of ordinary differential equations (continuous time, continuous populations). This model is suitable for large populations but leads to the 'atto-fox' problem: populations never become extinct but reach extremely small quantities that are somehow able to bounce back. One resolution is to turn the populations into discrete quantities (you can't have half a fox, or $10^{-18}$ foxes – but you can rescale units to talk in terms of 1000s of foxes, for example). Having done that we can also add uncertainty by introducing a transition function between states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lv_transition_fun(t, n):\n",
    "    x, y = n\n",
    "    s = .001\n",
    "    return {\n",
    "        (x + s, y): 2/3 * x,\n",
    "#         (x - s, y): x * x / 20,\n",
    "        (x - 4/3*s, y+s): x * y,\n",
    "        (x, y - s): y\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact()\n",
    "def lv_demo():\n",
    "    ts, history = evolution(lv_transition_fun, (0, 25000), [1, 0.75],\n",
    "                        t_eval=np.linspace(0, 25000, 101),\n",
    "                        maxrate=10)\n",
    "    prey, predators = history[:, 0], history[:, 1]\n",
    "    fig = plt.figure(figsize=(8, 6))\n",
    "    plt.plot(ts, prey, 'g-',\n",
    "             ts, predators, 'r-')\n",
    "    plt.grid()\n",
    "    ax = plt.gca()\n",
    "    ax.set_ylim([0, None])\n",
    "    ax.legend(['prey', 'predators'])\n",
    "    plt.show()\n",
    "\n",
    "    fig = plt.figure(figsize=(15, 4))\n",
    "    plt.plot(prey, predators, 'k')\n",
    "    plt.grid()\n",
    "    ax = plt.gca()\n",
    "    ax.set_aspect('equal')\n",
    "    ax.set_xlabel('prey')\n",
    "    ax.set_ylabel('predators')\n",
    "    ax.set_xlim([0, None])\n",
    "    ax.set_ylim([0, None])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While the basic, deterministic Lotka-Volterra system has solutions that are closed loops, the realisations of the stochastic system also exhibit periodic-like behaviour, but with some variation and therefore not closed loops.\n",
    "\n",
    "If we calculate the _ensemble average_ of several realisations (perhaps with different initial conditions), we will find that the average also executes very periodic behaviour."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact()\n",
    "def lv_ensemble_demo():\n",
    "    # Same initial conditions\n",
    "    n_states = 5\n",
    "    initial_states = [(1, 1) for _ in range(n_states)]\n",
    "\n",
    "\n",
    "    # # Or different initial conditions\n",
    "    # initial_states = [(1 + r*np.cos(th), 1 + r*np.sin(th)) \n",
    "    #                   for r in np.linspace(.1, .4, 4)\n",
    "    #                   for th in np.linspace(0, 2*np.pi, 9)[:-1]]\n",
    "\n",
    "    tmax = 25000\n",
    "    ts = np.linspace(0, tmax, 201)\n",
    "\n",
    "    histories = [evolution(lv_transition_fun, (0, tmax), s, t_eval=ts,\n",
    "                          maxrate=20)[1] for s in initial_states]\n",
    "\n",
    "    average_x = np.average([h[:, 0] for h in histories], axis=0)\n",
    "    average_y = np.average([h[:, 1] for h in histories], axis=0)\n",
    "\n",
    "    @interact(step=widgets.IntSlider(min=0, max=len(ts)-1, value=0, continuous_update=False))\n",
    "    def plot_evolution(step):\n",
    "        final_states = [h[1, :] for h in histories]\n",
    "        plt.figure(figsize=[8,8])\n",
    "        plt.plot(\n",
    "            average_x, average_y, 'r-',\n",
    "            [x[step][0] for x in histories], \n",
    "            [x[step][1] for x in histories],\n",
    "            'k.',\n",
    "            average_x[step], average_y[step], 'ro'\n",
    "        )\n",
    "        plt.gca().set_xlim([-.1, 2])\n",
    "        plt.gca().set_ylim([-.1, 2])\n",
    "        plt.gca().set_aspect('equal')\n",
    "        plt.gca().grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stochastic harmonic oscillator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sho_transition_fun(t, s):\n",
    "    x, y = s\n",
    "    return {\n",
    "        (x + 1, y): y,\n",
    "        (x - 1, y): -y,\n",
    "        (x, y + 1): -x,\n",
    "        (x, y - 1): x\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact()\n",
    "def sho_demo():\n",
    "    tmax = 100\n",
    "    ts, history = evolution(sho_transition_fun, (0, tmax), [40, 40],\n",
    "                        t_eval=np.linspace(0, tmax, 121))\n",
    "    plt.plot(ts, history)\n",
    "    plt.show()\n",
    "\n",
    "    plt.plot(history[:, 0], history[:, 1])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although each individual realisation is noisy, the ensemble average again shows remarkable periodicity. However, unlike the Lotka-Volterra system, here there is quite a lot of variance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact()\n",
    "def sho_ensemble_demo():\n",
    "    initial_states = [(30 + r*np.cos(th), 40 + r*np.sin(th)) \n",
    "                  for r in np.linspace(1, 5, 5)\n",
    "                  for th in np.linspace(0, 2*np.pi, 9)]\n",
    "\n",
    "    histories = [evolution(sho_transition_fun, (0, 25), s, t_eval=np.linspace(0, 25, 161))[1]\n",
    "                 for s in initial_states]\n",
    "    average_x = np.average([h[:, 0] for h in histories], axis=0)\n",
    "    average_y = np.average([h[:, 1] for h in histories], axis=0)\n",
    "\n",
    "    @interact(step=widgets.IntSlider(min=0, max=len(histories[0])-1, value=0, continuous_update=False))\n",
    "    def plot_evolution(step):\n",
    "        final_states = [h[step][1] for h in histories]\n",
    "\n",
    "        plt.figure(figsize=[8,8])\n",
    "        plt.plot(\n",
    "            average_x, average_y, 'r-',\n",
    "            [x[step, 0] for x in histories],\n",
    "            [x[step, 1] for x in histories],\n",
    "                'k.',\n",
    "            average_x[step], average_y[step], 'ro',)\n",
    "        plt.gca().set_xlim([-60, 60])\n",
    "        plt.gca().set_ylim([-60, 60])\n",
    "        plt.gca().set_aspect('equal')\n",
    "        plt.gca().grid()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decay with regeneration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transition_fun(t, s):\n",
    "    x, y = s\n",
    "    return {\n",
    "        (x + 1, y): 1 - 0.05*x + y,\n",
    "        (x - 1, y): 1 + 0.05*x - y,\n",
    "        (x, y + 1): 1 - 0.05*y - x,\n",
    "        (x, y - 1): 1 + 0.05*y + x\n",
    "    }\n",
    "\n",
    "ts, history = evolution(transition_fun, (0, 100), [10, 10],\n",
    "                    t_eval=np.linspace(0, 100, 151))\n",
    "\n",
    "plt.figure(figsize=(14, 5))\n",
    "plt.plot(ts, history)\n",
    "plt.show()\n",
    "\n",
    "plt.plot(history[:, 0], history[:, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Warning: This is slow.\n",
    "\n",
    "initial_states = [(60 + r*np.cos(th), 60 + r*np.sin(th)) \n",
    "                  for r in np.linspace(1, 5, 5)\n",
    "                  for th in np.linspace(0, 2*np.pi, 9)[:-1]]\n",
    "\n",
    "histories = [evolution(transition_fun, (0, 100), s, t_eval=np.linspace(0, 100, 251))[1] for s in initial_states]\n",
    "average_x = np.average([h[:, 0] for h in histories], axis=0)\n",
    "average_y = np.average([h[:, 1] for h in histories], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(step=widgets.IntSlider(min=0, max=len(histories[0])-1, value=0, continuous_update=False))\n",
    "def plot_evolution(step):\n",
    "    final_states = [h[step][1] for h in histories]\n",
    "\n",
    "    plt.figure(figsize=[8,8])\n",
    "    plt.plot(\n",
    "        average_x, average_y, 'r-',\n",
    "        [x[step, 0] for x in histories],\n",
    "        [x[step, 1] for x in histories],\n",
    "            'k.',\n",
    "        average_x[step], average_y[step], 'ro',\n",
    "    )\n",
    "    plt.gca().set_xlim([-120, 120])\n",
    "    plt.gca().set_ylim([-120, 120])\n",
    "    plt.gca().set_aspect('equal')\n",
    "    plt.gca().grid()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A singular process\n",
    "\n",
    "This is the stochastic version of the dynamical system\n",
    "$$ \\dot{s} = s^2, $$\n",
    "which is a classic example of a system with a 'finite-time singularity', meaning that its solutions, which take the form\n",
    "$$ s = \\frac{1}{s_0^{-1} - t}, $$\n",
    "reach infinity in a finite amount of time. Contrast this to $\\dot{s} = s$, with solution $s = s_0 \\exp t$, which diverges but only as $t \\rightarrow \\infty$.\n",
    "\n",
    "The stochastic version also exhibits singular behaviour."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact()\n",
    "def singular_demo():\n",
    "    def transition_fun(t, s):\n",
    "        return {\n",
    "            s + 0.01: 100 * s ** 2\n",
    "        }\n",
    "\n",
    "    s0 = 1\n",
    "    ts, history = evolution(transition_fun, (0, 2), s0,\n",
    "                       maxrate=1e8)\n",
    "\n",
    "    plt.plot(ts, history,\n",
    "            ts, 1/(1/s0 - ts), 'k--')\n",
    "    ax = plt.gca()\n",
    "    ax.set_yscale('log')\n",
    "    ax.legend(['stochastic solution', 'continuous solution'])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Queueing process with multiple servers\n",
    "\n",
    "Inspired when I was living at Minack with 40 other people with three showers between us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(\n",
    "    in_rate=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    out_rate=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    n_servers=widgets.IntSlider(min=0, max=12, value=3),\n",
    ")\n",
    "def queueing_process(in_rate, out_rate, n_servers):\n",
    "    def transition_fun(t, s):\n",
    "        todo, inprog, done = s\n",
    "        outcomes = {}\n",
    "        if inprog == n_servers:  # max capacity already, so new tickets go to the queue\n",
    "            outcomes[todo + 1, inprog, done] = in_rate\n",
    "        else:\n",
    "            outcomes[todo, inprog + 1, done] = in_rate\n",
    "\n",
    "        if todo > 0:\n",
    "            outcomes[todo - 1, inprog, done + 1] = out_rate * inprog\n",
    "        else:\n",
    "            outcomes[0, inprog - 1, done + 1] = out_rate * inprog\n",
    "\n",
    "\n",
    "        return outcomes\n",
    "\n",
    "    ts, history = evolution(transition_fun, (0, 100), (0, 0, 0),\n",
    "                        t_eval=np.linspace(0, 100, 101))\n",
    "\n",
    "    ax = plt.gca()\n",
    "    ax.plot(\n",
    "        ts, history[:, 0], 'r-',\n",
    "        ts, history[:, 1], 'k:',\n",
    "#         ts, history[:, 2], 'k-'\n",
    "    )\n",
    "    ax.set_ylim([0, None])\n",
    "    ax.grid()\n",
    "    ax.legend(['queueing', 'in progress', 'done'])\n",
    "#     return ts, history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Immigration-death system (Mathematical Biology Sheet 2 Question 4)\n",
    "\n",
    "Rates:\n",
    "* $\\lambda_1, \\lambda_2$ - birth of one or two offspring\n",
    "* $\\beta n$ - deaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@interact(\n",
    "    l_1=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    l_2=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    beta=widgets.FloatSlider(min=0, max=1, value=0.1, continuous_update=False),\n",
    ")\n",
    "def immigration_death_demo(l_1, l_2, beta):\n",
    "    def transition_fun(t, s, l_1, l_2, beta):\n",
    "        return {\n",
    "            s + 1: l_1,\n",
    "            s + 2: l_2,\n",
    "            s - 1: beta * s\n",
    "        }\n",
    "\n",
    "    tmax = 500\n",
    "    ts, history = evolution(transition_fun, (0, tmax), 0,\n",
    "                        args=(l_1, l_2, beta),\n",
    "                        t_eval=np.arange(0, tmax, 1))\n",
    "    \n",
    "    mean = 1/beta * (l_1 + 2*l_2)\n",
    "    sigma = (1/beta * (l_1 + 3*l_2))**0.5\n",
    "    plt.plot(ts, history, 'r-',\n",
    "             ts, [mean] * len(ts), 'k--',\n",
    "             ts, [mean + sigma] * len(ts), 'k:',\n",
    "             ts, [mean - sigma] * len(ts), 'k:',\n",
    "    )\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Immigration-death system (Mathematical Biology Sheet 2 Question 5)\n",
    "\n",
    "Similar to the above, except now we introduce two or three, rather than one or two, at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@interact(\n",
    "    l_2=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    l_3=widgets.FloatSlider(min=0, max=5, value=1, continuous_update=False),\n",
    "    beta=widgets.FloatSlider(min=0, max=1, value=0.1, continuous_update=False),\n",
    ")\n",
    "def immigration_2_or_3_death_demo(l_2, l_3, beta):\n",
    "    def transition_fun(t, s, l_2, l_3, beta):\n",
    "        return {\n",
    "            s + 2: l_2,\n",
    "            s + 3: l_3,\n",
    "            s - 1: beta * s\n",
    "        }\n",
    "\n",
    "    tmax = 500\n",
    "    ts, history = evolution(transition_fun, (0, tmax), 0,\n",
    "                        args=(l_2, l_3, beta),\n",
    "                        t_eval=np.arange(0, tmax, 1))\n",
    "    \n",
    "    mean = 1/beta * (2*l_2 + 3*l_3)\n",
    "#     sigma = 0  # FIXME\n",
    "    sigma = ((3 * mean - 5 * l_2 / beta) / 3) ** 0.5\n",
    "    plt.plot(ts, history, 'r-',\n",
    "             ts, [mean] * len(ts), 'k--',\n",
    "             ts, [mean + sigma] * len(ts), 'k:',\n",
    "             ts, [mean - sigma] * len(ts), 'k:',\n",
    "    )\n",
    "    plt.title(f'$\\\\mu = {mean:.2f}$, $\\\\sigma = {sigma:.2f}$')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Immigration-death system with quadratic death rate (Mathematical Biology sheet 2 question 6)\n",
    "\n",
    "Throughout we normalize such that $\\beta = 1$, so that we can work with $r = \\lambda / \\beta$ instead.\n",
    "\n",
    "In the steady state, the generating function satisfies the equation\n",
    "$$\n",
    "s \\frac{\\partial^2\\phi}{\\partial s^2} = r \\phi, \\qquad \\phi(s=1) = 1, \\text{with $\\phi(s=0)$ bounded}.\n",
    "$$\n",
    "This equation does not have a simple closed-form solution, although the solution can be written as:\n",
    "$$\n",
    "\\phi = \\frac{ s^{1/2} I_1\\left(2(rs)^{1/2}\\right) }{I_1\\left(2r^{1/2}\\right)}\n",
    "$$\n",
    "where $I_n$ is a [modified Bessel function of the first kind](https://mathworld.wolfram.com/ModifiedBesselFunctionoftheFirstKind.html), the solution to the differential equation\n",
    "$$\n",
    "x^2 y'' + x y' - (x^2 + n^2) y = 0\n",
    "$$\n",
    "that is regular at $x = 0$. Then (according to Wolfram Alpha)...\n",
    "$$\n",
    "\\mu = \\frac{\\partial\\phi}{\\partial s}(s=1) = \\frac{1}{2} \\left[ \n",
    "\\frac{r^{1/2}}{I_1(2 r^{1/2})} \\left(\n",
    "I_0(2r^{1/2}) + I_2(2r^{1/2})\n",
    "\\right) + 1\\right].\n",
    "$$\n",
    "We also know that\n",
    "$$\n",
    "\\sigma = \\sqrt{r + \\mu - \\mu^2}.\n",
    "$$\n",
    "\n",
    "It can be shown ([source](https://dlmf.nist.gov/10.40)) that, as $z \\rightarrow \\infty$, then for any $n$,\n",
    "$$\n",
    "I_n(z) \\sim \\frac{\\mathrm{e}^z}{(2\\pi z)^{1/2}},\n",
    "$$\n",
    "and therefore (at leading order):\n",
    "$$ \\mu \\sim r^{1/2}, \\quad \\sigma \\sim r^{1/4}. $$\n",
    "This is consistent with what you might expect from a scaling (dimensional analytical) argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.special import iv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(\n",
    "    r=widgets.FloatSlider(min=0, max=200, value=100, step=0.1, continuous_update=False),\n",
    ")\n",
    "def immigration_nonlinear_death_demo(r):\n",
    "    def f(s):\n",
    "        return s * (s - 1)\n",
    "    \n",
    "    def transition_fun(t, s, r):\n",
    "        return {\n",
    "            s + 1: r,\n",
    "            s - 1: f(s)\n",
    "        }\n",
    "\n",
    "    tmax = 500\n",
    "    ts, history = evolution(transition_fun, (0, tmax), 0,\n",
    "                        args=(r,),\n",
    "                        t_eval=np.arange(0, tmax, 1))\n",
    "\n",
    "    z = 2*r**0.5\n",
    "    mean = 0.5 * (r**0.5/iv(1, z) * (iv(0, z) + iv(2, z)) + 1 )\n",
    "    sigma = (r + mean - mean**2) ** 0.5\n",
    "    plt.plot(ts, history, 'r-',\n",
    "             ts, [mean] * len(ts), 'k--',\n",
    "             ts, [mean + sigma] * len(ts), 'k:',\n",
    "             ts, [mean - sigma] * len(ts), 'k:',\n",
    "    )\n",
    "    plt.show()"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}