feat(model): SIR model

docs/references/models/sir.md

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

docs/tutorials/sir.py

@@ -0,0 +1,55 @@
+# ---
+# 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.03,
+        "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"],
+)
+
+# %%

hamilflow/models/sir.py

@@ -0,0 +1,169 @@
+"""Main module for Brownian motion."""
+
+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 `gamma`.
+    :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):
+            results.append(
+                {
+                    "t": t_i * self.system.delta_t,
+                    "S": susceptible,
+                    "I": infected,
+                    "R": recovered,
+                },
+            )
+
+            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)
+
+        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
@@ -97,6 +101,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

@@ -94,6 +94,7 @@ ignore = [
 "tests/**/*.py" = [
     "S101",  # assert: Fine in tests
     "SLF001",  # private-member-access: find in tests
+    "PLR2004", # magic value is fine in tests
 ]
 
 

tests/test_models/test_sir.py

@@ -0,0 +1,80 @@
+"""Tests for the SIR model."""
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from hamilflow.models.sir import SIR
+
+
+@pytest.fixture
+def sir_system() -> dict[str, float]:
+    """Params for the SIR system."""
+    return {
+        "beta": 0.3,
+        "alpha": 0.1,
+        "delta_t": 1.0,
+    }
+
+
+@pytest.fixture
+def sir_initial_condition() -> dict[str, int]:
+    """Generate initial conditions for the SIR model."""
+    return {
+        "susceptible_0": 999,
+        "infected_0": 1,
+        "recovered_0": 0,
+    }
+
+
+@pytest.fixture
+def sir_model(sir_system: dict, sir_initial_condition: dict) -> SIR:
+    """SIR model."""
+    return SIR(sir_system, sir_initial_condition)
+
+
+def test_sir_initialization(sir_model: SIR) -> None:
+    """Test the initialization of the SIR model."""
+    assert sir_model.system.beta == 0.3
+    assert sir_model.system.alpha == 0.1
+    assert sir_model.system.delta_t == 1.0
+    assert sir_model.initial_condition.susceptible_0 == 999
+    assert sir_model.initial_condition.infected_0 == 1
+    assert sir_model.initial_condition.recovered_0 == 0
+
+
+def test_sir_definition(sir_model: SIR) -> None:
+    """Test the definition of the SIR model."""
+    definition = sir_model.definition
+    assert definition["system"]["beta"] == 0.3
+    assert definition["system"]["alpha"] == 0.1
+    assert definition["system"]["delta_t"] == 1.0
+    assert definition["initial_condition"]["susceptible_0"] == 999
+    assert definition["initial_condition"]["infected_0"] == 1
+    assert definition["initial_condition"]["recovered_0"] == 0
+
+
+def test_sir_generate_from(sir_model: SIR) -> None:
+    """Test the data generation of the SIR model."""
+    n_steps = 10
+    result = sir_model.generate_from(n_steps)
+    assert isinstance(result, pd.DataFrame)
+    assert len(result) == n_steps
+    assert all(col in result.columns for col in ["t", "S", "I", "R"])
+
+
+def test_sir_step(sir_model: SIR) -> None:
+    """Test the step function of the SIR model."""
+    delta_s, delta_i, delta_r = sir_model._step(999, 1)
+    assert delta_s < 0
+    assert delta_i > 0
+    assert delta_r == 0
+
+
+def test_sir_call(sir_model: SIR) -> None:
+    """Test the call function of the SIR model."""
+    t = np.arange(1, 10)
+    result = sir_model(t)
+    assert isinstance(result, pd.DataFrame)
+    assert len(result) == len(t) + 1
+    assert all(col in result.columns for col in ["t", "S", "I", "R"])