From 9c4013a0915187e28ae3ecfd851ecd7e65efd20a Mon Sep 17 00:00:00 2001 From: Katherine Rosenfeld Date: Tue, 14 Apr 2026 14:01:39 -0700 Subject: [PATCH 1/5] 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/5] 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 From f7782a0740aea7a7065a107a753102bb20b585d2 Mon Sep 17 00:00:00 2001 From: Jonathan Bloedow Date: Fri, 31 Jul 2026 09:47:35 -0700 Subject: [PATCH 3/5] Fix CI: extend coverage-omit list to measles templates + noqa E741 Two CI failures on PR #24: - Lint (Python 3.12): E741 on measles_plot.py:60 (I = np.array(...)). Kept the SEIR-standard variable name with an inline noqa noting the domain convention. - Tests (3.10 + 3.14): coverage 89.36% < 90% threshold, because measles.py and measles_plot.py are template files (copied verbatim to user output dirs, not exercised by laser-init's own tests). Existing templates (si/sir/seir/plot.py) are already excluded from coverage; extended the same omit list to the two measles templates. --- pyproject.toml | 2 ++ src/laser/init/models/measles_plot.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d393c30..3c547f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,6 +176,8 @@ omit = [ "**/models/si.py", "**/models/sir.py", "**/models/seir.py", + "**/models/measles.py", + "**/models/measles_plot.py", ] [tool.coverage.report] diff --git a/src/laser/init/models/measles_plot.py b/src/laser/init/models/measles_plot.py index 728b2a0..d5155bc 100644 --- a/src/laser/init/models/measles_plot.py +++ b/src/laser/init/models/measles_plot.py @@ -57,7 +57,7 @@ def global_seir_fractions(model, scenario: pl.DataFrame) -> plt.Figure: S = np.array(global_tracker.S) / total_pop E = np.array(global_tracker.E) / total_pop - I = np.array(global_tracker.I) / total_pop + I = np.array(global_tracker.I) / total_pop # noqa: E741 — SEIR domain notation R = np.array(global_tracker.R) / total_pop fig, ax = plt.subplots(figsize=(10, 6)) From 404b67379b0f2ac0e573c1cadfc1d1675f69d843 Mon Sep 17 00:00:00 2001 From: Jonathan Bloedow Date: Fri, 31 Jul 2026 10:00:10 -0700 Subject: [PATCH 4/5] Apply ruff format to measles templates + loaders/tests --- src/laser/init/loaders/abm.py | 4 +++- src/laser/init/models/measles.py | 20 ++++++++++++-------- src/laser/init/models/measles_plot.py | 4 +++- tests/test_loaders.py | 4 +--- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/laser/init/loaders/abm.py b/src/laser/init/loaders/abm.py index b60ba32..f128b02 100644 --- a/src/laser/init/loaders/abm.py +++ b/src/laser/init/loaders/abm.py @@ -98,7 +98,9 @@ def emit_script( 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) + self._emit_generic( + model, shape_filename, cxr_filename, pop_filename, exp_filename, output_dir + ) return diff --git a/src/laser/init/models/measles.py b/src/laser/init/models/measles.py index f125de3..58ab8c5 100644 --- a/src/laser/init/models/measles.py +++ b/src/laser/init/models/measles.py @@ -10,7 +10,9 @@ 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") +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) @@ -60,13 +62,15 @@ def main(config_file: Path, data_dir: Path) -> None: # 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)), - }) + 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"] diff --git a/src/laser/init/models/measles_plot.py b/src/laser/init/models/measles_plot.py index d5155bc..25f5157 100644 --- a/src/laser/init/models/measles_plot.py +++ b/src/laser/init/models/measles_plot.py @@ -6,7 +6,9 @@ from matplotlib.backends.backend_pdf import PdfPages -def show_plots(model, scenario: pl.DataFrame, output_dir: Path | None, name: str = "measles") -> Path: +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 diff --git a/tests/test_loaders.py b/tests/test_loaders.py index e83d55f..6720369 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -301,9 +301,7 @@ def test_abm_emit_script_measles_model(self, tmp_path): ) 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 / "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): From e24559e2d4ce449a12fe451974ec2055ae384323 Mon Sep 17 00:00:00 2001 From: Jonathan Bloedow Date: Fri, 31 Jul 2026 10:24:28 -0700 Subject: [PATCH 5/5] Address Copilot review comments on PR #24 - measles.py: add path_type=Path to both click.Path options so main() receives Path objects, matching the pattern in seir.py/sir.py/si.py. - measles_plot.py: correct return type to Path | None; close figures in the no-output-dir path (moved plt.close outside the if/else so every figure is closed regardless of whether it was written to PDF). - loaders/abm.py: docstring now notes that pop_filename/exp_filename are accepted for signature compatibility but ignored for MEASLES. --- src/laser/init/loaders/abm.py | 8 +++++--- src/laser/init/models/measles.py | 4 ++-- src/laser/init/models/measles_plot.py | 5 +++-- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/laser/init/loaders/abm.py b/src/laser/init/loaders/abm.py index f128b02..f630c8a 100644 --- a/src/laser/init/loaders/abm.py +++ b/src/laser/init/loaders/abm.py @@ -78,15 +78,17 @@ def emit_script( Creates a YAML configuration file with data file paths and simulation parameters, 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. + laser-generic and only consumes shape and crude-rate data; ``pop_filename`` and + ``exp_filename`` are accepted for signature compatibility with the other model + types but are ignored for MEASLES. Args: mode: Model mode (must be "ABM"). 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. - exp_filename: Path to the CSV file with life expectancy data. + pop_filename: Path to the CSV file with age distribution. Ignored for MEASLES. + exp_filename: Path to the CSV file with life expectancy data. Ignored for MEASLES. output_dir: Directory where model script and config will be written. Raises: diff --git a/src/laser/init/models/measles.py b/src/laser/init/models/measles.py index 58ab8c5..de589c2 100644 --- a/src/laser/init/models/measles.py +++ b/src/laser/init/models/measles.py @@ -22,14 +22,14 @@ "-c", "--config", "config_file", - type=click.Path(exists=True), + type=click.Path(exists=True, path_type=Path), default=Path(__file__).parent / "config.yaml", help="Path to the configuration YAML file.", ) @click.option( "-d", "--data-dir", - type=click.Path(exists=True), + type=click.Path(exists=True, path_type=Path), default=None, help="Path to the data directory.", ) diff --git a/src/laser/init/models/measles_plot.py b/src/laser/init/models/measles_plot.py index 25f5157..79a66d4 100644 --- a/src/laser/init/models/measles_plot.py +++ b/src/laser/init/models/measles_plot.py @@ -8,7 +8,7 @@ def show_plots( model, scenario: pl.DataFrame, output_dir: Path | None, name: str = "measles" -) -> Path: +) -> Path | None: """Generate visualization plots for measles ABM model output. Creates plots for analyzing measles ABM simulation results including @@ -36,9 +36,10 @@ def show_plots( with PdfPages(pdf_path) as pdf: for fig in figs: pdf.savefig(fig) - plt.close(fig) else: pdf_path = None + for fig in figs: + plt.close(fig) return pdf_path