{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69c73ab37081b35c",
   "metadata": {},
   "source": [
    "# Simulating Discrete models"
   ]
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:23.406045233Z",
     "start_time": "2026-03-02T18:10:23.372200854Z"
    }
   },
   "source": [
    "from epimodels.discrete import models as DM\n",
    "from IPython.display import Markdown as md"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import matplotlib.pyplot as P"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "id": "24940ecc8c6fa07f",
   "metadata": {},
   "source": [
    "## SIS model\n",
    "\n",
    "The SIS (Susceptible-Infectious-Susceptible) model for diseases that do not confer immunity after recovery."
   ]
  },
  {
   "cell_type": "code",
   "id": "a41cc87db7594492",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:26.394027058Z",
     "start_time": "2026-03-02T18:10:26.322478242Z"
    }
   },
   "source": [
    "modelsis = DM.SIS()\n",
    "md(str(modelsis))"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "28afc637fa80a992",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:35.485227234Z",
     "start_time": "2026-03-02T18:10:35.062285661Z"
    }
   },
   "source": [
    "modelsis([0, 1, 1000], [0,50], 1001, {'beta': 2, 'gamma': 1})\n",
    "modelsis.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "id": "f163062e-23ac-4a90-a2d8-df773a41b1c0",
   "metadata": {},
   "source": [
    "## SIR Model\n",
    "\n",
    "The SIR (Susceptible-Infectious-Removed) model is a classic compartmental model."
   ]
  },
  {
   "cell_type": "code",
   "id": "7f98fd94-0ab9-4445-b419-c445588640aa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:42.321785138Z",
     "start_time": "2026-03-02T18:10:42.262256090Z"
    }
   },
   "source": [
    "sirmodel = DM.SIR()\n",
    "md(str(sirmodel))"
   ],
   "outputs": [],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "4d1f7ade-f222-4e7d-a116-95dddfee3a86",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:51.703505745Z",
     "start_time": "2026-03-02T18:10:51.362072061Z"
    }
   },
   "source": [
    "sirmodel([1000, 1, 0], [0,500], 1001, {'beta': .2, 'gamma': .1})\n",
    "sirmodel.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SEIS Model\n",
    "\n",
    "The SEIS (Susceptible-Exposed-Infectious-Susceptible) model includes an exposed compartment but no permanent immunity."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SEIS()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 0, 1], [0, 100], 1001, {'beta': 0.5, 'e': 0.3, 'r': 0.1, 'b': 0})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SEIR Model\n",
    "\n",
    "The SEIR (Susceptible-Exposed-Infectious-Removed) model includes an exposed compartment for diseases with incubation period."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SEIR()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 0, 1, 0], [0, 100], 1001, {'beta': 0.5, 'e': 0.3, 'r': 0.1, 'b': 0, 'alpha': 0.01})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SIRS Model\n",
    "\n",
    "The SIRS (Susceptible-Infectious-Removed-Susceptible) model accounts for waning immunity."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SIRS()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 1, 0], [0, 200], 1001, {'beta': 0.3, 'r': 0.1, 'b': 0, 'w': 0.02})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "id": "57b72ff6f1cdd923",
   "metadata": {},
   "source": [
    "## SIpRpS model\n",
    "\n",
    "This model assumes that a fraction $\\delta$ of infectious individuals acquire full immunity while the remaining $(1-\\delta)$ returns to the susceptible stage."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SIpRpS()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 1, 0], [0, 100], 1001, {'beta': 0.5, 'r': 0.2, 'b': 0, 'delta': 0.7})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SEIpRpS model\n",
    "\n",
    "Similar to SIpRpS but with an exposed compartment. A fraction $\\delta$ of infectious individuals acquire full immunity while the remaining $(1-\\delta)$ returns to the susceptible stage."
   ]
  },
  {
   "cell_type": "code",
   "id": "d36dbcc5272d073f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:53.492103624Z",
     "start_time": "2026-03-02T18:10:53.471813966Z"
    }
   },
   "source": [
    "model = DM.SEIpRpS()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "id": "b88858349a3f722a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:56.206953412Z",
     "start_time": "2026-03-02T18:10:55.891541369Z"
    }
   },
   "source": [
    "model([1000, 0, 1, 0], [0, 50], 1001, {'beta': 1, 'r': .3, 'e': 1, 'b': 0, 'delta': 0.6})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 16
  },
  {
   "cell_type": "code",
   "id": "6c29d9d975e474d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-02T18:10:57.177794117Z",
     "start_time": "2026-03-02T18:10:57.107906939Z"
    }
   },
   "source": [
    "model.parameter_table()"
   ],
   "outputs": [],
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SIpR model\n",
    "\n",
    "This model allows recovered individuals to be reinfected at a reduced rate $p\\beta$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SIpR()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 18
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 1, 0], [0, 150], 1001, {'beta': 0.3, 'r': 0.1, 'b': 0, 'p': 0.3})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 19
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SEIpR model\n",
    "\n",
    "Similar to SIpR but with an exposed compartment. Recovered individuals can be reinfected at a reduced rate $p\\beta$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SEIpR()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 20
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model([1000, 0, 1, 0], [0, 150], 1001, {'beta': 0.3, 'e': 0.2, 'r': 0.1, 'b': 0, 'alpha': 0.01, 'p': 0.3})\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 21
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SEQIAHR model\n",
    "\n",
    "A complex model with asymptomatic and hospitalized compartments, suitable for modeling diseases like COVID-19."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.SEQIAHR()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 22
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "params = {\n",
    "    'chi': 0.3,    # Quarantine reduction factor\n",
    "    'phi': 0.1,    # Hospitalization rate from I\n",
    "    'beta': 0.5,   # Transmission rate\n",
    "    'rho': 0.1,    # Recovery rate from H\n",
    "    'delta': 0.2,  # Recovery rate from I\n",
    "    'gamma': 0.1,  # Recovery rate from A\n",
    "    'alpha': 0.3,  # Incubation rate\n",
    "    'mu': 0.01,    # Mortality rate in H\n",
    "    'p': 0.5,      # Proportion asymptomatic\n",
    "    'q': 20,       # Quarantine start day\n",
    "    'r': 30        # Quarantine duration\n",
    "}\n",
    "inits = [10000, 0, 1, 0, 0, 0, 0, 0]\n",
    "model(inits, [0, 100], 10001, params)\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 23
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Influenza model\n",
    "\n",
    "A complex age-structured model for influenza with four age classes (0-2, 3-14, 15-59, >60 years). Each age class has compartments for susceptible, incubating, subclinical, clinical, and complicated individuals."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model = DM.Influenza()\n",
    "md(str(model))"
   ],
   "outputs": [],
   "execution_count": 24
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "inits = [\n",
    "    100, 0, 0, 0, 0,  # Age group 1 (0-2 years): S1, E1, Is1, Ic1, Ig1\n",
    "    200, 0, 0, 0, 0,  # Age group 2 (3-14 years): S2, E2, Is2, Ic2, Ig2\n",
    "    500, 0, 1, 0, 0,  # Age group 3 (15-59 years): S3, E3, Is3, Ic3, Ig3\n",
    "    200, 0, 0, 0, 0   # Age group 4 (>60 years): S4, E4, Is4, Ic4, Ig4\n",
    "]\n",
    "params = {\n",
    "    'beta': 0.0003,  # Transmission rate\n",
    "    'r': 0.3,        # Recovery rate\n",
    "    'e': 0.5,        # Incubation rate\n",
    "    'c': 0.5,        # Clinical progression rate\n",
    "    'g': 0.2,        # Complication rate\n",
    "    'd': 0.1,        # Death/recovery rate from complications\n",
    "    'pc1': 0.3,      # Clinical proportion age group 1\n",
    "    'pc2': 0.3,      # Clinical proportion age group 2\n",
    "    'pc3': 0.4,      # Clinical proportion age group 3\n",
    "    'pc4': 0.6,      # Clinical proportion age group 4\n",
    "    'pp1': 0.05,     # Complication proportion age group 1\n",
    "    'pp2': 0.05,     # Complication proportion age group 2\n",
    "    'pp3': 0.1,      # Complication proportion age group 3\n",
    "    'pp4': 0.3,      # Complication proportion age group 4\n",
    "    'b': 0           # Birth rate\n",
    "}\n",
    "model(inits, [0, 100], 1001, params)\n",
    "model.plot_traces()"
   ],
   "outputs": [],
   "execution_count": 25
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}