patchsim 0.1.0__py3-none-any.whl

patchsim/cli.py

@@ -0,0 +1,393 @@
+import argparse
+import json
+import logging
+import shutil
+import sys
+import textwrap
+from importlib import resources
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from patchsim import __version__
+from patchsim.calibration import get_calibration_plan, run_calibration
+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,
+)
+from patchsim.sensitivity import get_sensitivity_plan, run_sensitivity
+from patchsim.utils.geo import generate_contacts
+
+
+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)
+    sensitivity = get_sensitivity_plan(config, list(net.all_compartments), required=False)
+    calibration = get_calibration_plan(config, net, y0, required=False)
+    if not json_output:
+        print(f"Configuration is valid: {config_path}")
+        if sensitivity:
+            print(f"Planned sensitivity evaluations: {sensitivity.evaluation_count}")
+        if calibration:
+            print(
+                f"Calibration observations: {calibration.n}; fitted variables: {calibration.p}; "
+                f"starts: {calibration.start_count}"
+            )
+            print(f"Maximum forward simulations: {calibration.max_forward_simulations}")
+            for warning in calibration.warnings:
+                print(f"Warning: {warning}", file=sys.stderr)
+        if net.groups:
+            diagnostics = net.interaction_diagnostics
+            print(
+                "Interaction diagnostics: "
+                f"units={diagnostics['units']}, "
+                f"max reciprocity residual={diagnostics['max_local_reciprocity_residual']:.6g}"
+            )
+    if json_output:
+        result = {
+            "ok": True,
+            "config": config_path,
+            "model_name": config.get("ModelName"),
+            "num_patches": num_patches,
+            "patches": patches,
+            "solver": config["Solver"],
+            "time_step": config["TimeStep"],
+        }
+        if net.groups:
+            result.update(
+                {
+                    "num_groups": net.num_groups,
+                    "groups": net.groups,
+                    "interaction": net.interaction_diagnostics,
+                }
+            )
+        if sensitivity:
+            result["sensitivity"] = {
+                "name": sensitivity.name,
+                "method": "sobol",
+                "num_parameters": len(sensitivity.parameters),
+                "evaluation_count": sensitivity.evaluation_count,
+            }
+        if calibration:
+            result["calibration"] = {
+                "name": calibration.name,
+                "method": "least_squares",
+                "n": calibration.n,
+                "p": calibration.p,
+                "start_count": calibration.start_count,
+                "max_forward_simulations": calibration.max_forward_simulations,
+                "warnings": list(calibration.warnings),
+            }
+        return result
+    return None
+
+
+def _cmd_sensitivity(config_path: str, *, json_output: bool = False) -> dict[str, Any]:
+    summary = run_sensitivity(config_path, progress=lambda message: print(message, file=sys.stderr))
+    result = {"ok": True, "config": config_path, **summary}
+    if not json_output:
+        action = "Reused" if summary["reused"] else "Completed"
+        print(f"{action} sensitivity study: {summary['output_dir']}")
+        print(f"Samples: {summary['samples_path']}")
+        print(f"Responses: {summary['responses_path']}")
+        print(f"Indices: {summary['indices_path']}")
+        print(f"Manifest: {summary['manifest_path']}")
+    return result
+
+
+def _cmd_calibrate(config_path: str, *, json_output: bool = False) -> dict[str, Any]:
+    summary = run_calibration(config_path, progress=lambda message: print(message, file=sys.stderr))
+    result = {"ok": True, "config": config_path, **summary}
+    if not json_output:
+        action = "Reused" if summary["reused"] else "Completed"
+        print(f"{action} calibration study: {summary['output_dir']}")
+        print(f"Estimates: {summary['estimates_path']}")
+        print(f"Fitted seeds: {summary['fitted_seeds_path']}")
+        print(f"Attempts: {summary['attempts_path']}")
+        print(f"Residuals: {summary['residuals_path']}")
+        print(f"Manifest: {summary['manifest_path']}")
+    return result
+
+
+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 _cmd_generate_contacts(args: argparse.Namespace) -> None:
+    output_path, report_path, _report = generate_contacts(
+        args.source,
+        args.output,
+        id_column=args.id_column,
+        population_column=args.population_column,
+        kernel=args.kernel,
+        decay=args.decay,
+        scale=args.scale,
+        min_distance_km=args.min_distance_km,
+        normalize=args.normalize,
+        self_weight=args.self_weight,
+        self_share=args.self_share,
+        centroid_crs=args.centroid_crs,
+        force=args.force,
+    )
+    print(f"Wrote contacts to: {output_path}")
+    print(f"Wrote validation report to: {report_path}")
+    if args.normalize == "none":
+        unit = "scale * population_i * population_j / km**decay" if args.kernel == "gravity" else "scale / km**decay"
+        print(f"Warning: unnormalized weights use raw kernel units ({unit}); they are not probabilities.")
+
+
+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 sensitivity -c my-project/config.yaml
+    uv run patchsim calibrate -c my-project/config.yaml
+    uv run patchsim list-models
+    uv run patchsim generate-contacts centroids.csv contacts.csv --id-column id \
+        --kernel distance --decay 2 --min-distance-km 0.001 --normalize row --self-share 0.9
+        """
+        ),
+    )
+    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")
+
+    sensitivity_p = subparsers.add_parser("sensitivity", help="Run or reuse a Sobol sensitivity study")
+    sensitivity_p.add_argument("-c", "--config", required=True, help="Path to simulation config YAML")
+    sensitivity_p.add_argument("--json", action="store_true", help="Emit machine-readable JSON summary")
+
+    calibration_p = subparsers.add_parser("calibrate", help="Run or reuse a bounded calibration study")
+    calibration_p.add_argument("-c", "--config", required=True, help="Path to simulation config YAML")
+    calibration_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")
+
+    contacts_p = subparsers.add_parser(
+        "generate-contacts",
+        help="Generate a validated spatial contact network",
+    )
+    contacts_p.add_argument("source", help="Centroid CSV, GeoJSON, JSON, or Shapefile input")
+    contacts_p.add_argument("output", help="Output edge-list CSV")
+    contacts_p.add_argument("--id-column", required=True, help="Identifier column in the source")
+    contacts_p.add_argument("--population-column", help="Positive population column required by gravity")
+    contacts_p.add_argument("--kernel", required=True, choices=["distance", "gravity"])
+    contacts_p.add_argument("--decay", required=True, type=float, help="Dimensionless distance-decay exponent")
+    contacts_p.add_argument(
+        "--min-distance-km",
+        required=True,
+        type=float,
+        help="Positive distance floor in kilometres",
+    )
+    contacts_p.add_argument("--normalize", required=True, choices=["none", "row"])
+    contacts_p.add_argument("--scale", type=float, help="Positive raw-kernel scale; only for normalization none")
+    diagonal = contacts_p.add_mutually_exclusive_group()
+    diagonal.add_argument("--self-weight", type=float, help="Raw diagonal weight for normalization none")
+    diagonal.add_argument("--self-share", type=float, help="Diagonal share in [0, 1) for row normalization")
+    contacts_p.add_argument("--centroid-crs", help="Projected CRS used to centroid polygon vector input")
+    contacts_p.add_argument("--force", action="store_true", help="Replace both output artifacts")
+
+    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 == "sensitivity":
+            result = _cmd_sensitivity(args.config, json_output=args.json)
+            if args.json:
+                _emit_json(result)
+        elif args.command == "calibrate":
+            result = _cmd_calibrate(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})
+        elif args.command == "generate-contacts":
+            _cmd_generate_contacts(args)
+        else:
+            parser.print_help()
+            raise SystemExit(2)
+    except Exception as e:
+        logging.error(f"Command failed: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

patchsim/core/expressions.py

@@ -0,0 +1,144 @@
+"""Evaluate transition rate expressions from user config files.
+
+Evaluation is restricted to arithmetic over named values so that a config file cannot
+execute arbitrary code. Any other construct is rejected.
+"""
+
+import ast
+import math
+import numbers
+import operator
+from functools import lru_cache
+from typing import Any, Mapping
+
+_MAX_DEPTH = 64  # maximum expression nesting depth
+_MAX_LENGTH = 1000  # maximum expression length in characters
+
+_BINARY_OPS = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.Pow: operator.pow,
+}
+
+_UNARY_OPS = {
+    ast.UAdd: operator.pos,
+    ast.USub: operator.neg,
+}
+
+
+@lru_cache(maxsize=256)
+def _parse(expression: str) -> ast.AST:
+    if len(expression) > _MAX_LENGTH:
+        raise ValueError(f"is too long to evaluate ({len(expression)} characters, limit {_MAX_LENGTH})")
+
+    try:
+        tree = ast.parse(expression, mode="eval")
+    except SyntaxError as exc:
+        raise ValueError(f"could not be parsed: {exc.msg}") from exc
+    except (RecursionError, MemoryError) as exc:
+        raise ValueError("is too deeply nested to parse") from exc
+
+    return tree.body
+
+
+def evaluate(expression: str, scope: Mapping[str, Any]) -> float:
+    """Evaluate an arithmetic rate expression against a scope of named values.
+
+    Args:
+        expression: Arithmetic expression over parameter and compartment names.
+        scope: Mapping of names available to the expression.
+
+    Returns:
+        The numeric value of the expression.
+
+    Raises:
+        ValueError: If the expression uses an unsupported construct, references an
+            unknown name, or cannot be evaluated numerically.
+    """
+    return _evaluate_node(_parse(expression), scope, depth=0)
+
+
+def _evaluate_node(node: ast.AST, scope: Mapping[str, Any], depth: int) -> float:
+    if depth > _MAX_DEPTH:
+        raise ValueError(f"is too deeply nested to evaluate (limit {_MAX_DEPTH} levels)")
+
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, bool) or not isinstance(node.value, numbers.Real):
+            raise ValueError(f"{type(node.value).__name__} literals are not allowed in rate expressions")
+        return _coerce(node.value, "a numeric literal")
+
+    if isinstance(node, ast.Name):
+        if node.id not in scope:
+            raise ValueError(f"unknown name '{node.id}'; expected a parameter or compartment")
+        return _coerce(scope[node.id], f"name '{node.id}'")
+
+    if isinstance(node, ast.BinOp) and type(node.op) in _BINARY_OPS:
+        left = _evaluate_node(node.left, scope, depth + 1)
+        right = _evaluate_node(node.right, scope, depth + 1)
+        return _apply(_BINARY_OPS[type(node.op)], left, right)
+
+    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPS:
+        return _apply(_UNARY_OPS[type(node.op)], _evaluate_node(node.operand, scope, depth + 1))
+
+    raise ValueError(
+        f"{_describe(node)} is not allowed in rate expressions, "
+        "which may only combine parameters and compartments arithmetically"
+    )
+
+
+def _coerce(value: Any, description: str) -> float:
+    """Convert a value to a finite real float, or raise ValueError.
+
+    Accepts ``numbers.Real`` rather than ``float`` because solver state can be a numpy
+    scalar and only ``numpy.float64`` subclasses ``float``. Complex results (from a
+    negative base to a fractional power) and non-finite results (from overflow) are
+    rejected.
+    """
+    if isinstance(value, complex):
+        raise ValueError(f"{description} is a complex number; rates must be real")
+    if isinstance(value, bool) or not isinstance(value, numbers.Real):
+        raise ValueError(f"{description} is not a real number")
+    try:
+        result = float(value)
+    except (OverflowError, ValueError) as exc:
+        raise ValueError(f"{description} is too large to represent as a floating point number") from exc
+    if not math.isfinite(result):
+        raise ValueError(f"{description} is {result}; rates must be finite")
+    return result
+
+
+def _apply(op, *operands: float) -> float:
+    """Apply an arithmetic operator and validate its result."""
+    try:
+        value = op(*operands)
+    except ZeroDivisionError as exc:
+        raise ValueError("attempted division by zero") from exc
+    except OverflowError as exc:
+        raise ValueError("overflowed to a value that is not finite") from exc
+    except ValueError as exc:
+        raise ValueError(f"could not be evaluated: {exc}") from exc
+
+    return _coerce(value, "the expression")
+
+
+def _describe(node: ast.AST) -> str:
+    """Return a readable name for a disallowed construct, for the error message."""
+    descriptions = {
+        ast.Attribute: "attribute access",
+        ast.Call: "function calls",
+        ast.Subscript: "indexing",
+        ast.ListComp: "comprehensions",
+        ast.GeneratorExp: "comprehensions",
+        ast.DictComp: "comprehensions",
+        ast.SetComp: "comprehensions",
+        ast.Lambda: "lambdas",
+        ast.List: "list literals",
+        ast.Tuple: "tuple literals",
+        ast.Dict: "dict literals",
+    }
+    for node_type, description in descriptions.items():
+        if isinstance(node, node_type):
+            return description
+    return f"{type(node).__name__} expressions"