docs/examples/08-strain-stratification.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "obvious-radio",
   "metadata": {},
   "source": [
    "# Strain stratification\n",
    "\n",
    "Summer provides a special disease-strain stratification class ([StrainStratification](../api/stratification.rst#summer2.stratification.StrainStratification)) that treats each stratum as a separate strain of the infection. This allows you to model multiple strains of an infection (or multiple infectious diseases with similar compartmental structures) which are able to infect people separately.\n",
    "\n",
    "They key difference between StrainStratification vs. a normal Stratification is that the strain stratification adjusts the count of infectious people per strain. For example, in a normal stratification with an age-based split into 'young' and 'old' you will have frequency-dependent infection flows calculated as follows:\n",
    "\n",
    "```python\n",
    "# Find a common force of infection for both young and old\n",
    "num_infected = num_old_infected + num_young_infected\n",
    "force_of_infection = contact_rate * num_infected / num_pop\n",
    "\n",
    "# Use that common force of infection to get flow rates for young/old infection \n",
    "infect_rate_young = force_of_infection * num_young_susceptible\n",
    "infect_rate_old = force_of_infection * num_old_susceptible\n",
    "```\n",
    "\n",
    "Consider now how this will be calculated for two strains (\"mild\" and \"wild\") when applied to an unstratified susceptible compartment:\n",
    "\n",
    "```python\n",
    "# Find a different force of infection for mild and wild\n",
    "force_of_infection_mild = contact_rate * num_infected_mild / num_pop\n",
    "force_of_infection_wild = contact_rate * num_infected_wild / num_pop\n",
    "\n",
    "# Use the different force of infection values to get flow rates for mild/wild infection \n",
    "infect_rate_mild = force_of_infection_mild * num_susceptible\n",
    "infect_rate_wild = force_of_infection_wild * num_susceptible\n",
    "```\n",
    "\n",
    "Let's work through a code example. For starters, let's create an SIR model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "british-puzzle",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import CompartmentalModel\n",
    "import pandas as pd\n",
    "\n",
    "pd.options.plotting.backend = \"plotly\"\n",
    "\n",
    "def build_model():\n",
    "    \"\"\"Returns a model for the stratification examples\"\"\"\n",
    "    model = CompartmentalModel(\n",
    "        times=[1990, 2010],\n",
    "        compartments=[\"S\", \"I\", \"R\"],\n",
    "        infectious_compartments=[\"I\"],\n",
    "        timestep=0.1,\n",
    "    )\n",
    "\n",
    "    # Add people to the model\n",
    "    model.set_initial_population(distribution={\"S\": 990, \"I\": 10})\n",
    "\n",
    "    # Susceptible people can get infected.\n",
    "    model.add_infection_frequency_flow(name=\"infection\", contact_rate=2, source=\"S\", dest=\"I\")\n",
    "\n",
    "    # Infectious people take 3 years, on average, to recover.\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",
    "    model.add_death_flow(name=\"infection_death\", death_rate=0.05, source=\"I\")\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "regulated-blend",
   "metadata": {},
   "source": [
    "Lets see what this model looks like without any stratifications:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fluid-gnome",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "major-detective",
   "metadata": {},
   "source": [
    "Now we can add a strain stratification to the infected (I) and recovered (R) compartments. We will assume immunity to one strain gives you immunity to the other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "changing-defendant",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import StrainStratification\n",
    "\n",
    "strata = ['mild', 'wild']\n",
    "strat = StrainStratification(name=\"strain\", strata=strata, compartments=['I', 'R'])\n",
    "\n",
    "# At the start of the simulation, 20% of infected people have wild strain.\n",
    "strat.set_population_split({'mild': 0.8, 'wild': 0.2})\n",
    "\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "violent-infection",
   "metadata": {},
   "source": [
    "Note that despite the stratification, the model results are the same in aggregate, because we have not applied any adjustments to the flows or strata infectiousness yet. Let's do that: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "informational-ratio",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import StrainStratification, Multiply\n",
    "\n",
    "strata = ['mild', 'wild']\n",
    "strat = StrainStratification(name=\"strain\", strata=strata, compartments=['I', 'R'])\n",
    "\n",
    "# Again, 20% of infected people have wild strain at the start.\n",
    "strat.set_population_split({'mild': 0.8, 'wild': 0.2})\n",
    "\n",
    "# The wild strain kills at 1.2x the rate as the mild strain does.\n",
    "strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"mild\": None,  \n",
    "    \"wild\": Multiply(1.2),\n",
    "})\n",
    "\n",
    "\n",
    "# Wild strain is twice as infectious than the mild strain (or equivalently, people are twice as susceptible to it).\n",
    "strat.set_flow_adjustments(\"infection\", {\n",
    "    \"mild\": None,  \n",
    "    \"wild\": Multiply(2),\n",
    "})\n",
    "\n",
    "\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "potential-cache",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "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": 5
}