From 9c4013a0915187e28ae3ecfd851ecd7e65efd20a Mon Sep 17 00:00:00 2001 From: Katherine Rosenfeld Date: Tue, 14 Apr 2026 14:01:39 -0700 Subject: [PATCH 1/2] initial laser-measles template --- src/laser/init/cli.py | 7 +- src/laser/init/loaders/abm.py | 83 +++++++++++++++- src/laser/init/models/__init__.py | 2 +- src/laser/init/models/measles.py | 129 +++++++++++++++++++++++++ src/laser/init/models/measles_plot.py | 133 ++++++++++++++++++++++++++ 5 files changed, 346 insertions(+), 8 deletions(-) create mode 100644 src/laser/init/models/measles.py create mode 100644 src/laser/init/models/measles_plot.py diff --git a/src/laser/init/cli.py b/src/laser/init/cli.py index 6656aae..48be85d 100644 --- a/src/laser/init/cli.py +++ b/src/laser/init/cli.py @@ -56,7 +56,7 @@ ) @click.option( "--model", - type=click.Choice(["SI", "SIR", "SEIR"], case_sensitive=False), + type=click.Choice(["SI", "SIR", "SEIR", "MEASLES"], case_sensitive=False), default="SEIR", help="Select the type of epidemiological model to prepare data for (default: SEIR)", ) @@ -106,7 +106,7 @@ def cli( end_year: End year for simulation (1950-2100, must be >= start_year). output_dir: Output directory path. If None, defaults to "./ISOCODE/start_year". mode: Modeling mode, either "ABM" (agent-based model) or "MPM" (metapopulation model). - model: Epidemiological model type - "SI", "SIR", or "SEIR". + model: Epidemiological model type - "SI", "SIR", "SEIR", or "MEASLES". shape_source: Administrative boundary data source - "unocha", "geoboundaries", or "gadm". If None, uses config value or defaults to "unocha". raster_source: Population raster data source - currently only "worldpop" supported. @@ -436,7 +436,7 @@ def emit_model_script( Args: mode: Model mode ("ABM" or "MPM"). - model: Model type ("SI", "SIR", or "SEIR"). + model: Model type ("SI", "SIR", "SEIR", or "MEASLES"). shapes_filename: Path to the administrative boundaries GeoPackage. cxr_filename: Path to the crude birth/death rate CSV. pop_filename: Path to the age distribution CSV. @@ -458,6 +458,7 @@ def emit_model_script( "ABM/SI": abm.AbmLoader, "ABM/SIR": abm.AbmLoader, "ABM/SEIR": abm.AbmLoader, + "ABM/MEASLES": abm.AbmLoader, "MPM/SI": mpm.MpmLoader, "MPM/SIR": mpm.MpmLoader, "MPM/SEIR": mpm.MpmLoader, diff --git a/src/laser/init/loaders/abm.py b/src/laser/init/loaders/abm.py index 782b308..b60ba32 100644 --- a/src/laser/init/loaders/abm.py +++ b/src/laser/init/loaders/abm.py @@ -1,6 +1,8 @@ import shutil from pathlib import Path +from laser.init.logger import logger + __yaml__ = """ data_dir: %%data_dir%% @@ -23,6 +25,25 @@ naive_population: true """ +__measles_yaml__ = """ +data_dir: %%data_dir%% + +datafiles: + shape_data: %%shape_data%% + cxr_data: %%cxr_data%% + +simulation: + nyears: 2 + seed: 42 + start_time: "%%start_time%%" + beta: 20.0 + seasonality: 0.0 + distance_exponent: 2.0 + mixing_scale: 0.01 + initial_infections: 50 + naive_population: true +""" + class AbmLoader: def __init__(self) -> None: @@ -55,12 +76,13 @@ def emit_script( """Generate ABM model script and configuration files. Creates a YAML configuration file with data file paths and simulation parameters, - then copies the appropriate model script (SI, SIR, or SEIR) and plotting utilities - to the output directory. + then copies the appropriate model script (SI, SIR, SEIR, or MEASLES) and plotting + utilities to the output directory. The MEASLES model uses laser-measles instead of + laser-generic and requires only shape and crude rate data files. Args: mode: Model mode (must be "ABM"). - model: Model type ("SI", "SIR", or "SEIR"). + model: Model type ("SI", "SIR", "SEIR", or "MEASLES"). shape_filename: Path to the GeoPackage file with administrative boundaries. cxr_filename: Path to the CSV file with crude birth/death rates. pop_filename: Path to the CSV file with age distribution. @@ -73,6 +95,34 @@ def emit_script( assert mode.upper() == "ABM", f"AbmLoader only supports ABM mode, got {mode}" + if model.upper() == "MEASLES": + self._emit_measles(shape_filename, cxr_filename, output_dir) + else: + self._emit_generic(model, shape_filename, cxr_filename, pop_filename, exp_filename, output_dir) + + return + + def _emit_generic( + self, + model: str, + shape_filename: Path, + cxr_filename: Path, + pop_filename: Path, + exp_filename: Path, + output_dir: Path, + ) -> None: + """Generate a laser-generic model script and configuration. + + Args: + model: Model type ("SI", "SIR", or "SEIR"). + shape_filename: Path to the GeoPackage file with administrative boundaries. + cxr_filename: Path to the CSV file with crude birth/death rates. + pop_filename: Path to the CSV file with age distribution. + exp_filename: Path to the CSV file with life expectancy data. + output_dir: Directory where model script and config will be written. + """ + logger.info("Emitting laser-generic %s model script to %s", model, output_dir) + yaml = __yaml__.replace("%%data_dir%%", str(output_dir.absolute())) yaml = yaml.replace("%%shape_data%%", str(shape_filename.name)) yaml = yaml.replace("%%cxr_data%%", str(cxr_filename.name)) @@ -84,4 +134,29 @@ def emit_script( shutil.copy2(source_dir / f"{model.lower()}.py", Path(output_dir) / f"{model.lower()}.py") shutil.copy2(source_dir / "plot.py", Path(output_dir) / "plot.py") - return + def _emit_measles( + self, + shape_filename: Path, + cxr_filename: Path, + output_dir: Path, + start_time: str = "2000-01", + ) -> None: + """Generate a laser-measles ABM model script and configuration. + + Args: + shape_filename: Path to the GeoPackage file with administrative boundaries. + cxr_filename: Path to the CSV file with crude birth/death rates. + output_dir: Directory where model script and config will be written. + start_time: Simulation start time in YYYY-MM format. + """ + logger.info("Emitting laser-measles ABM model script to %s", output_dir) + + yaml = __measles_yaml__.replace("%%data_dir%%", str(output_dir.absolute())) + yaml = yaml.replace("%%shape_data%%", str(shape_filename.name)) + yaml = yaml.replace("%%cxr_data%%", str(cxr_filename.name)) + yaml = yaml.replace("%%start_time%%", start_time) + (Path(output_dir) / "config.yaml").write_text(yaml) + + source_dir = Path(__file__).parent.parent / "models" + shutil.copy2(source_dir / "measles.py", Path(output_dir) / "measles.py") + shutil.copy2(source_dir / "measles_plot.py", Path(output_dir) / "measles_plot.py") diff --git a/src/laser/init/models/__init__.py b/src/laser/init/models/__init__.py index ef0a257..283381e 100644 --- a/src/laser/init/models/__init__.py +++ b/src/laser/init/models/__init__.py @@ -1 +1 @@ -"""Epidemiological model templates (SI, SIR, SEIR) and plotting utilities.""" +"""Epidemiological model templates (SI, SIR, SEIR, Measles) and plotting utilities.""" diff --git a/src/laser/init/models/measles.py b/src/laser/init/models/measles.py new file mode 100644 index 0000000..f125de3 --- /dev/null +++ b/src/laser/init/models/measles.py @@ -0,0 +1,129 @@ +import importlib.util +from pathlib import Path + +import click +import geopandas as gpd +import numpy as np +import pandas as pd +import polars as pl +import yaml +from laser.measles.abm import ABMModel, ABMParams, components +from laser.measles.components import create_component + +spec = importlib.util.spec_from_file_location("module_name", Path(__file__).parent / "measles_plot.py") +plot = importlib.util.module_from_spec(spec) +spec.loader.exec_module(plot) + + +@click.command() +@click.option( + "-c", + "--config", + "config_file", + type=click.Path(exists=True), + default=Path(__file__).parent / "config.yaml", + help="Path to the configuration YAML file.", +) +@click.option( + "-d", + "--data-dir", + type=click.Path(exists=True), + default=None, + help="Path to the data directory.", +) +def main(config_file: Path, data_dir: Path) -> None: + """Run a measles ABM (Agent-Based Model) epidemiological simulation. + + Loads configuration and data files, constructs a spatial scenario from + administrative boundaries, sets up the ABM with vital dynamics, disease + transmission, and state tracking components, then runs the simulation + and generates output plots. + + Args: + config_file: Path to the YAML configuration file. + data_dir: Path to the data directory, or None to use config value. + + Returns: + None + + Raises: + click.exceptions.ClickException: If config_file or data_dir paths are invalid. + KeyError: If required configuration keys are missing. + FileNotFoundError: If data files specified in config cannot be found. + """ + config = yaml.safe_load(Path(config_file).read_text()) + + data_dir = Path(data_dir or config["data_dir"]) + datafiles = config["datafiles"] + gdf = gpd.read_file(data_dir / datafiles["shape_data"]) + cxr_df = pd.read_csv(data_dir / datafiles["cxr_data"]) + + # Build the scenario Polars DataFrame from the GeoPackage + centroids = gdf.geometry.centroid + scenario = pl.DataFrame({ + "id": [f"patch_{i}" for i in range(len(gdf))], + "lat": centroids.y.to_numpy(), + "lon": centroids.x.to_numpy(), + "pop": gdf["population"].to_numpy().astype(np.int64), + "mcv1": np.zeros(len(gdf)), + }) + + sim = config["simulation"] + + # Configure model parameters + params = ABMParams( + num_ticks=sim["nyears"] * 365, + seed=sim.get("seed", 42), + start_time=sim.get("start_time", "2000-01"), + ) + + model = ABMModel(scenario=scenario, params=params) + + # Vital dynamics + cbr = cxr_df["CBR"].iloc[0] if "CBR" in cxr_df.columns else 30.0 + cdr = cxr_df["CDR"].iloc[0] if "CDR" in cxr_df.columns else 10.0 + vd_params = components.VitalDynamicsParams( + crude_birth_rate=float(cbr), + crude_death_rate=float(cdr), + ) + + # Infection seeding — seed the largest population patch + largest_patch = scenario.sort("pop", descending=True)["id"][0] + seeding_params = components.InfectionSeedingParams( + target_patches=[largest_patch], + infections_per_patch=sim.get("initial_infections", 50), + ) + + # Infection process + infection_params = components.InfectionParams( + beta=sim.get("beta", 20.0), + seasonality=sim.get("seasonality", 0.0), + distance_exponent=sim.get("distance_exponent", 2.0), + mixing_scale=sim.get("mixing_scale", 0.01), + ) + + # Initialize equilibrium states if not a naive population + if not sim.get("naive_population", True): + model.add_component(components.InitializeEquilibriumStatesProcess) + + # Assemble components + model.components = [ + create_component(components.VitalDynamicsProcess, vd_params), + create_component(components.InfectionSeedingProcess, seeding_params), + create_component(components.InfectionProcess, infection_params), + components.StateTracker, + create_component( + components.StateTracker, + components.StateTrackerParams(aggregation_level=0), + ), + ] + + model.run() + + plot.show_plots(model, scenario, output_dir=Path(__file__).parent, name="measles") + + return + + +if __name__ == "__main__": + main() diff --git a/src/laser/init/models/measles_plot.py b/src/laser/init/models/measles_plot.py new file mode 100644 index 0000000..728b2a0 --- /dev/null +++ b/src/laser/init/models/measles_plot.py @@ -0,0 +1,133 @@ +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import polars as pl +from matplotlib.backends.backend_pdf import PdfPages + + +def show_plots(model, scenario: pl.DataFrame, output_dir: Path | None, name: str = "measles") -> Path: + """Generate visualization plots for measles ABM model output. + + Creates plots for analyzing measles ABM simulation results including + global SEIR dynamics, spatial attack rates, and infectious spread. + If output_dir is provided, saves all plots to a single PDF file. + + Args: + model: laser-measles ABM model instance with completed simulation results. + scenario: Polars DataFrame with scenario data (id, lat, lon, pop, mcv1). + output_dir: Directory where output PDF will be saved, or None to skip saving. + name: Base name for the output PDF file. + + Returns: + Path to the saved PDF file, or None if output_dir is None. + """ + + plots = [ + global_seir_fractions, + spatial_attack_rate, + infectious_over_time, + ] + figs = [plot_func(model, scenario) for plot_func in plots] + if output_dir: + pdf_path = Path(output_dir) / f"{name}_output.pdf" + with PdfPages(pdf_path) as pdf: + for fig in figs: + pdf.savefig(fig) + plt.close(fig) + else: + pdf_path = None + + return pdf_path + + +def global_seir_fractions(model, scenario: pl.DataFrame) -> plt.Figure: + """Plot global SEIR fractions over time. + + Args: + model: laser-measles ABM model instance with completed simulation. + scenario: Polars DataFrame with scenario data. + + Returns: + Matplotlib Figure with stacked SEIR fraction plot. + """ + global_tracker = model.get_instance("StateTracker")[0] + total_pop = scenario["pop"].sum() + ticks = np.arange(model.params.num_ticks) + + S = np.array(global_tracker.S) / total_pop + E = np.array(global_tracker.E) / total_pop + I = np.array(global_tracker.I) / total_pop + R = np.array(global_tracker.R) / total_pop + + fig, ax = plt.subplots(figsize=(10, 6)) + ax.stackplot(ticks, S, E, I, R, labels=["S", "E", "I", "R"], alpha=0.8) + ax.set_xlabel("Day") + ax.set_ylabel("Fraction of Population") + ax.set_title("Global SEIR Fractions Over Time") + ax.legend(loc="center right") + ax.set_xlim(0, len(ticks) - 1) + ax.set_ylim(0, 1) + fig.tight_layout() + + return fig + + +def spatial_attack_rate(model, scenario: pl.DataFrame) -> plt.Figure: + """Plot spatial attack rate by patch. + + Args: + model: laser-measles ABM model instance with completed simulation. + scenario: Polars DataFrame with scenario data. + + Returns: + Matplotlib Figure with attack rate scatter plot. + """ + patch_tracker = model.get_instance("StateTracker")[1] + pops = scenario["pop"].to_numpy() + + # Cumulative recovered at end as proxy for attack rate + final_R = np.array(patch_tracker.R)[-1] + attack_rate = final_R / pops + + fig, ax = plt.subplots(figsize=(10, 6)) + scatter = ax.scatter( + scenario["lon"].to_numpy(), + scenario["lat"].to_numpy(), + c=attack_rate, + s=pops / pops.max() * 200, + cmap="YlOrRd", + alpha=0.7, + edgecolors="black", + ) + plt.colorbar(scatter, ax=ax, label="Attack Rate") + ax.set_xlabel("Longitude") + ax.set_ylabel("Latitude") + ax.set_title("Spatial Attack Rate") + fig.tight_layout() + + return fig + + +def infectious_over_time(model, scenario: pl.DataFrame) -> plt.Figure: + """Plot infectious count per patch over time as a heatmap. + + Args: + model: laser-measles ABM model instance with completed simulation. + scenario: Polars DataFrame with scenario data. + + Returns: + Matplotlib Figure with infectious heatmap. + """ + patch_tracker = model.get_instance("StateTracker")[1] + I_data = np.array(patch_tracker.I) + + fig, ax = plt.subplots(figsize=(12, 6)) + im = ax.imshow(I_data.T, aspect="auto", cmap="hot", interpolation="nearest") + plt.colorbar(im, ax=ax, label="Infectious Count") + ax.set_xlabel("Day") + ax.set_ylabel("Patch Index") + ax.set_title("Infectious Agents per Patch Over Time") + fig.tight_layout() + + return fig From 7ab090df062af6f8187bb0daa21fd0b4a4a0771c Mon Sep 17 00:00:00 2001 From: Katherine Rosenfeld Date: Wed, 15 Apr 2026 11:54:22 -0700 Subject: [PATCH 2/2] Add measles ABM support dependencies and tests --- CHANGELOG.md | 7 +++ pyproject.toml | 2 + tests/test_loaders.py | 120 ++++++++++++++++++++++++++++++++++++++++++ tests/test_models.py | 27 ++++++++-- 4 files changed, 153 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index edf0e0f..a4f19d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file. ## [Unreleased] ### Added +- Measles ABM template support across the project + - Added a measles-specific ABM template and plotting helpers using `laser-measles` + - Added `MEASLES` CLI and loader support for emitting `measles.py`, `measles_plot.py`, and a measles-specific `config.yaml` + - Added tests covering measles model module imports and ABM loader emission for the `MEASLES` model +- Declared new runtime dependencies for the measles workflow in `pyproject.toml` + - Added `laser-measles>=0.10.0` + - Added `polars>=1.0.0` - Comprehensive documentation overhaul - Updated pyproject.toml with proper package description - Completely rewrote README.md with installation instructions, prerequisites, troubleshooting, advanced usage, and comprehensive examples diff --git a/pyproject.toml b/pyproject.toml index 5ace62e..eff4b8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,9 +22,11 @@ dependencies = [ "click>=8.3.1", "geopandas>=1.1.2", "laser-generic>=1.0.1", + "laser-measles>=0.10.0", "matplotlib>=3.10.8", "openai>=2.24.0", "pandas>=2.3.3", + "polars>=1.0.0", "pycountry>=26.2.16", "pyogrio>=0.12.1", "pyyaml>=6.0.3", diff --git a/tests/test_loaders.py b/tests/test_loaders.py index 29092d5..6c9da28 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -269,6 +269,126 @@ def test_abm_config_file_contains_paths(self, tmp_path): assert config is not None assert isinstance(config, dict) + def test_abm_emit_script_measles_model(self, tmp_path): + """Test that emit_script works for MEASLES model. + + Given MEASLES model type + When emit_script() is called + Then it should create measles.py, measles_plot.py, and config.yaml + + Failure indicates MEASLES model generation is broken. + """ + loader = abm.AbmLoader() + output_dir = tmp_path / "output" + output_dir.mkdir() + + # Create dummy input files + shape_file = tmp_path / "shapes.gpkg" + cxr_file = tmp_path / "cxr.csv" + pop_file = tmp_path / "pop.csv" + exp_file = tmp_path / "exp.csv" + + for f in [shape_file, cxr_file, pop_file, exp_file]: + f.touch() + + loader.emit_script( + mode="ABM", + model="MEASLES", + shape_filename=shape_file, + cxr_filename=cxr_file, + pop_filename=pop_file, + exp_filename=exp_file, + output_dir=output_dir, + ) + + assert (output_dir / "measles.py").exists(), "measles.py should be created" + assert ( + output_dir / "measles_plot.py" + ).exists(), "measles_plot.py should be created" + assert (output_dir / "config.yaml").exists(), "config.yaml should be created" + + def test_abm_measles_config_has_correct_keys(self, tmp_path): + """Test that measles config.yaml contains expected keys. + + Given MEASLES model type + When emit_script() generates config.yaml + Then it should contain measles-specific simulation parameters + + Failure indicates measles config template is malformed. + """ + import yaml + + loader = abm.AbmLoader() + output_dir = tmp_path / "output" + output_dir.mkdir() + + shape_file = tmp_path / "shapes.gpkg" + cxr_file = tmp_path / "cxr.csv" + pop_file = tmp_path / "pop.csv" + exp_file = tmp_path / "exp.csv" + + for f in [shape_file, cxr_file, pop_file, exp_file]: + f.touch() + + loader.emit_script( + mode="ABM", + model="MEASLES", + shape_filename=shape_file, + cxr_filename=cxr_file, + pop_filename=pop_file, + exp_filename=exp_file, + output_dir=output_dir, + ) + + config = yaml.safe_load((output_dir / "config.yaml").read_text()) + assert "data_dir" in config + assert "datafiles" in config + assert "shape_data" in config["datafiles"] + assert "cxr_data" in config["datafiles"] + sim = config["simulation"] + assert "beta" in sim + assert "seasonality" in sim + assert "distance_exponent" in sim + assert "mixing_scale" in sim + assert "initial_infections" in sim + assert "naive_population" in sim + + def test_abm_measles_does_not_create_generic_files(self, tmp_path): + """Test that MEASLES model does not create generic model files. + + Given MEASLES model type + When emit_script() is called + Then it should NOT create plot.py or any generic model scripts + + Failure indicates MEASLES routing is leaking to generic path. + """ + loader = abm.AbmLoader() + output_dir = tmp_path / "output" + output_dir.mkdir() + + shape_file = tmp_path / "shapes.gpkg" + cxr_file = tmp_path / "cxr.csv" + pop_file = tmp_path / "pop.csv" + exp_file = tmp_path / "exp.csv" + + for f in [shape_file, cxr_file, pop_file, exp_file]: + f.touch() + + loader.emit_script( + mode="ABM", + model="MEASLES", + shape_filename=shape_file, + cxr_filename=cxr_file, + pop_filename=pop_file, + exp_filename=exp_file, + output_dir=output_dir, + ) + + assert not (output_dir / "plot.py").exists(), ( + "generic plot.py should not be created for MEASLES" + ) + assert not (output_dir / "seir.py").exists(), "seir.py should not be created for MEASLES" + class TestMpmLoader: """Test suite for MPM loader.""" diff --git a/tests/test_models.py b/tests/test_models.py index d2f3b51..665d92f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,11 +1,10 @@ """Tests for laser.init.models modules. -This module tests model template generation for SI, SIR, and SEIR +This module tests model template generation for SI, SIR, SEIR, and Measles epidemiological models. """ -import pytest -from laser.init.models import si, sir, seir, plot +from laser.init.models import si, sir, seir, plot, measles, measles_plot class TestModelModulesExist: @@ -54,3 +53,25 @@ def test_plot_module_exists(self): Failure indicates plot module has been removed or renamed. """ assert plot is not None + + def test_measles_model_module_exists(self): + """Test that measles model module exists. + + Given the models package + When checking for measles module + Then it should be available + + Failure indicates measles model has been removed or renamed. + """ + assert measles is not None + + def test_measles_plot_module_exists(self): + """Test that measles_plot utilities module exists. + + Given the models package + When checking for measles_plot module + Then it should be available + + Failure indicates measles_plot module has been removed or renamed. + """ + assert measles_plot is not None