feat(model): SIR model

.github/workflows/docs-preview.yaml

@@ -11,23 +11,15 @@ on:
 concurrency: preview-${{ github.ref }}
 
 jobs:
-  pre-commit:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v4.1.6
-    - uses: actions/setup-python@v5.1.0
-      with:
-        python-version: "3.12"
-    - uses: pre-commit/action@v3.0.1
   deploy-preview:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v5.1.0
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.12"
       - name: Setup Poetry
-        uses: abatilo/actions-poetry@v3.0.0
+        uses: abatilo/actions-poetry@v4
         with:
           poetry-version: 1.8.3
       - name: Install Packages

.github/workflows/docs-publish.yaml

@@ -12,23 +12,15 @@ on:
 
 
 jobs:
-  pre-commit:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v4.1.6
-    - uses: actions/setup-python@v5.1.0
-      with:
-        python-version: "3.12"
-    - uses: pre-commit/action@v3.0.1
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v5.1.0
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.12"
       - name: Setup Poetry
-        uses: abatilo/actions-poetry@v3.0.0
+        uses: abatilo/actions-poetry@v4
         with:
           poetry-version: 1.8.3
       - name: Install Packages

.github/workflows/publish.yaml

@@ -14,13 +14,13 @@ jobs:
     runs-on: ${{ matrix.os }}
     steps:
       - name: Checkout
-        uses: actions/checkout@v4.1.6
+        uses: actions/checkout@v6
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5.1.0
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
       - name: Setup Poetry
-        uses: abatilo/actions-poetry@v3.0.0
+        uses: abatilo/actions-poetry@v4
         with:
           poetry-version: 1.8.3
       - name: Install Packages
@@ -31,15 +31,15 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4.1.6
+      - uses: actions/checkout@v6
         with:
           fetch-depth: 0
       - name: Set up Python
-        uses: actions/setup-python@v5.1.0
+        uses: actions/setup-python@v6
         with:
           python-version: "3.12"
       - name: Setup Poetry
-        uses: abatilo/actions-poetry@v3.0.0
+        uses: abatilo/actions-poetry@v4
         with:
           poetry-version: 1.8.3
       - name: Install dependencies

.github/workflows/tests.yaml

@@ -20,13 +20,13 @@ jobs:
     runs-on: ${{ matrix.os }}
     steps:
       - name: Checkout
-        uses: actions/checkout@v4.1.6
+        uses: actions/checkout@v6
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5.1.0
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
       - name: Setup Poetry
-        uses: abatilo/actions-poetry@v3.0.0
+        uses: abatilo/actions-poetry@v4
         with:
           poetry-version: 1.8.3
       - name: Install Packages
@@ -35,7 +35,7 @@ jobs:
         run: |
           poetry run pytest --cache-clear --junitxml=pytest.xml --cov=hamilflow tests/ > pytest-coverage.txt
       - name: Pytest coverage comment
-        uses: MishaKav/pytest-coverage-comment@v1.1.51
+        uses: MishaKav/pytest-coverage-comment@v1.7.1
         with:
           pytest-coverage-path: ./pytest-coverage.txt
           junitxml-path: ./pytest.xml

docs/references/models/sir.md

@@ -0,0 +1,3 @@
+# SIR
+
+::: hamilflow.models.sir

docs/tutorials/sir.py

@@ -0,0 +1,58 @@
+# ---
+# jupyter:
+#   jupytext:
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.16.4
+#   kernelspec:
+#     display_name: .venv
+#     language: python
+#     name: python3
+# ---
+
+# %% [markdown]
+# # SIR Model
+#
+# In this tutorial, we will learn how to use the SIR model.
+
+# %%
+import plotly.express as px
+
+from hamilflow.models.sir import SIR
+
+# %% [markdown]
+# ## Model
+
+# %%
+sir_1 = SIR(
+    system={
+        "beta": 0.3,
+        "alpha": 0.1,
+        "delta_t": 0.1,
+    },
+    initial_condition={
+        "susceptible_0": 999,
+        "infected_0": 1,
+        "recovered_0": 0,
+    },
+)
+
+# %%
+n_steps = 100
+
+sir_1_results = sir_1.generate_from(n_steps=n_steps)
+sir_1_results.head()
+
+# %%
+px.line(
+    sir_1_results,
+    x="t",
+    y=["S", "I", "R"],
+)
+
+# %%
+sir_1._step(999, 1)  # noqa: SLF001
+
+# %%

