Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 161 additions & 0 deletions examples/er_simulation/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,161 @@
# Emergency Room Simulation

This example models a busy emergency room using Mesa's [Action system](https://github.com/mesa/mesa/pull/3461).

In plain language, patients show up over time, get triaged, and then wait for treatment. Some cases are minor and leave after triage. Others need treatment, and critical patients are always prioritized. If a doctor is already treating a moderate patient and a critical case arrives, that ongoing treatment can be interrupted and resumed later.

The goal of this example is not medical realism. It is a compact, easy-to-read model for exploring timed actions, interruption, resumption, and event-driven scheduling in Mesa.

## Why this example is interesting

This model is a nice fit for Mesa's Actions API because an emergency room is full of things that take time:

- triage takes time
- treatment takes time
- patients arrive at random times
- some work can be interrupted
- some work cannot

That makes it a good teaching example for discrete-event behavior inside an agent-based model.

## What the model does

Each run follows the same high-level flow:

1. Patients arrive at random intervals.
2. A free doctor triages the patient, or the patient waits in the triage queue.
3. Triage assigns one of three severities: `critical`, `moderate`, or `minor`.
4. Minor patients leave directly after triage.
5. Critical and moderate patients wait for treatment.
6. Critical patients are treated first.
7. Moderate treatments can be interrupted if a critical patient needs immediate attention.
8. Interrupted moderate treatments are resumed later.

## Mesa features demonstrated

| Feature | How it appears here |
|---------|----------------------|
| `Action` subclassing | `Triage` and `Treat` are both custom actions |
| Timed actions | Triage and treatment both have durations |
| Interruptible actions | Moderate treatment can be bumped by critical arrivals |
| Resumable actions | Interrupted treatment continues later instead of restarting |
| `schedule_event` | Patient arrivals are scheduled using an exponential distribution |
| Priority queues | Critical patients move ahead of moderate patients |
| Data collection | Queue lengths, throughput, interruptions, and wait time are tracked |

## Project structure

| File | Purpose |
|------|---------|
| `er_simulation/model.py` | Main ER model, queues, scheduling, and metrics |
| `er_simulation/agents.py` | Doctor and patient agents plus the `Triage` and `Treat` actions |
| `render_visuals.py` | Generates a PNG or GIF dashboard from a simulation run |
| `tests.py` | Basic checks for processing, waiting times, triage, and interruptions |

## Quick start

From inside this folder:

```bash
pip install -r requirements.txt
```

Then run a simple simulation:

```bash
python -c "
from er_simulation.model import ERModel
from er_simulation.agents import Patient, Doctor

model = ERModel(n_doctors=4, arrival_rate=1.5, seed=11)
for _ in range(40):
model.step()

patients = list(model.agents_by_type[Patient])
discharged = [p for p in patients if p.status == 'discharged']

print(f'Total patients: {len(patients)}')
print(f'Discharged: {len(discharged)}')
print(f'Triage queue: {model.triage_queue_length}')
print(f'Treatment queue: {model.treatment_queue_length}')
print(f'Interruptions: {model.total_interruptions}')
if discharged:
avg_wait = sum(p.wait_time for p in discharged) / len(discharged)
print(f'Avg wait time: {avg_wait:.1f}')
for doc in model.agents_by_type[Doctor]:
print(f'Doctor {doc.unique_id}: {doc.patients_treated} treated')
"
```

## How to make a visual

If you want something you can attach to a PR, use the renderer:

```bash
python render_visuals.py --steps 40 --n-doctors 4 --arrival-rate 1.5 --seed 11 --output er_snapshot.png
python render_visuals.py --steps 40 --n-doctors 4 --arrival-rate 1.5 --seed 11 --output er_animation.gif --fps 4
```

This produces:

- a static dashboard image showing the final state and recent trends
- an animated GIF showing how queues, throughput, wait time, and doctor activity evolve over time

## Parameters you can play with

These are the main knobs worth changing:

| Parameter | Meaning |
|-----------|---------|
| `n_doctors` | Number of doctors available for triage and treatment |
| `arrival_rate` | How quickly patients arrive on average |
| `seed` | Random seed for reproducible runs |
| `steps` | How long the simulation runs |

Good starter combinations:

- `n_doctors=2, arrival_rate=3.0` for heavy overload
- `n_doctors=4, arrival_rate=1.5` for a balanced run with visible interruptions
- `n_doctors=6, arrival_rate=1.0` for a better-staffed system

## What to look for

Here are a few behaviors newcomers usually find interesting:

- If you keep staffing fixed and raise `arrival_rate`, the triage queue grows fast.
- If you keep `arrival_rate` fixed and add doctors, more patients make it through treatment and the system feels less congested.
- Minor patients help reduce pressure because they leave directly after triage.
- Interruptions are easiest to notice under moderate pressure, where doctors have enough time to begin moderate treatments before a critical case arrives.
- Under extreme overload, the system can become so busy triaging incoming patients that interruptions actually become less visible than you might expect.

## Questions worth exploring

If you are experimenting with the model, these are good prompts:

- When does the main bottleneck shift from triage to treatment?
- How many doctors are needed before queues stop growing rapidly?
- Do interruptions clearly help critical patients, or do they mostly delay moderate ones?
- Which parameter change helps more: lowering arrivals or adding one more doctor?

## Example output

One run with `n_doctors=4`, `arrival_rate=1.5`, `seed=11`, and `40` steps produced:

```text
Total patients: 63
Discharged: 22
Triage queue: 26
Treatment queue: 11
Interruptions: 3
Avg wait time: 16.6
Doctor 1: 2 treated
Doctor 2: 3 treated
Doctor 3: 1 treated
Doctor 4: 2 treated
```

