docs/examples/10-derived-outputs-stratified.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "norman-accused",
   "metadata": {},
   "source": [
    "# Derived outputs in stratified models\n",
    "\n",
    "In a [previous example](03-derived-outputs.ipynb), we saw how to track derived outputs for a simple, unstratified model. In this example we'll look into tracking derived outputs for a stratified model.\n",
    "\n",
    "We'll be looking at:\n",
    "\n",
    "- [Tracking cumulative deaths by age group](#Tracking-cumulative-deaths-by-age-group)\n",
    "- [Tracking incidence for different strata](#Tracking-disease-incidence)\n",
    "- [Tracking daily notifications](#Tracking-daily-notifications)\n",
    "- [Tracking hospital occupancy](#Tracking-hospital-occupancy)\n",
    "\n",
    "Let's start by defining a reasonably complicated compartmental SEIR model stratified by age and clinical status. This model will have:\n",
    "\n",
    "An age stratification with strata:\n",
    "\n",
    "- young\n",
    "- old\n",
    "\n",
    "A \"clinical status\"\" stratification with strata:\n",
    "\n",
    "- asymptomatic: has disease, doesn't know it, no symptoms\n",
    "- symptomatic: has disease, doesn't know it (or not tested), symptoms\n",
    "- isolated: has disease, knows it (tested), is isolated at home, symptoms\n",
    "- hospital: has disease, knows it (tested), severse symptoms and in hospital\n",
    "\n",
    "We'll use this model as a basis for defining our derived outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naval-chick",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import CompartmentalModel, Stratification, Multiply, Overwrite\n",
    "import pandas as pd\n",
    "pd.options.plotting.backend = \"plotly\"\n",
    "\n",
    "def build_model():\n",
    "    \"\"\"Returns a new SIR model\"\"\"\n",
    "    # Create basic SEIR model.\n",
    "    model = CompartmentalModel(\n",
    "        times=[0, 20],\n",
    "        compartments=[\"S\", \"E\", \"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=\"E\")\n",
    "    model.add_transition_flow(name=\"incidence\", fractional_rate=0.5, source=\"E\", dest=\"I\")\n",
    "    model.add_transition_flow(name=\"recovery\", fractional_rate=0.3, source=\"I\", dest=\"R\")\n",
    "    model.add_death_flow(name=\"infection_death\", death_rate=0.05, source=\"I\")\n",
    "\n",
    "    # Stratify by age.\n",
    "    age_strat = Stratification('age', ['young', 'old'], ['S', 'E', 'I', 'R'])\n",
    "    age_strat.set_population_split({'young': 0.5, 'old': 0.5})\n",
    "    model.stratify_with(age_strat)\n",
    "    \n",
    "    # Stratify by clinical status for infected people.\n",
    "    strata = ['asymptomatic', 'symptomatic', 'isolated', 'hospital']\n",
    "    clinical_strat = Stratification('clinical', strata, ['I'])\n",
    "    clinical_strat.set_population_split({'asymptomatic': 1, 'symptomatic': 0, 'isolated': 0, 'hospital': 0})\n",
    "    \n",
    "    # Half of young people become asymptomatic.\n",
    "    young_incidence_adjustments = {\n",
    "        \"asymptomatic\": Multiply(0.4),\n",
    "        \"symptomatic\": Multiply(0.3),\n",
    "        \"isolated\": Multiply(0.2),\n",
    "        \"hospital\": Multiply(0.1),\n",
    "    }\n",
    "    clinical_strat.set_flow_adjustments(\n",
    "        \"incidence\",\n",
    "        young_incidence_adjustments,\n",
    "        source_strata={'age': 'young'}\n",
    "    )\n",
    "\n",
    "    # A higher proporiton of old people go to hospital\n",
    "    old_incidence_adjustments = {\n",
    "        \"asymptomatic\": Multiply(0.3),\n",
    "        \"symptomatic\": Multiply(0.3),\n",
    "        \"isolated\": Multiply(0.2),\n",
    "        \"hospital\": Multiply(0.3),\n",
    "    }\n",
    "    clinical_strat.set_flow_adjustments(\n",
    "        \"incidence\",\n",
    "        old_incidence_adjustments,\n",
    "        source_strata={'age': 'old'}\n",
    "    )\n",
    "   \n",
    "    # Adjust risk of dying given clinical status\n",
    "    clinical_strat.set_flow_adjustments(\"infection_death\", {\n",
    "        \"asymptomatic\": Overwrite(0),  # Can't die if no symptoms\n",
    "        \"symptomatic\": None,\n",
    "        \"isolated\": None,\n",
    "        \"hospital\": Multiply(2),  # Severe cases go to hospital\n",
    "    })\n",
    "\n",
    "    # Adjust infectiousness given clinical status\n",
    "    clinical_strat.add_infectiousness_adjustments(\"I\", {\n",
    "        \"asymptomatic\": Multiply(0.5),\n",
    "        \"symptomatic\": None,\n",
    "        \"isolated\": Multiply(0.2),\n",
    "        \"hospital\": Multiply(0.1),\n",
    "    })    \n",
    "    model.stratify_with(clinical_strat)\n",
    "    \n",
    "    return model\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "local-addition",
   "metadata": {},
   "source": [
    "Now that we have a model that we can inspect, let's use the derived output requests from our [previous example](03-derived-outputs.ipynb) to calculate some quantities of interest.\n",
    "\n",
    "## Tracking cumulative deaths by age group\n",
    "\n",
    "We can use a [flow output](../api/model.rst#summer2.model.CompartmentalModel.request_output_for_flow) plus a [cumulative output](../api/model.rst#summer2.model.CompartmentalModel.request_cumulative_output) to track the cumulative number of people who have died from the disease, segmented by age group:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "surprised-porter",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "age_strata = [\"young\", \"old\"]\n",
    "for age_stratum in age_strata:\n",
    "    output_name = f\"death_{age_stratum}\"\n",
    "    model.request_output_for_flow(\n",
    "        output_name,\n",
    "        flow_name=\"infection_death\",\n",
    "        source_strata={\"age\": age_stratum},\n",
    "        save_results=False\n",
    "    )\n",
    "    model.request_cumulative_output(f\"accum_{output_name}\", output_name)\n",
    "    \n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "binding-provider",
   "metadata": {},
   "source": [
    "## Tracking disease incidence\n",
    "\n",
    "We can use a [flow output](../api/model.rst#summer2.model.CompartmentalModel.request_output_for_flow) to track the number of people who progress from exposed to infected per day ('incidence') and break it down by either age group, severity, or both.\n",
    "\n",
    "To start, let's look at incidence (daily flow from E to I), by age group:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "driven-newton",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "age_strata = [\"young\", \"old\"]\n",
    "for age_stratum in age_strata:\n",
    "    model.request_output_for_flow(\n",
    "        f\"incidence_{age_stratum}\",\n",
    "        flow_name=\"incidence\",\n",
    "        source_strata={\"age\": age_stratum}\n",
    "    )\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "continued-candidate",
   "metadata": {},
   "source": [
    "We can also inspect incidence by clinical status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "joined-prospect",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "clinical_strata = [\"asymptomatic\", \"symptomatic\", \"isolated\", \"hospital\"]\n",
    "output_names = [f\"incidence_{s}\" for s in clinical_strata]\n",
    "for clinical_stratum, output_name in zip(clinical_strata, output_names):\n",
    "    model.request_output_for_flow(\n",
    "        output_name,\n",
    "        flow_name=\"incidence\",\n",
    "        dest_strata={\"clinical\": clinical_stratum}\n",
    "    )\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "associate-printing",
   "metadata": {},
   "source": [
    "Finally, we can get a break down of incidence by both age AND clinical status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pending-right",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "model = build_model()\n",
    "\n",
    "age_strata = [\"young\", \"old\"]\n",
    "clinical_strata = [\"asymptomatic\", \"symptomatic\", \"isolated\", \"hospital\"]\n",
    "strata = list(product(age_strata, clinical_strata))\n",
    "output_names = [f\"incidence_{a}_{c}\" for a, c in strata]\n",
    "\n",
    "for (age, clinical), output_name in zip(strata, output_names):\n",
    "    model.request_output_for_flow(\n",
    "        output_name,\n",
    "        flow_name=\"incidence\",\n",
    "        dest_strata={\"clinical\": clinical, \"age\": age}\n",
    "    )\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "surprised-intro",
   "metadata": {},
   "source": [
    "## Tracking daily notifications\n",
    "\n",
    "We could use the same approach to track \"notifications\": disease incidence that is detected via testing. In our (simplified) clinical status strata, we could define notifications as:\n",
    "\n",
    "- asymptomatic: no notification\n",
    "- symptomatic: no notification\n",
    "- isolated: notification\n",
    "- hospital: notification\n",
    "\n",
    "Given this definition, notifications can be calculated as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "extra-showcase",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "# Step 1: Get model to track the incidence flows for the notification strata\n",
    "notify_strata = [\"isolated\", \"hospital\"]\n",
    "output_names = [f\"incidence_{s}\" for s in notify_strata]\n",
    "for notify_strata, output_name in zip(notify_strata, output_names):\n",
    "    model.request_output_for_flow(\n",
    "        output_name,\n",
    "        flow_name=\"incidence\",\n",
    "        dest_strata={\"clinical\": notify_strata}\n",
    "    )\n",
    "    \n",
    "# Step 2: Aggregate the notification strata\n",
    "model.request_aggregate_output(\n",
    "    name=\"notifications\",\n",
    "    sources=output_names,\n",
    ")\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "informative-moisture",
   "metadata": {},
   "source": [
    "## Tracking hospital occupancy\n",
    "\n",
    "We can use a [compartmental output](../api/model.rst#summer2.model.CompartmentalModel.request_output_for_compartments) to track the number of infected people who are in the hospital strata per timestep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rocky-injection",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "model.request_output_for_compartments(\n",
    "    'hospital_occupancy',\n",
    "    compartments=[\"I\"],\n",
    "    strata={\"clinical\": \"hospital\"}\n",
    ")\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9775d6df",
   "metadata": {},
   "source": [
    "## Using DerivedOutput GraphObjects\n",
    "\n",
    "Previously requested derived outputs are available as GraphObjects via the DerivedOutput constructor. <br>These can be operated on in exactly the same way as all other GraphObjects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d9f238d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2.parameters import DerivedOutput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b2dc6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "\n",
    "model.request_output_for_compartments(\n",
    "    'hospital_occupancy',\n",
    "    compartments=[\"I\"],\n",
    "    strata={\"clinical\": \"hospital\"},\n",
    "    save_results=False\n",
    ")\n",
    "\n",
    "model.request_output_for_compartments(\n",
    "    'all_infected',\n",
    "    compartments=[\"I\"],\n",
    "    save_results=False\n",
    ")\n",
    "\n",
    "model.request_function_output(\"prop_infected_hosp\", DerivedOutput(\"hospital_occupancy\") / DerivedOutput(\"all_infected\") )\n",
    "\n",
    "model.run()\n",
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "careful-gather",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Now you know how to request derived ouputs that target specific strata in a stratified model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "forty-inclusion",
   "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
}