docs/examples/09-mixing-matrices.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "inner-diploma",
   "metadata": {},
   "source": [
    "# Mixing matrices\n",
    "\n",
    "By default the Summer compartmental model assumes that each person in the model comes into contact with every other person at the same rate (homogeneous mixing). This isn't always true. It may be the case, for example, that children tend to interact more with other children, and less with the elderly (with-like or assortative mixing). This difference in social mixing can be expressed and modelled using a \"mixing matrix\". A NxN matrix which defines how each N strata of a stratification interact and may infect each other.\n",
    "\n",
    "\n",
    "For example, in a 'child'/'adult' stratification, we could have the following mixing matrix (arbitrary numbers):\n",
    "\n",
    "| | child | adult |\n",
    "| --- | --- | --- |\n",
    "| child | 0.2 | 0.3 |\n",
    "| adult | 0.5 | 0.7 |\n",
    "\n",
    "\n",
    "In this mixing matrix, the columns are the infectors and the rows the infected. So the above matrix represents the following infector -> infected relationships:\n",
    "\n",
    "| | child | adult |\n",
    "| --- | --- | --- |\n",
    "| **child** | child -> child | adult -> child |\n",
    "| **adult** | child -> adult | adult -> adult |\n",
    "\n",
    "\n",
    "This worked example may clarify. We will calculate the frequency-dependent infection rates for adults and children using the above mixing matrix. Assume the following scenario:\n",
    "\n",
    "- 1000 people, 10 infected, 990 susceptible\n",
    "- 80% of the population is adults, 20% are children\n",
    "- So, 800 adults with 792 susceptible and 8 infected; and\n",
    "- 200 children with 198 susceptible and 2 infected\n",
    "\n",
    "\n",
    "```python\n",
    "child_force_of_inf = 0.2 * 2 / 200 + 0.3 * 8 / 800\n",
    "adult_force_of_inf = 0.5 * 2 / 200 + 0.7 * 8 / 800\n",
    "\n",
    "# Here we assume the contact rate accounts for the use of a mixing matrix .\n",
    "child_inf_rate = contact_rate * child_force_of_inf * 198\n",
    "adult_inf_rate = contact_rate * adult_force_of_inf * 792\n",
    "```\n",
    "\n",
    "## Using a mixing matrix\n",
    "\n",
    "Let's start by defining some code to build a model and plot the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "current-banking",
   "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 mixing matrix examples\"\"\"\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=2, 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",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "innocent-mortgage",
   "metadata": {},
   "source": [
    "For starters let's see what our 'vanilla' SIR model looks like when not stratified at all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seeing-diversity",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build and run model without any stratification\n",
    "model = build_model()\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "amino-frederick",
   "metadata": {},
   "source": [
    "A mixing matrix can be added to any \"full\" stratificiation: which is one that affects all compartments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chinese-cookie",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "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",
    "# Implicitly there is a 50-50 split between young and old.\n",
    "strata = [\"young\", \"old\"]\n",
    "strat = Stratification(name=\"age\", strata=strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "\n",
    "# Define a NxN (in this case 2x2) mixing matrix.\n",
    "# The order of the rows/columns will be the same as \n",
    "# the order of the strata passed to the Stratification\n",
    "age_mixing_matrix = np.array([\n",
    "    [0.2, 0.3],\n",
    "    [0.5, 0.7],\n",
    "]) \n",
    "\n",
    "# Add the mixing matrix to the stratification\n",
    "strat.set_mixing_matrix(age_mixing_matrix)\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": "explicit-indian",
   "metadata": {},
   "source": [
    "## Time-varying mixing matrices\n",
    "\n",
    "You can specify a mixing matrix that varies over time, by using a function of time that returns a mixing matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caroline-latitude",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from summer2 import Stratification\n",
    "from summer2.parameters import Function\n",
    "from summer2.functions import get_piecewise_scalar_function\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",
    "# An age mixing matrix for 'normal' social mixing\n",
    "normal_mixing_matrix = np.array([\n",
    "    [0.2, 0.3],\n",
    "    [0.5, 0.7],\n",
    "]) \n",
    "\n",
    "# An age mixing matrix for 'age lockdown': \n",
    "#  - no interaction between young and old\n",
    "#  - increased interaction inside each group\n",
    "lockdown_mixing_matrix = np.array([\n",
    "    [0.3, 0.0],\n",
    "    [0.0, 0.8],\n",
    "]) \n",
    "\n",
    "start_lockdown = 3  # Day 3\n",
    "end_lockdown = 8  # Day 8\n",
    "\n",
    "# Here we build a Function that returns the appropropriate matrix based on the current timestep\n",
    "# The same machinery could be used to have 'variable strength' lockdowns, or other changes over time\n",
    "\n",
    "def mix(lvalue, rvalue, balance):\n",
    "    #Return a linear combination of the 2 arrays lvalue and rvalue\n",
    "    return (1.0 - balance)*lvalue + balance*rvalue\n",
    "\n",
    "# This could easily be extended to multiple lockdowns using exactly the same functions\n",
    "is_lockdown = get_piecewise_scalar_function([start_lockdown,end_lockdown],[0.0,1.0,0.0])\n",
    "final_mm = Function(mix, [normal_mixing_matrix, lockdown_mixing_matrix, is_lockdown])\n",
    "\n",
    "# Add the mixing matrix to the stratification\n",
    "strat.set_mixing_matrix(final_mm)\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": "stunning-camera",
   "metadata": {},
   "source": [
    "## Multiple mixing matrices for multiple stratifications\n",
    "\n",
    "You can set mixing matrices for multiple stratifications. You also have the option not to. A key assumption is that these two types of mixing are independent. For example, below we assume that there is heterogeneous mixing between young and old agegroup as well as mixing between urban and rural residents, but the model will assume that the way that young and old people mix in rural areas is the same as the way that young and old people mix in urban areas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mineral-section",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, Multiply\n",
    "from summer2.parameters import Data\n",
    "\n",
    "# Age stratification with young/old mixing\n",
    "age_strata = [\"young\", \"old\"]\n",
    "age_strat = Stratification(name=\"age\", strata=age_strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "age_mixing_matrix = Data(np.array([\n",
    "    [0.2, 0.3],\n",
    "    [0.5, 0.7],\n",
    "]) )\n",
    "age_mixing_matrix.node_name = \"age_mixing\"\n",
    "age_strat.set_mixing_matrix(age_mixing_matrix)\n",
    "\n",
    "\n",
    "# Location stratification with urban/rural mixing\n",
    "loc_strata = [\"urban\", \"rural\"]\n",
    "loc_strat = Stratification(name=\"location\", strata=loc_strata, compartments=[\"S\", \"I\", \"R\"])\n",
    "# Rural people have worse health care, higher mortality rates,\n",
    "loc_strat.set_flow_adjustments(\"infection_death\", {\n",
    "    \"urban\": None,\n",
    "    \"rural\": Multiply(3),\n",
    "})\n",
    "loc_strat.set_population_split({\"urban\": 0.7, \"rural\": 0.3})\n",
    "loc_mixing_matrix = Data(np.array([\n",
    "    [0.8, 0.2],\n",
    "    [0.2, 0.8],\n",
    "]))\n",
    "loc_mixing_matrix.node_name = \"loc_mixing\"\n",
    "loc_strat.set_mixing_matrix(loc_mixing_matrix)\n",
    "\n",
    "\n",
    "# Build and run model with the stratifications we just defined\n",
    "model = build_model()\n",
    "# Apply age, then location stratifications \n",
    "model.stratify_with(age_strat)\n",
    "model.stratify_with(loc_strat)\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "central-variable",
   "metadata": {},
   "source": [
    "The combined age and location matrix is the Kronecker product of the age and location matrices. We can visualize it here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa1ad98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query the model's graph so we are just looking at the relevant nodes\n",
    "mixing_graph = model.graph.filter(targets=\"mixing_matrix\")\n",
    "mixing_graph.draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40afffa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a callable function for this graph, that returns all its nodes values\n",
    "mixing_data = mixing_graph.get_callable(output_all=True)()\n",
    "mixing_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "strong-northwest",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import itertools\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(15, 5), dpi=100)\n",
    "combined_strata = [\n",
    "    f\"{age[0]} + {loc[:3]}\" for age, loc \n",
    "    in itertools.product(age_strata, loc_strata)\n",
    "]\n",
    "\n",
    "plots = [\n",
    "    ['Age based mixing', age_strata, mixing_data[\"age_mixing\"]],\n",
    "    ['Location based mixing', loc_strata, mixing_data[\"loc_mixing\"]],\n",
    "    ['Age and location based mixing', combined_strata, mixing_data[\"mixing_matrix\"]],\n",
    "\n",
    "]\n",
    "for ax, (title, strata, matrix) in zip(axes, plots):\n",
    "    ax.set_title(title)\n",
    "    ax.set_xticks(np.arange(len(strata)))\n",
    "    ax.set_yticks(np.arange(len(strata)))\n",
    "    ax.set_xticklabels(strata)\n",
    "    ax.set_yticklabels(strata)\n",
    "    ax.imshow(matrix, cmap=\"Greys\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "boolean-forge",
   "metadata": {},
   "source": [
    "## Prem et al. mixing matrices based on age and location\n",
    "\n",
    "You can obtain estimated mixing matrices from [Projecting social contact matrices in 152 countries using contact surveys and demographic data](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1005697#sec020) by Prem et al. in PLOS Computational Biology in 2017.\n",
    "\n",
    "This paper is accompanied by age and location specific mixing matrices for 152 countries. You can download the matrices as Excel spreadsheets [here](https://doi.org/10.1371/journal.pcbi.1005697.s002). The paper provides mixing matrices for 5 location types:\n",
    "\n",
    "- home\n",
    "- school\n",
    "- work\n",
    "- other_locations\n",
    "- all_locations\n",
    "\n",
    "The rows and columns indices of each matrix represent a 5 year age bracket from 0-80, giving us a 16x16 matrix. \n",
    "\n",
    "A more recent version of these social mixing matrices can be obtained from [Kiesha Prem's GitHub](https://github.com/kieshaprem/synthetic-contact-matrices)"
   ]
  }
 ],
 "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
}