docs/examples/07-age-stratification.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "regulated-ireland",
   "metadata": {},
   "source": [
    "# Age stratification\n",
    "\n",
    "Summer provides a special age stratification class ([AgeStratification](../api/stratification.rst#summer2.stratification.AgeStratification)) that automatically adds ageing flows to the model. See the linked documentation for more details on the flows.\n",
    "\n",
    "`AgeStratification` is used similarly to the more general-purpose `Stratification` class, except it makes some assumptions of the meaning of the strata in order to implement ageing flows.\n",
    "\n",
    "**Important caveat**: The AgeStratification assumes that your time unit is years.\n",
    "\n",
    "Let's start by creating a vanilla SIR model. You can see that there are no flows by default:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "changing-insight",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import CompartmentalModel, AgeStratification\n",
    "import pandas as pd\n",
    "pd.options.plotting.backend = \"plotly\"\n",
    "\n",
    "model = CompartmentalModel(\n",
    "    times=[1990, 2010],\n",
    "    compartments=[\"S\", \"I\", \"R\"],\n",
    "    infectious_compartments=[\"I\"],\n",
    ")\n",
    "model.set_initial_population(distribution={\"S\": 900, \"I\": 100, \"R\": 10})\n",
    "\n",
    "model.flows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "extended-parks",
   "metadata": {},
   "source": [
    "Next, lets add an age stratification with the following age strata:\n",
    "\n",
    "- 0: 0-19 year olds\n",
    "- 20: 20-39 year olds\n",
    "- 40: 40-59 year olds\n",
    "- 60: ages 60+\n",
    "\n",
    "Now you can see that 9 ageing flows have been added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "appointed-network",
   "metadata": {},
   "outputs": [],
   "source": [
    "strata = [0, 20, 40, 60]\n",
    "strat = AgeStratification(name=\"age\", strata=strata, compartments=['S', 'I', 'R'])\n",
    "strat.set_population_split({'0': 0.4, '20': 0.3, '40': 0.2, '60': 0.1})\n",
    "model.stratify_with(strat)\n",
    "model.flows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "outstanding-wisconsin",
   "metadata": {},
   "source": [
    "The rate of ageing is the reciprocal of the width of the age bracket: that is - people in the 0-19 age bracket will leave the bracket at a rate of 5% per year. People in the oldest age bracket will not leave their age bracket through an ageing process (although other flows may cause them to leave the model or transition between compartments within their bracket)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "subjective-average",
   "metadata": {},
   "source": [
    "Let's visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ordered-florida",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.run()\n",
    "outputs = model.get_outputs_df()\n",
    "outputs.plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "simple-husband",
   "metadata": {},
   "source": [
    "As you can see above, these ageing flows add new dynamics to the model. They are easier to visualise if we consider the age strata for a single compartment type:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f96b5a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs[model.query_compartments({\"name\": \"S\"})].plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "meaning-vector",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Now you know how to use an AgeStratification to automatically generate ageing flows. See the [API reference](../api/stratification.rst#summer2.stratification.AgeStratification) for more details."
   ]
  }
 ],
 "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
}