hamilflow/models/sir.py

@@ -0,0 +1,169 @@
+"""Main module for SIR model in epidemiology."""
+
+from collections.abc import Mapping
+from functools import cached_property
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from pydantic import BaseModel, Field, computed_field
+
+from hamilflow.models.utils.typing import TypeTime
+
+
+class SIRSystem(BaseModel):
+    """Definition of the SIR system.
+
+    :cvar beta: Transmission rate
+    :cvar alpha: Recovery rate
+    :cvar delta_t: Time granularity of the simulation
+    """
+
+    beta: float = Field(default=0.3, gt=0, description="Transmission rate", frozen=True)
+    alpha: float = Field(default=0.1, gt=0, description="Recovery rate", frozen=True)
+    delta_t: float = Field(ge=0.0, default=1.0)
+
+
+class SIRIC(BaseModel):
+    """The initial condition for an SIR model simulation.
+
+    :cvar susceptible_0: Initial number of susceptible individuals
+    :cvar infected_0: Initial number of infectious individuals
+    :cvar recovered_0: Initial number of recovered individuals
+    """
+
+    susceptible_0: int = Field(
+        default=999,
+        ge=0,
+        description="Initial susceptible population",
+    )
+    infected_0: int = Field(
+        default=1,
+        ge=0,
+        description="Initial infectious population",
+    )
+    recovered_0: int = Field(
+        default=0,
+        ge=0,
+        description="Initial recovered population",
+    )
+
+    @computed_field  # type: ignore[misc]
+    @cached_property
+    def n(self) -> int:
+        """Total population in the simulation."""
+        return self.susceptible_0 + self.infected_0 + self.recovered_0
+
+
+class SIR:
+    r"""SIR model simulation.
+
+    The SIR model divides a population into three compartments:
+
+    - $S$ (Susceptible): Individuals who can be infected.
+    - $I$ (Infectious): Individuals who are currently infected and can transmit the disease.
+    - $R$ (Recovered): Individuals who have recovered and are assumed to have immunity.
+    - $N$(Total population): $N = S + I + R$.
+
+    The dynamics of the compartments are governed by the following system of ordinary differential equations:
+
+    $$
+    \begin{split}
+    \frac{dS(t)}{dt} &= -\beta I(t) S(t) \\
+    \frac{dI(t)}{dt} &= \beta S(t) I(t) - \alpha I(t) \\
+    \frac{dR(t)}{dt} &= \alpha I(t),
+    \end{split}
+    $$
+
+    with the constraint
+
+    $$
+    N = S(t) + I(t) + R(t).
+    $$
+
+    Where:
+    - $\beta$ is the transmission rate (probability of infection per contact per unit time).
+    - $\alpha$ is the recovery rate (rate at which infected individuals recover per unit time).
+
+    :param system: The parameters of the SIR system, including `beta` and `alpha`.
+    :param initial_condition: The initial state of the population, including `S0`, `I0`, and `R0`.
+    """
+
+    def __init__(
+        self,
+        system: Mapping[str, float],
+        initial_condition: Mapping[str, int],
+    ) -> None:
+        self.system = SIRSystem(**system)
+        self.initial_condition = SIRIC(**initial_condition)
+
+    @cached_property
+    def definition(self) -> dict[str, dict[str, Any]]:
+        """Model params and initial conditions defined as a dictionary."""
+        return {
+            "system": self.system.model_dump(),
+            "initial_condition": self.initial_condition.model_dump(),
+        }
+
+    def generate_from(self, n_steps: int) -> pd.DataFrame:
+        """Simulate the SIR model and return time series data.
+
+        :param n_steps: Number of steps to simulate
+        :return: DataFrame with time, S, I, R columns
+        """
+        time_steps = np.arange(1, n_steps) * self.system.delta_t
+
+        return self(time_steps)
+
+    def _step(
+        self,
+        susceptible: float,
+        infected: float,
+    ) -> tuple[int, int, int]:
+        """Calculate changes in S, I, R populations for one time step.
+
+        :param susceptible: Current susceptible population
+        :param infected: Current infected population
+        :param recovered: Current recovered population
+        :return: tuple of (dS, dI, dR) changes
+        """
+        delta_s = -self.system.beta * susceptible * infected * self.system.delta_t
+        delta_i = (
+            self.system.beta * susceptible * infected - self.system.alpha * infected
+        ) * self.system.delta_t
+        delta_r = self.system.alpha * infected * self.system.delta_t
+
+        return int(delta_s), int(delta_i), int(delta_r)
+
+    def __call__(self, t: TypeTime) -> pd.DataFrame:
+        """Generate the SIR model simulation based on the given time array."""
+        susceptible = self.initial_condition.susceptible_0
+        infected = self.initial_condition.infected_0
+        recovered = self.initial_condition.recovered_0
+
+        results = [
+            {
+                "t": 0,
+                "S": susceptible,
+                "I": infected,
+                "R": recovered,
+            },
+        ]
+
+        for t_i in np.array(t):
+            delta_s, delta_i, delta_r = self._step(susceptible, infected)
+
+            susceptible = max(susceptible + delta_s, 0)
+            infected = max(infected + delta_i, 0)
+            recovered = max(recovered + delta_r, 0)
+
+            results.append(
+                {
+                    "t": t_i,
+                    "S": susceptible,
+                    "I": infected,
+                    "R": recovered,
+                },
+            )
+
+        return pd.DataFrame(results)

