examples: add SIR epidemic model with Mesa 4.0 discrete space API

examples/sir_epidemic/Readme.md

@@ -0,0 +1,82 @@
+# SIR Epidemic Model
+
+## Abstract
+
+This model simulates the spread of an infectious disease through a population using the classic Susceptible-Infected-Recovered (SIR) compartmental framework. Agents move randomly on a 2D grid, and infected agents can transmit the disease to susceptible neighbors with a configurable probability. After a set recovery period, infected agents become recovered and immune. The model demonstrates how epidemic dynamics emerge from simple local interaction rules.
+
+## Background
+
+The SIR model is one of the foundational frameworks in mathematical epidemiology, originally formulated by Kermack and McKendrick (1927). It divides a population into three compartments:
+
+- **Susceptible (S):** Individuals who can contract the disease.
+- **Infected (I):** Individuals who have the disease and can spread it.
+- **Recovered (R):** Individuals who have recovered and are immune.
+
+While the original SIR model uses differential equations to describe population-level dynamics, agent-based implementations allow exploration of spatial effects, stochastic variation, and heterogeneous contact patterns that differential equation models cannot capture.
+
+## Model Description
+
+### Agents
+
+Each agent (`Person`) is a `CellAgent` on an `OrthogonalMooreGrid`. Agents have:
+- A `state` attribute: "Susceptible", "Infected", or "Recovered"
+- An `infection_timer` tracking how long they have been infected
+- A `recovery_time` threshold for transitioning to "Recovered"
+
+### Rules
+
+At each step, every agent:
+1. **Moves** to a random neighboring cell (Moore neighborhood).
+2. **Spreads infection** (if infected): Each susceptible agent in the new cell's neighborhood is infected with probability `infection_rate`.
+3. **Recovers** (if infected): After `recovery_time` steps, transitions to "Recovered".
+
+### Space
+
+A toroidal `OrthogonalMooreGrid` of configurable width and height. Multiple agents can occupy the same cell.
+
+### Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `num_agents` | 100 | Total number of agents |
+| `width` | 20 | Grid width |
+| `height` | 20 | Grid height |
+| `infection_rate` | 0.3 | Probability of transmission per contact |
+| `recovery_time` | 10 | Steps until recovery |
+| `initial_infected` | 3 | Number of initially infected agents |
+
+## How to Run
+
+### Install dependencies
+
+```bash
+pip install -U "mesa[rec]"
+```
+
+### Run the visualization
+
+```bash
+solara run app.py
+```
+
+### Run headless (no visualization)
+
+```bash
+python model.py
+```
+
+## Results and Discussion
+
+With default parameters (100 agents, 30% infection rate, recovery time of 10 steps), the epidemic typically follows an S-shaped curve:
+
+1. **Early phase (steps 1-10):** Slow spread as few agents are infected.
+2. **Exponential growth (steps 10-25):** Rapid increase in infections as the disease spreads through the susceptible population.
+3. **Peak and decline (steps 25-40):** Infections peak and decline as the susceptible pool is depleted.
+4. **Equilibrium (steps 40+):** Nearly all agents have recovered; the epidemic ends.
+
+The spatial nature of the model means that local clusters of infection can form and spread outward, creating wave-like patterns visible in the grid visualization.
+
+## References
+
+- Kermack, W. O., & McKendrick, A. G. (1927). A contribution to the mathematical theory of epidemics. *Proceedings of the Royal Society of London. Series A*, 115(772), 700-721.
+- Wilensky, U. (1998). NetLogo Virus model. Center for Connected Learning and Computer-Based Modeling, Northwestern University.

examples/sir_epidemic/app.py

@@ -0,0 +1,76 @@
+"""SIR Epidemic Model - Visualization.
+
+Run with: solara run app.py
+"""
+
+from mesa.visualization import SolaraViz, make_plot_component, make_space_component
+from model import Person, SIRModel
+
+# Color mapping for agent states
+STATE_COLORS = {
+    "Susceptible": "tab:blue",
+    "Infected": "tab:red",
+    "Recovered": "tab:green",
+}
+
+
+def agent_portrayal(agent):
+    """Define how agents are displayed on the grid."""
+    if not isinstance(agent, Person):
+        return None
+    return {
+        "color": STATE_COLORS.get(agent.state, "tab:gray"),
+        "marker": "o",
+        "size": 30,
+    }
+
+
+model_params = {
+    "num_agents": {
+        "type": "SliderInt",
+        "value": 100,
+        "label": "Number of Agents",
+        "min": 10,
+        "max": 300,
+        "step": 10,
+    },
+    "width": 20,
+    "height": 20,
+    "infection_rate": {
+        "type": "SliderFloat",
+        "value": 0.3,
+        "label": "Infection Rate",
+        "min": 0.0,
+        "max": 1.0,
+        "step": 0.05,
+    },
+    "recovery_time": {
+        "type": "SliderInt",
+        "value": 10,
+        "label": "Recovery Time (steps)",
+        "min": 1,
+        "max": 30,
+        "step": 1,
+    },
+    "initial_infected": {
+        "type": "SliderInt",
+        "value": 3,
+        "label": "Initial Infected",
+        "min": 1,
+        "max": 20,
+        "step": 1,
+    },
+}
+
+# Create visualization
+page = SolaraViz(
+    SIRModel,
+    components=[
+        make_space_component(agent_portrayal),
+        make_plot_component(["Susceptible", "Infected", "Recovered"]),
+    ],
+    model_params=model_params,
+    name="SIR Epidemic Model",
+)
+
+page  # noqa: B018

