Add measles ABM template support

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

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",

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,

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")

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."""

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()