diff --git a/examples/sir_epidemic/Readme.md b/examples/sir_epidemic/Readme.md new file mode 100644 index 00000000..4ad21ced --- /dev/null +++ b/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. diff --git a/examples/sir_epidemic/app.py b/examples/sir_epidemic/app.py new file mode 100644 index 00000000..4cb60786 --- /dev/null +++ b/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 diff --git a/examples/sir_epidemic/metadata.toml b/examples/sir_epidemic/metadata.toml new file mode 100644 index 00000000..b7e0d9f8 --- /dev/null +++ b/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 "] + +[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"] diff --git a/examples/sir_epidemic/model.py b/examples/sir_epidemic/model.py new file mode 100644 index 00000000..3e67f2ff --- /dev/null +++ b/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)