Mesa Wrapper

Thin Mesa wrapper around the existing Simulation for SolaraViz.

This module defines the HospitalABM class, which serves as a wrapper around the existing AMR Hub simulation to enable integration with SolaraViz for visualization purposes. The functionality mirrors the existing simulation, defined in run.py, but is adapted to fit the Mesa framework for easier visualization.

HospitalABM

Bases: Model

Mesa model wrapper around the AMR Hub simulation for visualization.

Source code in src/amr_hub_abm/mesa_wrapper.py

class HospitalABM(mesa.Model):
    """Mesa model wrapper around the AMR Hub simulation for visualization."""

    def __init__(
        self,
        config_path: str = "tests/inputs/simulation_config.yml",
    ) -> None:
        """Initialize the model and seed infections for demo purposes."""
        super().__init__()

        self.config_path = config_path

        self.simulation = create_simulation(sim_config)

        # seed infections for visualization demo
        self.simulation.agents[0].infection_status = InfectionStatus.INFECTED
        self.simulation.agents[1].infection_status = InfectionStatus.EXPOSED

    def create_new_simulation(self) -> None:
        """Create a new simulation."""
        self.simulation = create_simulation(sim_config)

    def step(self) -> None:
        """Advance the wrapped simulation by one time step."""
        if self.simulation.time >= self.simulation.total_simulation_time:
            self.create_new_simulation()
            return

        self.simulation.step(record=True)

__init__(config_path='tests/inputs/simulation_config.yml')

Initialize the model and seed infections for demo purposes.

Source code in src/amr_hub_abm/mesa_wrapper.py

def __init__(
    self,
    config_path: str = "tests/inputs/simulation_config.yml",
) -> None:
    """Initialize the model and seed infections for demo purposes."""
    super().__init__()

    self.config_path = config_path

    self.simulation = create_simulation(sim_config)

    # seed infections for visualization demo
    self.simulation.agents[0].infection_status = InfectionStatus.INFECTED
    self.simulation.agents[1].infection_status = InfectionStatus.EXPOSED

create_new_simulation()

Create a new simulation.

Source code in src/amr_hub_abm/mesa_wrapper.py

def create_new_simulation(self) -> None:
    """Create a new simulation."""
    self.simulation = create_simulation(sim_config)

step()

Advance the wrapped simulation by one time step.

Source code in src/amr_hub_abm/mesa_wrapper.py

def step(self) -> None:
    """Advance the wrapped simulation by one time step."""
    if self.simulation.time >= self.simulation.total_simulation_time:
        self.create_new_simulation()
        return

    self.simulation.step(record=True)