docs/examples/11-flows-between-strata.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "spiritual-guinea",
   "metadata": {},
   "source": [
    "# Flows between strata\n",
    "\n",
    "In the [previous introduction to flows](04-flow-types.ipynb) and [introduction to stratification](06-stratification-introduction.ipynb) we saw a workflow where:\n",
    "\n",
    "- An unstratified model is defined\n",
    "- Flows are added to that model\n",
    "- The model is then stratified, splitting the flows between new strata\n",
    "\n",
    "This approach works fine for many workflows, but in some cases, we want to define flows that move people between strata. For example, we might want to model people migrating from a rural location to an urban location over time.\n",
    "\n",
    "This example will show you how to implement flows between strata. Let's start with a baseline model, stratified by location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "superb-celebration",
   "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=[1990, 2020],\n",
    "        compartments=[\"pop\"],\n",
    "        infectious_compartments=[],\n",
    "        timestep=0.1,\n",
    "    )\n",
    "    model.set_initial_population(distribution={\"pop\": 20e6})    \n",
    "    model.add_crude_birth_flow(\"birth\", 0.02, \"pop\")\n",
    "    model.add_death_flow(\"death\", 0.01, \"pop\")\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "atmospheric-funds",
   "metadata": {},
   "source": [
    "## Unstratified model\n",
    "\n",
    "In our example model, there is only one compartment with a birth and death rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "juvenile-wellington",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model()\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "embedded-exhaust",
   "metadata": {},
   "source": [
    "## Stratified model\n",
    "\n",
    "Next lets split the population into urban and rural."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collect-summer",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification\n",
    "\n",
    "model = build_model()\n",
    "\n",
    "strat = Stratification('location', ['urban', 'rural'], ['pop'])\n",
    "strat.set_population_split({'rural': 0.7, 'urban': 0.3})\n",
    "model.stratify_with(strat)\n",
    "\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "spread-lyric",
   "metadata": {},
   "source": [
    "Note that, by default, 50% of the births, which are based on total population, are born into the urban/rural stratum respectively. This isn't physically realistic but we'll ignore it for simplicity's sake."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "based-scott",
   "metadata": {},
   "source": [
    "## Stratified model with migration\n",
    "\n",
    "Now we can add a [transition flow](../api/model.rst#summer2.model.CompartmentalModel.add_transition_flow) where 2% of the rural population migrates to the urban compartment per year."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "christian-donna",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification\n",
    "\n",
    "model = build_model()\n",
    "\n",
    "# Add an urban/rural stratification with an inter-location migration flow.\n",
    "strat = Stratification('location', ['urban', 'rural'], ['pop'])\n",
    "strat.set_population_split({'rural': 0.7, 'urban': 0.3})\n",
    "model.stratify_with(strat)\n",
    "model.add_transition_flow(\n",
    "    'migration',\n",
    "    fractional_rate=0.02,\n",
    "    source='pop',\n",
    "    dest='pop',\n",
    "    source_strata={'location': 'rural'},\n",
    "    dest_strata={'location': 'urban'},\n",
    "    # Expected flow count can be used as a sanity check,\n",
    "    # to assert that the expected number of flows was added.\n",
    "    expected_flow_count=1\n",
    ")\n",
    "\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "english-means",
   "metadata": {},
   "source": [
    "## Stratified model with age-based migration\n",
    "\n",
    "We can take this example one step further with the observation that:\n",
    "\n",
    "- people aged 0-19 are unlikely to migrate\n",
    "- people aged 20-39 are likely to migrate\n",
    "- people aged 40+ are less likely to migrate\n",
    "\n",
    "We can use an [age stratification](07-age-stratification.ipynb) to model the age strata and ageing flows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "physical-burning",
   "metadata": {},
   "outputs": [],
   "source": [
    "from summer2 import Stratification, AgeStratification, Overwrite\n",
    "\n",
    "model = build_model()\n",
    "\n",
    "# Add an urban/rural stratification with an inter-location migration flow.\n",
    "strat = Stratification('location', ['urban', 'rural'], ['pop'])\n",
    "strat.set_population_split({'rural': 0.7, 'urban': 0.3})\n",
    "model.stratify_with(strat)\n",
    "model.add_transition_flow(\n",
    "    'migration',\n",
    "    fractional_rate=0,  # To be overwritten\n",
    "    source='pop',\n",
    "    dest='pop',\n",
    "    source_strata={'location': 'rural'},\n",
    "    dest_strata={'location': 'urban'},\n",
    "    # Expected flow count can be used as a sanity check,\n",
    "    # to assert that the expected number of flows was added.\n",
    "    expected_flow_count=1\n",
    ")\n",
    "\n",
    "# Set age-specific migration rates.\n",
    "age_strat = AgeStratification('age', [0, 20, 40], ['pop']) \n",
    "age_strat.set_population_split({'0': 0.2, '20': 0.4, '40': 0.4})\n",
    "age_strat.set_flow_adjustments(\"migration\", {\n",
    "    '0': Overwrite(0),  # No migration\n",
    "    '20': Overwrite(0.05),  # 5% of 20-39 year olds per year\n",
    "    '40': Overwrite(0.01),  # 1% of 40+ year olds per year\n",
    "})\n",
    "model.stratify_with(age_strat)\n",
    "\n",
    "# Track urban and rural populations\n",
    "model.request_output_for_compartments(\n",
    "    'urban_pop',\n",
    "    compartments=[\"pop\"],\n",
    "    strata={\"location\": \"urban\"}\n",
    ")\n",
    "model.request_output_for_compartments(\n",
    "    'rural_pop',\n",
    "    compartments=[\"pop\"],\n",
    "    strata={\"location\": \"rural\"}\n",
    ")\n",
    "\n",
    "model.run()\n",
    "model.get_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b6a1f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_derived_outputs_df().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "certified-field",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Now you know how to add flows between strata after a model has been stratified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collective-america",
   "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
}