docs/examples/01-basic-model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic model introduction\n",
    "\n",
    "This page introduces the processes for building and running a simple compartmental disease model with Summer.\n",
    "In the following example, we will create an SIR compartmental model for a general, unspecified emerging infectious disease spreading through a fully susceptible population. In this model there will be:\n",
    "\n",
    "- three compartments: susceptible (S), infected (I) and recovered (R)\n",
    "- a starting population of 1000 people, with 10 of them infected (and infectious)\n",
    "- an evaluation timespan from day zero to 20 in 0.1 day steps\n",
    "- inter-compartmental flows for infection, deaths and recovery\n",
    "\n",
    "First, let's look at a complete example of this model in action, and then examine the details of each step. This is the complete example model that we will be working with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import CompartmentalModel\n",
    "\n",
    "import pandas as pd\n",
    "pd.options.plotting.backend = \"plotly\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model\n",
    "model = CompartmentalModel(\n",
    "    times=[0, 20],\n",
    "    compartments=[\"S\", \"I\", \"R\"],\n",
    "    infectious_compartments=[\"I\"],\n",
    "    timestep = 0.1\n",
    ")\n",
    "model.set_initial_population(distribution={\"S\": 990, \"I\": 10})\n",
    "model.add_infection_frequency_flow(name=\"infection\", contact_rate=1, source=\"S\", dest=\"I\")\n",
    "model.add_transition_flow(name=\"recovery\", fractional_rate=1/3, source=\"I\", dest=\"R\")\n",
    "model.add_death_flow(name=\"infection_death\", death_rate=0.05, source=\"I\")\n",
    "\n",
    "# Run the model\n",
    "model.run()\n",
    "\n",
    "# Plot the outputs\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Now let's inspect each step of the example in more details. To start, here's how to create a new model: let's import the summer2 library and create a new [CompartmentalModel](../api/model.rst) object. You can see that our model has an attribute called `compartments`, which contains a description of each modelled compartment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CompartmentalModel(\n",
    "    times=[0, 20],\n",
    "    compartments=[\"S\", \"I\", \"R\"],\n",
    "    infectious_compartments=[\"I\"],\n",
    "    timestep=0.1\n",
    ")\n",
    "\n",
    "# View a description of the model compartments\n",
    "model.compartments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding a population \n",
    "\n",
    "Initially the model compartments are all empty. Let's add:\n",
    "\n",
    "- 990 people to the susceptible (S) compartment, plus\n",
    "- 10 in the infectious (I) compartment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add people to the model\n",
    "model.set_initial_population(distribution={\"S\": 990, \"I\": 10})\n",
    "\n",
    "# View the initial population\n",
    "\n",
    "# This will be performed automatically when a model is run, \n",
    "# but must be run manually to inspect the value interactively\n",
    "from summer2.population import calculate_initial_population\n",
    "\n",
    "calculate_initial_population(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding inter-compartmental flows \n",
    "\n",
    "Now, let's add some flows for people to transition between the compartments. These flows will define the dynamics of our infection. We will add:\n",
    "\n",
    "- an infection flow from S to I (using frequency-dependent transmission)\n",
    "- a recovery flow from I to R\n",
    "- an infection death flow, that depletes people from the I compartment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Susceptible people can get infected.\n",
    "model.add_infection_frequency_flow(name=\"infection\", contact_rate=1.0, source=\"S\", dest=\"I\")\n",
    "\n",
    "# Infectious people take 3 days, on average, to recover.\n",
    "# If the model was run at this stage of construction, \n",
    "# then the basic reproduction number (R0) of this infection would be 3.\n",
    "model.add_transition_flow(name=\"recovery\", fractional_rate=1/3, source=\"I\", dest=\"R\")\n",
    "\n",
    "# Add an infection-specific death flow to the I compartment.\n",
    "# This now slightly reduces the actual sojourn time in the I compartment\n",
    "# from the original request of 3 days, and so slightly reduces R0 as well.\n",
    "model.add_death_flow(name=\"infection_death\", death_rate=0.05, source=\"I\")\n",
    "\n",
    "# Inspect the new flows, which we just added to the model.\n",
    "model.flows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the model\n",
    "\n",
    "Now we can calculate the outputs for the model over the requested time period. \n",
    "The model calculates the compartment sizes by solving a system of differential equations (defined by the flows we just added) over the requested time period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View the model outputs\n",
    "\n",
    "The model's results are available as a Pandas DataFrame, via the get_outputs_df() method \n",
    "This is available after the model has been run. Let's have a look at what's inside:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_outputs_df()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the outputs\n",
    "\n",
    "You can get a better idea of what is going on inside the model by visualizing how the compartment sizes change over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_outputs_df().plot(title = \"SIR Model Outputs\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "That's it for now, now you know how to:\n",
    "\n",
    "- Create a model\n",
    "- Add a population\n",
    "- Add flows\n",
    "- Run the model\n",
    "- Access and visualize the outputs\n",
    "\n",
    "A detailed API reference for the CompartmentalModel class can be found [here](../api/model.rst)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus: how the model works inside\n",
    "\n",
    "This section presents a code snippet that shows an approximation of what is happening inside the model we just built and ran.\n",
    "\n",
    "In the example code below we use the [Euler method](https://en.wikipedia.org/wiki/Euler_method) to solve an ordinary differential equation (ODE) which is defined by the model's flows.  Euler's method is easy to reason about, but inappropriate for most problems; summer2 uses an adaptive [Runge-Kutta method](https://en.wikipedia.org/wiki/Runge-Kutta_methods) as its default solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "TIMESTEP = 0.1\n",
    "START_TIME = 0\n",
    "END_TIME = 20\n",
    "\n",
    "# Get times\n",
    "time_period = END_TIME - START_TIME + 1\n",
    "num_steps = time_period / TIMESTEP\n",
    "times = np.linspace(START_TIME, END_TIME, num=int(num_steps))\n",
    "\n",
    "# Define initial conditions\n",
    "initial_conditions = np.array([990.0, 10.0, 0.0])  # S, I, R\n",
    "\n",
    "# Define outputs\n",
    "outputs = np.zeros((int(num_steps), 3))\n",
    "outputs[0] = initial_conditions\n",
    "\n",
    "# Model parameters\n",
    "contact_rate = 1.0\n",
    "sojourn_time = 3.0\n",
    "death_rate = 0.05\n",
    "\n",
    "# Calculate outputs for each timestep\n",
    "for t_idx, t in enumerate(times):\n",
    "    if t_idx == 0:\n",
    "        continue\n",
    "\n",
    "    flow_rates = np.zeros(3)\n",
    "    compartment_sizes = outputs[t_idx - 1 ]\n",
    "\n",
    "    # Susceptible people can get infected (frequency-dependent).\n",
    "    num_sus = compartment_sizes[0]\n",
    "    num_inf = compartment_sizes[1]\n",
    "    num_pop = compartment_sizes.sum()\n",
    "    force_of_infection = contact_rate * num_inf / num_pop\n",
    "    infection_flow_rate = force_of_infection * num_sus\n",
    "    flow_rates[0] -= infection_flow_rate\n",
    "    flow_rates[1] += infection_flow_rate\n",
    "\n",
    "    # Infectious take some time to recover.\n",
    "    num_inf = compartment_sizes[1]\n",
    "    recovery_flow_rate = num_inf / sojourn_time\n",
    "    flow_rates[1] -= recovery_flow_rate\n",
    "    flow_rates[2] += recovery_flow_rate\n",
    "    \n",
    "    # Add an infection-specific death flow to the I compartment.\n",
    "    num_inf = compartment_sizes[1]\n",
    "    recovery_flow_rate = num_inf * death_rate\n",
    "    flow_rates[1] -= recovery_flow_rate\n",
    "    \n",
    "    # Calculate compartment sizes at next timestep given flowrates.\n",
    "    outputs[t_idx] = compartment_sizes + flow_rates * TIMESTEP  \n",
    "    \n",
    "# Plot the results as a function of time for S, I, R respectively.\n",
    "pd.DataFrame(outputs,columns=[\"S\",\"I\",\"R\"]).plot()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('jaxtest')",
   "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.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "79d70c6600ff0f236e7435c5ae4904b6925bd23ac9f566aeb17edb414d890ad2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}