diff --git a/examples/er_simulation/README.md b/examples/er_simulation/README.md new file mode 100644 index 00000000..af574225 --- /dev/null +++ b/examples/er_simulation/README.md @@ -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. diff --git a/examples/er_simulation/er_simulation/__init__.py b/examples/er_simulation/er_simulation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/er_simulation/er_simulation/agents.py b/examples/er_simulation/er_simulation/agents.py new file mode 100644 index 00000000..723377bf --- /dev/null +++ b/examples/er_simulation/er_simulation/agents.py @@ -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 diff --git a/examples/er_simulation/er_simulation/model.py b/examples/er_simulation/er_simulation/model.py new file mode 100644 index 00000000..7b347bf4 --- /dev/null +++ b/examples/er_simulation/er_simulation/model.py @@ -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() diff --git a/examples/er_simulation/requirements.txt b/examples/er_simulation/requirements.txt new file mode 100644 index 00000000..a8ac5c11 --- /dev/null +++ b/examples/er_simulation/requirements.txt @@ -0,0 +1,2 @@ +mesa +numpy diff --git a/examples/er_simulation/tests.py b/examples/er_simulation/tests.py new file mode 100644 index 00000000..8c62ae1f --- /dev/null +++ b/examples/er_simulation/tests.py @@ -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