examples/sir_epidemic/metadata.toml

@@ -0,0 +1,14 @@
+[model]
+title = "SIR Epidemic Model"
+abstract = "A Susceptible-Infected-Recovered epidemic model demonstrating disease spread on a 2D grid using Mesa's discrete space API. Agents move randomly, infected agents spread disease to susceptible neighbors, and recover after a fixed period."
+authors = ["Ved Prakash <vedprakash9162@gmail.com>"]
+
+[model.classification]
+domain = ["epidemiology", "public-health"]
+space = "Grid"
+time = "Discrete"
+complexity = "Basic"
+
+[model.mesa]
+mesa_version_min = "4.0"
+keywords = ["SIR", "epidemic", "infection", "grid", "CellAgent", "disease-spread"]

examples/sir_epidemic/model.py

@@ -0,0 +1,114 @@
+"""SIR Epidemic Model.
+
+A Susceptible-Infected-Recovered epidemic model demonstrating disease spread
+on a 2D grid using Mesa's discrete space API. Agents move randomly, infected
+agents can spread the disease to susceptible neighbors, and infected agents
+recover after a set number of steps.
+"""
+
+import mesa
+from mesa.discrete_space import CellAgent, OrthogonalMooreGrid
+
+
+class Person(CellAgent):
+    """An agent representing a person in the SIR model.
+
+    Attributes:
+        state: Current health state - "Susceptible", "Infected", or "Recovered".
+        infection_timer: Number of steps since infection.
+        recovery_time: Steps required to recover from infection.
+    """
+
+    def __init__(self, model, initial_state="Susceptible"):
+        super().__init__(model)
+        self.state = initial_state
+        self.infection_timer = 0
+        self.recovery_time = model.recovery_time
+
+    def step(self):
+        """Move randomly, spread infection, and check recovery."""
+        # Move to a random neighboring cell
+        neighborhood = self.cell.neighborhood
+        self.cell = neighborhood.select_random_cell()
+
+        # If infected, try to spread to susceptible neighbors
+        if self.state == "Infected":
+            for neighbor in self.cell.neighborhood.agents:
+                if (
+                    neighbor.state == "Susceptible"
+                    and self.random.random() < self.model.infection_rate
+                ):
+                    neighbor.state = "Infected"
+
+            # Check if agent should recover
+            self.infection_timer += 1
+            if self.infection_timer >= self.recovery_time:
+                self.state = "Recovered"
+
+
+class SIRModel(mesa.Model):
+    """A simple SIR epidemic model.
+
+    Args:
+        num_agents: Total number of agents in the simulation.
+        width: Width of the grid.
+        height: Height of the grid.
+        infection_rate: Probability of infection spreading per contact.
+        recovery_time: Number of steps before an infected agent recovers.
+        initial_infected: Number of agents initially infected.
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        num_agents=100,
+        width=20,
+        height=20,
+        infection_rate=0.3,
+        recovery_time=10,
+        initial_infected=3,
+    ):
+        super().__init__()
+        self.infection_rate = infection_rate
+        self.recovery_time = recovery_time
+
+        # Create grid
+        self.grid = OrthogonalMooreGrid((width, height), torus=True, random=self.random)
+
+        # Create DataCollector
+        self.datacollector = mesa.DataCollector(
+            model_reporters={
+                "Susceptible": lambda m: sum(
+                    1 for a in m.agents if a.state == "Susceptible"
+                ),
+                "Infected": lambda m: sum(1 for a in m.agents if a.state == "Infected"),
+                "Recovered": lambda m: sum(
+                    1 for a in m.agents if a.state == "Recovered"
+                ),
+            }
+        )
+
+        # Create agents and place on random cells
+        for _ in range(num_agents):
+            agent = Person(self, initial_state="Susceptible")
+            agent.cell = self.grid.all_cells.select_random_cell()
+
+        # Infect initial agents
+        infected_agents = self.random.sample(list(self.agents), initial_infected)
+        for agent in infected_agents:
+            agent.state = "Infected"
+
+        self.datacollector.collect(self)
+
+    def step(self):
+        """Advance the model by one step."""
+        self.agents.shuffle_do("step")
+        self.datacollector.collect(self)
+
+
+if __name__ == "__main__":
+    model = SIRModel()
+    for _ in range(50):
+        model.step()
+    data = model.datacollector.get_model_vars_dataframe()
+    print(data)