## Notes

- This is a teaching model, not a hospital operations model.
- Severity assignment and service durations are simplified on purpose.
- The same doctors perform both triage and treatment, which creates useful queueing behavior for demonstration.
Empty file.
88 changes: 88 additions & 0 deletions examples/er_simulation/er_simulation/agents.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
"""Actions and agents for the ER simulation."""

import mesa
from mesa.experimental.actions import Action


class Triage(Action):
"""Assess a patient's severity and route them accordingly."""

def __init__(self, nurse, patient):
dur = max(0.5, nurse.model.rng.normal(1.5, 0.3))
super().__init__(nurse, duration=dur)
self.patient = patient

def on_complete(self):
severity = self.agent.model.rng.choice(
["critical", "moderate", "minor"], p=[0.15, 0.45, 0.4]
)
self.patient.severity = severity
self.patient.status = "triaged"

if severity == "minor":
self.patient.status = "discharged"
self.patient.discharge_time = self.agent.model.time
return

self.agent.model.enqueue_patient(self.patient)


class Treat(Action):
"""Treat a patient; duration depends on severity."""

DURATIONS = {"critical": 12.0, "moderate": 6.0, "minor": 2.0}

def __init__(self, doctor, patient):
base = self.DURATIONS.get(patient.severity, 4.0)
dur = max(1.0, doctor.model.rng.normal(base, base * 0.2))
super().__init__(
doctor,
duration=dur,
interruptible=(patient.severity != "critical"),
)
self.patient = patient

def on_start(self):
self.agent.current_patient = self.patient
self.patient.status = "being_treated"

def on_complete(self):
self.patient.status = "discharged"
self.patient.discharge_time = self.agent.model.time
self.agent.patients_treated += 1
self.agent.current_patient = None
self.agent.model.try_assign(self.agent)

def on_interrupt(self, progress):
self.patient.status = "waiting"
self.agent.current_patient = None
self.agent.model.enqueue_patient(self.patient)

def on_resume(self):
self.agent.current_patient = self.patient
self.patient.status = "being_treated"


class Patient(mesa.Agent):
"""A patient arriving in the ER."""

def __init__(self, model, arrival_time):
super().__init__(model)
self.arrival_time = arrival_time
self.discharge_time = None
self.severity = None
self.status = "arrived"

@property
def wait_time(self):
end = self.discharge_time if self.discharge_time else self.model.time
return end - self.arrival_time


class Doctor(mesa.Agent):
"""A doctor who treats patients."""

def __init__(self, model):
super().__init__(model)
self.patients_treated = 0
self.current_patient = None
65 changes: 65 additions & 0 deletions examples/er_simulation/er_simulation/model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
"""Emergency room simulation using Mesa's Action system."""

import mesa
import numpy as np

from .agents import Doctor, Patient, Treat, Triage


class ERModel(mesa.Model):
"""Discrete-event ER simulation with stochastic patient arrivals."""

def __init__(self, n_doctors=3, arrival_rate=1.5, seed=42):
super().__init__()
self.rng = np.random.default_rng(seed)
self.arrival_rate = arrival_rate
self._queue = []

for _ in range(n_doctors):
Doctor(self)

self._schedule_arrival()

def _schedule_arrival(self):
gap = self.rng.exponential(1.0 / self.arrival_rate)
self.schedule_event(self._on_arrival, at=self.time + gap)

def _on_arrival(self):
patient = Patient(self, arrival_time=self.time)

nurse = self._get_free_doctor()
if nurse:
nurse.start_action(Triage(nurse, patient))
else:
self._queue.append(patient)

self._schedule_arrival()

def enqueue_patient(self, patient):
self._queue.append(patient)
self._queue.sort(
key=lambda p: (
{"critical": 0, "moderate": 1, "minor": 2}.get(p.severity, 3),
p.arrival_time,
)
)
self._assign_waiting()

def try_assign(self, doctor):
if self._queue and not doctor.is_busy:
patient = self._queue.pop(0)
doctor.start_action(Treat(doctor, patient))

def _assign_waiting(self):
for doc in self.agents_by_type[Doctor]:
if not doc.is_busy and self._queue:
patient = self._queue.pop(0)
doc.start_action(Treat(doc, patient))

def _get_free_doctor(self):
free = [d for d in self.agents_by_type[Doctor] if not d.is_busy]
return self.rng.choice(free) if free else None

def step(self):
self.run_for(1.0)
self._assign_waiting()
2 changes: 2 additions & 0 deletions examples/er_simulation/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
mesa
numpy
33 changes: 33 additions & 0 deletions examples/er_simulation/tests.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
from er_simulation.agents import Doctor, Patient
from er_simulation.model import ERModel


def test_patients_are_processed():
model = ERModel(n_doctors=3, arrival_rate=1.5, seed=42)
for _ in range(20):
model.step()

patients = list(model.agents_by_type[Patient])
assert len(patients) > 0

discharged = [p for p in patients if p.status == "discharged"]
assert len(discharged) > 0


def test_all_doctors_work():
model = ERModel(n_doctors=3, arrival_rate=2.0, seed=42)
for _ in range(30):
model.step()

for doc in model.agents_by_type[Doctor]:
assert doc.patients_treated > 0


def test_wait_times_are_positive():
model = ERModel(n_doctors=2, arrival_rate=1.5, seed=42)
for _ in range(25):
model.step()

discharged = [p for p in model.agents_by_type[Patient] if p.status == "discharged"]
for p in discharged:
assert p.wait_time >= 0
Loading