diff --git a/src/patchsim/cli.py b/src/patchsim/cli.py index 8ba1268..75b022d 100644 --- a/src/patchsim/cli.py +++ b/src/patchsim/cli.py @@ -45,11 +45,13 @@ def _emit_json(payload: dict[str, Any]) -> None: def _cmd_run(config_path: str, *, json_output: bool = False) -> dict[str, Any]: - logging.info("Starting PatchSim simulation...") + if not json_output: + print("Starting PatchSim simulation...") config = load_config(config_path) net, y0, patches, num_patches = setup_simulation(config) summary = run_simulation(config, config["ModelName"], net, y0, patches, num_patches) - logging.info("Simulation completed successfully.") + if not json_output: + print("Simulation completed successfully.") return {"ok": True, "config": config_path, **summary} if json_output else summary @@ -59,7 +61,8 @@ def _cmd_validate(config_path: str, *, json_output: bool = False, schema: bool = config = load_config(config_path) _net, _y0, patches, num_patches = setup_simulation(config) - logging.info("Configuration is valid: %s", config_path) + if not json_output: + print(f"Configuration is valid: {config_path}") if json_output: return { "ok": True, @@ -82,9 +85,33 @@ def _copy_template_tree(template_node, target_path: Path) -> None: target_path.write_bytes(template_node.read_bytes()) -def _render_template_config(project_name: str, template_name: str) -> str: - config = get_init_template_config(template_name, project_name) - return yaml.safe_dump(config, sort_keys=False) +def _write_seed_for_template(project_dir: Path, config: dict[str, Any]) -> None: + """Write a seed CSV whose columns match the template's compartments. + + The scaffold ships a single patch-population file; seed every patch fully + susceptible and place one infectious individual in the first patch. + """ + import pandas as pd + + compartments = list(config.get("compartments") or ["S", "I", "R"]) + patch_df = pd.read_csv(project_dir / config["PatchFile"]) + patch_col = next(c for c in patch_df.columns if c.lower() == "patch") + pop_col = next(c for c in patch_df.columns if c.lower() == "population") + + susceptible = "S" if "S" in compartments else compartments[0] + infectious = "I" if "I" in compartments else compartments[-1] + + rows = [] + for idx, record in patch_df.iterrows(): + seeded = 1 if idx == 0 else 0 + row = {"patch": record[patch_col], **{c: 0 for c in compartments}} + row[infectious] = seeded + row[susceptible] = int(record[pop_col]) - seeded + rows.append(row) + + seed_path = project_dir / config["SeedFile"] + seed_path.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame(rows)[["patch", *compartments]].to_csv(seed_path, index=False) def _cmd_init(name: str, force: bool = False, template: str = "sir") -> None: @@ -114,10 +141,12 @@ def _cmd_init(name: str, force: bool = False, template: str = "sir") -> None: template_root = resources.files("patchsim").joinpath("templates", "project") _copy_template_tree(template_root, project_dir) + template_config = get_init_template_config(template, project_dir.name) config_path = project_dir / "config.yaml" - config_path.write_text(_render_template_config(project_dir.name, template), encoding="utf-8") + config_path.write_text(yaml.safe_dump(template_config, sort_keys=False), encoding="utf-8") + _write_seed_for_template(project_dir, template_config) - logging.info("Created project scaffold at: %s", project_dir) + print(f"Created project scaffold at: {project_dir}") def _list_builtin_models() -> list[dict[str, str]]: @@ -127,15 +156,16 @@ def _list_builtin_models() -> list[dict[str, str]]: def _cmd_list_models(*, json_output: bool = False) -> list[dict[str, str]]: models = _list_builtin_models() if not models: - logging.info("No built-in models found.") + if not json_output: + print("No built-in models found.") return [] if json_output: return models - logging.info("Built-in models and templates:") + print("Built-in models and templates:") for model in models: - logging.info("- %s (%s)", model["name"], model["kind"]) + print(f"- {model['name']} ({model['kind']})") return models diff --git a/src/patchsim/core/model_runner.py b/src/patchsim/core/model_runner.py index e66febd..4cf1dc5 100644 --- a/src/patchsim/core/model_runner.py +++ b/src/patchsim/core/model_runner.py @@ -59,4 +59,4 @@ def solve(self, y0, t_range): def visualize(self, t, results, patches, outdir, model_name): from patchsim.utils.viz import plot_patch_subplots - plot_patch_subplots(t, results, patches, outdir, model_name) + plot_patch_subplots(t, results, patches, outdir, model_name, compartments=self.compartments) diff --git a/src/patchsim/core/simulation.py b/src/patchsim/core/simulation.py index 56eb725..48ddc3a 100644 --- a/src/patchsim/core/simulation.py +++ b/src/patchsim/core/simulation.py @@ -101,12 +101,7 @@ def get_config_schema() -> dict[str, Any]: def get_model_catalog() -> list[dict[str, str]]: """Return built-in model references and YAML templates for the CLI.""" - catalog: list[dict[str, str]] = [ - {"name": "ka_fmd_sirsv_discrete", "kind": "python"}, - ] - for template_name in sorted(MODEL_TEMPLATE_CONFIGS): - catalog.append({"name": template_name, "kind": "yaml-template"}) - return catalog + return [{"name": name, "kind": "yaml-template"} for name in sorted(MODEL_TEMPLATE_CONFIGS)] def get_available_template_names() -> list[str]: @@ -178,10 +173,17 @@ def load_config(config_path: str) -> dict[str, Any]: def setup_simulation(config: dict[str, Any]) -> tuple[NetworkModel, dict[str, float], list, int]: """Set up the simulation model and initial conditions.""" - # Load patch data + # Load patch data (accept either case for the 'patch'/'population' columns) patch_df = pd.read_csv(config["PatchFile"]) - patches = patch_df["patch"].tolist() - populations = patch_df.set_index("patch")["Population"].to_dict() + patch_col = next((c for c in patch_df.columns if c.lower() == "patch"), None) + pop_col = next((c for c in patch_df.columns if c.lower() == "population"), None) + if patch_col is None or pop_col is None: + raise ValueError( + f"PatchFile ({config['PatchFile']}) must have 'patch' and 'population' columns.\n" + f"Found columns: {list(patch_df.columns)}" + ) + patches = patch_df[patch_col].tolist() + populations = patch_df.set_index(patch_col)[pop_col].to_dict() for p, pop in populations.items(): if pop <= 0: diff --git a/src/patchsim/utils/viz.py b/src/patchsim/utils/viz.py index cea71e5..025fc84 100644 --- a/src/patchsim/utils/viz.py +++ b/src/patchsim/utils/viz.py @@ -4,7 +4,7 @@ import matplotlib.pyplot as plt -def plot_patch_subplots(t_range, out_ode, patches, output_dir, model_name, patch_parameters=None): +def plot_patch_subplots(t_range, out_ode, patches, output_dir, model_name, patch_parameters=None, compartments=None): """ Plots all patches as subplots in a single figure and saves the figure. """ @@ -17,7 +17,9 @@ def plot_patch_subplots(t_range, out_ode, patches, output_dir, model_name, patch axes = axes.flatten() if n > 1 else [axes] for i, patch in enumerate(patches): ax = axes[i] - for c in ["S", "I", "R"]: + # Plot the model's actual compartments; fall back to those present for this patch. + comps = compartments or [k[: -len(f"_{i}")] for k in out_ode if k.endswith(f"_{i}")] + for c in comps: ax.plot(t_range, out_ode[f"{c}_{i}"], label=c) title = f"Patch {patch} (ODE)" if patch_parameters and patch in patch_parameters: @@ -33,5 +35,5 @@ def plot_patch_subplots(t_range, out_ode, patches, output_dir, model_name, patch fig.delaxes(axes[j]) plt.tight_layout() os.makedirs(output_dir, exist_ok=True) - plt.savefig(os.path.join(output_dir, f"all_patches_{model_name}_ode.png")) + plt.savefig(os.path.join(output_dir, f"patch_timeseries_{model_name}_ode.png")) plt.close() diff --git a/tests/test_cli_e2e.py b/tests/test_cli_e2e.py new file mode 100644 index 0000000..ef07443 --- /dev/null +++ b/tests/test_cli_e2e.py @@ -0,0 +1,47 @@ +import json +import subprocess +from pathlib import Path + +import pytest + +TEMPLATES = ["sir", "sis", "sirs", "seir"] + + +def _patchsim(*args, cwd=None): + return subprocess.run(["patchsim", *args], capture_output=True, text=True, cwd=cwd) + + +@pytest.mark.parametrize("template", TEMPLATES) +def test_init_then_run_produces_existing_artifacts(tmp_path, template): + """Every built-in template must scaffold AND run end-to-end, and the artifact + paths reported by `run --json` must actually exist on disk.""" + proj = tmp_path / f"{template}-proj" + init = _patchsim("init", str(proj), "--template", template) + assert init.returncode == 0, init.stderr + + run = _patchsim("run", "-c", "config.yaml", "--json", cwd=str(proj)) + assert run.returncode == 0, f"`run` failed for template {template}:\n{run.stderr}" + + summary = json.loads(run.stdout) + assert Path(summary["csv_path"]).is_file(), f"reported csv missing: {summary['csv_path']}" + assert Path(summary["plot_path"]).is_file(), f"reported plot missing: {summary['plot_path']}" + + +def test_list_models_output_is_clean(): + """Human-facing `list-models` output matches the documented format and does not + leak logging decoration (timestamps / source-file references).""" + res = _patchsim("list-models") + assert res.returncode == 0 + assert "- sir (yaml-template)" in res.stdout + assert "- sis (yaml-template)" in res.stdout + combined = res.stdout + res.stderr + assert "cli.py" not in combined + assert "INFO" not in combined + + +def test_list_models_excludes_unimplemented_models(): + """The catalog only advertises models that actually exist (no dangling entries).""" + res = _patchsim("list-models", "--json") + assert res.returncode == 0 + names = {m["name"] for m in json.loads(res.stdout)["models"]} + assert names == {"seir", "sir", "sirs", "sis"} diff --git a/tests/test_simulation_sanity.py b/tests/test_simulation_sanity.py new file mode 100644 index 0000000..41aa507 --- /dev/null +++ b/tests/test_simulation_sanity.py @@ -0,0 +1,97 @@ +"""Empirical sanity checks on simulation dynamics. + +These lock in epidemiologically-correct behaviour for a multi-patch SIRS model +(waning immunity) so future changes can't silently break the maths. +""" + +import numpy as np +import pandas as pd +import yaml + +from patchsim.core.model_runner import Model +from patchsim.core.simulation import load_config, setup_simulation + + +def _solve_sirs(tmp_path, *, waning=0.02, tmax=160, npatches=5, seed_patch=0): + """Build a multi-patch SIRS scenario with a ring network and solve it.""" + data = tmp_path / "data" + for sub in ("patch", "seeds", "networks"): + (data / sub).mkdir(parents=True, exist_ok=True) + + patches = [f"P{i}" for i in range(npatches)] + pops = [1000 * (i + 1) for i in range(npatches)] + pd.DataFrame({"patch": patches, "Population": pops}).to_csv(data / "patch" / "p.csv", index=False) + + infected = [5 if i == seed_patch else 0 for i in range(npatches)] + pd.DataFrame( + { + "patch": patches, + "S": [p - inf for p, inf in zip(pops, infected, strict=True)], + "I": infected, + "R": [0] * npatches, + } + ).to_csv(data / "seeds" / "s.csv", index=False) + + rows = [] + for i, src in enumerate(patches): + rows.append((0, src, src, 0.8)) + rows.append((0, src, patches[(i + 1) % npatches], 0.1)) + rows.append((0, src, patches[(i - 1) % npatches], 0.1)) + pd.DataFrame(rows, columns=["day", "source", "target", "weight"]).to_csv(data / "networks" / "n.csv", index=False) + + cfg = { + "ModelName": "sanity", + "PatchFile": str(data / "patch" / "p.csv"), + "SeedFile": str(data / "seeds" / "s.csv"), + "NetworkFile": str(data / "networks" / "n.csv"), + "OutputDir": str(tmp_path / "out"), + "TMax": tmax, + "compartments": ["S", "I", "R"], + "Parameters": {"beta": 0.35, "gamma": 0.1, "waning": waning}, + "Transitions": {"S -> I": "beta", "I -> R": "gamma * I", "R -> S": "waning * R"}, + } + cfg_path = tmp_path / "config.yaml" + cfg_path.write_text(yaml.safe_dump(cfg)) + + config = load_config(str(cfg_path)) + net, y0, patches_out, num_patches = setup_simulation(config) + model = Model(net, compartments=list(net.base_model.compartments)) + t = np.arange(tmax, dtype=float) + return model.solve(y0, t), num_patches, t + + +def _total(out, num_patches, comp): + return np.sum([out[f"{comp}_{i}"] for i in range(num_patches)], axis=0) + + +def test_population_is_conserved(tmp_path): + out, n, t = _solve_sirs(tmp_path) + total = sum(_total(out, n, c) for c in ("S", "I", "R")) + assert np.max(np.abs(total - total[0])) < 1e-3 * total[0] + + +def test_no_compartment_goes_negative(tmp_path): + out, n, t = _solve_sirs(tmp_path) + for key, series in out.items(): + assert np.min(series) > -1e-6, f"{key} went negative" + + +def test_epidemic_grows_from_seed(tmp_path): + out, n, t = _solve_sirs(tmp_path) + total_I = _total(out, n, "I") + assert total_I.max() > 5 * total_I[0] + + +def test_infection_spreads_through_network_to_unseeded_patch(tmp_path): + out, n, t = _solve_sirs(tmp_path, seed_patch=0) + assert out["I_2"][0] == 0 # P2 starts with no infections + assert out["I_2"].max() > 1.0 # but gets infected via the network + + +def test_waning_immunity_sustains_endemic_infection(tmp_path): + with_waning, n, t = _solve_sirs(tmp_path, waning=0.05, tmax=200) + without_waning, n2, t2 = _solve_sirs(tmp_path, waning=0.0, tmax=200) + final_with = sum(with_waning[f"I_{i}"][-1] for i in range(n)) + final_without = sum(without_waning[f"I_{i}"][-1] for i in range(n2)) + assert final_with > final_without # waning replenishes S -> infection persists + assert final_with > 1.0 diff --git a/tests/test_smoke.py b/tests/test_smoke.py index b2f6a83..2054f01 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -51,7 +51,6 @@ def test_cli_list_models(): ) assert result.returncode == 0 output = (result.stdout + result.stderr).lower() - assert "ka_fmd_sirsv_discrete" in output assert "sir (yaml-template)" in output assert "sis (yaml-template)" in output