From dbb310d178cbde25d9a12921ae418fb1021f55d9 Mon Sep 17 00:00:00 2001
From: jtrauer <james.trauer@monash.edu>
Date: Fri, 4 Oct 2024 21:53:49 +1000
Subject: [PATCH] Add notebook illustrating force of infection generation
 without use of built-in versions

---
 .../jtrauer/construct_force_infection.ipynb   | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 notebooks/user/jtrauer/construct_force_infection.ipynb

diff --git a/notebooks/user/jtrauer/construct_force_infection.ipynb b/notebooks/user/jtrauer/construct_force_infection.ipynb
new file mode 100644
index 0000000000..1798eda862
--- /dev/null
+++ b/notebooks/user/jtrauer/construct_force_infection.ipynb
@@ -0,0 +1,95 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "90759a2a-40e5-431a-a8bd-26e51a650aff",
+   "metadata": {},
+   "source": [
+    "Illustration of how to construct the force of infection\n",
+    "(frequency-dependent) without using `summer2`'s built-in\n",
+    "`add_infection_frequency_flow`.\n",
+    "(I think this is correct, but may need David to confirm\n",
+    "and haven't tested in detail.)\n",
+    "Also, note that this is not an epidemiologically sensible model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7f879ec-74d9-4e2c-975f-a5cb3c670a8e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from summer2 import CompartmentalModel\n",
+    "from summer2.parameters import Parameter, Function, CompartmentValues"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "758c8abb-8e5f-43e4-a76d-47a505e6f932",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "comps = [\"susceptible\", \"infectious\", \"recovered\"]\n",
+    "epi_model = CompartmentalModel([0.0, 100.0], comps, [\"infectious\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13199b5d-d8b5-4c1f-8506-2e313fae2e43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_force_infection(comp_vals):\n",
+    "    \"\"\"Manually construct the force of infection calculation.\n",
+    "    \"\"\"\n",
+    "    infect_comps = epi_model.query_compartments({\"name\": \"infectious\"}, as_idx=True)\n",
+    "    infect_pop = comp_vals[infect_comps]\n",
+    "    all_comps = epi_model.query_compartments({\"name\": comps}, as_idx=True)\n",
+    "    total_pop = comp_vals[all_comps].sum()\n",
+    "    return infect_pop / total_pop\n",
+    "\n",
+    "epi_model.add_transition_flow(\n",
+    "    \"infection\",    \n",
+    "    Function(get_force_infection, [CompartmentValues]) * Parameter(\"beta\"),\n",
+    "    \"susceptible\",\n",
+    "    \"infectious\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0fd8b13c-af6b-4170-9e61-95ec7cbaa566",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epi_model.set_initial_population({\"susceptible\": 100.0})\n",
+    "epi_model.run({\"beta\": 1.0})"
+   ]
+  }
+ ],
+ "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.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}