docs/examples/06-stratification-introduction.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "juvenile-dakota",
   "metadata": {},
   "source": [
    "# Stratification introduction\n",
    "\n",
    "So far we've looked at how to create a compartmental model, add flows, request derived outputs and use different solvers. Now we'll look into stratifying a model using the [Stratification](../api/stratification.rst) class.\n",
    "\n",
    "Consider the following scenario, you are modelling some infectious disease that has different susceptibility and infection mortality across different age groups. We can use stratifications to model this.\n",
    "\n",
    "In this example we'll cover:\n",
    "\n",
    "- [No stratification](#No-stratification)\n",
    "- [Minimal stratification](#Minimal-stratification)\n",
    "- [Population distribution](#Population-distribution)\n",
    "- [Flow adjustments](#Flow-adjustments)\n",
    "- [Infectiousness adjustments](#Infectiousness-adjustments)\n",
    "- [Partial stratifications](#Partial-stratifications)\n",
    "- [Multiple stratifications](#Multiple-stratifications)\n",
    "- [Multiple interdependent stratifications](#Multiple-interdependent-stratifications)\n",
    "\n",
    "First, let's import the summer library and create a new [CompartmentalModel](../api/model.rst) with S, I, R compartments plus infection, death and recovery flows. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decimal-mayor",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import CompartmentalModel\n",
    "import pandas as pd\n",
    "pd.options.plotting.backend = \"plotly\"\n",
    "\n",
    "def build_model():\n",
    "    \"\"\"Returns a model for the stratification examples\"\"\"\n",
    "    model = CompartmentalModel(\n",
    "        times=[0, 20],\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": "acknowledged-aside",
   "metadata": {},
   "source": [
    "## No stratification\n",
    "\n",
    "With no stratification, this is just a regular SIR model: there are 3 compartments where susceptible people get infected/infectious, some of them die, and some of them recover."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "municipal-archives",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and run model with no stratifications\n",
    "model = build_model()\n",
    "model.run()\n",
    "model.get_outputs_df().plot()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "natural-shoulder",
   "metadata": {},
   "source": [
    "## Minimal stratification\n",
    "\n",
    "Next, let's try a simple stratification where we split the population into 'young' (say, 0 to 39 years old) and 'old' (aged 40 and above). Notice the following changes to the model outputs:\n",
    "\n",
    "- There are now 6 compartments instead of 3: each original compartment has been split into an \"old\" and \"young\" compartment, with the original population evenly divided between them (by default)\n",
    "- The model dynamics haven't changed otherwise: we will get the same results as before if we add the old and young compartments back together. This is because there is homogeneous mixing between strata and no demographic processes, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "temporal-variation",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments, which\n",
    "# splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62077e4d",
   "metadata": {},
   "source": [
    "## Importation flows and stratification\n",
    "\n",
    "We build a model with the same simple stratifaction as above; 2 age compartments, \"young\" and \"old\", but then add an importation flow to the model.  Note the following important details:\n",
    "\n",
    "- In addition to the existing (transition) infections, there are new importatation infections for both young and old, each at half the total rate specified.  This is because split_imports is set to True, and therefore evenly divides its total amongst the target compartments.  The increase in infections compared to the previous run is consistent with this.\n",
    "- The importation flow is added to the model directly, but only _after_ the Stratification has been applied.  This is because split_imports uses the model state at the time it is called in order to determine its splitting values.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf6f86f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments, which\n",
    "# splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "\n",
    "# Stratify the model first\n",
    "model.stratify_with(strat)\n",
    "\n",
    "# Now the following call is aware of the changes made by the Stratification\n",
    "model.add_importation_flow(\"infection_imports\", 10, \"I\", split_imports=True)\n",
    "\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "approximate-organization",
   "metadata": {},
   "source": [
    "## Population distribution\n",
    "\n",
    "We may not always wish to split the population evenly between strata. For example, we might know that 60% of the population is 'young' while 40% is 'old'. Notice that\n",
    "\n",
    "- The stratified compartments are now split according to a 60:40 ratio into young and old respectively\n",
    "- The overall model dynamics still haven't changed otherwise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sitting-capitol",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments, which\n",
    "# splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Set a population distribution\n",
    "strat.set_population_split({\"young\": 0.6, \"old\": 0.4})\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spanish-screen",
   "metadata": {},
   "source": [
    "## Flow adjustments\n",
    "\n",
    "As noted so far, we've been successful in subdividing the population, but haven't actually changed our model dynamics, which is kind of boring. Next let's look at how we can adjust the flow rates based on strata. Let's assume two new facts about our disease:\n",
    "\n",
    "- young people are twice as susceptible to infection\n",
    "- old people are three times as likely to die from the infectious disease, while younger people are half as likely as under the original parameters we requested\n",
    "- younger people take twice as long to recover\n",
    "\n",
    "These inter-strata differences can be modelled using flow adjustments. Now we're seeing some genuinely new model dynamics. Note how there are fewer recovered 'old' people at the end of the model run, because of their higher death rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seeing-referral",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments, which\n",
    "# splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Set a population distribution\n",
    "strat.set_population_split({\"young\": 0.6, \"old\": 0.4})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow\n",
    "strat.set_flow_adjustments(\"infection\", {\n",
    "    \"old\": None,  # No adjustment for old people, use baseline requested value\n",
    "    \"young\": Multiply(2),  # Young people are twice as susceptible to infection\n",
    "})\n",
    "\n",
    "# Add an adjustment to the 'infection_death' flow\n",
    "strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"old\": Multiply(3),  # Older people die at three times the rate requested under the original parameters\n",
    "    \"young\": Multiply(0.5), # Younger people die at half the rate requested under the original parameters\n",
    "})\n",
    "\n",
    "# Add an adjustment to the 'recovery' flow\n",
    "strat.set_flow_adjustments(\"recovery\", {\n",
    "    \"old\": None,  # No adjustment for old people, use baseline\n",
    "    \"young\": Multiply(0.5),  # Young people take twice as long to recover\n",
    "})\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "driven-acquisition",
   "metadata": {},
   "source": [
    "## Infectiousness adjustments\n",
    "\n",
    "In addition to adjusting flow rates for each strata, you can also set an infectiousness level for a given strata. This affects how likely an infectious person in that stratum is to infect someone else. For example we could consider the following:\n",
    "\n",
    "- young people are 1.2 times as infectious, because they're not wearing face masks as much\n",
    "- young people are twice as susceptible to the disease, because some of them have immature immune systems\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ultimate-notice",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply, Overwrite\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments, which\n",
    "# splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Set a population distribution\n",
    "strat.set_population_split({\"young\": 0.5, \"old\": 0.5})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow\n",
    "strat.set_flow_adjustments(\"infection\", {\n",
    "    \"old\": None,  # No adjustment for old people, use baseline\n",
    "    \"young\": Multiply(2),  # Young people twice as susceptible\n",
    "})\n",
    "\n",
    "# Add an adjustment to infectiousness levels for young people in the 'I' compartment\n",
    "strat.add_infectiousness_adjustments(\"I\", {\n",
    "    \"old\": None,  # No adjustment for old people, use baseline\n",
    "    \"young\": Multiply(1.2),  # Young people 1.2x more infectious\n",
    "})\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "solved-campus",
   "metadata": {},
   "source": [
    "## Partial stratifications\n",
    "\n",
    "So far we've been stratifying all compartments, but Summer allows only some of the compartments to be stratified. For example, we can stratify only the infectious compartment to model three different levels of disease severity: asymptomatic, mild and severe.\n",
    "\n",
    "When you do a partial stratification, flow rates into that stratified compartment will automatically be adjusted with an even split to conserve the behaviour by default, e.g. a flow rate of 3 from a source will be evenly split into (1, 1, 1) across the three destinations. This behaviour can be manually overriden with a flow adjustment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "juvenile-reduction",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply, Overwrite\n",
    "\n",
    "# Create a stratification named 'severity', applying to the infectious, which\n",
    "# splits that compartment into 'asymptomatic', 'mild' and 'severe'.\n",
    "strata = [\"asymptomatic\", \"mild\", \"severe\"]\n",
    "strat = Stratification(name=\"severity\", strata=strata, compartments=['I'])\n",
    "\n",
    "# Set a population distribution - everyone starts out asymptomatic.\n",
    "strat.set_population_split({\"asymptomatic\": 1.0, \"mild\": 0, \"severe\": 0})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow, overriding default split.\n",
    "strat.set_flow_adjustments(\"infection\", {\n",
    "    \"asymptomatic\": Multiply(0.3),  # 30% of incident cases are asymptomatic\n",
    "    \"mild\": Multiply(0.5),  # 50% of incident cases are mild\n",
    "    \"severe\": Multiply(0.2),  # 20% of incident cases are severe\n",
    "})\n",
    "\n",
    "# Add an adjustment to the 'infection_death' flow\n",
    "strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "strat.add_infectiousness_adjustments(\"I\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "# Build and run model with the stratification we just defined\n",
    "model = build_model()\n",
    "model.stratify_with(strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "external-imperial",
   "metadata": {},
   "source": [
    "## Multiple stratifications\n",
    "\n",
    "A model can have multiple stratifications applied in series. For example, we can add an 'age' stratification, followed by a 'severity' one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "practical-blink",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply, Overwrite\n",
    "\n",
    "### Age stratification\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments,\n",
    "# which splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "age_strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "age_strat.set_population_split({\"young\": 0.6, \"old\": 0.4})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow\n",
    "age_strat.set_flow_adjustments(\"infection\", {\n",
    "    \"old\": None,  # No adjustment for old people, use unstratified parameter value\n",
    "    \"young\": Multiply(2),  # Young people are twice as susceptible\n",
    "})\n",
    "\n",
    "# Add an adjustment to infectiousness levels for young people the 'I' compartment\n",
    "age_strat.add_infectiousness_adjustments(\"I\", {\n",
    "    \"old\": None,  # No adjustment for old people, use unstratified parameter value\n",
    "    \"young\": Multiply(1.2),  # Young people are 5x more infectious\n",
    "})\n",
    "\n",
    "\n",
    "### Disease severity stratification\n",
    "\n",
    "# Create a stratification named 'severity', applying to the infectious compartment, which\n",
    "# splits that compartment into 'asymptomatic', 'mild' and 'severe'.\n",
    "strata = [\"asymptomatic\", \"mild\", \"severe\"]\n",
    "severity_strat = Stratification(name=\"severity\", strata=strata, compartments=[\"I\"])\n",
    "severity_strat.set_population_split({\"asymptomatic\": 1.0, \"mild\": 0, \"severe\": 0})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow (overriding the default split of one third to each stratum)\n",
    "severity_strat.set_flow_adjustments(\"infection\", {\n",
    "    \"asymptomatic\": Multiply(0.3),  # 30% of cases are asympt.\n",
    "    \"mild\": Multiply(0.5),  # 50% of cases are mild.\n",
    "    \"severe\": Multiply(0.2),  # 20% of cases are severse.\n",
    "})\n",
    "\n",
    "# Add an adjustment to the 'infection_death' flow\n",
    "severity_strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "severity_strat.add_infectiousness_adjustments(\"I\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "\n",
    "# Build and run model with the stratifications we just defined\n",
    "model = build_model()\n",
    "# Apply age, then severity stratifications \n",
    "model.stratify_with(age_strat)\n",
    "model.stratify_with(severity_strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "honest-instrument",
   "metadata": {},
   "source": [
    "## Multiple interdependent stratifications\n",
    "\n",
    "In the previous example we assumed that the age and severity stratifications were independent. For example, we assumed that the proportion of infected people who have a disease severity of asymptomatic, mild and severe is the same for both young and old people. Perhaps, for a given disease, this is not true! it's easy to imagine an infection for which younger people tend towards being more asymptomatic, and older people tend towards having a more severe infection.\n",
    "\n",
    "This interdependency between stratifications can be modelled using Summer, where a flow adjustment for a stratification can selectively refer to strata used for previous stratifications. You can refer to the API reference for [set_flow_adjustments](../api/stratification.rst#summer2.stratification.Stratification.set_flow_adjustments) for more details.\n",
    "\n",
    "To clarify, let's consider the example described above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "contained-discussion",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply, Overwrite\n",
    "\n",
    "### Age stratification\n",
    "\n",
    "# Create a stratification named 'age', applying to all compartments,\n",
    "# which splits the population into 'young' and 'old'.\n",
    "strata = [\"young\", \"old\"]\n",
    "age_strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "age_strat.set_population_split({\"young\": 0.6, \"old\": 0.4})\n",
    "\n",
    "### Disease severity stratification (depends on the age stratification)\n",
    "\n",
    "# Create a stratification named 'severity', applying to the infectious, which\n",
    "# splits that compartment into 'asymptomatic', 'mild' and 'severe'.\n",
    "strata = [\"asymptomatic\", \"mild\", \"severe\"]\n",
    "severity_strat = Stratification(name=\"severity\", strata=strata, compartments=[\"I\"])\n",
    "severity_strat.set_population_split({\"asymptomatic\": 1.0, \"mild\": 0, \"severe\": 0})\n",
    "\n",
    "# Add an adjustment to the 'infection' flow for young people\n",
    "# where younger people tend towards asymptomatic infection\n",
    "young_infection_adjustments = {\n",
    "    \"asymptomatic\": Multiply(0.5),  # 50% of cases are asympt.\n",
    "    \"mild\": Multiply(0.4),  # 40% of cases are mild.\n",
    "    \"severe\": Multiply(0.1),  # 10% of cases are severe.\n",
    "}\n",
    "severity_strat.set_flow_adjustments(\n",
    "    \"infection\",\n",
    "    young_infection_adjustments,\n",
    "    source_strata={'age': 'young'}  # Only apply this adjustment to flows of young people\n",
    ")\n",
    "\n",
    "# Add an adjustment to the 'infection' flow for old people\n",
    "# where older people tend towards severe infection\n",
    "old_infection_adjustments = {\n",
    "    \"asymptomatic\": Multiply(0.1),  # 10% of cases are asympt.\n",
    "    \"mild\": Multiply(0.4),  # 40% of cases are mild.\n",
    "    \"severe\": Multiply(0.5),  # 50% of cases are severe.\n",
    "}\n",
    "severity_strat.set_flow_adjustments(\n",
    "    \"infection\",\n",
    "    old_infection_adjustments,\n",
    "    source_strata={'age': 'old'}   # Only apply this adjustment to flows of old people\n",
    ")\n",
    "\n",
    "# Add an adjustment to the 'infection_death' flow (for all age groups)\n",
    "severity_strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "# Adjust infectiousness levels (for all age groups)\n",
    "severity_strat.add_infectiousness_adjustments(\"I\", {\n",
    "    \"asymptomatic\": Multiply(0.5),\n",
    "    \"mild\": None,\n",
    "    \"severe\": Multiply(1.5),\n",
    "})\n",
    "\n",
    "\n",
    "# Build and run model with the stratifications we just defined\n",
    "model = build_model()\n",
    "# Apply age, then severity stratifications \n",
    "model.stratify_with(age_strat)\n",
    "model.stratify_with(severity_strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "impressed-jamaica",
   "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
}