patchsim 0.1.0b1__py3-none-any.whl

patchsim/__init__.py

@@ -0,0 +1,24 @@
+"""PatchSim public SDK interface."""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+from patchsim.core.model import CompartmentalModel, NetworkModel
+from patchsim.core.simulation import load_config, run_simulation, setup_simulation
+from patchsim.utils.viz import plot_patch_subplots
+
+try:
+    __version__ = version("patchsim")
+except PackageNotFoundError:  # pragma: no cover - local editable fallback
+    __version__ = "0.1.0b1"
+
+__all__ = [
+    "CompartmentalModel",
+    "NetworkModel",
+    "__version__",
+    "load_config",
+    "plot_patch_subplots",
+    "run_simulation",
+    "setup_simulation",
+]

patchsim/cli.py

@@ -0,0 +1,246 @@
+import argparse
+import json
+import logging
+import shutil
+import textwrap
+from importlib import resources
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from patchsim import __version__
+from patchsim.core.simulation import (
+    get_available_template_names,
+    get_config_schema,
+    get_init_template_config,
+    get_model_catalog,
+    load_config,
+    run_simulation,
+    setup_simulation,
+)
+
+
+def _configure_logging(*, json_output: bool = False) -> None:
+    """Configure CLI logging with rich handler when available."""
+    if json_output:
+        logging.basicConfig(level=logging.WARNING)
+        return
+
+    try:
+        from rich.logging import RichHandler
+
+        logging.basicConfig(
+            level=logging.INFO,
+            format="%(message)s",
+            datefmt="[%X]",
+            handlers=[RichHandler(rich_tracebacks=True)],
+        )
+    except Exception:
+        logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
+
+
+def _emit_json(payload: dict[str, Any]) -> None:
+    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+def _cmd_run(config_path: str, *, json_output: bool = False) -> dict[str, Any]:
+    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)
+    if not json_output:
+        print("Simulation completed successfully.")
+    return {"ok": True, "config": config_path, **summary} if json_output else summary
+
+
+def _cmd_validate(config_path: str, *, json_output: bool = False, schema: bool = False) -> dict[str, Any] | None:
+    if schema:
+        return get_config_schema()
+
+    config = load_config(config_path)
+    _net, _y0, patches, num_patches = setup_simulation(config)
+    if not json_output:
+        print(f"Configuration is valid: {config_path}")
+    if json_output:
+        return {
+            "ok": True,
+            "config": config_path,
+            "model_name": config.get("ModelName"),
+            "num_patches": num_patches,
+            "patches": patches,
+        }
+    return None
+
+
+def _copy_template_tree(template_node, target_path: Path) -> None:
+    if template_node.is_dir():
+        target_path.mkdir(parents=True, exist_ok=True)
+        for child in template_node.iterdir():
+            _copy_template_tree(child, target_path / child.name)
+        return
+
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+    target_path.write_bytes(template_node.read_bytes())
+
+
+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:
+    project_dir = Path(name)
+
+    # Safety checks: prevent deleting cwd, parent dirs, root, or non-directories
+    resolved = project_dir.resolve()
+    cwd = Path.cwd().resolve()
+
+    if project_dir.exists() and not project_dir.is_dir():
+        raise NotADirectoryError(f"Target exists and is not a directory: {project_dir}")
+
+    # Block deletion of root, cwd, or any ancestor of cwd
+    if force and (
+        resolved == Path(resolved.anchor)  # Filesystem root (/ or C:\)
+        or resolved == cwd  # Current working directory
+        or cwd in resolved.parents  # resolved is ancestor of cwd (e.g., ..)
+    ):
+        raise ValueError(f"Refusing to overwrite unsafe target: {resolved}")
+
+    if project_dir.exists() and any(project_dir.iterdir()) and not force:
+        raise FileExistsError(f"Refusing to overwrite existing directory: {project_dir}. Use --force to overwrite.")
+
+    if project_dir.exists() and force:
+        shutil.rmtree(project_dir)
+
+    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(yaml.safe_dump(template_config, sort_keys=False), encoding="utf-8")
+    _write_seed_for_template(project_dir, template_config)
+
+    print(f"Created project scaffold at: {project_dir}")
+
+
+def _list_builtin_models() -> list[dict[str, str]]:
+    return get_model_catalog()
+
+
+def _cmd_list_models(*, json_output: bool = False) -> list[dict[str, str]]:
+    models = _list_builtin_models()
+    if not models:
+        if not json_output:
+            print("No built-in models found.")
+        return []
+
+    if json_output:
+        return models
+
+    print("Built-in models and templates:")
+    for model in models:
+        print(f"- {model['name']} ({model['kind']})")
+    return models
+
+
+def main() -> None:
+    """Command-line interface for running the PatchSim simulation."""
+    parser = argparse.ArgumentParser(
+        description=(
+            "PatchSim: A modular metapopulation simulation framework for multi-disease epidemiological modelling."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=textwrap.dedent(
+            """
+Examples:
+    uv run patchsim init my-project
+    uv run patchsim init my-project --template seir
+    uv run patchsim run -c my-project/config.yaml
+    uv run patchsim validate -c my-project/config.yaml
+    uv run patchsim list-models
+        """
+        ),
+    )
+    parser.add_argument("--version", action="version", version=f"patchsim {__version__}")
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    init_p = subparsers.add_parser("init", help="Scaffold a new PatchSim project")
+    init_p.add_argument("name", help="Directory name for the new project")
+    init_p.add_argument("--force", action="store_true", help="Overwrite target directory if it already exists")
+    init_p.add_argument(
+        "--template",
+        choices=get_available_template_names(),
+        default="sir",
+        help="Starter template to use for config.yaml",
+    )
+
+    run_p = subparsers.add_parser("run", help="Run a simulation")
+    run_p.add_argument("-c", "--config", required=True, help="Path to simulation config YAML")
+    run_p.add_argument("--json", action="store_true", help="Emit machine-readable JSON summary")
+
+    validate_p = subparsers.add_parser("validate", help="Validate config and inputs")
+    validate_p.add_argument("-c", "--config", help="Path to simulation config YAML")
+    validate_p.add_argument("--schema", action="store_true", help="Print the configuration JSON Schema")
+    validate_p.add_argument("--json", action="store_true", help="Emit machine-readable JSON summary")
+
+    list_p = subparsers.add_parser("list-models", help="List available built-in models")
+    list_p.add_argument("--json", action="store_true", help="Emit machine-readable JSON list")
+
+    args = parser.parse_args()
+    json_mode = bool(getattr(args, "json", False) or getattr(args, "schema", False))
+    _configure_logging(json_output=json_mode)
+
+    try:
+        if args.command == "init":
+            _cmd_init(args.name, force=args.force, template=args.template)
+        elif args.command == "run":
+            result = _cmd_run(args.config, json_output=args.json)
+            if args.json:
+                _emit_json(result)
+        elif args.command == "validate":
+            if not args.schema and not args.config:
+                parser.error("the following arguments are required: -c/--config")
+            result = _cmd_validate(args.config, json_output=args.json, schema=args.schema)
+            if result is not None:
+                _emit_json(result)
+        elif args.command == "list-models":
+            result = _cmd_list_models(json_output=args.json)
+            if args.json:
+                _emit_json({"models": result})
+        else:
+            parser.print_help()
+            raise SystemExit(2)
+    except Exception as e:
+        logging.error(f"Simulation failed: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

patchsim/core/model.py

@@ -0,0 +1,237 @@
+"""
+Core model implementation for compartmental models.
+"""
+
+import re
+from typing import Any, Callable, Dict
+
+from scipy.integrate import odeint
+
+
+class CompartmentalModel:
+    """Base class for compartmental models."""
+
+    def __init__(self, compartments: list[str], parameters: dict[str, float], transitions: list[dict[str, Any]]):
+        """Initialize the model with compartments, parameters, and transitions."""
+        self.compartments = compartments
+        self.parameters = parameters
+        self.transitions = transitions
+
+    def compute_rates(self, state: dict[str, float], parameters: dict[str, float] | None = None) -> dict[str, float]:
+        """Compute transition rates for each compartment.
+
+        Args:
+            state: Current compartment state
+            parameters: Optional parameter override (defaults to self.parameters)
+        """
+        params = parameters if parameters is not None else self.parameters
+        rates = {}
+        for transition in self.transitions:
+            transition_label = transition["transition"]
+            source, target = [p.strip() for p in transition_label.split("->")]
+            rate = transition["rate"]
+            rate_expr = rate
+            # Handle rate expressions
+            if isinstance(rate_expr, str):
+                # Safe evaluation: build scope from parameters and state, disable builtins
+                scope = {**params, **state}
+                try:
+                    rate_val = eval(rate_expr, {"__builtins__": {}}, scope)
+                except (KeyError, NameError, ValueError, SyntaxError, TypeError, ZeroDivisionError) as e:
+                    msg = f"Invalid rate expression '{rate_expr}' in transition '{transition_label}': {e}"
+                    raise ValueError(msg) from e
+            else:
+                rate_val = rate_expr
+
+            # If expression already includes the source compartment, don't multiply again.
+            if isinstance(rate, str) and re.search(rf"\b{re.escape(source)}\b", rate):
+                flow = rate_val
+            else:
+                flow = rate_val * state[source]
+
+            rates[transition_label] = flow
+        return rates
+
+
+class NetworkModel:
+    """Network model for multi-patch simulations."""
+
+    def __init__(self, base_model: CompartmentalModel, num_patches: int, network_matrix: list[list[float]]):
+        """Initialize the network model."""
+        self.base_model = base_model
+        self.num_patches = num_patches
+        self.network = network_matrix
+        self.all_compartments = [f"{c}_{i}" for i in range(num_patches) for c in base_model.compartments]
+
+    def get_patch_state(self, full_state: Dict[str, float], patch_idx: int) -> Dict[str, float]:
+        """Get state for a specific patch."""
+        return {c: full_state[f"{c}_{patch_idx}"] for c in self.base_model.compartments}
+
+    def get_patch_population(self, state: Dict[str, float]) -> float:
+        """Get total population for a patch."""
+        return sum(state[c] for c in self.base_model.compartments)
+
+    def compute_force_of_infection(self, full_state: dict[str, float], infected_compartment: str = "I") -> list[float]:
+        """Compute force of infection for each patch (per-capita rate, before beta scaling).
+
+        Args:
+            full_state: Current state of all compartments
+            infected_compartment: Name of the compartment representing infected individuals
+
+        Returns:
+            List of per-capita forces of infection (model_runner applies beta * FOI * S)
+        """
+        lambdas = []
+        for i in range(self.num_patches):
+            if self.num_patches == 1:
+                # Single patch case: infected proportion
+                patch_state = self.get_patch_state(full_state, 0)
+                infected = patch_state[infected_compartment]
+                total_pop = self.get_patch_population(patch_state)
+                force = infected / total_pop if total_pop > 0 else 0
+            else:
+                # Multi-patch case: network-weighted infected proportion
+                force = 0
+                for j in range(self.num_patches):
+                    patch_state_j = self.get_patch_state(full_state, j)
+                    infected_j = patch_state_j[infected_compartment]
+                    pop_j = self.get_patch_population(patch_state_j)
+                    force += self.network[i][j] * (infected_j / pop_j if pop_j > 0 else 0)
+            lambdas.append(force)
+        return lambdas
+
+    def _adjust_infection_rate(
+        self,
+        patch_params: dict[str, float],
+        original_rate_expr: Any,
+        rate: float,
+        patch_state: dict[str, float],
+        lambdas: list[float],
+        patch_idx: int,
+        is_infection_transition: bool,
+        has_network: bool,
+    ) -> float:
+        """Adjust infection rate for network-mediated FOI.
+
+        Args:
+            patch_params: Parameters for the current patch
+            original_rate_expr: Original rate expression from transition definition
+            rate: Computed rate from base model
+            patch_state: Current state for the patch
+            lambdas: Force of infection for each patch
+            patch_idx: Current patch index
+            is_infection_transition: Whether this is an infection transition
+            has_network: Whether multi-patch network exists
+
+        Returns:
+            Adjusted rate incorporating network FOI if applicable
+        """
+        if is_infection_transition and has_network:
+            # Network case: Apply network FOI (lambdas already computed)
+            # Check if original expression includes beta term
+            beta = patch_params.get("beta", 1.0)
+            if isinstance(original_rate_expr, str) and re.search(r"\bbeta\b", original_rate_expr):
+                # Rate expression includes beta; apply network FOI correction
+                adjusted_rate = beta * patch_state["S"] * lambdas[patch_idx]
+            else:
+                # Rate is already computed; apply FOI scaling
+                adjusted_rate = rate * lambdas[patch_idx] if patch_state["S"] > 0 else 0
+        else:
+            # Single patch or non-infection transition: use rate as-is
+            adjusted_rate = rate
+        return adjusted_rate
+
+    def compute_derivatives(self, state: dict[str, float]) -> dict[str, float]:
+        """Compute derivatives for all compartments based on transitions, incorporating network-mediated FOI."""
+        derivatives = {c: 0.0 for c in self.all_compartments}
+
+        # Compute network-mediated force of infection for each patch
+        lambdas = self.compute_force_of_infection(state)
+
+        # Process each patch
+        for i in range(self.num_patches):
+            # Get state for this patch
+            patch_state = self.get_patch_state(state, i)
+
+            # Resolve patch parameters using canonical patch ordering when available.
+            if hasattr(self, "patch_parameters"):
+                patch_name = None
+                if hasattr(self, "patch_names") and i < len(self.patch_names):
+                    patch_name = self.patch_names[i]
+                elif self.patch_parameters:
+                    import logging
+
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        "patch_parameters defined but patch_names not set; "
+                        "patch-specific parameters will be ignored for patch %d",
+                        i,
+                    )
+                patch_params = {**self.base_model.parameters, **self.patch_parameters.get(patch_name, {})}
+            else:
+                patch_params = self.base_model.parameters
+
+            # Compute rates with patch-specific parameters without mutating shared state
+            rates = self.base_model.compute_rates(patch_state, parameters=patch_params)
+
+            # Update derivatives based on rates, applying network-mediated FOI to infection transitions
+            for transition in self.base_model.transitions:
+                transition_label = transition["transition"]
+                source, target = [p.strip() for p in transition_label.split("->")]
+                rate = rates[transition_label]
+                original_rate_expr = transition.get("rate", "")
+
+                # Apply network-mediated FOI to susceptible-to-infection transitions.
+                # Allow model-level override via `infection_compartments` attribute.
+                infection_compartments = set(getattr(self, "infection_compartments", {"I", "E"}))
+                is_infection_transition = source == "S" and target in infection_compartments
+
+                has_network = self.num_patches > 1 and self.network is not None
+
+                # Use shared helper to adjust infection rate for network FOI
+                adjusted_rate = self._adjust_infection_rate(
+                    patch_params,
+                    original_rate_expr,
+                    rate,
+                    patch_state,
+                    lambdas,
+                    i,
+                    is_infection_transition,
+                    has_network,
+                )
+
+                # Decrease source compartment
+                derivatives[f"{source}_{i}"] -= adjusted_rate
+                # Increase target compartment
+                derivatives[f"{target}_{i}"] += adjusted_rate
+
+        return derivatives
+
+    def simulate_discrete(self, y0_dict: dict[str, float], t_range: list[float]) -> dict[str, list[float]]:
+        """Run discrete-time simulation."""
+        state = y0_dict.copy()
+        history = {c: [state[c]] for c in self.all_compartments}
+
+        for _ in t_range[1:]:
+            derivatives = self.compute_derivatives(state)
+            new_state = {c: state[c] + derivatives[c] for c in self.all_compartments}
+            state = new_state
+            for c in self.all_compartments:
+                history[c].append(state[c])
+
+        return history
+
+    def simulate_ode(
+        self, y0_dict: dict[str, float], t_range: list[float], integrator: Callable = odeint
+    ) -> tuple[list[float], dict[str, list[float]]]:
+        """Run ODE simulation."""
+        y0 = [y0_dict[c] for c in self.all_compartments]
+
+        def rhs(y, t):
+            state = {c: y[i] for i, c in enumerate(self.all_compartments)}
+            derivatives = self.compute_derivatives(state)
+            return [derivatives[c] for c in self.all_compartments]
+
+        sol = integrator(rhs, y0, t_range)
+        out = {c: sol[:, i] for i, c in enumerate(self.all_compartments)}
+        return t_range, out

patchsim/core/model_runner.py

@@ -0,0 +1,62 @@
+from scipy.integrate import odeint
+
+
+class Model:
+    """
+    High-level simulation model.
+    Owns the Network and builds/solves the ODE.
+    """
+
+    def __init__(self, network_model, compartments):
+        self.network = network_model
+        self.compartments = compartments
+        self.all_vars = self.network.all_compartments
+
+    def construct_ode(self):
+        def rhs(y, t):
+            state = {v: y[i] for i, v in enumerate(self.all_vars)}
+            dydt = {v: 0.0 for v in self.all_vars}
+
+            # Compute network-based force of infection (per-capita, without beta)
+            lambdas = self.network.compute_force_of_infection(state)
+
+            # Apply transitions for all patches
+            for i in range(self.network.num_patches):
+                patch_state = {c: state[f"{c}_{i}"] for c in self.compartments}
+                rates = self.network.base_model.compute_rates(patch_state)
+                patch_params = self.network.base_model.parameters
+
+                for key, rate in rates.items():
+                    src, tgt = [p.strip() for p in key.split("->")]
+
+                    # Determine if this is an infection transition
+                    infection_compartments = set(getattr(self.network, "infection_compartments", {"I", "E"}))
+                    is_infection_transition = src == "S" and tgt in infection_compartments
+
+                    # Use network helper to apply FOI adjustment consistently
+                    has_network = self.network.num_patches > 1 and self.network.network is not None
+                    adjusted_rate = self.network._adjust_infection_rate(
+                        patch_params, key, rate, patch_state, lambdas, i, is_infection_transition, has_network
+                    )
+
+                    dydt[f"{src}_{i}"] -= adjusted_rate
+                    dydt[f"{tgt}_{i}"] += adjusted_rate
+
+            return [dydt[v] for v in self.all_vars]
+
+        return rhs
+
+    def solve(self, y0, t_range):
+        rhs = self.construct_ode()
+        # Validate all required variables are present in y0
+        missing = [v for v in self.all_vars if v not in y0]
+        if missing:
+            raise ValueError(f"Missing initial values for: {missing}")
+        y0_vec = [y0[v] for v in self.all_vars]
+        sol = odeint(rhs, y0_vec, t_range)
+        return {v: sol[:, i] for i, v in enumerate(self.all_vars)}
+
+    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, compartments=self.compartments)