hamilflow/models/utils/typing.py

@@ -5,4 +5,4 @@
 
 from numpy.typing import ArrayLike
 
-TypeTime = TypeVar("TypeTime", Sequence[float], Sequence[int], ArrayLike)
+TypeTime = TypeVar("TypeTime", bound=Sequence[float] | Sequence[int] | ArrayLike)

mkdocs.yml

@@ -71,6 +71,10 @@ plugins:
       allow_errors: false
       include_requirejs: true
 
+watch:
+  - docs
+  - hamilflow
+
 extra_javascript:
   - javascripts/mathjax.js
   - https://polyfill.io/v3/polyfill.min.js?features=es6
@@ -86,6 +90,7 @@ nav:
     - "Pendulum": tutorials/pendulum.py
     - "Kepler Problem": tutorials/kepler_problem.py
     - "Harmonic Oscillator Chain": tutorials/harmonic_oscillator_chain.py
+    - "SIR": tutorials/sir.py
   - References:
     - "Introduction": references/index.md
     - "Models":
@@ -97,6 +102,7 @@ nav:
         - Dynamics: references/models/kepler_problem/dynamics.md
         - Numerics: references/models/kepler_problem/numerics.md
       - "Harmonic Oscillator Chain": references/models/harmonic_oscillator_chain.md
+      - "SIR": references/models/sir.md
     - Mathematics:
       - Trigonometrics: references/maths/trigonometrics.md
   - "Changelog": changelog.md

pyproject.toml

@@ -75,6 +75,8 @@ target-version = "py310"
 select = ["ALL"]
 ignore = [
     "D107",  # undocumented-public-init: we document the class instead
+    "D203",  # one-blank-line-before-class: we use D211 instead
+    "D213",  # multi-line-summary-second-line: we use D212 instead
     "E501",  # line-too-long: we use black instead
     "TID252",  # relative-imports
     "PD901",  # pandas-df-variable-name
@@ -94,6 +96,7 @@ ignore = [
 "tests/**/*.py" = [
     "S101",  # assert: Fine in tests
     "SLF001",  # private-member-access: find in tests
+    "PLR2004", # magic value is fine in tests
 ]