diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2d98693..b3e0976 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,6 +18,8 @@ on: branches: [main] concurrency: + # One cancelable group per workflow+ref; covers Ubuntu matrix, macOS smoke, + # notebooks, and docs (no per-OS duplicate concurrency keys). group: ci-${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true @@ -90,6 +92,10 @@ jobs: fh.write(f"{name}={text}\n") print(f"{name}={text}") + # Path gates (expensive jobs skip when false unless workflow=true): + # - test → Ubuntu pytest matrix AND macOS core smoke (test-macos) + # - notebooks → Ubuntu nbmake only (never macOS) + # - docs / lint → as named lint = hit("src/", "tests/", "pyproject.toml", ".pre-commit-config.yaml") test = hit("src/", "tests/", "scripts/", "pyproject.toml", "uv.lock") notebooks = hit("examples/", "src/", "pyproject.toml", "uv.lock") @@ -122,6 +128,8 @@ jobs: uvx ruff format --check src/ tests/ test: + # Ubuntu full suite runs unmarked + marked tests (Ray when extras installed). + # macOS smoke (TASK-2142) uses: -m "not slow and not distributed and not ray" name: test (${{ matrix.python-version }}) needs: changes if: needs.changes.outputs.test == 'true' || needs.changes.outputs.workflow == 'true' @@ -132,6 +140,8 @@ jobs: include: - python-version: "3.10" coverage: "false" + - python-version: "3.11" + coverage: "false" - python-version: "3.12" coverage: "true" @@ -187,7 +197,44 @@ jobs: files: ./coverage.xml fail_ci_if_error: false + test-macos: + # Core smoke on macOS (Q6/Q7/R8): Python 3.12, CPU torch, hard-fail. + # Shares the ``test`` path filter with Ubuntu pytest (not ``notebooks``). + # Marker expression from TASK-2141; no Ray / Dask extras (ray tests deselected). + # nbmake never runs on macOS — only notebook-plan / notebooks on Ubuntu. + name: test (macos-3.12) + needs: changes + if: needs.changes.outputs.test == 'true' || needs.changes.outputs.workflow == 'true' + runs-on: macos-latest + steps: + - uses: actions/checkout@v7 + + - name: Install uv + uses: astral-sh/setup-uv@v9.0.0 + with: + python-version: "3.12" + enable-cache: true + cache-dependency-glob: "uv.lock" + + - name: Install dependencies + run: > + uv sync --frozen + --extra dev --extra rl --extra mpc --extra symmetry + --extra lightning --extra equivariance + --extra baselines-ode --extra msm + + - name: Run macOS core smoke tests + env: + OMP_NUM_THREADS: "1" + MKL_NUM_THREADS: "1" + OPENBLAS_NUM_THREADS: "1" + TORCH_NUM_THREADS: "1" + run: > + uv run --no-sync pytest tests/ -n auto + -m "not slow and not distributed and not ray" + notebook-plan: + # Ubuntu-only nbmake planning; macOS never executes tutorial notebooks. name: notebook-plan needs: changes if: needs.changes.outputs.notebooks == 'true' || needs.changes.outputs.workflow == 'true' @@ -241,6 +288,14 @@ jobs: "examples/1[4-9]_*.ipynb examples/2[0-6]_*.ipynb" ), }, + { + "shard": "27-46", + "pattern": ( + "examples/2[7-9]_*.ipynb " + "examples/3[0-9]_*.ipynb " + "examples/4[0-6]_*.ipynb" + ), + }, ] } matrix = smoke if os.environ["EVENT_NAME"] == "pull_request" else full @@ -305,7 +360,7 @@ jobs: # by path filters; this aggregator still reports success when skips are expected. ci: name: ci - needs: [changes, lint, test, notebook-plan, notebooks, docs] + needs: [changes, lint, test, test-macos, notebook-plan, notebooks, docs] if: always() runs-on: ubuntu-latest steps: @@ -321,7 +376,14 @@ jobs: needs = json.loads(os.environ["NEEDS_JSON"]) if needs["changes"]["result"] != "success": sys.exit("changes job did not succeed") - for job in ("lint", "test", "notebook-plan", "notebooks", "docs"): + for job in ( + "lint", + "test", + "test-macos", + "notebook-plan", + "notebooks", + "docs", + ): result = needs[job]["result"] if result not in ("success", "skipped"): sys.exit(f"{job} failed with result={result}") diff --git a/CHANGELOG.md b/CHANGELOG.md index 635420c..5af5b10 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,62 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.13.0] - 2026-08-12 + +Ops-and-trust release on the 0.12 stack: default safetensors checkpoints, +fit callbacks and tracking adapters, a config-driven CLI, thin Ray Tune +helpers, homogeneous representation-level explainability, and a broader CI +matrix. Homogeneous scientific defaults and the linear latent operator +contract are unchanged. + +### Added + +- Default checkpoint container ``safetensors_v1`` (directory or ``.kgckpt`` / + ``.zip`` bundle) with ``meta.json`` / ``config.json`` / + ``model.safetensors``. Core dependency ``safetensors>=0.4``. Load + auto-detects the safe layout and legacy format-1 ``.pt`` pickles. See + ``SECURITY.md`` and Sphinx serialization docs. +- Optional ``.kgckpt`` zip packaging of the safetensors directory layout. +- Fit callback protocol (``FitCallback`` / ``NoOpFitCallback``) wired through + ``run_fit_loop`` and ``GraphKoopmanModel.fit(..., callbacks=)``. In-tree + ``CsvFitLogger`` and optional ``TensorBoardFitLogger``; W&B / MLflow remain + examples-only DIY (no cloud SDK pins). +- Console script ``koopman-graph`` with ``train`` / ``predict`` subcommands, + JSON configs in core and optional YAML via ``[cli]`` (PyYAML). Sphinx + ``cli`` page and ``examples/cli/`` smoke config. +- Power-user ``koopman_graph.tuning``: ``fit_history_metrics``, lazy + ``run_ray_tune``, and example search-space scaffolds. Ray Tune example + script migrated; Optuna stays examples-only. +- Homogeneous ``explain_representation`` / + ``RepresentationExplanation`` (GNNExplainer by default; optional Captum + integrated gradients via ``[explain]``). Interpretive / non-causal masks; + not ResDMD and not ``ModeEnergyAttribution``. Tutorial + ``examples/46_representation_explain.ipynb``. Bib keys + ``Ying2019GNNExplainer``, ``Sundararajan2017IntegratedGradients``. +- CI: Ubuntu pytest on Python 3.10, 3.11, and 3.12 (coverage on 3.12); + macOS core smoke on 3.12 with + ``-m "not slow and not distributed and not ray"`` (hard-fail). Notebook + push shards extended through ``46_*``. + +### Changed + +- **Behavior change:** ``Model.save`` / the default checkpoint path write + ``safetensors_v1`` instead of a pickle-only format-1 ``.pt``. Pass + ``format="legacy_pt"`` only when you intentionally need the legacy pickle + layout. Prefer loading trusted ``legacy_pt`` files only (see + ``SECURITY.md``). + +### Notes + +- Portable inference export (TorchScript / ``torch.export`` / ONNX) is + **not** shipped in 0.13.0 (Phase 83 cut-candidate). Training checkpoints + (``safetensors_v1``) remain the supported weight-sharing path; ONNX stays + deferred. +- Windows is best-effort community support and is not in the CI matrix. +- Representation explain remains a homogeneous MVP (rejects hetero / + RelGraph / hypergraph / adaptive topology / delays / controls). +- Homogeneous training defaults are unchanged relative to 0.12.0. + ## [0.12.0] - 2026-08-05 Scale-and-honesty release on the 0.11 stack: Kronecker-sum exact spectrum diff --git a/CITATION.cff b/CITATION.cff index 12f8425..f6fa421 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -13,8 +13,8 @@ authors: repository-code: "https://github.com/tjkessler/KoopmanGraph" url: "https://koopmangraph.readthedocs.io/" license: Apache-2.0 -version: 0.12.0 -date-released: 2026-08-05 +version: 0.13.0 +date-released: 2026-08-12 doi: 10.5281/zenodo.21811756 identifiers: - type: doi diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9f1dfd1..48845fa 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -153,6 +153,9 @@ ruff format --check src/ tests/ # Tests with coverage gate (matches the CI 3.12 job; use -n auto for speed) pytest tests/ -n auto --cov=koopman_graph --cov-report=term-missing --cov-fail-under=90 + +# Optional: same marker expression planned for macOS smoke (TASK-2142) +# pytest tests/ -n auto -m "not slow and not distributed and not ray" ``` With uv after `uv sync --extra dev`: @@ -238,15 +241,25 @@ Branch protection on `main` is configured through a GitHub ruleset Mark the aggregator check from [`ci.yml`](.github/workflows/ci.yml) as **required** before merge: -- `ci` — succeeds when path-selected jobs (`lint`, `test`, `notebooks`, - `docs`) succeed or are intentionally skipped; fails if any selected job fails +- `ci` — succeeds when path-selected jobs (`lint`, `test`, `test-macos`, + `notebooks`, `docs`) succeed or are intentionally skipped; fails if any + selected job fails Individual jobs still appear on the workflow run for diagnosis: - `lint` — Ruff via `uvx` (no full torch install) -- `test (3.10)`, `test (3.12)` — pytest with `pytest-xdist`; coverage gate on 3.12 only +- `test (3.10)`, `test (3.11)`, `test (3.12)` — pytest with `pytest-xdist`; + coverage gate on 3.12 only. Ubuntu runs the full default suite (including + `@pytest.mark.ray` when `[ray]` / `[distributed]` extras are installed). + Opt-in multi-process DDP uses `@pytest.mark.distributed` plus + `KOOPMAN_GRAPH_DISTRIBUTED_TESTS=1`. +- `test (macos-3.12)` — macOS core smoke on `macos-latest` with + `-m "not slow and not distributed and not ray"` (hard-fail; no Ray/Dask + extras). Shares the same path filter as Ubuntu `test` (not the notebooks + filter). Windows is not in CI (best-effort community). - `notebooks (smoke)` on pull requests — 10 representative tutorials; on pushes - to `main`, `notebooks (01-13)` / `notebooks (14-26)` run the full suite + to `main`, `notebooks (01-13)` / `notebooks (14-26)` / `notebooks (27-46)` + run the full suite on **Ubuntu only** (macOS does not run nbmake) - `docs` — Sphinx documentation build (`sphinx-build -W`, warnings as errors) The **Draft paper** workflow ([`draft-pdf.yml`](.github/workflows/draft-pdf.yml)) diff --git a/README.md b/README.md index e3fbd99..c0dd2ee 100644 --- a/README.md +++ b/README.md @@ -135,11 +135,14 @@ Featured tutorials: [01 synthetic](https://github.com/tjkessler/KoopmanGraph/blo - [Scope and limitations](https://koopmangraph.readthedocs.io/en/latest/limitations.html) — when not to use; measured boundaries - [Architecture](https://koopmangraph.readthedocs.io/en/latest/architecture.html) — public vs power-user API layers - [FAQ / troubleshooting](https://koopmangraph.readthedocs.io/en/latest/faq.html) — install, imports, checkpoints -- [Installation](https://koopmangraph.readthedocs.io/en/latest/installation.html) — dependencies and install paths -- What’s new in 0.12.0: Kronecker-sum exact spectrum for eligible discrete - and continuous graph operators (three-path spectrum contract), Scale / - architecture honesty for residual dense inverse and eig-reg ceilings, and - docs-only topology-criticality honesty — see [CHANGELOG.md](CHANGELOG.md). +- [Installation](https://koopmangraph.readthedocs.io/en/latest/installation.html) — dependencies, install paths, and CI platforms +- [CLI](https://koopmangraph.readthedocs.io/en/latest/cli.html) — `koopman-graph train` / `predict` config workflow +- [SECURITY.md](SECURITY.md) — supported versions and checkpoint trust boundaries +- What’s new in 0.13.0: default `safetensors_v1` checkpoints (legacy pickle via + `format="legacy_pt"`), fit callbacks / CSV–TensorBoard tracking, config-driven + CLI, thin Ray Tune helpers, homogeneous `explain_representation` MVP, and + Ubuntu 3.10–3.12 plus macOS core-smoke CI — see [CHANGELOG.md](CHANGELOG.md). + Portable TorchScript / ONNX export is not shipped in this release. ## Related software @@ -164,7 +167,7 @@ If you use KoopmanGraph in research, please cite: publisher = {Zenodo}, doi = {10.5281/zenodo.21811756}, url = {https://github.com/tjkessler/KoopmanGraph}, - version = {0.12.0}, + version = {0.13.0}, } ``` diff --git a/SECURITY.md b/SECURITY.md index 9566bed..29a0f89 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -7,7 +7,7 @@ Security updates are provided for the **latest published release** of | Version lineage | Security updates | | --------------- | ---------------- | -| Latest PyPI release (currently the 0.8.x line) | Supported | +| Latest PyPI release (currently the 0.12.x line) | Supported | | Older published majors/minors | Not regularly backported | Check the installed version with: @@ -45,14 +45,38 @@ This channel is the same maintainer contact used for ## Loading checkpoints and dataset caches -Model checkpoints, conformal calibration files, hierarchical wrappers, and -on-disk teaching caches (``*.pt``) are loaded with PyTorch +Model checkpoints use one of two on-disk containers (logical architecture +schema remains format-1 either way): + +### Preferred: `safetensors_v1` directory or `.kgckpt` zip + +Default `GraphKoopmanModel.save` / `save_checkpoint` writes a **directory** +containing `meta.json`, `config.json` (architecture metadata), and +`model.safetensors` (weights). When the destination path ends in `.kgckpt` +or `.zip`, the same three files are packed into a zip bundle (still no +pickle). Weight tensors load through the safetensors library rather than +pickle. Prefer this container when sharing trained models. + +The JSON sidecar files are not executable code, but an untrusted checkpoint can +still carry misleading architecture metadata or adversarial weight values that +affect numerical results. Treat third-party directories or `.kgckpt` files as +untrusted inputs; inspect `config.json` before load when provenance is unclear. + +### Legacy: `legacy_pt` pickle file + +Single ``*.pt`` / ``*.pth`` files (and some training `checkpoint_path` +writers that still pin this format) are loaded with PyTorch ``torch.load(..., weights_only=False)``. That API deserializes Python objects -and can execute code from a malicious file. - -**Trust boundary:** only load checkpoints and cache artifacts that you created -yourself or obtained from a source you trust. Do not load untrusted ``.pt`` / -checkpoint files shared over the network or from unknown mirrors. Prefer -rebuilding teaching caches with the repository download scripts (SHA256-verified -where digests are pinned) rather than accepting third-party ``traffic.pt`` / -``contact.pt`` blobs. +and can execute code from a malicious file. Pass ``format="legacy_pt"`` only +when you need a single pickle file. + +**Trust boundary for `legacy_pt`:** only load pickle checkpoints that you +created yourself or obtained from a source you trust. Do not load untrusted +``.pt`` / checkpoint files shared over the network or from unknown mirrors. + +### Teaching caches and other ``*.pt`` artifacts + +Conformal calibration files, hierarchical wrappers, and on-disk teaching +caches (``*.pt``) remain pickle-based. Prefer rebuilding teaching caches with +the repository download scripts (SHA256-verified where digests are pinned) +rather than accepting third-party ``traffic.pt`` / ``contact.pt`` blobs. diff --git a/docs/source/api.rst b/docs/source/api.rst index 2153eef..eb764c1 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -239,14 +239,19 @@ the root façade. Specialized helpers (``compute_generator_spectrum``, ``koopman_spectral_clustering``, ``ClusteringResult``, ``estimate_coupling_from_snapshots``, ``CouplingEstimate``, ``attribute_mode_energy``, ``ModeEnergyAttribution``, -``evaluate_topology_transfer``, ``TopologyTransferReport``) are imported +``evaluate_topology_transfer``, ``TopologyTransferReport``, +``explain_representation``, ``RepresentationExplanation``) are imported from :mod:`koopman_graph.analysis` only. The helpers live in the ``spectrum`` / ``similarity`` / ``anomaly`` / ``plotting`` / ``residuals`` / -``sindy`` / ``clustering`` / ``topology_estimation`` / ``transfer`` -submodules. +``sindy`` / ``clustering`` / ``topology_estimation`` / ``transfer`` / +``explain`` submodules. ``attribute_mode_energy`` is an interpretive diagnostic on assembled :math:`K_{\mathrm{eff}}` (not a causal claim; not a ResDMD residual). -Finite-dictionary ResDMD +:func:`~koopman_graph.analysis.explain_representation` / +:class:`~koopman_graph.analysis.RepresentationExplanation` provide +homogeneous MVP node / edge / feature masks (GNNExplainer or optional +Captum integrated gradients) — **interpretive** / **non-causal**; not +``ModeEnergyAttribution`` and not ResDMD. Finite-dictionary ResDMD (:func:`~koopman_graph.analysis.resdmd`, :class:`~koopman_graph.analysis.ResDMDReport`) and the finite-matrix resolvent-norm grid @@ -265,7 +270,7 @@ advantage allowed) — see :doc:`limitations`. .. automodule:: koopman_graph.analysis :members: - :exclude-members: KoopmanSpectrum, SpectralResidualReport, spectral_residuals, ResDMDReport, resdmd, ResolventNormGrid, resolvent_norm_grid, TopologyTransferReport, evaluate_topology_transfer + :exclude-members: KoopmanSpectrum, SpectralResidualReport, spectral_residuals, ResDMDReport, resdmd, ResolventNormGrid, resolvent_norm_grid, TopologyTransferReport, evaluate_topology_transfer, RepresentationExplanation, explain_representation :show-inheritance: .. automodule:: koopman_graph.analysis.residuals @@ -284,6 +289,10 @@ advantage allowed) — see :doc:`limitations`. :members: :show-inheritance: +.. automodule:: koopman_graph.analysis.explain + :members: + :show-inheritance: + Long-horizon statistics (power-user) ------------------------------------ @@ -487,6 +496,38 @@ export. :exclude-members: GraphSnapshotSequence :show-inheritance: +Experiment tracking adapters (power-user) +----------------------------------------- + +Optional :class:`~koopman_graph.FitCallback` adapters live under +:mod:`koopman_graph.tracking` (not on root ``__all__``). Prefer +``from koopman_graph.tracking import CsvFitLogger, TensorBoardFitLogger``. +``TensorBoardFitLogger`` needs a peer ``tensorboard`` install. Cloud SDKs +(W&B, MLflow) are not pinned; see :doc:`faq` and +``examples/tracking/wandb_mlflow_callback.py``. + +.. automodule:: koopman_graph.tracking + :members: + :imported-members: + :show-inheritance: + +Hyperparameter search helpers (power-user) +------------------------------------------ + +Optional HPO helpers live under :mod:`koopman_graph.tuning` (not on root +``__all__``). Prefer ``from koopman_graph.tuning import fit_history_metrics, +run_ray_tune``. This package is **not** an AutoML product: search spaces +remain caller-owned; ``example_*`` scaffolds are smoke-only. +``run_ray_tune`` and the example spaces require +``pip install "koopman-graph[ray]"`` (lazy import). Optuna is +examples-only (no library Optuna API). See +``examples/scripts/ray_tune_koopman_example.py``. + +.. automodule:: koopman_graph.tuning + :members: + :imported-members: + :show-inheritance: + Distributed trainers and Dask prep (power-user) ------------------------------------------------ diff --git a/docs/source/architecture.rst b/docs/source/architecture.rst index f74b9c9..bc8d46b 100644 --- a/docs/source/architecture.rst +++ b/docs/source/architecture.rst @@ -168,6 +168,9 @@ leading-underscore helpers across peer modules. optimizer is constructed so orbit ``K_self`` parameters are trained. Prefer calling it only via ``run_fit_loop``; it is not a :mod:`koopman_graph.training` package ``__all__`` export. +* ``callbacks`` — ``FitCallback`` protocol and ``NoOpFitCallback`` (also + root-exported). Observe-only hooks invoked from ``run_fit_loop``; + adapters live in :mod:`koopman_graph.tracking` Do not invent a ``training/objectives/`` or ``training/loop/`` subtree. Prefer ``from koopman_graph.training import …``. Do not import @@ -547,7 +550,7 @@ Encoder and decoder remain peers: both import from ``nn.gnn``; neither imports the other. ``koopman_graph.analysis`` package layout (spectrum / similarity / anomaly / -plotting / residuals / topology estimation / SINDy / clustering): +plotting / residuals / topology estimation / SINDy / clustering / explain): * ``spectrum`` — re-exports neutral-leaf ``compute_spectrum``, ``compute_generator_spectrum``, and ``discrete_spectrum_at_delta_t``, plus @@ -582,6 +585,13 @@ plotting / residuals / topology estimation / SINDy / clustering): leading Koopman mode embeddings; sign/ordering canonicalized by magnitude-sorted ``|λ|`` and largest-magnitude positive pivot). Cluster quality inherits operator / spectrum quality. +* ``explain`` — + :func:`~koopman_graph.analysis.explain_representation` and frozen + :class:`~koopman_graph.analysis.RepresentationExplanation` + (homogeneous MVP node / edge / feature masks via GNNExplainer or optional + Captum integrated gradients; **interpretive** / **non-causal**; distinct + from :func:`~koopman_graph.analysis.attribute_mode_energy` and from + ResDMD — see :doc:`limitations`) * ``resdmd`` — :func:`~koopman_graph.analysis.resdmd` and frozen :class:`~koopman_graph.analysis.ResDMDReport` (finite-dictionary ResDMD @@ -805,7 +815,7 @@ code stays valid; capability packages alternate import paths. Do not silently demote either without a separately versioned breaking migration. -**Keep in** ``koopman_graph.__all__`` (core workflow; exactly these 35): +**Keep in** ``koopman_graph.__all__`` (core workflow; exactly these 37): * :class:`~koopman_graph.model.GraphKoopmanModel` * :class:`~koopman_graph.nn.encoder.GNNEncoder`, @@ -850,6 +860,8 @@ versioned breaking migration. * Primary spectrum entrypoints: :class:`~koopman_graph.spectrum_types.KoopmanSpectrum`, :func:`~koopman_graph.analysis.compute_spectrum` +* Fit observe hooks (Q11): :class:`~koopman_graph.training.FitCallback`, + :class:`~koopman_graph.training.NoOpFitCallback` * ``__version__`` Dataset generators remain via :mod:`koopman_graph.datasets` (not root @@ -865,6 +877,7 @@ Dataset generators remain via :mod:`koopman_graph.datasets` (not root ``KoopmanSparsityLoss``, ``WorstCaseReconstructionLoss``, ``LieConsistencyLoss``, ``PDEResidualLoss`` * :mod:`koopman_graph.training` — ``FitHistory``, ``LossWeights`` + (``FitCallback`` / ``NoOpFitCallback`` stay on root ``__all__`` as well) * :mod:`koopman_graph.adaptation` — ``RecursiveKoopmanAdapter``, ``AdaptationStepResult``, ``KoopmanObserver``, ``FilterResult`` * :mod:`koopman_graph.env` — ``GraphKoopmanEnv`` @@ -878,7 +891,8 @@ Dataset generators remain via :mod:`koopman_graph.datasets` (not root ``detect_anomaly``, ``calibrate_anomaly_threshold``, ``AnomalyDetectionResult``, ``estimate_coupling_from_snapshots``, ``CouplingEstimate``, ``identify_sparse_dynamics``, ``SINDyReport``, - ``koopman_spectral_clustering``, ``ClusteringResult`` + ``koopman_spectral_clustering``, ``ClusteringResult``, + ``explain_representation``, ``RepresentationExplanation`` * :mod:`koopman_graph.observables` — ``graph_laplacian_features`` * :mod:`koopman_graph.uq` — ``EnsembleGraphKoopmanModel``, ``LatentGaussianKoopmanUQ``, ``ConformalKoopmanUQ`` @@ -913,7 +927,21 @@ change. Examples: * :mod:`koopman_graph.training` — loss schedules, ``train_one_epoch``, - ``run_fit_loop``, ``resolve_device``, and related helpers + ``run_fit_loop``, ``FitCallback`` / ``NoOpFitCallback``, + ``resolve_device``, and related helpers (``FitCallback`` / + ``NoOpFitCallback`` are also root-exported) +* :mod:`koopman_graph.tracking` — optional CSV / TensorBoard fit adapters + implementing ``FitCallback`` (not on root ``__all__``; may import + ``training`` types; ``training`` must not import tracking adapters) +* :mod:`koopman_graph.tuning` — hyperparameter-search helpers + (``fit_history_metrics``, lazy ``run_ray_tune``, example-only search + scaffolds). Not an AutoML product; search spaces stay caller-owned. + Off root ``__all__``. May import ``training`` types; ``training`` must + not import ``tuning``. Importing the package does not import Ray +* :mod:`koopman_graph.cli` — config-driven ``koopman-graph`` console script + (``train`` / ``predict``; JSON in core, YAML via ``[cli]``). Off root + ``__all__``. Leaf façade: may import ``model`` / ``serialization`` / + ``data``; other library packages must **not** import ``cli`` * :mod:`koopman_graph.distributed` — optional DDP / Fabric trainer adapters, ``KoopmanLightningModule`` Trainer sugar, ``run_ray_train_fit_loop``, ``fit_ensemble_with_ray``, rank helpers, and @@ -924,7 +952,12 @@ Examples: * :mod:`koopman_graph.datasets.molecular` — synthetic MD / contact-graph teaching oracles and optional ``[md]`` I/O stubs * :mod:`koopman_graph.serialization` — checkpoint build/load internals behind - ``GraphKoopmanModel.save`` / ``load`` + ``GraphKoopmanModel.save`` / ``load``. Default on-disk container is the + ``safetensors_v1`` layout (``meta.json``, ``config.json``, + ``model.safetensors``) as a directory or a ``.kgckpt`` / ``.zip`` bundle; + ``format="legacy_pt"`` retains a single pickle ``.pt`` file. Logical + schema remains ``FORMAT_VERSION`` 1. Trust boundaries are documented in + the repository ``SECURITY.md``. * :mod:`koopman_graph.datasets.dynamics` — Laplacian diffusion primitives and shared generation validators used by benchmark generators (dense ``I - alpha * L_sym`` step operators; shares pseudoinverse-normalized @@ -1727,8 +1760,27 @@ Façade vs functional training ownership * **Functional loop** — :func:`~koopman_graph.training.run_fit_loop` owns device placement, optimizer / scheduler setup, epoch iteration (``train_one_epoch`` / ``train_windowed_epoch`` / ``eval_one_epoch``), early - stopping, best-weight tracking, optional checkpoint writes, and history - assembly. Prefer extending that helper over growing the model class. + stopping, best-weight tracking, optional checkpoint writes, history + assembly, and observe-only ``callbacks`` hooks + (``on_fit_start`` / ``on_epoch_end`` / ``on_fit_end``). Prefer extending + that helper over growing the model class. +* **Tracking adapters** — :mod:`koopman_graph.tracking` implements + :class:`~koopman_graph.FitCallback` for CSV and TensorBoard. Cloud SDKs + (W&B, MLflow, …) stay out of core; DIY callbacks are documented in + :doc:`faq` / ``examples/tracking/``. Layer rule: ``tracking`` → + ``training`` types only; never ``training`` → tracking adapters. +* **HPO helpers** — :mod:`koopman_graph.tuning` flattens + :class:`~koopman_graph.training.FitHistory` for trial reporters and + wraps Ray Tune behind a lazy façade (``[ray]``). It does **not** fork + loss mathematics or invent search spaces. Example-only + ``example_*`` scaffolds are smoke ranges, not scientific defaults. + Optuna remains examples-only (no library Optuna API). Layer rule: + ``tuning`` → ``training`` types only; never ``training`` → ``tuning``. +* **Distributed / Lightning logging** — ``fit(..., callbacks=...)`` is + single-process only (``strategy="ddp"`` raises if callbacks are set). + For :class:`~koopman_graph.distributed.KoopmanLightningModule`, attach + Lightning loggers to ``Trainer``; do not expect a parallel FitCallback + stack on the Lightning path in this release. * **Neighbor-subgraph training** — pass :class:`~koopman_graph.data.NeighborWindowSampler` via ``fit(..., sampler=...)`` (mutually exclusive with ``window_length``). Training losses (rollout, @@ -2100,6 +2152,11 @@ payloads are **deprecated** and rejected with a clear re-save error (no silent migration). Other ``format_version`` values remain unsupported. Formal multi-version checkpoint tracking and migration branches begin at 1.0. +On disk, that schema is stored in the default ``safetensors_v1`` container +(directory or ``.kgckpt`` / ``.zip`` bundle) or an explicit ``legacy_pt`` +pickle file (see :mod:`koopman_graph.serialization` and the repository +``SECURITY.md``). The container choice does not bump ``FORMAT_VERSION``. + v0.5.0 capability architecture ------------------------------ @@ -2650,6 +2707,58 @@ TDL parity, PyEMMA-scale MD, guaranteed sample-efficiency) live in - Power-user - Exact Aut only; self-block MVP; no sample-efficiency guarantee +v0.13.0 capability architecture +------------------------------- + +Ops-and-trust surfaces preserve the existing façade / package / power-user +tiers. Do **not** promote ``cli``, ``tracking``, ``tuning``, or explain +helpers onto root ``__all__`` beyond the curated ``FitCallback`` / +``NoOpFitCallback`` pair (Q11). Portable TorchScript / ONNX export remains +cut for this release (see :doc:`limitations`). + +**Shipped capability list (0.13 peers)** + +* Default ``safetensors_v1`` checkpoint container behind + :mod:`koopman_graph.serialization` (``legacy_pt`` escape hatch) +* :class:`~koopman_graph.FitCallback` / :class:`~koopman_graph.NoOpFitCallback` + on the root façade; :mod:`koopman_graph.tracking` CSV / TensorBoard adapters +* Console script package :mod:`koopman_graph.cli` (leaf; off root ``__all__``) +* :mod:`koopman_graph.tuning` Ray Tune helpers (off root ``__all__``; Optuna + examples-only) +* :mod:`koopman_graph.analysis.explain` homogeneous + ``explain_representation`` MVP (package ``analysis.__all__`` only) +* CI: Ubuntu 3.10–3.12 + macOS core smoke (not Windows) + +**Capability and API-tier map (0.13 peers)** + +.. list-table:: + :header-rows: 1 + :widths: 24 18 58 + + * - Capability home + - API tier + - Contracts to preserve + * - :mod:`koopman_graph.serialization` (``safetensors_v1``) + - Package (behind ``Model.save`` / ``load``) + - Default safe container; ``FORMAT_VERSION`` 1 unchanged; trust + boundaries in ``SECURITY.md`` + * - :mod:`koopman_graph.training.callbacks` + - Public root façade (``FitCallback`` / ``NoOpFitCallback``) + - Observe-only; single-process ``fit``; no ``training`` → ``tracking`` + * - :mod:`koopman_graph.tracking` + - Power-user + - ``tracking`` → ``training`` types only; no cloud SDK pins + * - :mod:`koopman_graph.cli` + - Power-user / console script + - Leaf package; no reverse imports into library cores + * - :mod:`koopman_graph.tuning` + - Power-user + - Lazy Ray; caller-owned search spaces; never ``training`` → ``tuning`` + * - :mod:`koopman_graph.analysis.explain` + - Package (``analysis.__all__``) + - Homogeneous MVP; interpretive / non-causal; ≠ ResDMD / + ``ModeEnergyAttribution`` + Related documentation --------------------- @@ -2657,6 +2766,7 @@ Related documentation * :doc:`capabilities` — feature inventory * :doc:`limitations` — scope boundaries and when not to use * :doc:`quickstart` — end-user training and prediction walkthrough +* :doc:`cli` — config-driven console script * :doc:`faq` — install / import / checkpoint troubleshooting and support routing * Repository ``CONTRIBUTING.md`` — development setup and contribution workflow * Repository ``SECURITY.md`` — supported versions and private vulnerability reporting diff --git a/docs/source/capabilities.rst b/docs/source/capabilities.rst index c9fd568..db02fe4 100644 --- a/docs/source/capabilities.rst +++ b/docs/source/capabilities.rst @@ -241,9 +241,12 @@ inverse / spectrum; see :doc:`faq` and :doc:`limitations`). opt-in (``KOOPMAN_GRAPH_DISTRIBUTED_TESTS=1``); default PR CI does not require multi-node hardware. -Ray Tune is **examples-only**: see -``examples/scripts/ray_tune_koopman_example.py`` (search space stays in the -script; the library does not expose a Tune / AutoML API). Optional +Ray Tune HPO is supported via power-user helpers under +:mod:`koopman_graph.tuning` (``fit_history_metrics``, ``run_ray_tune``, +optional ``example_*`` smoke scaffolds) and the examples-only script +``examples/scripts/ray_tune_koopman_example.py``. Search spaces remain +**user-defined / script-owned**; the library is not an AutoML product. +Optuna is **examples-only** (no library Optuna API). Optional ``pip install "koopman-graph[dask]"`` activates :mod:`koopman_graph.distributed.dask_prep` helpers (``materialize_sequences``, ``materialize_window_index_list``) for offline @@ -263,6 +266,16 @@ Analysis * ``attribute_mode_energy`` / ``ModeEnergyAttribution`` — interpretive type / relation energy fractions on assembled ``K_eff`` (not causal; not a ResDMD residual on relation-attributed modes) +* ``explain_representation`` / ``RepresentationExplanation`` — + representation-level node / edge / feature masks on a homogeneous + snapshot for targets ``latent``, ``one_step_forecast``, or + ``reconstruction``. Default algorithm ``gnn_explainer`` uses PyG + GNNExplainer (``Ying2019GNNExplainer``); optional + ``algorithm="integrated_gradients"`` uses Captum via the ``[explain]`` + extra (``Sundararajan2017IntegratedGradients``). Masks are + **interpretive** and **non-causal**; they are **not** + ``ModeEnergyAttribution`` and **not** ResDMD residuals. Homogeneous MVP + only — see :doc:`limitations` * ``spectral_residuals`` / ``SpectralResidualReport`` — held-out data-driven residuals and ``trustworthy_mask()``; optional ``plot_spectrum(..., annotate_untrustworthy=True)``. Diagnostic in the @@ -364,7 +377,51 @@ Research tooling ``koopman_graph.baselines.gnn`` (``ForecasterProtocol`` deviation tables; not dedicated-library SOTA or leaderboard-matched protocols) * Benchmark datasets and Jupyter tutorials under ``examples/`` -* Model ``save`` / ``load`` checkpoints; ≥90% coverage enforced in CI +* Model ``save`` / ``load`` checkpoints (default ``safetensors_v1`` + directory or ``.kgckpt`` / ``.zip`` bundle; ``legacy_pt`` pickle escape + hatch); ≥90% coverage enforced in CI. See the repository + ``SECURITY.md`` for load trust boundaries. + +Experiment tracking (fit callbacks) +----------------------------------- + +Observe-only hooks on the functional fit path (single-process +``GraphKoopmanModel.fit`` / :func:`~koopman_graph.training.run_fit_loop`): + +* :class:`~koopman_graph.FitCallback` protocol and + :class:`~koopman_graph.NoOpFitCallback` (root-exported) +* In-tree adapters under :mod:`koopman_graph.tracking` (power-user; not on + root ``__all__``): :class:`~koopman_graph.tracking.CsvFitLogger` (stdlib + CSV) and :class:`~koopman_graph.tracking.TensorBoardFitLogger` (requires + peer ``pip install tensorboard``) +* Callbacks must not mutate model parameters; default ``callbacks=None`` + preserves prior ``FitHistory`` behavior +* ``fit(..., strategy="ddp")`` does not accept ``callbacks`` yet — use + single-process fit or Lightning loggers for distributed runs +* Lightning ``Trainer`` users should attach Lightning loggers to + :class:`~koopman_graph.distributed.KoopmanLightningModule`; this release + does not add a second FitCallback stack on Trainer +* Weights & Biases / MLflow are intentionally not core dependencies — + see :doc:`faq` and ``examples/tracking/wandb_mlflow_callback.py`` + +Hyperparameter search helpers +----------------------------- + +Power-user package :mod:`koopman_graph.tuning` (not on root ``__all__``): + +* :func:`~koopman_graph.tuning.fit_history_metrics` — flatten a + :class:`~koopman_graph.training.FitHistory` into scalar floats for + Tune / Optuna reporters +* :func:`~koopman_graph.tuning.run_ray_tune` — thin Ray Tune façade + (lazy import; requires ``pip install "koopman-graph[ray]"``) +* :func:`~koopman_graph.tuning.example_lr_loguniform_space` / + :func:`~koopman_graph.tuning.example_lr_latent_dim_space` — *example* + smoke scaffolds only (not scientific defaults) + +This is **not** an AutoML product. Searchable hyperparameters and ranges +remain caller-owned. Optuna is **examples-only** (no library Optuna API). +See ``examples/scripts/ray_tune_koopman_example.py`` and +``examples/scripts/README.md``. Stability mode selection ------------------------ diff --git a/docs/source/cli.rst b/docs/source/cli.rst new file mode 100644 index 0000000..7fd7b03 --- /dev/null +++ b/docs/source/cli.rst @@ -0,0 +1,116 @@ +Command-line interface +====================== + +The ``koopman-graph`` console script is a config-driven façade over the library +``fit`` / ``predict`` / ``save`` / ``load`` path. It does not fork training +mathematics. Prefer the Python API when you need custom encoders, distributed +strategies, or experiment callbacks; use the CLI for reproducible train and +predict smokes from JSON (or YAML) configs. + +Install the package as usual (:doc:`installation`). The console script is then +available as ``koopman-graph``. YAML configs additionally need +``pip install "koopman-graph[cli]"`` (PyYAML); JSON configs do not. + +Subcommands +----------- + ++-------------------+----------------------------------------------------------+ +| Command | Role | ++===================+==========================================================+ +| ``version`` | Print ``koopman-graph {version}`` and exit 0. | +| | Equivalent to ``--version``. | ++-------------------+----------------------------------------------------------+ +| ``train`` | Load a train config, build an allowlisted model and | +| | sequence, call ``fit``, write a checkpoint | +| | (default ``safetensors_v1``). | ++-------------------+----------------------------------------------------------+ +| ``predict`` | Load a checkpoint, take the first snapshot from | +| | ``--data``, roll out ``--steps``, write a ``.pt`` | +| | forecast payload. | ++-------------------+----------------------------------------------------------+ + +Global flags include ``--help`` and ``--version``. + +Train +~~~~~ + +.. code-block:: bash + + koopman-graph train \ + --config examples/cli/synthetic_train.json \ + --out /tmp/kg-cli + +* ``--config`` (required) — JSON always; ``.yaml`` / ``.yml`` require + ``pip install "koopman-graph[cli]"`` (PyYAML). +* ``--out`` (optional) — when set, relative ``checkpoint.path`` values resolve + under this directory. + +On success, stdout reports ``wrote checkpoint: …`` and the process exits 0. + +Predict +~~~~~~~ + +.. code-block:: bash + + koopman-graph predict \ + --checkpoint /tmp/kg-cli/model.kgckpt \ + --data examples/cli/synthetic_train.json \ + --steps 5 \ + --out /tmp/kg-cli/forecast.pt + +* ``--checkpoint`` — ``safetensors_v1`` directory or ``.kgckpt`` zip, or a + legacy ``.pt`` checkpoint (auto-detected). +* ``--data`` — trusted ``.pt`` ``GraphSnapshotSequence`` (or list of + ``Data``), or JSON/YAML with a ``data`` section (or a bare data mapping with + ``kind``). +* ``--steps`` — autoregressive horizon (default 5). +* ``--out`` — destination ``.pt`` file. + +The forecast payload is a ``torch.save`` dict with keys ``steps``, +``forecasts`` (list of PyG ``Data``), and ``summary``. + +Exit codes +---------- + +* ``0`` — success. +* ``1`` — validation, I/O, or runtime error; a short ``error: …`` message is + written to stderr (unknown config keys use dotted paths such as + ``model.not_allowed``). + +Train config schema (MVP) +------------------------- + +Top-level sections: ``model``, ``data``, optional ``fit``, optional +``checkpoint``. Unknown keys are rejected. + +**model.** Required: ``encoder``, ``in_channels``, ``hidden_channels``, +``latent_dim``. Encoder kinds in the CLI MVP: ``gcn``, ``gat``, ``sage`` +(matched decoder peers). Optional passthrough includes ``num_layers``, +``time_step``, ``dynamics_mode``, and other allowlisted factory kwargs. + +**data.** Required ``kind``: + +* ``synthetic_path`` — seeded path-graph decay trajectory + (``num_nodes``, ``num_timesteps``, ``seed``, ``feature_dim`` / ``in_channels``). +* ``cached_sequence`` — requires ``path`` to a trusted pickle ``.pt`` sequence + (same trust boundary as other teaching caches; see the security notes in the + repository). + +**fit.** Allowlisted ``GraphKoopmanModel.fit`` kwargs (for example ``epochs``, +``lr``, ``device``). + +**checkpoint.** ``path`` (file or directory); optional ``format`` +(``safetensors_v1`` default, or ``legacy_pt``). + +A minimal runnable example ships at ``examples/cli/synthetic_train.json`` +(see ``examples/cli/README.md``). + +Limitations +----------- + +* Heterogeneous, sheaf, cell, and control-heavy stacks are not exposed in the + CLI MVP; use the Python API. +* The CLI package is a power-user façade: other library modules must not import + ``koopman_graph.cli``. +* Experiment tracking and HPO remain separate surfaces (callbacks / ``tuning``); + the CLI does not pin cloud SDKs. diff --git a/docs/source/faq.rst b/docs/source/faq.rst index b1c849c..19bf19a 100644 --- a/docs/source/faq.rst +++ b/docs/source/faq.rst @@ -24,6 +24,13 @@ that the installed ``torch`` version and CUDA tag match the PyG wheel index you used. With uv, ``uv pip install torch --torch-backend=auto`` (or a specific backend) helps pick a matching PyTorch index; see :doc:`installation`. +Which platforms does CI cover? +------------------------------ + +CI runs the full Ubuntu test matrix on Python 3.10–3.12 and a macOS core smoke +job on Python 3.12. Windows is best-effort community support (not in CI). See +:doc:`installation` (Supported platforms / CI). + CUDA vs CPU mismatches ---------------------- @@ -273,9 +280,11 @@ Use the power-user :mod:`koopman_graph.distributed` helpers (not root Prefer a picklable (ideally module-level) factory. This does **not** change UQ coverage guarantees and does **not** shard one model across GPUs — that is Ray Train / DDP / Fabric. -* **Ray Tune HPO** (examples-only) — - ``examples/scripts/ray_tune_koopman_example.py``. The search space stays - in the script; KoopmanGraph does not expose a Tune / AutoML API. +* **Ray Tune HPO** — power-user helpers in :mod:`koopman_graph.tuning` + (``fit_history_metrics``, ``run_ray_tune``, optional ``example_*`` + smoke scaffolds) plus ``examples/scripts/ray_tune_koopman_example.py``. + The search configuration stays script-/caller-owned; KoopmanGraph is + not an AutoML product. Optuna is examples-only (no library Optuna API). Default ``fit`` / ``run_fit_loop`` remain single-process when ``strategy`` is unset. Distributed training does **not** reduce dense @@ -355,6 +364,38 @@ Typical pattern:: Do **not** replace :func:`~koopman_graph.training.run_fit_loop` / DDP / Fabric with a Dask-worker training loop. +How do I log fits to CSV, TensorBoard, W&B, or MLflow? +------------------------------------------------------- + +Use observe-only :class:`~koopman_graph.FitCallback` hooks on +single-process ``fit`` / :func:`~koopman_graph.training.run_fit_loop`: + +* **CSV / TensorBoard (in-tree)** — + :class:`~koopman_graph.tracking.CsvFitLogger` and + :class:`~koopman_graph.tracking.TensorBoardFitLogger`:: + + from koopman_graph.tracking import CsvFitLogger, TensorBoardFitLogger + + model.fit( + sequence, + epochs=20, + callbacks=[ + CsvFitLogger("runs/fit.csv"), + TensorBoardFitLogger("runs/tb"), # needs: pip install tensorboard + ], + ) + +* **W&B / MLflow (DIY, no library pin)** — implement ``FitCallback`` and call + ``wandb.log`` / ``mlflow.log_metrics`` from ``on_epoch_end`` using + ``train_breakdown.to_floats()``. Sketch classes live in + ``examples/tracking/wandb_mlflow_callback.py``. Install those SDKs + yourself; KoopmanGraph does not declare them as dependencies. +* **Lightning Trainer** — attach Lightning loggers to + :class:`~koopman_graph.distributed.KoopmanLightningModule` / + ``Trainer``; do not expect ``fit(..., callbacks=)`` on the Lightning + path. Native ``strategy="ddp"`` rejects non-None ``callbacks`` until + that path is wired. + What is the difference between ``spectral_residuals`` and ResDMD? ----------------------------------------------------------------- diff --git a/docs/source/index.rst b/docs/source/index.rst index bab34d4..69a64eb 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -25,6 +25,7 @@ for the public vs power-user API contract. installation quickstart + cli faq capabilities limitations diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 49a2e6f..4d23573 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -26,6 +26,19 @@ With uv you can also pin an interpreter for the project: uv python pin 3.12 +Supported platforms / CI +~~~~~~~~~~~~~~~~~~~~~~~~ + +Continuous integration verifies: + +* **Ubuntu** — full pytest matrix on Python **3.10**, **3.11**, and **3.12** + (coverage gate on 3.12). +* **macOS** — core smoke on Python **3.12** + (``-m "not slow and not distributed and not ray"``). + +**Windows** is best-effort community support and is **not** in the CI matrix. +Package classifiers list Python 3.10–3.12 to match the tested interpreters. + PyTorch ~~~~~~~ @@ -95,6 +108,20 @@ Optional extras for symmetry-adapted operators (node-orbit partitions via pip install -e ".[symmetry]" +Optional extras for the config-driven ``koopman-graph`` CLI when using YAML +configs (JSON configs need no extra; PyYAML only): + +.. code-block:: bash + + pip install -e ".[cli]" + +Optional extras for Captum-backed representation attribution +(``algorithm="integrated_gradients"``; GNNExplainer needs no extra): + +.. code-block:: bash + + pip install -e ".[explain]" + Optional extras for Koopman model-predictive control (OSQP QP solver): .. code-block:: bash @@ -146,7 +173,12 @@ encode to invariant latents; latent :math:`K` remains non-equivariant). ``[baselines-graphcast]`` is a reserved no-op — the GraphCast teaching path is pure PyTorch. Exact-automorphism isotypic mode (``koopman_symmetry="isotypic"``) needs optional ``pynauty`` separately -from ``[symmetry]`` (``networkx`` approximate orbits). +from ``[symmetry]`` (``networkx`` approximate orbits). The ``[cli]`` extra +pins PyYAML so ``koopman-graph train --config *.yaml`` can load; the +console script itself ships with the core package (see below). The +``[explain]`` extra pins Captum for +``explain_representation(..., algorithm="integrated_gradients")``; the +PyG GNNExplainer path needs no extra. uv (project sync) ~~~~~~~~~~~~~~~~~ @@ -165,13 +197,14 @@ need a non-default accelerator; otherwise sync is enough for CPU development: Documentation and optional capability extras (``uv.lock`` includes ``mpc`` / ``symmetry`` for frozen sync; add ``lightning`` / ``distributed`` when you -need Fabric locally): +need Fabric locally, or ``cli`` for YAML configs): .. code-block:: bash uv sync --extra docs uv sync --extra mpc --extra symmetry uv sync --extra lightning + uv sync --extra cli cd docs && make html pip-compatible uv installs (after creating a venv) mirror the pip commands: @@ -201,13 +234,38 @@ version with the tag you intend; see PyPI and ``CHANGELOG.md``): .. code-block:: bash - pip install koopman-graph==0.12.0 - # or: uv pip install koopman-graph==0.12.0 + pip install koopman-graph==0.13.0 + # or: uv pip install koopman-graph==0.13.0 + +YAML CLI configs from PyPI: + +.. code-block:: bash + + pip install "koopman-graph[cli]" + # or: uv pip install "koopman-graph[cli]" Releases are published automatically when a maintainer creates a GitHub Release (see ``CONTRIBUTING.md`` in the repository). For the latest in-tree development checkout, use the editable install from source above. +Console script +-------------- + +A normal install registers the ``koopman-graph`` console script (packaging +metadata: ``[project.scripts]`` → ``koopman_graph.cli:main``). After +``pip install`` / ``pip install -e .`` / ``uv sync``, the entry point should be +on your ``PATH``: + +.. code-block:: bash + + koopman-graph --version + koopman-graph train --help + koopman-graph predict --help + +JSON train configs work with the core install. YAML requires +``pip install "koopman-graph[cli]"`` (or ``-e ".[cli]"``). See :doc:`cli` for +subcommands and schema, and ``examples/cli/`` for a synthetic smoke config. + Verify ------ @@ -218,6 +276,13 @@ Confirm the package imports: python -c "import koopman_graph; print(koopman_graph.__version__)" # or: uv run python -c "import koopman_graph; print(koopman_graph.__version__)" +Confirm the console script: + +.. code-block:: bash + + koopman-graph --version + # or: uv run koopman-graph --version + For a full development check after ``pip install -e ".[dev]"`` or ``uv sync --extra dev``: @@ -229,4 +294,5 @@ For a full development check after ``pip install -e ".[dev]"`` or Next steps ---------- -See :doc:`quickstart` for a minimal train-and-predict workflow. +See :doc:`quickstart` for a minimal Python train-and-predict workflow, and +:doc:`cli` for the config-driven console script. diff --git a/docs/source/limitations.rst b/docs/source/limitations.rst index 24cdc19..d96a215 100644 --- a/docs/source/limitations.rst +++ b/docs/source/limitations.rst @@ -195,6 +195,19 @@ Spectral analysis (Colbrook & Townsend, *Commun. Pure Appl. Math.*, 2023/2024; ``ColbrookTownsend2023ResDMD``). See ``examples/40_resdmd_pseudospectra.ipynb``. +* **Representation explain masks are interpretive, not causal.** + :func:`~koopman_graph.analysis.explain_representation` returns + :class:`~koopman_graph.analysis.RepresentationExplanation` node / edge / + feature masks under GNNExplainer (``Ying2019GNNExplainer``) or optional + Captum integrated gradients (``Sundararajan2017IntegratedGradients``; + ``[explain]`` extra). Masks are **non-causal**: they are not + interventional topology recovery, not a ResDMD residual bound, and not + complementary operator-level + :class:`~koopman_graph.analysis.ModeEnergyAttribution` on assembled + ``K_eff``. The MVP is **homogeneous** ``Data`` only and rejects + ``HeteroData``, RelGraph / ``hetero_graph``, hypergraph encoders / + operators, adaptive / learned pairwise topology, ``n_delays != 1``, and + ``control_dim > 0``. * **Largest Lyapunov exponent (Rosenstein) is available but sensitive.** ``koopman_graph.statistics.largest_lyapunov_exponent`` implements the Rosenstein et al. (1993) small-data-set method @@ -433,6 +446,10 @@ When to use something else ResDMD / spectral-measure packages when you need residual-based certificates beyond the finite-dictionary MVP and held-out ``spectral_residuals`` diagnostics. +* **Portable inference export (TorchScript / ONNX).** 0.13.0 does **not** + ship a TorchScript / ``torch.export`` / ONNX inference-export API. + Share weights with the default ``safetensors_v1`` checkpoint path (or + explicit ``legacy_pt`` when required). ONNX remains deferred. * **Flat-vector deep Koopman without graph structure.** Packages such as PyKoopman and DLKoopman remain appropriate when the state is naturally a vector and edge-based coupling is not part of the model. diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst index 9ad0e87..d74c56e 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart.rst @@ -81,6 +81,28 @@ Roll out from the first snapshot in the sequence: print(f"Predicted {len(future_graphs)} future snapshots") print(f"First prediction shape: {future_graphs[0].x.shape}") +CLI path (same workflow) +------------------------ + +The console script runs the same library ``fit`` / ``predict`` / checkpoint +path from a JSON config. After installing the package (see +:doc:`installation`), a tiny CPU smoke looks like: + +.. code-block:: bash + + koopman-graph train \ + --config examples/cli/synthetic_train.json \ + --out /tmp/kg-cli + + koopman-graph predict \ + --checkpoint /tmp/kg-cli/model.kgckpt \ + --data examples/cli/synthetic_train.json \ + --steps 5 \ + --out /tmp/kg-cli/forecast.pt + +YAML configs require the optional ``[cli]`` extra. Subcommands, exit codes, +and the MVP config schema are documented in :doc:`cli`. + Inspect the spectrum -------------------- diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 40d2a69..c04737e 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -88,6 +88,8 @@ Analysis and stability - Koopman eigenvalue analysis + held-out ``spectral_residuals`` / spurious-mode filter * - `40_resdmd_pseudospectra.ipynb `_ - Finite-dictionary ResDMD MVP + resolvent-norm grid (≠ ``spectral_residuals``; not infinite-dim certificates) + * - `46_representation_explain.ipynb `_ + - Homogeneous ``explain_representation`` masks (GNNExplainer; optional Captum IG) — interpretive / non-causal; ≠ ResDMD / ``ModeEnergyAttribution`` * - `08_loss_stability.ipynb `_ - Loss weighting and soft stability regularization * - `26_sparse_interpretable_operator.ipynb `_ @@ -176,8 +178,12 @@ the script):: python examples/scripts/ray_tune_koopman_example.py --epochs 1 --num-samples 2 +For fit-callback tracking sketches (CSV / TensorBoard in-tree; optional +W&B / MLflow DIY without library SDK pins), see +``examples/tracking/wandb_mlflow_callback.py`` and :doc:`faq`. + See ``examples/scripts/README.md`` and :doc:`capabilities` (Distributed -training). +training; Experiment tracking). Heterogeneous (multi-relational) graphs --------------------------------------- diff --git a/examples/46_representation_explain.ipynb b/examples/46_representation_explain.ipynb new file mode 100644 index 0000000..245f433 --- /dev/null +++ b/examples/46_representation_explain.ipynb @@ -0,0 +1,567 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "6a382d75", + "metadata": {}, + "source": [ + "# Representation-level explanations for GraphKoopman models\n", + "\n", + "This notebook shows how to attribute a homogeneous `GraphKoopmanModel` quantity\n", + "with `explain_representation`, inspect GNNExplainer node/edge masks on a small\n", + "path graph, and read those masks under the package honesty contract.\n", + "\n", + "**Honesty banner.** Masks are **interpretive** attributions under a chosen\n", + "algorithm. They are **non-causal** (not interventional topology recovery),\n", + "**not** ResDMD residual bounds, and **not**\n", + "`ModeEnergyAttribution` (the complementary operator-level energy diagnostic on\n", + "assembled `K_eff`). The MVP is homogeneous `Data` only.\n" + ] + }, + { + "cell_type": "markdown", + "id": "3d632b98", + "metadata": {}, + "source": [ + "## Setup\n", + "\n", + "Imports, seeds, package/torch versions, and a non-interactive Matplotlib\n", + "backend under pytest / CI.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "c7b8d626", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "koopman_graph 0.12.0\n", + "torch 2.13.0\n", + "CI_lite=False fit_epochs=15 explain_epochs=40\n" + ] + } + ], + "source": [ + "import os\n", + "import random\n", + "import warnings\n", + "\n", + "from tqdm.std import TqdmWarning\n", + "\n", + "warnings.filterwarnings(\"ignore\", category=TqdmWarning)\n", + "\n", + "import matplotlib\n", + "\n", + "if os.environ.get(\"PYTEST_CURRENT_TEST\"):\n", + " matplotlib.use(\"Agg\")\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "from torch_geometric.data import Data\n", + "\n", + "try:\n", + " from IPython import get_ipython\n", + "\n", + " if get_ipython() is not None and not os.environ.get(\"PYTEST_CURRENT_TEST\"):\n", + " get_ipython().run_line_magic(\"matplotlib\", \"inline\")\n", + "except (ImportError, NameError):\n", + " pass\n", + "\n", + "import koopman_graph\n", + "from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel\n", + "from koopman_graph.analysis import explain_representation\n", + "\n", + "SEED = 0\n", + "random.seed(SEED)\n", + "np.random.seed(SEED)\n", + "torch.manual_seed(SEED)\n", + "\n", + "IS_CI = bool(os.environ.get(\"PYTEST_CURRENT_TEST\"))\n", + "NUM_NODES = 6\n", + "IN_CHANNELS = 3\n", + "HIDDEN = 8\n", + "LATENT = 4\n", + "EPOCHS = 2 if IS_CI else 15\n", + "EXPLAIN_EPOCHS = 5 if IS_CI else 40\n", + "LR = 0.05\n", + "\n", + "print(f\"koopman_graph {koopman_graph.__version__}\")\n", + "print(f\"torch {torch.__version__}\")\n", + "print(f\"CI_lite={IS_CI} fit_epochs={EPOCHS} explain_epochs={EXPLAIN_EPOCHS}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "130a4031", + "metadata": {}, + "source": [ + "## Motivation and background\n", + "\n", + "Graph neural encoders mix node features with neighborhood messages, so a\n", + "forecast can depend on both local features and edge structure. Readers often\n", + "want a **compact** view of which nodes and edges the model relied on for a\n", + "chosen quantity (latent state, reconstruction, or one-step forecast).\n", + "\n", + "`explain_representation` wraps PyG explainers for that purpose:\n", + "\n", + "- `algorithm=\"gnn_explainer\"` — GNNExplainer mask optimization (Ying et al.,\n", + " NeurIPS 2019; `Ying2019GNNExplainer`); core dependency, no extra install.\n", + "- `algorithm=\"integrated_gradients\"` — Captum integrated gradients\n", + " (Sundararajan, Taly & Yan, ICML 2017; `Sundararajan2017IntegratedGradients`);\n", + " requires `pip install \"koopman-graph[explain]\"`.\n", + "\n", + "These tools produce **attribution masks**, not causal discovery. For\n", + "operator-level type/relation energy fractions on assembled `K_eff`, use\n", + "`attribute_mode_energy` instead (see\n", + "[`docs/source/capabilities.rst`](../docs/source/capabilities.rst) and\n", + "[`docs/source/limitations.rst`](../docs/source/limitations.rst)).\n" + ] + }, + { + "cell_type": "markdown", + "id": "b8737452", + "metadata": {}, + "source": [ + "## Minimal example: path graph + GNNExplainer\n", + "\n", + "Build a tiny homogeneous path, fit a short encode → Koopman → decode model, then\n", + "attribute `one_step_forecast` with GNNExplainer.\n" + ] + }, + { + "cell_type": "markdown", + "id": "c65fa7b7", + "metadata": {}, + "source": [ + "### Synthetic path and model\n", + "\n", + "Undirected path with random node features and a small homogeneous GraphKoopman\n", + "stack.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "e72df4b6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "nodes=6 edges=10 latent_dim=4\n" + ] + } + ], + "source": [ + "def path_graph(num_nodes: int = NUM_NODES, in_channels: int = IN_CHANNELS) -> Data:\n", + " \"\"\"Undirected path with random node features.\"\"\"\n", + " sources = list(range(num_nodes - 1)) + list(range(1, num_nodes))\n", + " targets = list(range(1, num_nodes)) + list(range(num_nodes - 1))\n", + " edge_index = torch.tensor([sources, targets], dtype=torch.long)\n", + " x = torch.randn(num_nodes, in_channels)\n", + " return Data(x=x, edge_index=edge_index)\n", + "\n", + "\n", + "def tiny_model() -> GraphKoopmanModel:\n", + " return GraphKoopmanModel(\n", + " encoder=GNNEncoder(IN_CHANNELS, HIDDEN, LATENT, num_layers=2),\n", + " decoder=GNNDecoder(LATENT, HIDDEN, IN_CHANNELS, num_layers=2),\n", + " latent_dim=LATENT,\n", + " time_step=0.1,\n", + " )\n", + "\n", + "\n", + "def make_pair_sequence(data: Data, steps: int = 8):\n", + " \"\"\"Build a short GraphSnapshotSequence of near-identity diffusion steps.\"\"\"\n", + " from koopman_graph.data import GraphSnapshotSequence\n", + "\n", + " snaps = []\n", + " x = data.x.clone()\n", + " for _ in range(steps):\n", + " snaps.append(Data(x=x.clone(), edge_index=data.edge_index.clone()))\n", + " # Mild diffusion along the path (keeps topology fixed).\n", + " deg = torch.zeros(x.shape[0], 1)\n", + " src, dst = data.edge_index\n", + " ones = torch.ones(src.shape[0], 1)\n", + " deg.index_add_(0, src, ones)\n", + " agg = torch.zeros_like(x)\n", + " agg.index_add_(0, dst, x[src])\n", + " x = 0.85 * x + 0.15 * (agg / deg.clamp_min(1.0))\n", + " return GraphSnapshotSequence(snaps)\n", + "\n", + "\n", + "torch.manual_seed(SEED)\n", + "data0 = path_graph()\n", + "sequence = make_pair_sequence(data0)\n", + "model = tiny_model()\n", + "print(f\"nodes={data0.num_nodes} edges={data0.num_edges} latent_dim={LATENT}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "a8d3fec0", + "metadata": {}, + "source": [ + "### Short fit\n", + "\n", + "Train briefly so the encoder/decoder are not random weights. Under pytest /\n", + "CI the epoch count is reduced for smoke speed.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f653b896", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "final reconstruction_loss=0.1688\n" + ] + } + ], + "source": [ + "history = model.fit(sequence, epochs=EPOCHS, lr=LR)\n", + "print(f\"final reconstruction_loss={history.reconstruction_loss[-1]:.4f}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "1c205bff", + "metadata": {}, + "source": [ + "### Attribute `one_step_forecast`\n", + "\n", + "Call `explain_representation` on the first snapshot with the default\n", + "GNNExplainer algorithm.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "807d96fe", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "algorithm='gnn_explainer' target='one_step_forecast'\n", + "node_mask shape=(6, 1) edge_mask shape=(10,)\n" + ] + } + ], + "source": [ + "snapshot = sequence[0]\n", + "explanation = explain_representation(\n", + " model,\n", + " snapshot,\n", + " target=\"one_step_forecast\",\n", + " algorithm=\"gnn_explainer\",\n", + " epochs=EXPLAIN_EPOCHS,\n", + " lr=0.05,\n", + ")\n", + "print(\n", + " f\"algorithm={explanation.algorithm!r} target={explanation.target!r}\\n\"\n", + " f\"node_mask shape={tuple(explanation.node_mask.shape)} \"\n", + " f\"edge_mask shape={tuple(explanation.edge_mask.shape)}\"\n", + ")\n", + "assert explanation.node_mask is not None and explanation.edge_mask is not None\n", + "assert explanation.node_mask.shape[0] == snapshot.num_nodes\n", + "assert explanation.edge_mask.shape[0] == snapshot.num_edges\n" + ] + }, + { + "cell_type": "markdown", + "id": "1c85c4cb", + "metadata": {}, + "source": [ + "## Progressive deep dive: targets and MVP rejects\n", + "\n", + "The same wrapper can attribute `latent` (encoder output) or `reconstruction`\n", + "(encode → decode). Unsupported modes raise actionable errors — useful as a\n", + "guardrail when a notebook later grows hetero / delay / control complexity.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c7babf5c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "latent mean|node|=0.1309 mean|edge|=0.4081\n", + "reconstruction mean|node|=0.1413 mean|edge|=0.4981\n", + "one_step_forecast mean|node|=0.1418 mean|edge|=0.4989\n", + "expected reject: explain_representation supports n_delays=1 only (got n_delays=2)\n" + ] + } + ], + "source": [ + "for target in (\"latent\", \"reconstruction\", \"one_step_forecast\"):\n", + " result = explain_representation(\n", + " model,\n", + " snapshot,\n", + " target=target,\n", + " algorithm=\"gnn_explainer\",\n", + " epochs=EXPLAIN_EPOCHS,\n", + " lr=0.05,\n", + " )\n", + " node_abs = result.node_mask.detach().abs().mean().item()\n", + " edge_abs = result.edge_mask.detach().abs().mean().item()\n", + " print(\n", + " f\"{target:18s} mean|node|={node_abs:.4f} mean|edge|={edge_abs:.4f}\"\n", + " )\n", + "\n", + "# MVP reject: delay embedding is out of scope.\n", + "delay_model = GraphKoopmanModel(\n", + " encoder=GNNEncoder(IN_CHANNELS * 2, HIDDEN, LATENT, num_layers=2),\n", + " decoder=GNNDecoder(LATENT, HIDDEN, IN_CHANNELS, num_layers=2),\n", + " latent_dim=LATENT,\n", + " time_step=0.1,\n", + " n_delays=2,\n", + ")\n", + "try:\n", + " explain_representation(delay_model, snapshot, epochs=1)\n", + "except ValueError as exc:\n", + " print(f\"expected reject: {exc}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "15deaa95", + "metadata": {}, + "source": [ + "### Optional Captum integrated gradients\n", + "\n", + "`algorithm=\"integrated_gradients\"` typically returns `(N, F)` feature-attribute\n", + "masks with `edge_mask=None`. Skip this cell unless the `[explain]` extra is\n", + "installed (`pip install \"koopman-graph[explain]\"`).\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "c0d47d16", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "IG node_mask shape=(6, 3) feature_mask shape=(6, 3) edge_mask=None\n" + ] + } + ], + "source": [ + "try:\n", + " import captum # noqa: F401\n", + "except ImportError:\n", + " print('Captum not installed; skip IG demo. Install with: pip install \"koopman-graph[explain]\"')\n", + "else:\n", + " ig = explain_representation(\n", + " model,\n", + " snapshot,\n", + " target=\"one_step_forecast\",\n", + " algorithm=\"integrated_gradients\",\n", + " index=0,\n", + " )\n", + " print(\n", + " f\"IG node_mask shape={tuple(ig.node_mask.shape)} \"\n", + " f\"feature_mask shape={tuple(ig.feature_mask.shape)} \"\n", + " f\"edge_mask={ig.edge_mask}\"\n", + " )\n" + ] + }, + { + "cell_type": "markdown", + "id": "5cfe8d8c", + "metadata": {}, + "source": [ + "## Results\n", + "\n", + "Plot absolute GNNExplainer masks for the `one_step_forecast` explanation above.\n", + "On this under-trained path, expect **near-uniform** node `|mask|` bars and\n", + "**uneven** edge `|mask|` mass concentrated on a few directed edges (the printed\n", + "top indices name the peaks for this seeded run). Read those peaks as explainer\n", + "diagnostics — **relative** contrast only, not absolute importance scores and\n", + "not a causal ranking of nodes or edges.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "262647cd", + "metadata": {}, + "outputs": [ + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAA7QAAAE+CAYAAAC9c1w5AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjExLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlcelbwAAAAlwSFlzAAAPYQAAD2EBqD+naQAAabFJREFUeJzt3Qm8lOP///FPe9oXUtoVqbQQLUQpEhUlSkVkbUEkSwhps2UvS8gS36SQLJUUfZNEJCGUpU1FO+01/8f7+v3v+c5MM+fMzJmzzMzr+XjM45xzn3ube5br/lzX57qufD6fz2cAAAAAACSZ/Ll9AgAAAAAAxIOAFgAAAACQlAhoAQAAAABJiYAWAAAAAJCUCGgBAAAAAEmJgBYAAAAAkJQIaAEAAAAASYmAFgAAAACQlAhogST2yCOP2LvvvmupaPTo0TZz5syE7e/TTz+122+/3fr27Wt//fVXwvYLrn0k3333nY0aNYq3SAS7du1yn8dVq1ZxjZBnTZgwwSZOnGjp5sCBA3bTTTdRXiIpENAiJXz++efuxigWs2fPtpEjR1qyePjhh+29994LWqZgdtGiRZaKpk6dat98801C9rVw4UI7++yzrXjx4ta4cWMrUqSIpZtw75+ckGrX/t9//3XfNWvWrMl03f79+1vFihVz7fh53Z49e+zZZ5+1jRs3ZrrubbfdZj///HO2nUu011XreI9bbrnFxo8f77ZNpJdfftl/jOuvv97uvvtumzx5su3YsSPiuvq+DDVr1iz3vyeffDKufUfyzjvvuP1kp3Hjxtmbb76Z4TqxPvd4ffzxx/bJJ5/EtE3gdQ58DB8+3JJFgQIFLH/+/HbHHXfk9qkAmSKgRUr46aef3I1RLJYsWZJpgZmX6Cbiq6++Clo2aNAgO//883PtnJKFbkhOPvlkGzp0qLupKFWqlKWbcO+fnJBq116tivqu+fvvvzOtMNP30qWXXporx08lK1eutMcff9wqV66cbceI9rpqHa2rypkjjjjCBUz169e3rVu3Juxc5s6da3PmzHHHqFu3rgsqdJyjjjrKXnrppUPW1bIhQ4Ycsp97773X/e/999+Pa98ZVY7t27fPstMHH3xgn332WYbrxPrcc1LgdQ581KlTx5LJDTfc4ILz1atX5/apABkqmPG/geT04YcfunS/rl272ltvveVSZk4//XTr2LGjv0VXtbqqjfdadnv16mWnnXaa7d6921577TVbtmyZVahQwQWM9erVO2TfnTt3ttdff921Lqg22Vt+3nnn2bRp02zDhg3Wvn17a9euXabbyh9//GH/+c9/3Ha1a9d2N8Lezf8zzzzjbup0buvXr3fLHnzwQVuxYoV/32+//ba7gVZabWi61M6dO23AgAGZHidS6u8JJ5zg0o903aR379527LHHuuepVN7DDz/crrrqKne9PC+88IJ9+eWXli9fPvd/XdvAayHbt293qVx6HtWrV7cePXoE7SOQXkPVbmsfeh2j3VY3wlOmTHHb67XWut4N0Lfffuv+p9aJJk2a2MUXX2yFChUKeu4nnniiHTx40AVmul6RMgGi3VfBggVdbb9qv7t37+5uKgPF+vpEcz0jvX+038yOF3gNdIOpG+DLLrvMatWqlen5ZNe1/+9//+s+S/pf06ZN7YILLjjk2PPnz7cZM2a4922nTp3slFNO8f9PLVO6IddrUK1aNbf9McccE7S9Pkuq8FKgove/zk/rDx482P1f2R3ly5e3o48+2m699dZDjv/cc8/ZRRdd5H9O2o8+m2rt0PP54YcfXICmz02JEiX82/l8Plf5oOeo1uwOHTpYy5Yt/f+P9vhZvQb6bhwxYoSNGTPGta7L/v377brrrnMtk97rH+k6RXudozF9+nRr3bq1/zwC5cZ1PeOMM+zyyy93v/fp08cdTzf9AwcO9K/z559/unJE11H7uuSSS6xcuXL+/+s7Wf9fvny5Cyb1XVClShX//7Us8LtGAZoCzyuuuMJq1Kjhrofn1FNPdeXKvHnzXDknP/74o/uM6fty7969Qecfy75Dbd682b744ouIlcH6zOtzru98VerovXfhhRdao0aNgtbT+allWPtr2LChuz5e9oaupZ6PUtG98xw2bJgdeeSRhxwv1uee0XE9ynrSd1PJkiVddkk4mb2+4a5zvNcqs2N5+yldurT7XqxUqZJLFdZ7Xs9VWTJVq1Z130e63+jSpYv73nzqqafcNqGVbvfdd5/7LOs7Q/tVhaTuI9SaD+RVtNAiJS1evNgeeOAB94WstDbdVOpGywseVRjoJuSwww7z15zqJkaF3EknnWRvvPGG+79uxHUTqBuh0H0rqNLNvQrFwOUqAFWDr/8puH3iiScy3VY1uSpAfv/9d1fwqHBTrb+CDKlZs6a7mVP6one+ek6BKcdly5a1u+66y7+NqIAMrL3O7DjhKPBXAKPnUaZMGfccdI10PXXTruukG0ndWATePKgg1Hl6hbOCYJ2fR69L8+bN3Y2DzkU1wG3btg2bfqgbG918KmBTJUEs2+rGWzdCujkJrCHX81JBrW3U0qKbOr12ugkIfO5XX321C650c6LXIdI1imZf11xzjetTqZsIpU8qkAuslIjn9YnmekZ6/0RzPJ23goN77rnHXcOlS5e6bb7//vtMzyk7rr3ezz179nQVA3rPq+W3W7duQcdVIKL9/fPPP+5aK11Vn2mP3pNeC5WuvwJnBX8etWTr/+vWrXPvb6Uv6kZTGjRo4H7quWidcAGanoeupT4THp2LWvcULOgGU9dFN+56jQKftwIl3QTrvPV5OvPMM+2xxx7z/z+a4yfiGqil0muN9Cig1TLdYGd2naI5RrSUKu9VRobK6esaSu9fBSSqGPIo4NP+FFwrkFdFoK6BPmeiiphWrVq5IMErZ84991z79ddfMzyWKgh0TRXABNJnWUHJiy++6F/2/PPPuwAm2sqwSPsO13Kq9RQ0Rap8UKuePqf6zlm7dq37rCsjyqP3ia7xb7/95r6TVGmi73ev1VfXTK+Trq33fRWpq0Iszz2z43rnr8/ttm3bXBmtim59TwbK7PWNVjTXKppjaT96/bQfVWB6lU033nije8/r+1fvMT0vleNe6r6unb4XdJ/g0edU3/WBr68qoxM5ngWQLXxACpgwYYLuXPx/Dx8+3Fe4cGHfH3/84V92zz33+Bo1auT/+6GHHgr6W/r16+fr2LFj0LLx48f7qlWrFrTvAgUK+JYvXx60npbrHD777LOgbUuXLu3bunVrxG3379/v9v/8888H7a9Dhw6+G264wf/3qaee6p5DoFatWvnuvPNO9/vBgwd9VatW9T3++OP+/8+YMcNXqFAh319//RX1cUI1adLEd8opp/j/3rVrl69EiRK+Nm3a+Jdt377dXe+PPvoo4n6++OILdy7btm1zfy9evNiXL18+tz/Ppk2bfDt27PAfd/To0b7vv//eV6VKFd9NN93knmM024a67bbbfG3btvX/vWfPHl/lypV99957r3/ZmjVrfEWLFvW99tprQc+9YcOGvgMHDkR8XrHsq0WLFv7nII0bN/a/pvG+Pp7Mrkno+yfa4+m89b7avXt30DqdO3f2RSOR117vocMOO8y3evXqoOdYsmRJ35w5c9zf+vzpczh37lz/Orrmq1atiniOQ4YM8bVr187/99ChQ33t27cPWuf33393P/VZ0v6/+eabiPtbt26dW+fzzz/3L9M5a9n999/vX/bbb78F7Wv+/PnuNQzc9wsvvOArXry4b8OGDVEfPxHXQPvXPnQ8j95bWvbf//430+sUzTG2bNni9vfll19G3Mb7btG1Cicnr6toHZU3nh9//NGXP39+tz9PvXr1fA8//HDQdr169fL17t3b/b5y5Uq3n/Xr1wc9z82bN7vfL7vsMvfdHs71118fVB5p3bPPPtt9TxYrVsztR5+xI444wjdv3jx3XP0/cP1o9x1O9+7dDymHAumzXr9+fff94tHxBg4c6H8fHnvssb4bb7wx6DOscnLs2LFB3zHeNpHE8tyjOa7WOeaYY3x33HGHfx3dQ6jcuvLKK6N+fb1zq1Spku/aa68Nerz66qtRX6toj6X91KxZ07d3717/Mr3HdK8xa9Ys/7LJkye79513Drpe+gxMnz7dv46ee4MGDYKO99xzz/nKlCmT4WsB5DZSjpGyjjvuOFej6VGtZmYDRKimUy1BSqvTvYseaqlQC+GmTZtcK66oBjRcXxi1cgWm9akVU61MGtzIS+MK3Va1sdq/UuG0nndctV6pljhaSu1ViqlSk1TrK/pdrTRqQVDLarzHUWuGp2jRou55Bi5TDbBaRgL72agVQi3IOq6uoWqBVRPutdRoH2ohVyqZWi513UNTttQCpP5aSp+68847g65zZttmRDXUqg1X7btHLSVt2rRxfZ/UAuhRi7Bq6hOxr7POOsu9ToHvUe+aZfV9EOs1ieV4SpEPbCHR+8xLYY9VVq69Pp9KI73//vv95ys6t6+//tqlgqrPnK5rYNqkrrmuT2CaqtKslYatFj79VOqsR59PpePp86OWQbUWKeMgWt4AO4Epr57AtHuld+rzpPeAWo70/NW9Qb97dJ30HaL0faXJRiMR1yAamV2nRBxDLUNqLdW1ykhOXFePWoDVGqyMEZ2f9uWlbeo5qjVN312B5Yha1LzPlb4rle2ilG59V+v56Ts0Gmp1DDeAk56fMn4mTZrk9q2HWtZiGVsi0r4DW+j1fL307EjUOu6lnXtlr/c9p8++vgP0nvHoe0qvgV4nDaQWq2ieezTHVUrvL7/84tK/PbqHCMy0iOb19ejzH/iek8DPYGbXKpZjqWwJ7LKhTAi9nlruUcq/Mls8es8pu0Wt2/r8qsx+5ZVX7Oabbz7keei9DuRlBLRIWbq5D6RCIzC1JhwFXrqxPv7444OWa1nhwoX9fyvVMRwv4A0sCHRjFZgKG7qtNwiJCuRixYr5l+tvBaKx0I2V17dWQYL61Sr9KqvHCb2WCjLCLQu8vup7rP5hujnQDYcXmHgFo9LJlKqsgFWpt7oZVn80Ba5eAa+CXOmBSh0PFM22GfFeD+0ndL+h6b2RXut49pXRezKr74NYr0ksxwt9X+v/uqFSmnOsoxZn5drrnHWzGvr51DkrDV7UbSBcX7vAG1tdH22jAFgVTHoNFNh4FFQr1VZ9j5VurRRLpTZHG/h4FQnhBgrK6D3gpWAH0vXVjWc0aeeeRFyDaGR0nRJ1jIzSjXP6ugYGOQpUtK36w+t96gUT3udK79HAz42ugxe0Kr1UQZS6n6gCVGWE+kUq9T6zz5OeS+jn0aM+sBq7QJ8R/R6rjPbtBUk6V72uWXktIn3+VSkVr8yeezTH9dYJ933nieb1jbYPbbRlQjTHCv2uVIpx6Pe49q3rE0iVOkp/1/oqb3UN9F4MpO+xjN4XQF5AQIu0FdhSFlgAqQ9KrFMAeXQTpxpUb98qJDQQT+BgH+GOKc2aNQuqCY7mfEOpoFN/G9VCq/VE26gfbyzHSQTVIOtGVLXdGszHq20O5fX5VM2wBpdSS6AC8SuvvNL9X6/Dli1bXC22gjWvz3E022bEqyVXC6XXd07UD059SGORqH0l4vXJ6JqEvn9iOV7oCJd6rroRjGcKnqxcL52zKkTUAh2p1VzPV/38ItFgNjp39afzqC9dIF0rBWh6aEoW9TvTQCo658BWkEh0I6nzUOtK4MBD0Vyb0FFZVXGgh5dtEs33QCKugVeBF9gvPjRAz+g6RXOMzOh9rEFuwk3LEotEXddwg0Kp1UsBngbZUYWi97lSwJtRBYj+r9ZrlRdq7dVnVe+b0NaxQAp0Pvroo4iDNikjSNkses3UPzcWme1b9J2u5xTLtcro8x/Ymq/Pf2BGVazHyOy5R3Ncr5zW913giNqB31XRvr6JkJVj6fzV112vq1ehqconVXYFatGihRvg8dVXX3X9c3W/EBoI63tM4yYAeRmDQiFtqUZTwVIg3ZBoXsHAgXqUZqVUx2gogA0ceEUDbKggzahGW7WvCtQ0EIOCX49qSgOnLQh3vuHoOSig1UOjPHs1wNEeJxF0I+qN5OlRy3EgpX95AwspOFEtsQrw0Gkz1IqhlkYFtRrNMpZtI1Frka5F4AAoSo3WSJm6IY9FovaV1dcns2sS+v6J5XgKKLyBgLSuUvl0Ax6PrFwvpccp9T80MFKrnyqTRIMSaZCiwEFiFAR7rTB6b+o5eK0gCmpC0zI1CIyX0qfWNGUbKF1ey9Tyreub2WfxnHPOiXnuSmUi6OY58Pvm0UcfdQPY6MZTojl+Iq6BbvoVvCvY8oTOPZrRdYrmGJnRgHf6/vWee7wSdV3DUaDTr18/NyCPnq++75XuqhHZlWbt0X6994OCKO+1UOCmc1CadEbfXwrUNPCPWoU1ynQ4arVT2q1G0I80aFO8+46ltTwjquRQJY8qP7wuA0pDVwVM4Oc/2vIu2ucezXGV1aDXQmn0HgV5gXO9R/P6JkpWjuUNihb4mX366af9ZXMgVXhqwEx1EQrXuq1jadAyIC+jhRZpS31AVYir5lMFhwJBjcKrwEA1ovq/buh046Ha39C013BUw6uUOxWsarFQ/0QNm59RS5ZuaNSaoZpR9XtTa5lubNSiGViw6v+ad1atJApSQwPEwDRA3VwpKFeNe6zHSQRNLaCbavXfUR9eTZUQOOKoeCNI6kZEQY6CVd28ajTkUHqu2l6p37qJ1vOPdttwdGyN0KxCWkGgXn9Nb6IbU90IxCJR+8rq65PZ9Qz3/on2eOrjp9ZfraPPg4IW9dWNR1aul7IOdIOmFlpNx6XWf31etU+l14ueiwIn9fFVq4NuUtVHWGmxomv0yCOPuNFElQqvlmxlZQR2C9CNvb4D1GKs9GF9jpSG5/V9V7Cq89XNpq5NuOlddHxdLwWS0Y40q3PXtDFK01f/YW2rG2q15HlT1qgPXGbHT8Q10PG0X71/1OqvaxJ6M5zRdVJqY2bHyIwCKD3XaLoR5MR1jUTf+ZrzVGM0KCDU1Fne50pTySgA0fvU+87W87n22mvdMXWtFFipwkjva4/WV4aKrrkqSpXuqwBcLd6B08iFiibgjGff+l5Q5k1gn8x4KbDSfvTe0GdYx1VllSpgPSpv1QquChF13Yk0bU8szz2a4yrgVdmv7zu13GqU4dBrktnrG3qdAyllO3B07cxEe6xwAbzW0fE1S4PKAvUR1ndRaHaLPuO6Z9A2oVPraVR7jb6tmQ6AvCyfRobK7ZMAsko3BOqT5BUeuulWLXhgja++lHVTpdY+j1oxFHTqJl83MV7B5Q3EoIJUKZmB6Ufh9i0a4EM3YCokFyxY4G7c1EoWmN4UaVtRgaNz8aZQUU1x6LyL2l6BilKHVMDoBlK10SqgAymg1k2bUgFDC69ojhNIgYMC1MDBLRSka/CKwJRRtQirJVoFr+irRTW7miJBNwa6FurjpMLZS6XSDZVuLL0+v0p38watCHdctX6rkkFpfhltG0o3JbpxC73hUauRplfRQCgaqCowpTnSOUQSz74UnCuQVQpjvK9PoMyuSej7R4FtZsdT31RVTmggKLVU6P2koCDac8qOa++1UGhbfWa9/rOBvP6NXmt14BzF+mzos6NWD71n1Yc49LtB56frop9qzQ6cG1LXTNurVVh9yyK1LCuA0mdEFWWq4FJwqcqxwH5s+kwokA8c9Eivn75DlPar90bojXy0x8/qNfBav9UXXpUkSqtVn3y1zisAyuw6ZXYMtQyqFVkBRWjfRtFrrzlmdc0iyenrqkoBbRs6KKA+y3qfewMKefM265j6vtPnKrBiQ9+P+mwoONF5aJ9emrfKMm/wLG96Kl3bcAMRal19niO1oOm81HLs/T+WfQdSEKaKJz0yojJQr6XKTY8CZr1OgXO66j2hz78+y/rsh8ti0neVBq/TtqGvbzzPPdrj6nVUOa7KQb1uOg9VQugz5Mns9Q28zoFUfnndYqK9VpkdK9x+PKpM1vtM3+36/KriSen3oYGrXn9NdaT7mNAKcqUkq383kJcR0AIJ4gW0gSl6QLLzAloFFoiNWt0UUKoyALFRa5JG6lZwERrIIOep1VIVGRptF8lBXTlUweRVao4dO9ZlEOh7SSn2gUG0Kj8VMAdWAKl1XBVOyroIHbwKyGtIOQaAJKA+gKNGjYr4f/UrjbbVNJXPKS9R9gTBbHzUWqp0S4LZvEFp6tk9CBISS4MyqoVVlZLKRlOGjTIMvGBWlUV33HGHG3hNqfahU2OpVVqjIAPJgIAWSBClNcU6Qi4QLdWQZ5T6nNV+hpHohkepv3npnJD6dHOd2dyzyDmhqejI+5SmrX78SrPXd7W6JgUOmKWxPRTsquVd/ceBZEbKMQAAAAAgKTFtDwAAAAAgKRHQAgAAAACSEgEtAAAAACApEdACAAAAAJISAS0AAAAAICkR0AIAAAAAkhIBLQAAAAAgKRHQAgAAAACSEgEtAAAAACApEdACAAAAAJISAS0AAAAAICkR0AIAAAAAkhIBLQAAAAAgKRHQAgAAAACSEgEtAAAAACApEdACAAAAAJISAS0AAAAAICkR0AIAAAAAkhIBLQAAAAAgKRHQAgAAAACSEgEtAAAAACApEdACAAAAAJISAS2QzSZMmGBXXnlljm+bKt599127+OKLbefOnZaKxo4dawMGDMjSPj788EPr16+f+Xy+XDl+oCeeeMLGjBmTsP0BSE+UnelXHsaCshOB8vniuQMCIjh48KDNnj3b5s2bZ+vWrbPChQtblSpVrFGjRta+fXsrVKiQf93JkyfbW2+9ZWXLlnVfTPnzB9ev3HvvvbZ161Z77LHHErJNJIMHD7aTTjop217T6667ziZOnOjOKye3TRX333+/DRkyxLZs2WJlypSxVHPJJZfYJ598YmvWrIlr+927d9txxx1nV199td155505fvxQ8+fPt9atW9u3335r9evXT8g+gVRH2Xkoys70Kw9jQdmJQAWD/gKyYNmyZdajRw/bsGGD9erVy0499VS3fMmSJdanTx/Lly+fC3YV3MrSpUvtjTfecL+3aNHCevfuHbQ/rbt+/fqg4DQr26jV6KijjjrkvCtVqpRnX/crrrjC2rRpk9ungTzsueees7/++stuuOEGywtatmxpzZs3t6FDh2ZYkQTg/1B2Jh5lJzJD2ZlaCGiREL/99ptrlalTp47997//PaTmcPjw4da3b1934x2qRo0a7ua3W7duVrRo0aiOF8825557rmvJSiYnnniiewCRjBs3zi644AIrWbJknrlIl156qUtjVquvMjQAhEfZmT0oO5EZys7UQkCLhFAKzLZt2+z1118PmwajZf/5z39s+/bth/xv1KhR1rNnT9f37tZbb43qePFsk5k///zTpR83a9YsqLVr//79duONN1qBAgXskUcecT+V7vzDDz+4n2+++aZ99NFHLr1agUU0Lapvv/22v6VZLdelSpVyac9q4S5RokRQHyKlcL7wwgv+ZYHHnj59ur333nvunLp06WJnnXXWIcfS+U+ZMsVVNPzzzz9Wu3Zt12IeGGgE7lN9dNQns1y5cjZy5Miw5x/rOezYscNeeeUV++abb1w/zxNOOMG1rut5B1Lr/rPPPmu//vqr1apVy1WCRBLN88rs3LX9rFmz3HmoD6qOKVOnTnXL9Zqqpl/nG8/rd+DAAXvnnXfca6j3vvbfvXt3/3Ei+ffff+2OO+5w6cSjR492r0U4X3/9tf3000+uwihUtOcYKLP3crTP57zzznOvndL9Bw0alOFzBdIZZWd6lZ3R7jfW8lD3L+PHjw9a9/PPP3f3ZC+++KIVK1Ys5uOHouyk7MzLGBQKWaab7mnTprlUw+rVq0dcT4VP6dKlD1l+2mmnuRtgBambNm2K6pjxbJMZpR7ri13Bqwo6z+23325PP/20O54KP1FBoees/82cOdMVqBqkoW3bthkWZB61ZHfu3Nk9Onbs6Fqc1TdGwXTgYA+LFy92wVUg79gjRoxwwYWOrT627dq1c5UGgf7++2+3z4EDB9qRRx7p0sC//PJLq1evnivsQ/epIEr7bNiwoStMI4nlHFasWOH6UioNXC3kOvaTTz7plv3888/+9ZYvX24NGjRwQVXjxo1df2tdHz2HUNE+r4zO/a677rI5c+a4c1dafNOmTV1riQIwBbNNmjRxweIpp5ziUgLjef0uvPBC11JZvnx5957VZ+X888+3uXPnRjy/tWvX2umnn+7OURUrkYJZUd9X0XFDRXuOnmjey9E+H32W9F2Q0fME0h1lZ/qVndHuN5byUAG11tXzVeWrKgi6du3qxjJR8L93796Yjx8OZSdlZ56mQaGArFi2bJkGFvP1798/pu3uvPNOt93q1at9P/74o69AgQK+G2+80f//U0891VerVq2EbXPuuef6unfvfsjj33//9a974MAB31lnneUrU6aM79dff/VNnTrVbTtixIigffbq1ctXpEgR37Bhw4KW33bbbb78+fP7li5d6l82YMAAX+nSpTO9Hn/99Zfb58MPP5zhtt6xR48eHbS8VatWvmOPPTZoWdeuXd1zWbNmTdDyLl26+GrUqOHbt29f0D7vvfde/zo7d+6MeK6xnEPr1q195cuX9/3999/+ZZs3b/ZVqFDBd8opp/iXnX766b4qVar4tm/f7l/222+/+cqWLetegy1btsT8vDI694ceesi/bPfu3b5y5cr5mjVr5hs5cqR/+Z49e3xHHHGEr3fv3r5YX79Vq1a583711VeD1tu7d69v48aNQedTuXJl9/vXX3/tftd5rF+/PtNjXnnllb5ChQplul6kc4zlvRzt8/G0adPGV7NmzajPDUg3lJ3pV3ZGu99YysOWLVv6qlat6vvnn3/8y3SPpHKXsjM8ys7UQ8oxskwpK3LYYYeFbfX5/fff/X8rrSfcNDRqudNy9Wm4/vrr7eijj870uLFuoxrgcINCBY68rFGTlaKjWk61yK5atcr1vVXta6g9e/YcMt2JWncfeOABl2qpGtOMqAZZNbs6hvalVFydiwaxyozW79+/f9Ay1TJrlFvVUiu9aPPmzS49S+tVrlw5aF2NiKvn9dVXX7kBfMLtM9zrGes5aIAutSKq1VOteh6NUq3XTum0f/zxhxUpUsTVJqtfdGBfUNW+a59eipnE+rwinfu1117r/1vH1yBj77//vs2YMcO/XDXdaqFVqnSsr59+9wZCUzqdl/Kl5UccccQh+9O+NJiazl/p2dH0DVdte0YjXUb7HovmvRzr89HrvWjRokyfA5CuKDvTq+yMdr8q96ItD5VqrJbVu+++24oXL+5frvTh0JZnys7/oexMPQS0yLLDDz/c/2UZ6owzznDDy6ugUACjm+9I86oOGzbMBZMKHidNmhTVsWPZJtpBofR8FCQroNX5vvrqq+5GPpSCssAgTSpWrOj68agfS0YUyKkA1bDzCqR0HB3j008/dX2RM6MAIrT/qRdUKIhUcK90Xk0FsXDhQnccb4Yu/dRrIroh8Apl7S9cYJKVc1AKrxx77LFhU8dE6yigFKV8hzrmmGOC/o71eUU699BBlPRa6jUNDRC1fMGCBTG/fnovKLXtnnvucTdpSsnXwGkayCz0eW7cuNEFiUo11s1KuPdbOAq4A9PJYj3HWN7LsTwf7yZP5wcgPMrOyN83qVh2Rrtfb3k05aFXxoYblyF0GWXn/1B2ph4CWmSZCoAKFSrYF198ccj/zj77bPdT/VQiBbKBBdrNN99s9913n/sZjXi2yYwGvlF/T33h6bw10MPll18etv9TKBVWmd3Ia4AkzZerQRsUOIfW0kYj3P69IEjnENjyrD494Qa8uOyyy4L6XsY6Sm4s5xDuWnn9nbQfb71du3Ydsl7oslifVyznHmm593xiff1U0aJlGmhJtejqi61t1QJ78cUX+9fT50eDZOlmTUGjauajff9rcCYN8lGwYMG432PRvpejfT6ivu06PwDhUXZm/H2TamVntPv1ZoOIpjz0ng9lJ2VnuiOgRZapMFABo6BSg+kozSVet9xyixvVTz+zc5uMKJjQYDYaGOrll1926UEaICg0DUoFiAZj0GAKHqU87du375BRcUNb49Sq5s3H61GqkZeClggadEkpSArQQ4ONnFK3bl1X4IZLPdUyBWE6T/3UekolC6XrkpeeV6yvn2ruNSK3HholW+8Xjc4deu4a4EytpBppW7X1mjc5s5ZaDWqi2nwNWqWBQ+I9x1jey9E8H++cNHoogPAoO9Or7Ix2v2q5j7Y89MrYcN1iQpdRdv4PZWfqYZRjJIRabtTXUH0AA/sgelavXh3VfvRlrxpYpQ9pSpLs2iYS9aHUiIkaAVcpys8//7xVq1bNje6q2uFAGrFZ66pW2QsKNIWQRqXVdYhEo78qTUrH8mqE1RKsQDqWlN/MqA+mpoRQyrRGSgxtHdUQ/IEtj9lBr81VV13lUlS9EXlF/YNee+01NyWOrqPW07QBEydOdCMuejS1QOh7J7efV7Svn0apVD8vL31M1JKqR6SRi5VloGkmFCAqo0E3PhnRNBe6Kdbok/GcYyzv5Viej4JZpf+Fa4UA8D+UnelTdka731jKQ62rDDJlyQTe/yizLHD8kux8XtGi7KTszE600CIh1Afy448/doGl+tSpcFFtoOgLWDe4Sj/W/zKjAEgpv7qBjlY02yhYCJcapJYm9ZfVl/+ll17qBo/S8xCtr0JEaUAKMBSYedTfR3NwqrZYfVW+/fZbV8P81ltvHdIfMZBaIxUoqxDSNapZs6Y7bwUxap1LJG8wKwWOGjRC/W80GIIGYlKfzWj7ambFww8/7NJi1XKv66hjKj39oosuskcffdS/nlok16xZ46aDUd8opcHq2ujcVcGQV55XtK+f3jtq4ddgS+q7rcFG9LxVQaL3aiRekK/3pYJCDeoRKQ3PG/hD/cc1j26s5xjLezmW56N+wOqXSwstkDHKzvQqO6PdbyzloTJldJ+lRgU9VDGpdHZdJx3Dm24wO59XNCg7KTuzUz4NdZytR0DaUcGkVBeNvqc+I5qTUl/EoYPtfPfdd/b999+7eSxDRwVU+qNSkHQT3aFDh4RsE4lu4pW2o9pNDZqgVqXQQlXpsRqs4pxzznHBhgZ0UIujChz1FdQgDwo6NKdb4ATmov0qWFZhEUiFiJdSpAJLA0t88MEH7qcG3Im0rY6l1CsF4YF0fjpPzc0XOMG8/Pvvvy5VSQN3aaRn3QwErhNpn5HEcw4a7EI3LvrK0TWPNGex1tF+VCBrvR9//NEt07x6gSNSR/O8Yjl3BWcaFETvrUCqIdfrHBqcRfP6iY6l96V3k6H3WjTn470fle4eOhBIIM0dq/elbuxCB9+K5hwDj5/Zezma56MWWwXFPXr0cK2+AKJD2Zn6ZWe0+42nPPTW1WBS6iKluWbV11hp2qGBKmUnZWeqIaAF4hAY0AK5Ta20GlhKKWq5bfz48a4VYMWKFa7yBwA8lJ05R8G+gtlwfXHxfyg7Uwd9aAEgySkNTy3LeSHhRq0DH374IcEsAOQAZdZoIMvQMkGtx9dccw2vQQYoO1MHfWgBIMmpH6seeYHmngYA5Fw/7JEjR7oZGdQ9RWnH6oJy0003EdBmgrIzdeSJlGN1dtdIdUqZy58/ukZjjYq3YcMGN8+hRm4Lra0KHZFWfcLUTyERxwbi7TcDAEC6ouzMPitXrnSBrO5lNfaC7muBdJGrAa2GB1cNkuYQVQ2TBvl56qmn3BQpGX1gn3zySXv99dfd5NNKs2jdunXQOpoTVcOSBw7jXqdOHTd4SlaODQAAAADIO3K1SfKhhx5yc1FqpDe1kmoIcY2MmdGItNOnT3ejo2qKmIxo1E+NcOc9AoPZeI8NAAAAAMg7cjWg1XDimj/0+OOPd0OKK/+/Ro0abpTMSG688UbXsqo5DjOjlFClJifq2AAAAACAvCPXBoVSsKl5KTUJdCDNfah5H7NKE3TPmTPH/vnnH2vcuLGNHTvWTj755IQeW2nL69atc/OeZudk1ACA5KbePRrbQWM5pPt4DZSdAIBElp25FtCq/6uUL18+aPnhhx9uCxYsyNK+mzVrZgMGDHATSyug1e9nn322SyeuVKlS3Mfes2ePe3jWrl1r9erVy9K5AgDSx+rVq61KlSqWzlQRXLVq1dw+DQBAipSduRbQelH2/v37g5bv27fPChQokKV99+nTx/97iRIl3MBPU6dOdY/rrrsu7mOPHj3ahg0bFvYilypVKkvnDABIXdu3b3dBnDJ60p13DSg7kQx0r3j55Zfb0qVLbeLEidawYcOw6y1evNiNLuzdL37zzTd2wgknuL87dOgQcTsAWS87cy2g9aLs9evXBy3X35UrV07osTStj1pmNThUVo49ZMgQGzRo0CEXWcEsAS0AIDN0T/nfNaDsRDIYOnSo3Xzzze6hRpJw93s7d+60G264wR599FE788wz3ewZ5cqVc5mAp512msscBJB9ZWeuBbSKtE888UQ3+nD37t39LaQavXjgwIFBc8oqzTfcHLIZ9c8JzLNWetMff/xhtWrViunYofQFpQcAAABS2+zZs909Y6tWrYKW33vvvYesq9ZYTf2ollpvOkl1TTv//PPtvvvus2uuuSYh53TppZe6KStF97bRZjXqXK6//nr76aefrH379nb//fdboUKFEnJOQG7L1ZEpNFWO5ovVgE1ff/21S+nQB7Nfv37+dW677TZr166d/2/1iVVL65o1a/ytqvpbU++Igt8WLVrY5MmTXZ/ZDz/80Dp27GjVqlWzSy65JKZjAwAAIP2o1VVjsPz7778ugFXjyDPPPOO//8yIAlptoxbbt99+291rJor2uXDhQtuyZYsbMCeSZcuW2QMPPBDUHU/jvrzxxhvu/viJJ55I2DkBuS3XWmhFtVaTJk2yxx9/3H2wNIjTvHnz3OBMHv0emAY8Y8YMGzx4sPtd89Hefvvt/ul89FAL6vPPP+/mmR0+fLib3kdz0t56661B+dfRHBsAAADpR9l+PXr0CPu/0BZaBb9KK54yZYo/GzC7ePepmaVg1q5d26ZNm+a6xnXu3Nk+//xze//9912rrNKodf+sNGogFeTzZVS9gwypD23p0qVt27Zt9KEFAFBeRIGyE8nopJNOcg0mmgoynEWLFlnTpk3d72rJVQah7g/VeKIBRROVcuxRf15lJxYsWNBNOdmlS5dD1tm1a5drYf7qq69cQ87KlSvdck1d2aZNG1uxYkVCzwnIrfIiV1toAQAAgLxOAWnFihUj/t8LZgOpdXT69OkuGM5OjRo1cmnIodSNrm7dum4wVK9rnihdWRmMQKogoAUAAAAyEEsLa9++fXP0WhYuXPiQOTrVaqtZPtSPVunJGlxV/XnVkquW5tCBroBkRspxFpA2BSAaZw9/PyUu1MyhHXL7FJIW5QXXAkgE9X+dMGGCG6RKQWq3bt3skUceyXTWj/nz57t1Nbhq/fr17d1337UjjjiCFyU7jMm4f3O2uDk1e5CScpxkUuGGl5tdAACQlJIkCNFATtdee63/7+LFi0dcN3AKy5YtW7ogWAFCRn0RgWREyjEAAACQBMqUKeMe8SKYRSoioAWQY1IhE0HIRgAApFyLcQqnriK1/S8XAQAAAACAJEJACwAAAABISgS0AAAAAICkREALAAAAAEhKDAoF5AIGRwIAAACyjhZaAAAAAEBSIqAFAAAAACQlAloAAAAAQFIioAUAAAAAJCUCWgAAAABAUiKgBQAAAAAkJQJaAAAAAEBSIqAFAAAAACQlAloAAAAAQFIioAUAAAAAJKWCuX0CK1eutBdffNE2bNhgDRo0sGuuucYOO+ywDLf5888/7YUXXrDly5fbnXfeaXXr1g36//79+23q1Km2YMECK1iwoLVs2dI6d+5s+fLl868zYcIE+/jjj4O2q1q1qo0ePTrBzxAAAAAAkHIB7dKlS12w2aFDB2vevLkLbF977TWbP3++FS5cOOw2jz32mD3yyCNuG6171VVXBQW0Bw8etHr16lmTJk2sRYsWtnPnTuvXr59NmjTJ3njjDf96X3zxhf388892ww03+JeVKVMmm58xAKSPs4e/b6lg5tAOuX0KAAAgLwa0t99+uws6//Of/7i/e/ToYdWrV7eXX37Zrr766rDbdOzY0QYMGOBadJ955plD/q9W2FmzZlmNGjX8y0499VQ7/fTT7bbbbrMTTzzRv7xatWp2ySWXZMtzAwAAAACkaB/avXv32uzZs61bt27+ZRUqVLA2bdrYe++9F3G72rVrW6FChSL+XwFtYDArRx99tPv5999/By3/4YcfXIrzLbfcYu+++24Wng0AAAAAIG0C2lWrVtm+fftci2wg/f3rr78m9Fjjxo2zkiVLWtOmTf3L8ufP7/rsNm7c2KU39+7d27UQZ2TPnj22ffv2oAcAAAAAIM0C2l27drmfJUqUCFquwNP7XyK888479sADD9jTTz8d1Ed22LBhrk9t//79beTIkfbRRx+5vzNqHdaAUaVLl/Y/NIgUAADJ6qGHHrLKlSu7zKcTTjjBPv300wzXV6bT5Zdf7jKqVBlcq1Yttw8AANIuoFVAKFu2bAlavnnzZv//smrGjBl28cUX25gxY6xXr15B/zviiCOC/j755JNdn9pFixZF3N+QIUNs27Zt/sfq1asTcp4AAOS05557zlXuatYABarnnnuue/z+++8Rt9H4Fl9//bV99tln9u+//7qBGu+66y575ZVXcvTcAQDI9YBWrZtqMV22bFnQ8u+++86lAmfVzJkzrUuXLq5VdeDAgVFto8I5I0WKFLFSpUoFPQAASEaq7L3yyiutffv2riJ5xIgRVq5cubADLnqWLFliF1xwgR1zzDGuVbdTp0523HHHueUAAKRVQKvBm9R6qprhHTt2uGWaN1YtpD179vSvp/+rZTQWSh/WvLOjRo2ym2666ZD/q++u5qkN9NRTT7kaak0HBABAKlM2lKaua9WqVVC53Lp1a1cWR6KZAd588037/vvv7Z9//rG33nrLtehedNFFOXTmAADkoWl7FHC2a9fOjj/+eKtfv77997//dQGolnk+//xzW7hwoWtpFa3z7LPPuvllRf1fn3/+eVdjrIeC4/PPP9+1ni5evDhoWh7NWavCukCBAjZlyhS788473Ry2GqBqxYoVrp9ts2bNcuFKAACQc9avXx+2+436xmbU9eaee+5x5aXKbVE/WpXJmoIvowEV9fAwoCIAIGUC2rJly7pgVX1xNK/sI4884lKXAikIVWurR4NXKD1KFMAGTufjFa7qFxSON4iTRjjW3LeqVf7222/deTRs2DBo0CjkjLOHv5/0l3rmUFr1AaSGgwcPupbaSFRJ/NNPP7lp7zQlnjKiunfvbocddpj7GY4qpNVXFwCAlAtoRa2lp59+esT/N2/ePOhvFaDevLKR+rkGtspmRPPVhs5ZCwBAqqtUqZL7uXHjxqDlf/31l1WsWDHsNvqfZgNQlx1lN0nHjh2ta9eu9uSTT0YMaNVtaNCgQUEttMwSAABI+j60AAAgdygzSUHp3Llz/ct8Pp/7+5RTTglqsT1w4ID7XYNAeesF0jre/8JhQEUAQHYioAUAIA3dcsst9uKLL9q0adNcS+3gwYNd62m/fv3861xzzTXWqFEj97u65bRt29buvfdeN6qx1tV4FBokSq20AACkZcoxAADIeX369HEDKWowRo1joSnzZs2aFZQOrG5BBQv+71ZB409o3lmNbbFp0yarXr26PfjggzZgwABeQgBAriCgBQAgTd1www3uEYlGMA6kUZFDlwEAkJtIOQYAAAAAJCUCWgAAAABAUiKgBQAAAAAkJQJaAAAAAEBSIqAFAAAAACQlAloAAAAAQFIioAUAAAAAJCUCWgAAAABAUiKgBQAAAAAkJQJaAAAAAEBSKhjtirVr147rACtWrIhrOwAAQBkMAEBCAtrHHnss2lUBAEACUQYDAJDFgLZjx47RrgoAABKIMhgAgAT2oZ0+fXrE/40YMSKeXQIAAMpgAACyP6Dt1auXzZs375Dlw4cPt4cffjieXQIAAMpgAACyP6AdN26cnX/++bZkyRL/svvuu8/GjBljs2bNimeXAACAMhgAgOzpQxvokksusU2bNln79u1t/vz59tprr9mjjz5qH330kZ188skx72/Hjh1uf5UrV7ZChQpFvc3atWutWrVqVqxYsbj3G8+xAQDILYkugwEASMt5aAcOHGhXX321nXTSSW70xXgK0v3799u1115rhx9+uDVt2tSOPPJIVzBnZPny5da3b1+rWbOm1a1b1xYtWhTXfuM5NgAAeUEiymAAANKqhTZc39jSpUtbyZIlrWXLlvbpp5+6hwwePDiqfd5///329ttv23fffWfHHnusvfjii3bZZZdZgwYNrGHDhmG3+fjjj61Ro0auMK9Xr17c+43n2AAA5IbsKIMBAEirgHbixIlhl5cvX95+/PFH9/BEW5g+++yzdtVVV7mAUq644gp78MEHbfz48fbkk0+G3WbAgAHu55o1a7K033iODQBAbsiOMhgAgLQKaAMHgEqE9evXu6C0efPmQctbtGhhixcvztb9ZtexAQDIDokugwEASOs+tD6fz1atWuX/W78rhffNN9+Meh8a0ELUhzWQ/v7777/jOa2o9xvvsffs2WPbt28PegAAkJMSUQYDAJDWoxw//fTTLr1Jqbl79+611q1bW/78+e2vv/5yLZ833XRTpvsoUKCA+6ntQ4PGggXjOq2o9xvvsUePHm3Dhg2L+9wAAMiqRJTBAACkdQvtE0884QZlkrlz51qRIkXsp59+svfff98VtNHQNDn58uWzP//8M2i50oGrVKkSz2lFvd94jz1kyBDbtm2b/7F69eq4zxMAgNwqgwEASOuA9o8//nBBoVeYnnfeea7VU9MHZDRYUyCNzNikSRObMWOGf5lqmmfPnu1qmwODzN9++y3qc4tmv9EeO5RuGkqVKhX0AAAgJyWiDAYAIK0D2qOPPtpNebNjxw5744037KyzznLLV65c6f4XLaXvau7XMWPG2IIFC6xnz55WtGhRN8+s56677rJOnTr5/1bLqOai1bG8vkP6O7DvazT7jWYdAADymkSVwQAApG1Ae/fdd7s5W8uUKeNSdM844wy3/Pnnn3dT4UTr3HPPtXfeecdmzpxp/fr1cwHl/PnzrVy5cv51KlWqFFRAz5kzxzp37mzXXnut1alTx0aNGuX+njRpUkz7jWYdAADymkSVwQAApIK4Rl/q3r27m8h97dq1dsIJJ/gHWWrXrp21bds2pn117NjRPSIZPnx40N9dunRxj6zuN9p1AADISxJZBgMAkOziHk5Y/Xe8Pjyec845JxHnBAAAcqAM3rBhg73++uvuZ4MGDVywHM1MA/PmzXMZU8WKFbOLL77YqlWrxusFAMgV8c+PY2br1q1zfVj3798ftFw1xwAAIPtktQz+5Zdf7NRTT3WtvE2bNrV77rnHJkyY4LrieK2+oQ4ePOjSnbVOnz593IwBGpTqP//5j9WtWzchzwsAgGwPaDVdjWpkNZhSpEnfAQBA4iWqDL7tttvsuOOOsw8//NDNY3v11Vdb7dq1XYvtpZdeGnYbzX2rAam+/fZbq1Wrllt28803286dO7PwjAAAyOFBoW688UarUaOGfx7XLVu2uNraY445xsaNG5eF0wEAANldBu/bt88++OADN8K/gllR2rCmrtOAiZFo/7169fIHs1KiRAmrUKECLxoAIHkCWvWdeeihh6xixYr+wkyDUUycONEee+yxRJ8jAABIYBmsVOU9e/YEBaaiv5WKHM727dvt559/dmnKU6ZMsSFDhrgWW6U+Z0TH0baBDwAAcjWg1ZyvRx11lPtd09xs3LjR/V6/fn37/fffE3ZyAAAg8WWwlyJcsmTJoOWlSpWKmD7sBaIKptXXVtvOnj3bTaG3cOHCiMcaPXq0lS5d2v+oWrUqLykAIHcD2kAaTGLs2LG2detWe/rpp6169eqJOTMAAJAtZbAXyGq7QEpfDg1yQ7dRevH7779vd9xxh02bNs1OP/10u/POOyMeSy2527Zt8z/UBxgAgFwdFKpr167+30eOHGnnnnuujRo1yg477DB77bXXEnZyAAAg8WWw+ssqVXn58uXWvn17/3L9Xa9evbDbqHVVUwU1adIkaLn+fvXVVyMeq0iRIu4BAECeCWjVd8bTrFkzV9uqfjWqGS5btmwizw8AACS4DNZAUAqMX3rpJevbt68VLVrUli5dap999pm99dZb/vXeeOMNl8asEZGlR48eNnfuXDtw4ICb2kfT+Gg+2kaNGvEaAQCSM+VYNLF648aNCWYBAMhh8ZbB999/v+3atctOPvlk6927t7Vp08YuueQSO//88/3rfPTRR0Gtr0OHDnXBsFKdr7nmGtc6q0GhHnnkkYQ+JwAAsrWFVlSAffnll66/TajLL7883t0CAIAcKIM1SvKSJUvcPLQbNmxwAWrLli2D1tF8t5rKJ3DQKLXiKtD9448/rHPnzta2bVtSigEAyRXQqrb2qquussKFC7s+NaEIaAEAyB6JLIPV7/aCCy6I+P8zzzzzkGUFCxa0c845J4YzBgAgjwW0Gs1Qc8+pNhcAAOQcymAAABIwD22vXr3i2RQAAGQBZTAAAFkMaNXHZt68efFsCgAAsoAyGACALKYcjxs3zg0C0aVLF6tVq5bly5cv6P/XXXddPLsFAACUwQAAZG9AO3HiRFu1apVNmjTJypQpc8j/CWgBAMgelMEAAGQxoH3qqafcKIuarw4AAOQcymAAALLYh/bgwYMu3RgAAOQsymAAALIY0DZt2tRNqg4AAHIWZTAAAFlMOa5Ro4b17NnTLr30Uqtdu/Yhg0INHjw4nt0CAADKYAAAsjegXbhwoR177LH2xRdfuEeoWALa5cuX2/jx423Dhg3WoEEDGzBggJUoUSJL2yjQ3rVr1yHbnXnmmda3b1/3+3PPPWezZs0K+n+1atXskUceifrcAQDIaYksgwEASMuAdsmSJVGtt3Xr1rCjIHu+/vprO+2006xbt27u5/PPP+9GTlZhXaRIkbi36dq1q+3fv9+/zW+//Wa33nqrderUKWg/a9eutZtvvtm/rHTp0lE9LwAAckuiymAAANI2oI1W2bJlzefzRfz/kCFDrHXr1jZhwgT394UXXmhVq1Z1f3stqfFs07lz56Bt7rrrLhesXnTRRUHLK1eu7LYHACDVZFYGAwCQtoNCJcKePXtszpw5QQFl+fLlrW3btvbBBx8kbJsDBw7YSy+9ZL169bJixYoF/W/ZsmXWu3dvu/76623y5MkJe24AAAAAgBQOaFetWuXSgtVvNZD+/vXXXxO2zcyZM11q8dVXXx20vGDBgtasWTPX2luhQgXr379/plMRKaDevn170AMAAAAAkIIpx5kFhxLaaqrBnXbv3p2wbdTH9qSTTrLGjRsHLR8+fLhLx/J07NjRmjRpYtOmTbPzzz8/7L5Gjx5tw4YNi+r5AQAAAABStIXWG4Bpy5YtQcs3bdoUcRCLWLfZuHGjvffee4e0zkpgMCsnnHCCVa9e3b788suI56z+u9u2bfM/Vq9eneFzBAAAAACkYEBbpUoVF1QuXbo0aLn+btiwYUK2eeWVV6xw4cLWo0ePTM9HA2fs2LHDChQoEHEdjaJcqlSpoAcAAAAAIM0C2nz58rmBml544QU3tYB8+umnroX0kksu8a/3zDPP2KBBg2LaxvPiiy/axRdfbCVLlgxavm/fPnvttdeClo0ZM8Y2b94cNLUPAAAAACBN+9B+/vnnGf5/5MiRbj7YunXruocmiL/jjjusTZs2/nW++uorN8dsLNvIggUL7Mcff/RP7xNIrbCzZ8+2oUOHWp06dVzq8J9//ukCZfW3BQAg2WVWBgMAkNYB7UcffeRaQDW6sIJKeeKJJ9w0OF5/1ubNm2e4D6Xszp8/37WwbtiwwY4//nirWbNm0Dr9+vWz7t27x7SNlCtXzt5++203knGo/Pnzu0BXQazSlZXGXL9+fStevHi8lwMAgByTiDIYAIC0TTl+4403rGvXrnbEEUfYokWL/Mv37t1rDz74YEz7Uhpx06ZNXapvuMBUIw+fddZZMW0jxx13nHXu3DnDY1eqVMnOPvtsty+CWQBAMkhkGQwAQFoGtEr7nTx5sqsNDqR5XF999dVEnRsAAKAMBgAgsQHtL7/8Yq1atfK3lnoqVKjgpsoBAADZgzIYAIAsBrRKc1qxYsUhAe2cOXPcXK4AACB7UAYDAJDFgLZPnz52zTXX2JIlS1xAq8GZXn75Zbv66qvtqquuimeXAACAMhgAgOwf5VjT3fz9999uipsDBw5YxYoV3cjBGpF48ODB8ewSAABQBgMAkP0BbcGCBW3s2LF27733ulbagwcPWqNGjVxgCwAAsg9lMAAACZiH1uvHEzqlDgAAyH6JKoN3795tW7dudQM7KtsqWnv27HHzuWt+eM39DgBAng5ob7/99qh3ev/998d7PgAAIM7yuEiRInb00UfbeeedZ2XLls3wOiq76qabbrJnn33WbXfYYYfZU089ZRdeeGFU11/jZmiqvoEDB9pjjz3GawYAyNsBrVKLA2tlP/nkEytdurTVrVvXLfvhhx9s+/bt1rp16+w5UwAAYF988YUrg8uUKWPHHXecG5xRZfC2bdusefPm9ttvv9mgQYNs/vz5/jI6nIceeshee+01++qrr6x+/fr29NNPW48ePdw2+jsjEydOtJ9//jnT9QAAyG5R5xbNmDHD/1B/2csuu8zWrl1rn3/+uXusW7fOLWvcuHH2njEAAGlMZbBmFFC5q/J3wYIF7vcrrrjCBbSrVq2yzp07u6A2I+PGjXP7Of74411Q3L9/f6tRo4aNHz8+03lwb731VhfUqj8vAABJN23P1KlTXc1u8eLF/cv0+4MPPuj+BwAAsofK2dGjR7sUYU+xYsVcdx/9r3DhwjZs2DD78ssvI+5j48aNLvA95ZRTgpafeuqpGW63d+9eu/jii23EiBFWu3btBD0jAAByOKDdvHmzm7YnlJZt2rQpC6cDAAAyonI2XFkbWAYr2FW/2Ej++usv97N8+fJByw8//HD//8JRy6xacdUaHC11U1KXpMAHAAC5GtB26tTJ1dDOmTPHXzjpd/W9Of/88xN2cgAAIFjHjh1dGTx37lx/GazftaxDhw5unWnTprn1IvFGM96/f3/Q8n379lmBAgXCbqNy/uWXX7Y777zTfv/9d/dQi62Or98jUWuyxtzwHlWrVuUlBQAkTFydXzQi4oABA9x0ARol0Ssce/bs6eanBQAA2UNlcL9+/axt27bm8/ncMvWB7datmxvYSTTCsboGRVKlShX3c/369UHL9XflypXDbrNmzRoXkF5wwQX+Zeq7q+UKdleuXBk2GB4yZEhQf14FwAS1AIBcDWhVoGkwiIcffth++uknt0wjLR555JEJOzEAAHAoBauTJk2yRx991JXBCmaPPfZYq1Spkn+dLl26ZHjpSpYsaSeeeKLNnDnTunfv7m+d/fjjj900PB6lMCtl+KijjrLevXu7RyANBKnZDTKatkepzxmlPwMAkBVZGp6wYsWK7gEAAHKWAtjAIDZWd999t5tztkmTJtaiRQsbM2aMa2FV66/ntttus4ULF9qyZcsSdNYAACRW3AGtBo1Q2tOPP/7oUp7q1atnffv2dQNKAACAxLn99tujXlejHUdDY16opffxxx+3J554who0aGDz5s0LKsf1e6QUZFHLbbly5aI+NwAA8kRAq0nd27Vr59KeTjrpJJfu9Pzzz7sUZKUvNWvWLOEnCgBAulqyZIn/d6UAf/LJJ677T926dd2yH374wfVNVfpvLLp27eoe8QbHH3zwQUzHAwAgTwS0N998s/Xp08efniQHDhxwy/WYP39+os8TAIC0NWPGDP/vN954o1WvXt0NwujNB//vv/+6wRpV0QwAQDqJK6D96quv3JQAgaMZ6vehQ4dmmJoEAACyZurUqfb111/7g1nR7w8++KDLmtJgUQAApIu4AtoSJUrY2rVrD5mQXUP3a+TEWGkURfXJ1WTtRYsWzfI2Og9NMB9Ik8zXqVMnIccGACC3bN682ZVxRxxxRNByLVOZBgBAOvm/mdVjdNFFF7kJ3JUCpYJVjw8//NAN/a8RE6OlKQIuu+wyN6iE5rStUKGCvfjii1neZsSIEa4f0eWXX+5/aCL4rB4bAIDc1qlTJ1cGa+5X9ZvVQ7/36NHDDfQEAEA6iauFVoM/XX/99dahQwc7ePCgW5Y/f34XIKpfbbRGjhxps2bNcvPoqYX09ddft0svvdROOOEE98jKNmeeeaZNmTIloccGACC3aYYB9ZdVZWxgGdyzZ0/XrxYAgHQSVwut+uqoNXPdunVupMVPP/3U/a5lxYoVi3o/48ePt6uuusoFlKLCWGnBGjE5q9vs37/ffv75Z9uwYUPCjg0AQG7T6MYTJ050XX9UBuuhMvjVV1+1UqVK5fbpAQCQHPPQypFHHuke8fjzzz9dAdy0adOg5c2bN3eDXWR1m3fffdeWLl3q+sdq4nnVaJ9xxhlxHxsAgLykYsWK7gEAQDqLO6DVKMefffaZ6z8bKppWTm/gitCBpfR36IBOsW6j/rN33HGHVatWzfbu3WuDBg1y/Yq+++47N9VBPMf25v7Tw6N+SwAA5LSPPvrIZUX9+uuvbm54eeKJJ6x3795WpkwZXhAAQNqIK+V4yJAhbgCob7/91v75559DHtEoVKiQ+xkYIHp/e/+LdxsNlqFgVgoXLuyfwuCdd96J+9gyevRol+rlPapWrRrVcwUAIFHeeOMN69q1qxvleNGiRf7lqsDV1D0AAKSTuFpoVSs8c+ZMa9WqVdwH1ny1+fLlc+m/gZQKHClQjGcbUZCqUYxXr16dpf0okFdrb2ALLUEtACAnaVDDyZMnW/v27e3JJ5/0L+/SpYvLUBo1ahQvCAAgbcTVQnvgwAE3eXtWaC5b9WH94IMP/Mt2795tH3/8sbVt29a/TEGoRiKOdhufz+dqqQP9/vvv7nHcccfFdOxQRYoUcQNuBD4AAMhJv/zyi79CWZWzHlXcbty4kRcDAJBW4mqhVUGqFtoLLrggSwe/77773NQ/devWtRYtWthjjz1mJUuWtGuvvda/zrBhw2zhwoW2bNmyqLbR/LIKtq+77jqrX7++rVq1yu1D6/fq1SumYwMAkNco1XjFihXWoEGDoIBWc9FqnAgAANJJXAGt+qcqONTgE7Vr1w4qUGXw4MFR7addu3aulVQpUxqVWIXz/PnzXf/UwGNt3bo16m3UZ1YDVqnfrKYwKFu2rF155ZVu3tyiRYvGdGwAAPKaPn362DXXXGNPP/20K381Pd2MGTPslltuibr8BQAgrQPauXPnujlbNbKiN7pioFgKVE0Mr0ckd999d8zb1KxZ0432mNVjAwCQ1wwdOtSNyK9sJHUB0tQ9+fPnt379+hHQAgDSTlwB7ZIlSxJ/JgAAIFMFCxa0sWPH2r333uvK44MHD1qjRo2YkxYAkJbinocWAADkbl9asowAAOkurlGOAQAAAABImhZaDf4UD43ECAAA4kcZDABAFgNaTWsDAAByHmUwAABZDGg7duwY7aoAACCBKIMBAAiPPrQAAAAAgKREQAsAAAAASEoEtAAAAACApERACwAAAABISgS0AAAAAICkREALAAAAAEhKBLQAAAAAgKREQAsAAAAASEoFc/sEAABA7li+fLmNHz/eNmzYYA0aNLABAwZYiRIlIq6/Z88emzRpki1YsMAKFixoLVu2tO7du1v+/NSPAwByByUQAABp6Ouvv7YmTZrY5s2b7bTTTrMpU6a4nwpawzl48KDVq1fPPvnkE7fdMcccY7fccotdcMEF5vP5cvz8AQAQWmgBAEhDQ4YMsdatW9uECRPc3xdeeKFVrVrV/d23b99D1s+XL5/NmzfPKleu7F/WrFkzO+WUU2zx4sV20kkn5ej5AwAgtNACAJBm1Ao7Z84cF8R6ypcvb23btrUPPvgg7DYKaAODWalSpYr7uWXLlmw+YwAAwqOFFgCANLNq1Srbv3+/VatWLWi5/v7000+j3s+TTz5pZcqUcS21GQXPgWnM27dvj/OsAQA4FC20AACkGS/ALFasWNByDQi1e/fuqPYxefJkGzNmjD377LNWqlSpiOuNHj3aSpcu7X8orRkAgEQhoAUAIM0osAyXKrxp0ybX4pqZadOm2aWXXmpjx461bt26ZdpXd9u2bf7H6tWrs3j2AADkoZTjpUuXutpdb8qAgQMHZlqYZrbNzp077dVXXw2aVqB3795WoEAB/zrjxo07pJ9Q9erVXeEMAEAqU9/XsmXLuvL03HPP9S/X3w0bNsxw2+nTp7upeh599NGwg0eFKlKkiHsAAJByLbSLFi1y/W40FUCnTp1s5syZLvjctWtX3Ntoef369W3JkiVucIuTTz7Zhg0b5tbV/wIL7a1bt7rC2HuogAYAINVpgKdevXrZCy+84MpCUd/ZL7/80i655BL/es8884wNGjTI//f7779vF110kT3yyCPWv3//XDl3AADyTAut0pDatWtnTz/9tPv7/PPPd7XGKmCvu+66uLZRIf3FF19YhQoV/NuceOKJLgj+6quvrGnTpv7lFStWtI4dO2b78wQAIK8ZOXKkm4u2bt267qGy84477rA2bdr411G5uXDhQv9gTl27dnX9ZTVCsh6efv36uUpkAADSJqDVoBOqDVYg6lHasArEGTNmhA1oo9lGAW1gMOsFruFGVvz2229d3x/1JdJk8qqVzp+fbsUAgNSnwHT+/PmuVVZdeI4//nirWbNm0DoKVL3spaJFi9rEiRPD7qtWrVo5cs4AAOSZgFZTBhw4cMA/h51Hox/OnTs3YduIUqPKlSsXNK1AoUKFXC10q1atbO3atXbrrbfapEmTXDqVguJwmHoAAJBKVN4FZi6FatKkif/3woULB81bCwBAWge0e/fuDTtlgP72/peIbTQ41FNPPWVTpkyxkiVLBqVaBU4z0L59e2vcuLG9/fbbdsEFF0ScekD9cQEAAAAAuS/X8mu9UYk3b94c9ZQBsW7z5ptv2pVXXmnjx4+3zp07B/0vdM48jZasUY4XL14c8ZyZegAAAAAA8o5cC2grV65s5cuXd/1YA2l04kaNGmV5m6lTp7o+sRqhsU+fPpmej8/ncyM9KhU5Ek07oEA48AEAAAAASLOAVv12NCn7888/71pY5aOPPnIjLmrOWM+TTz5pAwYMiGkbpQ337NnTjYR8xRVXHHLsffv2uYGlFMR6Ro0a5SZ816jJAAAguYVmc0Wz/u+//+4euk8AACSHXB3Sd/jw4VajRg2rU6eOnXLKKXbeeefZfffdZ6effrp/HbXGamTjaLfRSMYakVGtp2+99Zablsd7zJo1y61ToEABNxWBBpM644wzrHbt2vbEE0+4/rYnnHBCLlwJAACQCCtWrLDjjjvOjbysx48//hhxXQ006dF4G61bt3b3Fz/99FNCX4yff/7ZZZN99913MW+rCnwNipkONN6JBh7T6Np6HWOhkbrPPfdcmzZtWradHw61f/9+GzFihMui1KCtaiCK1uGHH24FCxa05s2bJ/zSfvbZZ3bttdfa4MGDbeXKlVFv9+6777ruiprCTIPGIjnk6jy0JUqUcPPYLV261H0R1a9f34466qigdW644Yag1tfMtjnssMNcIBuO5tkTTc2j1lt9AJctW2Zly5a1Y4891k1JAAAAktftt9/ubkhvueUW1+3o5ptvtg8++OCQ9f766y83voYqu4sXL2533323e5x00kkJPycNKKn7jT/++MN1b8qIBrnUiNL6qXsg3dPovkVTKn344YcRxxlJdgpE9ejVq5d98803dvbZZ0cMRMaOHesy8HTPJ8q4u/HGG9193OrVq3P4zNOb5q/We1XTf6ny5ayzznKvXeAI6R5lWOo102ssuo/X50Kf10Sfkz5zauBSxYgavVQppAatUJMnT7YOHTq474Bnn33WTQPaqVMnN6WZpgVdvnx5Qs8NKRjQeho2bBjz/yItVx9YtcZGQ/1xNW0PAABIDboRVaW1XH755W5avoMHD4Zt5VTw2qNHD3vnnXeydR761157zXbv3m0VK1bMcD0FrC+99JKbRlDdoBQU6Lno/Nu1a+cCvssuu8xSkQIhr9vXmWee6TLn9LzDvS6qFNA4KWrRVXe0hx56yP2tbmjIWaeeeqp7iBqY9B4PnFUkkF5f3XdXq1bNTjvtNBdghgsys0oZGjNnznTvDbUgq5uhMjjVgBVK3wved0DXrl1dq66oMa106dJue7UiI2/jFQIAAClDwY7XiqkWO6UV//333y6dOJT6yq5bt86+//57N9tBTlLA+ttvvwUtq1SpkjsfpTtqqsCrr77aLddNv26svUyzrFKgqBTrV155xbWYqRU7dDaISDSDxCOPPOJaRTXGicY2SQRNwfjf//7X7rnnHvv1119t3LhxLpi96aabwqZqz5s3zx5//HFr0aKF/fPPP66VjYA2d9177712zjnnuKxHvUdUmRTuvafMA7XCZxcFovpMK+Pil19+cbOUKJh98MEHbdGiRYes//HHH7vpPIcOHepfpveWUt8JZpMDAS0AAEgZaiFS2unRRx/tAlkFShUqVHCDPQVSQHbxxRe7rks5HcyKBrRUsBYuFVr9fhXQeuc5aNAg10LbtGnThBxbrVEKpp977jmX1qvWTQWRuk6hFFjquGrN1jbvvfeeaz1Vi5dattT/8ZhjjknIean/slLGlYaq8VE0TorSU0PTtPX66vqpkuKqq65y3dHUqqu+ynq9jzzySLvooosSck7InIJUvZc1fo33vtXfClxDB17T2DVjxozJ9suq97I+33pP6H2uzAa1Cut7IZDex6oICexqoMBXnwsv0wN5HwEtAABIGQqC7rrrLhcYKfCK1PKom2ulPqoVR3bs2OH6AKo/oFpJ1S1JLabZRTf2Gigq0MaNG12Qphtw0bkoYNOAleGC33jpmlxwwQXu9xNPPNH1N1ZQEo4G4OzWrZvNnj3bqlevbi+//LJbvnPnTneNAmeMyAr1XVQ6qgJ3XQO1IGtQnnB9mtu3b+8G8mzcuLF7jXUuogBEaa/NmjVLyDkhc0qlV4WI3ieqePEo3Te0dVOzkGiQpjZt2mTrpdWAUHpfe5UaGuhJLbYKaEPp3BWEq2XZaz1WpYj60yJ55OooxwAAAImkUVbVYqeWPaXpPvzww2HXUzCmvpeBN9sKpNRic80111j//v0Tdk4aDEqDWSr9WUGsAuZw1J9PKbctW7Z0QZpussuVK+cPfhXwJkJgv1TNHqGWVrVs69ppJonAh4Jfna+uh7Zbv369W37EEUe4PodKL00EBUD16tVzQYf2rwAp0r7feOMNl2IsWk+ts3po9gq18qqiAjlD6bpTp0516cUKYPWINDir+qf26dPH/7da91UpsXjxYredV5GTVXv27HGVQBpT5/jjj7ciRYrYySefHHZdtRZ7n3W1zKqiRNkKGm1bj8wGcUPeQAstAABIGQpm47kx1iAwgbMqJJJaDtUCqWmENFCV0h/DtbiqNdJrkVTwqhbjTz75xD1k4MCBQQFBVqhl9bbbbnOB5AMPPOCWqc9gz549D2nJVlDrtb4pkNX5rFmzxq2v1lIF4FmlkXEV9Ku/rKZ/ySgoVR/JcBSYqHUNOUdTJYXO2xztQE9qSQ1s4U/UAFFqAda+1V9W6edqrY1E//coANZUnoG8kbSRtxHQAgAAZKP777/fPWKhVsrQlOREUQuWAmPNGKHUbI8GzgkdCVZTrWgKFAUJmjZR05sopVfp2Hokco5cjY6rFtd4Jaq1GNELl1ocrewY4TgwA0MZDrFQa64eSD4EtAAAIDmNyZfzx7w5MX1Gc9Prr7/uRitesGCBm6tXNAVOuP6qCma9wEMpwUrlVkvqrl27XIuW1xcXaSA3Pm8p8plD9iKgBQAASKMbfg2Wo365gSINgBXYiqZ+tuozqTRktdSqb2JcqIhAovBeAgEtAABA+vUz1iNeGqgKAPIKWmgBAABSVZK0GgNAvAhoAQAAACBVjUnt8QaYhxYAAAAAkJQIaAEAAAAASYmAFgAAAACQlAhoAQAAAABJiYAWAAAAAJCUCGgBAAAAAEmJgBYAAAAAkJQIaAEAAAAASYmAFgAAAACQlPJEQLt69Wr76quvbPv27QndJlHrAACQqrKrDAYAIOUD2t27d1vXrl2tTp06dumll1rFihXtySefzPI2iVoHAIBUlV1lMAAAOamg5aJhw4bZokWLbOXKlVapUiV75513rEuXLta0aVNr1qxZ3Nskah0AAFJVdpXBAACkTQvthAkT7KqrrnKFonTu3NmOP/54tzwr2yRqHQAAUlV2lcEAAKRFC+26detsw4YN1qRJk6DlquX95ptv4t4mUeuEs2fPHvfwbNu2zf1MRB+i/bt3WrKL5zrwvNPrNU+F11t43tFJ19c7o334fD7LC7KrDM7pstN2W87L6Lxz43yEc8raNRLeS7yXUvW9lFfPKYFlZ64FtJs3b3Y/y5cvH7Rcf3v/i2ebRK0TzujRo126VaiqVatG3CadlB5laSldn3c6P3eed3pJ5Ou9Y8cOK126tOW27CqD06LsvCv3X79DcE5cI95LfN5S+Hsps7Iz1wLaQoUK+QeYCLRr1y4rXLhw3Nskap1whgwZYoMGDfL/ffDgQVeIqzDPly+f5VWq3dCNg0alLFWqlKWTdH3uPG9e73SQTO9z1S6rQD7qqKMsL8iuMjgZys5ket/kJq4T14j3Ep+5ZCk7cy2gVWGSP39+W7t2bdBy/V2tWrW4t0nUOuEUKVLEPQKVKVPGkoUK7nQtvNP1ufO80wuvd96WF1pms7sMTqayM10/L7HiOnGNeC/xmcvrZWeuDQpVrFgxO+WUU+zdd9/1L/v3339t9uzZdtZZZ/mXrVixwt83J5ptErUOAACpKrvKYAAA0mranhEjRrhCUOlILVq0cHPZVahQwa655hr/Ovfff78tXLjQli1bFvU2iVoHAIBUlV1lMAAAaTNtT6tWrWzu3Ln2xx9/2OOPP27169e3+fPnW4kSJfzrHHPMMXbiiSfGtE2i1kkVSvW65557Dkn5Sgfp+tx53rze6SBd3+d5vQzO63jfcJ14L/GZy4v4bopfPl9emUMAAAAAAIBkaaEFAAAAACBeBLQAAAAAgKREQAsAAAAASEq5Osoxcsb69evdBPJHH320m8g+Xfzzzz/27bffWvXq1a1KlSqWTq+3Hnq902mORU0f8vPPP7v3eEZzYqYqPfeNGze6AXw0vUoqW7x4se3atStoWeXKla1mzZq5dk5IHr///rv9/fffdtxxxyXVYFY5afPmzW7gL5Wf5cqVy+3TydN27Njh7jUqVapktWrVyu3TybPvJ33ujj32WD5zEaxatcqV4UcddZR7IDa00KawAwcO2FVXXeVu8i6//HJ3w6cRQVPdmjVrbMCAAW50zjPOOMMmTpxo6UAjjzZt2tROOOEE93pXrFjRbrzxRkv1cd82bdpkV155pQti9X5v3LixNWnSxH766SdLp2D25JNPttNOO81+/fVXS3Xdu3d37/Hbb7/d/5g+fXpunxaSoNLr3HPPteOPP94uueQSO/LII+2FF17I7dPKUzTnsKZlUuBxxRVXWNWqVa1Xr162e/fu3D61PEvfRRr9+6GHHsrtU8lzVPHYp08f16igqb3q1atnTz31VG6fVp4LZFV+696tb9++rqJN9666t0H0CGhTmL40pk6d6moOv//+e5s9e7aNGjXK3n33XUtlK1eudF+aCmg0P2K6WLFihT399NP2559/2pIlS2zBggX2/PPPu2WpTM/37LPPtr/++su13Onv0qVL27XXXmvpYM+ePXbxxRe7gjCdDB482E0X4z1uuOGG3D4l5HGq+FC5oJai5cuX27PPPutusr05dmH2yy+/uOukFmwFtz/++KOrLL377ru5PGGMGzfOtT6qEhWHUqXIZ5995ipdv/rqK3efwjRrwW655Rbbv3+/a4zRNdL3kyqmhw8fzlsqBgS0KezFF190N7qqaZWWLVta69at3fJUpppStdCmU7qtXH311UGFqloqW7Ro4W72U5laW7p162b58//f15kKSz3vtWvXWroUho0aNbKLLrrI0oluIr/88kuXXg9kRjeMr7zyivXv398OP/xwt0yttEqpfemll7iA/5++S9u2beu/Hsp86dChQ8qXI/H47rvvbMSIEfbqq6/6yx/8zw8//GCTJk2yMWPG+Lt9FS5c2N2r4H9UGa/7tcMOO8z9rRT/unXruuWIHn1oU9S+fftcq+z1118ftFwpqemSgpvulOqjlod0KTyUibB161ZXiCqNUIVoqlOa7QcffOBaUtIpxVpGjx5tkydPdjX+uhlQUFK7du3cPi3kUWrx2L59+yEtaUr10+cH4anLijJfVHGI/9m5c6fr+vDYY4+l1Rgdsfj4449dAKsMKn3+NK6JvqNTfYyHWN15550urf/JJ5902YVffPGFLV261JXtiB4BbYratm2b60MbOgiU/lbLBlLfTTfd5Co2+vXrZ+lAhYEKAaU2qS+p+qCkMrVAq7LinXfesZIlS1o63gDoZmnLli12wQUXWNeuXd2Nd8GCFGs4lFfuhSsT9Z2B8B544AGXdvzyyy9ziQKosUANBGrRRnjr1q1z2RC9e/d2QZoCWQ00psrI0MaWdKb3UceOHd0YNzVq1HDBv8p2BbeIHjkSKapQoULuZ+hADmq1000gUtuwYcNcS7yCHY28mA7UX3jRokWuEM2XL5+1b98+pQfEUkXFKaec4lIplQ6oFmpRa5P6waUyDTLifY+VLVvW3SCpMkM33kA4lImxmzBhgus7q5TaBg0a8Mb6/2bNmuWyQ9RC6/Xh10jH6v6g3w8ePMi1+v+fOZXHCtJ+++03lzU4fvx4GzhwoOsriv/Ts2dP16dfg0N9/fXXbhyY999/n3EhYkRAm6I0KE6ZMmUO6Ueov9NxSpN0oj49Gm1RX4jqN51uVAusPtTq35TK/WjVsqQh/r1RftVCLY8++qi9+eablk40Wq2k8uuNrFFfWaFMjI76G2tgPf288MILefsF2Lt3rxu3YOTIkf7vX02NqFZI/a6B+mAukJXAARp79OjhMorok/1/lEn54Ycf2mWXXeafzkhlu65Tqg/gmmjkZqWwM8880/Wx02igopYcBTmqVURq0ijWaq3S66zBsdJlKo7ixYsHLVO/SqWeqmInlVtPAqnGW/0BdQOayv3d1HcttA+WWkzUKk+KFiJR6qP6WusmsXPnzm6Z0tXnzZtnTzzxBBcugFpklfKoNGMNLIlgSg/VI1Dz5s3d++uZZ57hcgXcg2qwLFUieXOE6zOnMvuII47gOplZgQIF3CBQGuE4kCpIuEaxIaBNYUoV0pesUhM1954KJ6UgDxo0yFKZ0qrVl86rSVWfDdUGKjWxfv36lqp0U6a+hWqhVTDn1YAqqEvldDG1RuvLv127di4rQYHd/fffbzfffHPa9S1NB0orV18j9cvSHJlqFdHrrTmXyT5BZhV+nTp1cu8TzfmogeNq1arl3kv4P2+//bZL6dd9gz5fXjlStGhRO+mkk7hMiJo+Z0ov1jzx6galimd95jQwlFepBHOpxSrDNDNHw4YNbeHCha7COtWnXEy0fL5U7mQG169OKYjKza9Tp47deuut/pqyVKXnqj4JoRTcP/zww5aq1BKvL8JQunHz0lFTkb7CpkyZYtOmTXNzJ6oQVRZC4NQT6UCjHOvGQa0rqf4Z13Q9Gsla/bI0wqimLFKfaSAzmlNVrWibNm1y34233XabfxofmCsjNfZCKI3FkG5dGWKhtNpjjjnGnxGH/6P+xJoqUuWzymqNMq4BK1X5jP+ZOnWqq0zasGGDVa5c2aUca3RoRI+AFgAAAACQlBgUCgAAAACQlAhoAQAAAABJiYAWAAAAAJCUCGgBAAAAAEmJgBYAAAAAkJQIaAEAAAAASYmAFgAAAACQlAhoARxiy5YtNmnSJNu/f3/cV2fjxo1uH5pMPbfPBQCAZCiLKDuB2BHQAjjEypUrrUePHrZ79+64r84PP/zg9nHgwIFcPxcAAJKhLKLsBGJHQAsgW1SoUMG6d+9u+fPzNQMAAGUnkD0KZtN+AWSzdevW2WeffWYXXXSRLVu2zH7//XerW7eu1apV65B1V6xYYd99952VK1fOWrRoYYULFz5kna+//trts169ehGPuWHDBlu0aJGVKFHCTjzxRCtdunTEdQ8//HDr3Lmz5cuXL+bzzcq5KB3srbfeskaNGlmdOnX863/yySd28OBBa9OmTcR9AgDSQzTlWTRl0eLFi+3PP/9065QtW9ZmzpxpF154oRUs+L9bbMpOIJv5ACSl6dOn+woVKuRr3769r1mzZr6zzz7bV7hwYd+4ceOC1rvuuut8xYsX97Vr185Xu3Zt39FHH+375Zdfgta57LLLfCVLlnT7qlGjhu+cc85Rx1ffjh07/Os88MADvtKlS7vjtGrVyle+fHnf+++/H/H85s6d6/axb9++mM43Eedy0003uW23bdvm/l6wYIGvYMGCvhkzZsR9vQEAqSGa8iyzsujgwYO+iy++mLITyAMIaIEkpQBRhevDDz/sX/bUU0+5QloFrbz33nsuiPzmm2/c33v37nWB7VlnneXfZtq0ab4iRYr4fvzxR/f3zp07fSeffHJQwT1r1ixfuXLlfCtXrvRvN3HiRN/hhx8eFGhmFtBmdr6JOpfdu3f7GjVq5OvVq5cLamvWrOkbOHBglq85ACC5RVOGRFMWvfXWW0Hr7Nq1y1XWUnYCOY/ObUASUzpv3759/X+3bt3atm3b5tKbRCMynnPOOda4cWP3d6FChezWW2+1jz76yDZv3uyWTZ482Tp16mTHHXec+/uwww6z66+/Pug4EyZMcOlUSr9688033TayadMml8qcqPNN1LkUKVLEXn/9dZd63LJlS5dS9sADD8RwZQEAqSiaMiSasmjKlClB6xQtWpSyE8gl9KEFkpj6whYvXtz/twI58UZh/OOPP6xJkyZB23h9VvU/9aldtWqVNWvWLGidmjVrBv2t/q6askAFeKBu3bq5IDlR55vIc9ENi0akfPHFF+2DDz7wHwsAkL6iKUOiKYtWr159yDo1atSI+VjRoOwEMkZAC6QwDczktcR6vL/1PylfvrwrcAOF/l2qVCmrUqWKa/HNTok8l2+++cYmTpzoAtuRI0dau3btrECBAtly3gCA5BBNGRJNWaQK4a1bt2a4DmUnkDNIOQZSmNJtZ8yYYTt37vQvU4qVapFVoHvrfPjhh7Znzx7/OkrVDdS+fXt7//33XY10II3sqL74iTzfRJzLrl27rFevXnb55ZfbnDlz7JdffrFRo0Yl7DwBAMkpmjIkmrLo1FNPdevs3bvXv2zatGkxHysRKDuR7vKpI21unwSA2L333ntuaoDASd41Pc8xxxxjv/32mwta9T+lHKsPaZ8+feynn36ysWPHuv5BmlJH/v33X2vYsKEdddRRdskll7h+RSq4//77b9uxY4fbVoW6+uJq+wEDBrjW3SVLltjs2bNt+fLlYeea1TQ5Z5xxhu3bt89NXxDN+SbqXPr3728ff/yx214pzrqh6NKli82fP9+aNm3K2w0A0lQ0ZUg0ZdE///xjxx9/vFWvXt169uxp3377rUst/uuvv9z/VPZQdgI5gxZaIElVrlzZ9cMJVLJkSevevbu/n6oGqfj888/d3K/6qforBXVeMCtad+HChW6AJs3Jp7lb1aqp/Xh9fNT/dNasWfbwww+7vkWad0/zvKoADxfMSoUKFdw+vP9Hc76JOBfVhCvtS4NCefvt0KGDjRgxwqZPn57AVwAAkGyiKc+iKYsU1H7xxRd22mmnuXVq165tzzzzjOvaokGkoj1WKMpOIHa00AIAAAAx0pgU6kvrufHGG+2zzz6zL7/8kmsJ5CAGhQIAAABipLEaTjjhBDd7wIIFC1xmkMapAJCzaKEFAAAAYqRRjl944QX74YcfXLcapSTXr1+f6wjkMAJaAAAAAEBSYlAoAAAAAEBSIqAFAAAAACQlAloAAAAAQFIioAUAAAAAJCUCWgAAAABAUiKgBQAAAAAkJQJaAAAAAEBSIqAFAAAAACQlAloAAAAAgCWj/weUl6+nKhhzjwAAAABJRU5ErkJggg==", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "top node by |mask|: 4 top edge by |mask|: 8 (4→3)\n", + "node |mask| range: [0.122, 0.144]; edge |mask| range: [0.078, 0.946]\n" + ] + } + ], + "source": [ + "node_w = explanation.node_mask.detach().cpu().squeeze(-1).numpy()\n", + "edge_w = explanation.edge_mask.detach().cpu().numpy()\n", + "src, dst = snapshot.edge_index.cpu().numpy()\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(9.5, 3.2))\n", + "\n", + "axes[0].bar(np.arange(len(node_w)), np.abs(node_w), color=\"steelblue\")\n", + "axes[0].set_xlabel(\"node index\")\n", + "axes[0].set_ylabel(\"|node_mask|\")\n", + "axes[0].set_title(\"GNNExplainer node mask (abs)\")\n", + "\n", + "axes[1].bar(np.arange(len(edge_w)), np.abs(edge_w), color=\"darkorange\")\n", + "axes[1].set_xlabel(\"edge index\")\n", + "axes[1].set_ylabel(\"|edge_mask|\")\n", + "axes[1].set_title(\"GNNExplainer edge mask (abs)\")\n", + "for i, (u, v) in enumerate(zip(src, dst)):\n", + " if i < 8: # annotate a few for readability\n", + " axes[1].text(i, abs(edge_w[i]), f\"{u}→{v}\", ha=\"center\", va=\"bottom\", fontsize=7)\n", + "\n", + "fig.suptitle(\n", + " \"Interpretive masks for one_step_forecast (not causal / not ResDMD / not ModeEnergy)\",\n", + " fontsize=10,\n", + ")\n", + "fig.tight_layout()\n", + "plt.show()\n", + "\n", + "top_node = int(np.argmax(np.abs(node_w)))\n", + "top_edge = int(np.argmax(np.abs(edge_w)))\n", + "print(\n", + " \"top node by |mask|:\",\n", + " top_node,\n", + " \"top edge by |mask|:\",\n", + " top_edge,\n", + " f\"({src[top_edge]}→{dst[top_edge]})\",\n", + ")\n", + "print(\n", + " f\"node |mask| range: [{np.abs(node_w).min():.3f}, {np.abs(node_w).max():.3f}]; \"\n", + " f\"edge |mask| range: [{np.abs(edge_w).min():.3f}, {np.abs(edge_w).max():.3f}]\"\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "id": "66f88324", + "metadata": {}, + "source": [ + "## Interpretation / discussion\n", + "\n", + "On this tiny path, masks fluctuate with initialization, explainer epochs, and\n", + "the short fit — treat them as a **diagnostic view** of the explainer's\n", + "optimization under the chosen target, not as ground-truth importance. Flat\n", + "node bars with a few large edge bars are common on under-trained synthetic\n", + "fixtures; that pattern is evidence about the explainer run, not a recovered\n", + "communication pathway.\n", + "\n", + "Complementary diagnostics:\n", + "\n", + "| Tool | Level | Claim type |\n", + "|------|-------|------------|\n", + "| `explain_representation` | node / edge / feature masks | interpretive, non-causal |\n", + "| `attribute_mode_energy` | type/relation energy on `K_eff` | interpretive operator energy |\n", + "| `spectral_residuals` / `resdmd` | residual diagnostics | data-driven / fixed-dictionary residuals — not this API |\n", + "\n", + "The optional Captum path above (`algorithm=\"integrated_gradients\"`) is for\n", + "feature-attribute masks when `[explain]` is installed; it does not replace\n", + "GNNExplainer edge masks or ModeEnergy / ResDMD diagnostics.\n" + ] + }, + { + "cell_type": "markdown", + "id": "36cd930b", + "metadata": {}, + "source": [ + "## Takeaways\n", + "\n", + "- Use `explain_representation` for **homogeneous** snapshot masks on\n", + " `latent` / `reconstruction` / `one_step_forecast`.\n", + "- Default `gnn_explainer` needs no extra; `integrated_gradients` needs\n", + " `[explain]` (Captum).\n", + "- Masks are **interpretive** and **non-causal** — not topology discovery, not\n", + " ResDMD, not `ModeEnergyAttribution`.\n", + "- MVP rejects hetero / RelGraph / hypergraph / adaptive topology / delays /\n", + " controls; expect actionable errors rather than silent fallbacks.\n", + "- Prefer short synthetic fixtures for teaching; do not over-claim mask contrast\n", + " on under-trained models.\n" + ] + }, + { + "cell_type": "markdown", + "id": "16ba1ea6", + "metadata": {}, + "source": [ + "## Further reading\n", + "\n", + "- Ying, Z., Bourgeois, D., You, J., Zitnik, M. & Leskovec, J. (2019).\n", + " GNNExplainer: Generating Explanations for Graph Neural Networks.\n", + " *Advances in Neural Information Processing Systems*, 32\n", + " (`Ying2019GNNExplainer`).\n", + "- Sundararajan, M., Taly, A. & Yan, Q. (2017). Axiomatic Attribution for Deep\n", + " Networks. *Proceedings of the 34th International Conference on Machine\n", + " Learning*, PMLR 70:3319–3328 (`Sundararajan2017IntegratedGradients`).\n", + "- Sphinx:\n", + " [`docs/source/capabilities.rst`](../docs/source/capabilities.rst)\n", + " (Analysis — `explain_representation`),\n", + " [`docs/source/limitations.rst`](../docs/source/limitations.rst)\n", + " (Spectral analysis — representation explain),\n", + " [`docs/source/api.rst`](../docs/source/api.rst)\n", + " (`explain_representation` / `RepresentationExplanation`),\n", + " [`docs/source/architecture.rst`](../docs/source/architecture.rst)\n", + " (`analysis.explain`).\n", + "- Related notebooks:\n", + " [`07_koopman_spectrum.ipynb`](07_koopman_spectrum.ipynb) (spectral residuals),\n", + " [`40_resdmd_pseudospectra.ipynb`](40_resdmd_pseudospectra.ipynb) (ResDMD MVP),\n", + " [`39_heterogeneous_relational_koopman.ipynb`](39_heterogeneous_relational_koopman.ipynb)\n", + " (mode-energy context on hetero operators).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "koopmangraph", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/cli/README.md b/examples/cli/README.md new file mode 100644 index 0000000..d438e74 --- /dev/null +++ b/examples/cli/README.md @@ -0,0 +1,41 @@ +# CLI examples + +Config-driven entry points for the `koopman-graph` console script. These +configs are small enough for local smoke runs and for the package CLI test +suite. They are **not** collected by `nbmake` CI. + +Install a normal editable or wheel build so the console script is on your +`PATH`. YAML configs additionally need the optional `[cli]` extra +(`pip install "koopman-graph[cli]"`). + +## Synthetic path train (`synthetic_train.json`) + +Tiny GCN train on a seeded path-graph decay trajectory (`data.kind = +synthetic_path`), three CPU epochs, writing a safetensors checkpoint +(`.kgckpt`): + +```bash +koopman-graph train \ + --config examples/cli/synthetic_train.json \ + --out /tmp/kg-cli +``` + +Relative `checkpoint.path` values resolve under `--out`, so this writes +`/tmp/kg-cli/model.kgckpt`. + +Forecast from the first snapshot of the same synthetic sequence: + +```bash +koopman-graph predict \ + --checkpoint /tmp/kg-cli/model.kgckpt \ + --data examples/cli/synthetic_train.json \ + --steps 5 \ + --out /tmp/kg-cli/forecast.pt +``` + +`--data` may be a train/data JSON (or YAML) file, or a trusted `.pt` +`GraphSnapshotSequence` cache. The forecast file is a `torch.save` payload +with `steps`, `forecasts`, and `summary`. + +See `koopman-graph train --help` and `koopman-graph predict --help` for +flags and exit-code conventions. diff --git a/examples/cli/synthetic_train.json b/examples/cli/synthetic_train.json new file mode 100644 index 0000000..b9409d7 --- /dev/null +++ b/examples/cli/synthetic_train.json @@ -0,0 +1,26 @@ +{ + "model": { + "encoder": "gcn", + "in_channels": 3, + "hidden_channels": 16, + "latent_dim": 8, + "num_layers": 2, + "dynamics_mode": "discrete", + "time_step": 0.1 + }, + "data": { + "kind": "synthetic_path", + "num_nodes": 8, + "num_timesteps": 20, + "seed": 0, + "feature_dim": 3 + }, + "fit": { + "epochs": 3, + "lr": 0.01, + "device": "cpu" + }, + "checkpoint": { + "path": "model.kgckpt" + } +} diff --git a/examples/scripts/README.md b/examples/scripts/README.md index e62c42f..7c9b796 100644 --- a/examples/scripts/README.md +++ b/examples/scripts/README.md @@ -34,9 +34,12 @@ or Ray). See `python examples/scripts/ddp_fit_torchrun.py --help` for flags. ## Ray Tune HPO (examples-only) -Tiny learning-rate search with Ray Tune. The **search space lives in the -script** — KoopmanGraph does not expose a Tune / AutoML API. Requires the -`[ray]` extra: +Tiny learning-rate search with Ray Tune. The script uses +`koopman_graph.tuning` helpers (`fit_history_metrics`, `run_ray_tune`, and an +*example-only* `example_lr_loguniform_space` scaffold). The **search +configuration still lives in the script** — those scaffolds are smoke ranges, +not scientific defaults, and KoopmanGraph is not an AutoML product. Requires +the `[ray]` extra: ```bash pip install "koopman-graph[ray]" diff --git a/examples/scripts/ray_tune_koopman_example.py b/examples/scripts/ray_tune_koopman_example.py index f56703b..786dc0f 100644 --- a/examples/scripts/ray_tune_koopman_example.py +++ b/examples/scripts/ray_tune_koopman_example.py @@ -1,15 +1,19 @@ #!/usr/bin/env python3 """Examples-only Ray Tune HPO pattern for KoopmanGraph. -Search space stays in this script — KoopmanGraph does **not** ship a Tune -search-space DSL or AutoML API. Requires ``pip install "koopman-graph[ray]"``. +This script owns the search configuration. It may call *example-only* +library scaffolds such as ``example_lr_loguniform_space``; those are smoke +ranges, not scientific defaults. KoopmanGraph is **not** an AutoML product. +Requires ``pip install "koopman-graph[ray]"``. Smoke run (two learning-rate trials on a tiny decay trajectory):: python examples/scripts/ray_tune_koopman_example.py --epochs 1 --num-samples 2 -Use native DDP / Fabric for multi-GPU *model* training; use Ray for trial / -ensemble-member parallelism. This script is intentionally outside ``nbmake`` CI. +Uses :func:`~koopman_graph.tuning.fit_history_metrics` and +:func:`~koopman_graph.tuning.run_ray_tune`. Prefer native DDP / Fabric for +multi-GPU *model* training; use Ray for trial / ensemble-member parallelism. +This script is intentionally outside ``nbmake`` CI. """ from __future__ import annotations @@ -23,6 +27,11 @@ from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.tuning import ( + example_lr_loguniform_space, + fit_history_metrics, + run_ray_tune, +) _RAY_INSTALL_HINT = 'pip install "koopman-graph[ray]"' @@ -54,7 +63,7 @@ def _build_model() -> GraphKoopmanModel: def _train_koopman(config: dict[str, Any]) -> None: - """Tune trainable: fit one model and report the final training loss.""" + """Tune trainable: fit one model and report metrics from FitHistory.""" from ray import tune torch.manual_seed(int(config["seed"])) @@ -66,8 +75,11 @@ def _train_koopman(config: dict[str, Any]) -> None: lr=float(config["lr"]), device="cpu", ) - final_loss = float(history.loss[-1]) if history.loss else float("nan") - tune.report({"loss": final_loss, "epochs": int(history.epochs)}) + metrics = fit_history_metrics(history) + # Ensure a finite loss key for Tune's metric="loss" even on empty history. + if "loss" not in metrics: + metrics["loss"] = float("nan") + tune.report(metrics) def parse_args() -> argparse.Namespace: @@ -75,7 +87,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Examples-only Ray Tune HPO for a tiny GraphKoopmanModel. " - "Search space lives in this script (not in the library)." + "Search configuration is owned by this script (not an AutoML API)." ), ) parser.add_argument( @@ -113,44 +125,31 @@ def main() -> None: if args.num_samples < 1: raise SystemExit(f"--num-samples must be >= 1; got {args.num_samples}") + # Script-owned search config: example_* scaffold is smoke-only, not science. + # Search space is owned by this script (not an AutoML API). try: - import ray - from ray import tune - from ray.train import RunConfig - from ray.tune import Tuner + param_space = example_lr_loguniform_space( + epochs=args.epochs, + seed=args.seed, + ) + results = run_ray_tune( + _train_koopman, + param_space, + num_samples=args.num_samples, + metric="loss", + mode="min", + storage_path=args.storage_path, + run_name="koopman_graph_tune_example", + ray_init_kwargs={ + "num_cpus": min(2, args.num_samples), + "ignore_reinit_error": True, + }, + ) except ImportError as exc: raise SystemExit( f"Ray Tune is required for this example; install with: {_RAY_INSTALL_HINT}" ) from exc - if not ray.is_initialized(): - ray.init(num_cpus=min(2, args.num_samples), ignore_reinit_error=True) - - # Search space is owned by this script (not by koopman_graph). - param_space = { - "lr": tune.loguniform(1e-3, 1e-1), - "epochs": args.epochs, - "seed": args.seed, - } - tune_config = tune.TuneConfig( - num_samples=args.num_samples, - metric="loss", - mode="min", - ) - run_config: RunConfig | None = None - if args.storage_path is not None: - run_config = RunConfig( - name="koopman_graph_tune_example", - storage_path=str(args.storage_path.resolve()), - ) - - tuner = Tuner( - _train_koopman, - param_space=param_space, - tune_config=tune_config, - run_config=run_config, - ) - results = tuner.fit() best = results.get_best_result(metric="loss", mode="min") best_loss = best.metrics.get("loss") if best.metrics else None print( diff --git a/examples/tracking/wandb_mlflow_callback.py b/examples/tracking/wandb_mlflow_callback.py new file mode 100644 index 0000000..fc5ae00 --- /dev/null +++ b/examples/tracking/wandb_mlflow_callback.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Reference FitCallback adapters for W&B and MLflow (peer installs). + +These classes are teaching sketches — not library APIs. Cloud SDKs are +intentionally not pinned by KoopmanGraph. Install peers yourself when +needed:: + + pip install wandb + pip install mlflow + pip install tensorboard # for TensorBoardFitLogger + +In-tree loggers (no cloud SDKs):: + + from koopman_graph.tracking import CsvFitLogger, TensorBoardFitLogger + +The ``__main__`` block below runs a tiny CSV-logged fit only (no W&B / +MLflow import), so the script stays usable without optional SDKs. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch +from torch_geometric.data import Data + +from koopman_graph import FitCallback, GNNDecoder, GNNEncoder, GraphKoopmanModel +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.tracking import CsvFitLogger +from koopman_graph.training.history import FitHistory, TrainingLossBreakdown + + +class WandbFitLogger: + """Sketch :class:`~koopman_graph.FitCallback` that calls ``wandb.log``. + + Requires a peer ``wandb`` install. Not part of ``koopman_graph.tracking``. + """ + + def __init__(self, *, project: str = "koopman-graph", **init_kwargs: Any) -> None: + self.project = project + self.init_kwargs = init_kwargs + self._wandb: Any | None = None + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + del model + import wandb + + self._wandb = wandb + self._wandb.init( + project=self.project, + config=dict(fit_kwargs), + **self.init_kwargs, + ) + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + del history_so_far + if self._wandb is None: + msg = "WandbFitLogger.on_epoch_end called before on_fit_start" + raise RuntimeError(msg) + payload = { + f"train/{key}": value for key, value in train_breakdown.to_floats().items() + } + if val_breakdown is not None: + val_floats = val_breakdown.to_floats() + payload.update({f"val/{key}": value for key, value in val_floats.items()}) + self._wandb.log(payload, step=epoch) + + def on_fit_end(self, *, history: FitHistory) -> None: + del history + if self._wandb is not None: + self._wandb.finish() + self._wandb = None + + +class MLflowFitLogger: + """Sketch :class:`~koopman_graph.FitCallback` that calls ``mlflow.log_metrics``. + + Requires a peer ``mlflow`` install. Not part of ``koopman_graph.tracking``. + """ + + def __init__(self, *, run_name: str | None = None) -> None: + self.run_name = run_name + self._mlflow: Any | None = None + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + del model + import mlflow + + self._mlflow = mlflow + self._mlflow.start_run(run_name=self.run_name) + for key, value in fit_kwargs.items(): + if isinstance(value, (bool, int, float, str)): + self._mlflow.log_param(key, value) + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + del history_so_far + if self._mlflow is None: + msg = "MLflowFitLogger.on_epoch_end called before on_fit_start" + raise RuntimeError(msg) + metrics = { + f"train_{key}": value for key, value in train_breakdown.to_floats().items() + } + if val_breakdown is not None: + metrics.update( + { + f"val_{key}": value + for key, value in val_breakdown.to_floats().items() + } + ) + self._mlflow.log_metrics(metrics, step=epoch) + + def on_fit_end(self, *, history: FitHistory) -> None: + del history + if self._mlflow is not None: + self._mlflow.end_run() + self._mlflow = None + + +def _assert_fit_callback(callback: FitCallback) -> None: + """Runtime-check that a sketch satisfies the FitCallback protocol.""" + assert isinstance(callback, FitCallback) + + +def _tiny_sequence() -> GraphSnapshotSequence: + """Deterministic five-step path graph for the CSV smoke path.""" + edge_index = torch.tensor( + [[0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3]], + dtype=torch.long, + ) + x0 = torch.ones(5, 3) + snapshots = [Data(x=x0 * (0.9**t), edge_index=edge_index) for t in range(5)] + return GraphSnapshotSequence(snapshots) + + +def main() -> None: + """Run a short CSV-logged fit (no W&B / MLflow required).""" + out_dir = Path("artifacts") / "tracking_demo" + out_dir.mkdir(parents=True, exist_ok=True) + csv_path = out_dir / "fit.csv" + logger = CsvFitLogger(csv_path) + _assert_fit_callback(logger) + _assert_fit_callback(WandbFitLogger()) + _assert_fit_callback(MLflowFitLogger()) + + model = GraphKoopmanModel( + encoder=GNNEncoder( + in_channels=3, hidden_channels=8, latent_dim=4, num_layers=2 + ), + decoder=GNNDecoder( + latent_dim=4, hidden_channels=8, out_channels=3, num_layers=2 + ), + latent_dim=4, + time_step=0.1, + ) + torch.manual_seed(0) + history = model.fit( + _tiny_sequence(), + epochs=2, + lr=1e-2, + device="cpu", + callbacks=[logger], + ) + print(f"wrote {csv_path} ({history.epochs} epochs)") + + +if __name__ == "__main__": + main() diff --git a/paper.bib b/paper.bib index 3738b4c..2e3eb7d 100644 --- a/paper.bib +++ b/paper.bib @@ -878,3 +878,35 @@ @article{Ghosh2025 publisher = {Elsevier BV} } +% Attribution / GNN explanation methods (explain_representation); verified +% from NeurIPS proceedings BibTeX and PMLR publisher record (TASK-2131). +% Pages omitted for Ying et al.: publisher BibTeX lists pages = {}. +@inproceedings{Ying2019GNNExplainer, + author = {Ying, Zhitao and Bourgeois, Dylan and You, Jiaxuan and + Zitnik, Marinka and Leskovec, Jure}, + title = {{GNNExplainer}: Generating Explanations for Graph Neural + Networks}, + booktitle = {Advances in Neural Information Processing Systems}, + volume = {32}, + year = {2019}, + publisher = {Curran Associates, Inc.}, + url = {https://proceedings.neurips.cc/paper_files/paper/2019/file/d80b7040b773199015de6d3b4293c8ff-Paper.pdf}, + eprint = {1903.03894}, + doi = {10.48550/arXiv.1903.03894} +} + +@inproceedings{Sundararajan2017IntegratedGradients, + author = {Sundararajan, Mukund and Taly, Ankur and Yan, Qiqi}, + title = {Axiomatic Attribution for Deep Networks}, + booktitle = {Proceedings of the 34th International Conference on Machine + Learning}, + pages = {3319--3328}, + year = {2017}, + editor = {Precup, Doina and Teh, Yee Whye}, + volume = {70}, + series = {Proceedings of Machine Learning Research}, + publisher = {PMLR}, + url = {https://proceedings.mlr.press/v70/sundararajan17a.html}, + eprint = {1703.01365} +} + diff --git a/paper.md b/paper.md index f2b79c2..38ed621 100644 --- a/paper.md +++ b/paper.md @@ -55,9 +55,9 @@ Tutorials under `examples/` cover synthetic dynamics, IEEE 118-bus, METR-LA, and # Research impact statement -`KoopmanGraph` is research software with no publications or downstream adoptions to cite at submission. Impact is shown through reproducible tutorials on synthetic, IEEE 118-bus, METR-LA, and epidemic graphs---per-horizon evaluation, topology-versus-vector baselines in the spirit of prior GNN-versus-Koopman power-grid studies [@Nandanoori2022], dynamic topology, spectral/stability analysis, and latent-space control---plus v0.5.0 tooling and notebooks for deep-ensemble / latent-Gaussian uncertainty, hierarchical multi-resolution forecasting, auxiliary-spectral continuous generators, and sparsity / worst-case reconstruction losses (with a GNN encoder zoo). Open artifacts include PyPI releases through v0.12.0, CI with coverage and notebook smoke tests, Sphinx docs, and a security policy. +`KoopmanGraph` is research software with no publications or downstream adoptions to cite at submission. Impact is shown through reproducible tutorials on synthetic, IEEE 118-bus, METR-LA, and epidemic graphs---per-horizon evaluation, topology-versus-vector baselines in the spirit of prior GNN-versus-Koopman power-grid studies [@Nandanoori2022], dynamic topology, spectral/stability analysis, and latent-space control---plus v0.5.0 tooling and notebooks for deep-ensemble / latent-Gaussian uncertainty, hierarchical multi-resolution forecasting, auxiliary-spectral continuous generators, and sparsity / worst-case reconstruction losses (with a GNN encoder zoo). Open artifacts include PyPI releases through v0.13.0, CI with coverage and notebook smoke tests, Sphinx docs, and a security policy. -The author uses the library to experiment with Koopman forecasting on networked systems and to lower the barrier for topology-aware Koopman models on standard PyG workflows. Zenodo archives are available through v0.12.0 (DOI: \url{https://doi.org/10.5281/zenodo.21811756}; earlier version DOIs are linked from that record). +The author uses the library to experiment with Koopman forecasting on networked systems and to lower the barrier for topology-aware Koopman models on standard PyG workflows. Zenodo archives are available through v0.12.0 (DOI: \url{https://doi.org/10.5281/zenodo.21811756}; earlier version DOIs are linked from that record; v0.13.0 follows GitHub Release ingest). # Conclusion diff --git a/pyproject.toml b/pyproject.toml index c0b8b88..8b68a7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,9 +33,16 @@ classifiers = [ dependencies = [ "torch>=2.8.0", "torch-geometric>=2.6.1", + "safetensors>=0.4", ] [project.optional-dependencies] +cli = [ + "pyyaml>=6", +] +explain = [ + "captum>=0.7", +] dev = [ "build>=1.2.0", "pytest>=8.4.0", @@ -106,6 +113,9 @@ distributed = [ "dask[distributed]>=2024.1.0", ] +[project.scripts] +koopman-graph = "koopman_graph.cli:main" + [project.urls] Homepage = "https://github.com/tjkessler/KoopmanGraph" Documentation = "https://koopmangraph.readthedocs.io/" @@ -129,6 +139,8 @@ filterwarnings = [ ] markers = [ "distributed: multi-process gloo/DDP smoke (opt-in via KOOPMAN_GRAPH_DISTRIBUTED_TESTS=1)", + "ray: starts a local Ray / Ray Train runtime (exclude on macOS smoke)", + "slow: long-running CPU tests reserved for full Ubuntu matrix (exclude on macOS smoke)", ] [tool.coverage.run] diff --git a/src/koopman_graph/__init__.py b/src/koopman_graph/__init__.py index aa859f6..d4ce623 100644 --- a/src/koopman_graph/__init__.py +++ b/src/koopman_graph/__init__.py @@ -63,18 +63,23 @@ Training customization, evaluation utilities, baselines, adaptation, and control live in capability modules: :mod:`koopman_graph.losses` (consistency / eigenvalue / physics residuals), :mod:`koopman_graph.training` -(``FitHistory``, ``LossWeights``), :mod:`koopman_graph.baselines`, +(``FitHistory``, ``LossWeights``, ``FitCallback``), :mod:`koopman_graph.baselines`, :mod:`koopman_graph.adaptation` (``RecursiveKoopmanAdapter`` and observers), :mod:`koopman_graph.env` (``GraphKoopmanEnv``), and :mod:`koopman_graph.metrics` (``evaluate_forecast``, ``EvaluationResult``, and low-level metrics). +``FitCallback`` / ``NoOpFitCallback`` are also re-exported at the package +root for ``fit(..., callbacks=)``. Physics-informed helpers such as ``graph_laplacian_features`` live in :mod:`koopman_graph.observables`. Power-user modules such as :mod:`koopman_graph.graph_utils`, :mod:`koopman_graph.nn`, -:mod:`koopman_graph.protocols`, :mod:`koopman_graph.spectrum_types`, and +:mod:`koopman_graph.protocols`, :mod:`koopman_graph.spectrum_types`, :mod:`koopman_graph.uq` (deep ensembles / optional ``predict_interval``), :mod:`koopman_graph.hierarchical` (TopK / SAG pool → coarse Koopman → unpool), -and :mod:`koopman_graph.mpc` (receding-horizon Koopman-MPC; ``[mpc]`` extra) +:mod:`koopman_graph.mpc` (receding-horizon Koopman-MPC; ``[mpc]`` extra), +:mod:`koopman_graph.tracking` (CSV / TensorBoard fit adapters), +:mod:`koopman_graph.cli` (``koopman-graph`` console script), and +:mod:`koopman_graph.tuning` (HPO metric helpers; not AutoML) are importable but intentionally omitted from ``__all__`` (encoder/decoder/delay and operator *classes*, including ``DelayEmbeddingEncoder`` and ``GraphKoopmanOperator``, remain root-stable). See the architecture docs for @@ -117,6 +122,7 @@ HypergraphKoopmanOperator, KoopmanOperator, ) +from koopman_graph.training import FitCallback, NoOpFitCallback __all__ = [ "CellComplexGNNDecoder", @@ -127,6 +133,7 @@ "DiffConvDecoder", "DiffConvEncoder", "E3EquivariantEncoder", + "FitCallback", "GATDecoder", "GATEncoder", "GNNDecoder", @@ -144,12 +151,11 @@ "KoopmanOperator", "KoopmanSpectrum", "MultiTrajectory", + "NoOpFitCallback", "RelGraphDecoder", "RelGraphEncoder", "SAGEDecoder", "SAGEEncoder", - "CellComplexGNNDecoder", - "CellComplexGNNEncoder", "SheafGNNDecoder", "SheafGNNEncoder", "SimplicialDecoder", @@ -157,4 +163,4 @@ "__version__", "compute_spectrum", ] -__version__ = "0.12.0" +__version__ = "0.13.0" diff --git a/src/koopman_graph/analysis/__init__.py b/src/koopman_graph/analysis/__init__.py index ee3bc15..a33a70b 100644 --- a/src/koopman_graph/analysis/__init__.py +++ b/src/koopman_graph/analysis/__init__.py @@ -74,6 +74,14 @@ :data:`~koopman_graph.analysis.TRANSFER_ADVANTAGE_EPSILON`). Measures zero-shot transfer against a mandatory ``pernode`` control; does **not** promise a transfer advantage (negative advantage is reportable). +``explain`` + Representation-level node / edge / feature masks + (:class:`~koopman_graph.analysis.RepresentationExplanation`, + :func:`~koopman_graph.analysis.explain_representation`). + Interpretive attributions under a chosen algorithm — **non-causal**, + **not** ResDMD, **not** + :class:`~koopman_graph.analysis.ModeEnergyAttribution`. Homogeneous + GNNExplainer MVP first. :class:`~koopman_graph.spectrum_types.KoopmanSpectrum` and the spectrum assembly helpers stay in the neutral top-level @@ -92,6 +100,10 @@ ClusteringResult, koopman_spectral_clustering, ) +from koopman_graph.analysis.explain import ( + RepresentationExplanation, + explain_representation, +) from koopman_graph.analysis.joint_stability import ( MAX_JOINT_LYAPUNOV_SIZE, MAX_JOINT_SCHUR_SIZE, @@ -157,6 +169,7 @@ "MAX_JOINT_SCHUR_SIZE", "ModeEnergyAttribution", "ResDMDReport", + "RepresentationExplanation", "ResolventNormGrid", "SINDyReport", "SpectralResidualReport", @@ -177,6 +190,7 @@ "empirical_spectral_measure", "estimate_coupling_from_snapshots", "evaluate_topology_transfer", + "explain_representation", "gershgorin_radius_bound", "identify_sparse_dynamics", "implied_timescales", diff --git a/src/koopman_graph/analysis/explain.py b/src/koopman_graph/analysis/explain.py new file mode 100644 index 0000000..77a6401 --- /dev/null +++ b/src/koopman_graph/analysis/explain.py @@ -0,0 +1,478 @@ +"""Representation-level node / edge / feature attribution masks. + +Honesty contract +---------------- +Masks produced by this module are **interpretive** attributions under a +chosen algorithm (for example PyG GNNExplainer or Captum integrated +gradients). They are **non-causal**: not causal or interventional topology +recovery, **not** certified ResDMD residuals, and **not** +:class:`~koopman_graph.analysis.ModeEnergyAttribution` (the complementary +operator-level energy diagnostic on assembled ``K_eff``). + +Homogeneous MVP first; hetero / adaptive / delay explanation is rejected +with actionable errors. ``algorithm="integrated_gradients"`` requires the +optional ``[explain]`` extra (``pip install "koopman-graph[explain]"``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +from torch import Tensor, nn +from torch_geometric.data import Data, HeteroData + +if TYPE_CHECKING: + from koopman_graph.model import GraphKoopmanModel + +ExplanationTarget = Literal["latent", "one_step_forecast", "reconstruction"] +ExplanationAlgorithm = Literal["gnn_explainer", "integrated_gradients"] + +_SUPPORTED_TARGETS: frozenset[str] = frozenset( + {"latent", "one_step_forecast", "reconstruction"} +) +_EXPLAIN_INSTALL_HINT = 'pip install "koopman-graph[explain]"' + + +@dataclass(frozen=True) +class RepresentationExplanation: + """Node / edge / feature masks for a chosen model quantity. + + Honesty contract + ---------------- + This result is an **interpretive** attribution under the recorded + ``algorithm``. It is **non-causal**: masks must not be read as causal + topology discovery or interventional importance. It is **not** a + ResDMD residual bound and is **not** + :class:`~koopman_graph.analysis.ModeEnergyAttribution`. + + Attributes + ---------- + target : {"latent", "one_step_forecast", "reconstruction"} + Model quantity the masks were requested for. + node_mask : Tensor or None + Optional per-node attribution mask. GNNExplainer object masks are + typically ``(num_nodes, 1)``; integrated-gradients attribute masks + are typically ``(num_nodes, num_features)``. + edge_mask : Tensor or None + Optional per-edge attribution mask, typically shape ``(num_edges,)`` + (often ``None`` for the Captum IG path). + feature_mask : Tensor or None + Optional per-feature attribution mask. For IG attribute masks, this + mirrors ``node_mask`` when shape is ``(num_nodes, num_features)``. + algorithm : str + Algorithm identifier (for example ``\"gnn_explainer\"`` or + ``\"integrated_gradients\"``). + index : int or None + Optional node index passed to the explainer when applicable. + """ + + target: ExplanationTarget + node_mask: Tensor | None + edge_mask: Tensor | None + feature_mask: Tensor | None + algorithm: str + index: int | None + + +class _HomogeneousExplainModule(nn.Module): + """Thin ``(x, edge_index)`` adapter for PyG ``Explainer``. + + Notes + ----- + Wraps :class:`~koopman_graph.model.GraphKoopmanModel` so PyG explainers + can call ``forward(x, edge_index)`` without a full ``Data`` batch. + """ + + def __init__( + self, + model: GraphKoopmanModel, + *, + target: ExplanationTarget, + reduce_to_node_scalar: bool = False, + ) -> None: + """Store the model and explanation target. + + Parameters + ---------- + model : GraphKoopmanModel + Homogeneous model whose encode / decode paths are attributed. + target : {"latent", "one_step_forecast", "reconstruction"} + Quantity exposed by :meth:`forward`. + reduce_to_node_scalar : bool, optional + When ``True``, sum the last feature dim so Captum IG sees a + per-node scalar (default ``False``). + """ + super().__init__() + self.model = model + self.target = target + self.reduce_to_node_scalar = reduce_to_node_scalar + + def forward(self, x: Tensor, edge_index: Tensor) -> Tensor: + """Return the selected homogeneous quantity for attribution. + + Parameters + ---------- + x : Tensor + Node features, shape ``(num_nodes, num_features)``. + edge_index : Tensor + COO edge index, shape ``(2, num_edges)``. + + Returns + ------- + Tensor + Encoded latent, reconstruction, or one-step forecast tensor + (optionally reduced to a per-node scalar). + """ + data = Data(x=x, edge_index=edge_index) + if self.target == "latent": + out = self.model.encode(data) + elif self.target == "reconstruction": + latent = self.model.encode(data) + out = self.model.decoder(latent, edge_index, None) + else: + # one_step_forecast: encode → advance → decode + decoded = self.model(data) + if not isinstance(decoded, Tensor): + msg = ( + "one_step_forecast explanation requires a tensor decode; " + f"got {type(decoded).__name__}" + ) + raise TypeError(msg) + out = decoded + if self.reduce_to_node_scalar: + # Captum IG needs a scalar output per selected node. + return out.sum(dim=-1, keepdim=True) + return out + + +def _reject_unsupported_model(model: GraphKoopmanModel) -> None: + """Raise if the model is outside the homogeneous explain MVP. + + Parameters + ---------- + model : GraphKoopmanModel + Candidate model to validate. + + Raises + ------ + ValueError + If delays, controls, hetero / hypergraph, or adaptive topology + modes are outside the MVP. + """ + if getattr(model, "n_delays", 1) != 1: + msg = ( + "explain_representation supports n_delays=1 only " + f"(got n_delays={model.n_delays})" + ) + raise ValueError(msg) + if int(getattr(model, "control_dim", 0) or 0) > 0: + raise ValueError( + "explain_representation does not support control_dim > 0 in the MVP" + ) + if model.uses_hetero_koopman or model._uses_relgraph_encode(): + raise ValueError( + "explain_representation homogeneous MVP does not support " + "RelGraph / hetero_graph models" + ) + if model.uses_hypergraph_koopman or model._uses_hypergraph_encode(): + raise ValueError( + "explain_representation homogeneous MVP does not support " + "hypergraph encoders / operators" + ) + if model.learns_pairwise_topology or model.adaptive_topology is not None: + raise ValueError( + "explain_representation homogeneous MVP does not support " + "adaptive / learned pairwise topology" + ) + + +def _reject_unsupported_data(data: Data | HeteroData) -> Data: + """Validate and return a homogeneous :class:`~torch_geometric.data.Data`. + + Parameters + ---------- + data : Data or HeteroData + Snapshot supplied by the caller. + + Returns + ------- + Data + The same homogeneous ``Data`` instance after validation. + + Raises + ------ + TypeError, ValueError + If ``data`` is hetero, not a ``Data``, or missing ``x`` / + ``edge_index``. + """ + if isinstance(data, HeteroData): + raise TypeError( + "explain_representation requires homogeneous torch_geometric.data.Data " + "(HeteroData is not supported in the MVP)" + ) + if not isinstance(data, Data): + raise TypeError( + "explain_representation requires torch_geometric.data.Data, " + f"got {type(data).__name__}" + ) + if data.x is None or data.edge_index is None: + raise ValueError("explain_representation requires Data.x and Data.edge_index") + return data + + +def _run_gnn_explainer( + model: GraphKoopmanModel, + data: Data, + *, + target: ExplanationTarget, + index: int | None, + epochs: int, + lr: float, +) -> RepresentationExplanation: + """Attribute ``target`` with PyG :class:`~torch_geometric.explain.GNNExplainer`. + + Parameters + ---------- + model : GraphKoopmanModel + Homogeneous model to explain. + data : Data + Snapshot with ``x`` and ``edge_index``. + target : {"latent", "one_step_forecast", "reconstruction"} + Quantity to attribute. + index : int or None + Optional node index forwarded to PyG ``Explainer``. + epochs : int + GNNExplainer mask-optimization epochs (must be >= 1). + lr : float + GNNExplainer learning rate (must be > 0). + + Returns + ------- + RepresentationExplanation + Node / edge masks and metadata for ``algorithm='gnn_explainer'``. + """ + from torch_geometric.explain import Explainer, GNNExplainer + from torch_geometric.explain.config import ModelConfig + + if epochs < 1: + raise ValueError(f"epochs must be >= 1, got {epochs}") + if lr <= 0: + raise ValueError(f"lr must be > 0, got {lr}") + + wrapper = _HomogeneousExplainModule(model, target=target) + wrapper.train() # GNNExplainer optimizes masks against a trainable model + explainer = Explainer( + model=wrapper, + algorithm=GNNExplainer(epochs=epochs, lr=lr), + explanation_type="model", + model_config=ModelConfig( + mode="regression", + task_level="node", + return_type="raw", + ), + node_mask_type="object", + edge_mask_type="object", + ) + explanation = explainer( + data.x, + data.edge_index, + index=index, + ) + feature_mask = getattr(explanation, "feature_mask", None) + return RepresentationExplanation( + target=target, + node_mask=explanation.node_mask, + edge_mask=explanation.edge_mask, + feature_mask=feature_mask if isinstance(feature_mask, Tensor) else None, + algorithm="gnn_explainer", + index=index, + ) + + +def _import_captum_stack() -> tuple[type, type, type]: + """Import Captum + PyG CaptumExplainer with an actionable install hint. + + Returns + ------- + CaptumExplainer, Explainer, ModelConfig + PyG / Captum classes needed for integrated gradients. + + Raises + ------ + ImportError + If Captum or PyG Captum helpers are not installed. + """ + try: + import captum # noqa: F401 + from torch_geometric.explain import CaptumExplainer, Explainer + from torch_geometric.explain.config import ModelConfig + except ImportError as exc: + msg = ( + "algorithm='integrated_gradients' requires Captum; " + f"install with: {_EXPLAIN_INSTALL_HINT}" + ) + raise ImportError(msg) from exc + return CaptumExplainer, Explainer, ModelConfig + + +def _run_integrated_gradients( + model: GraphKoopmanModel, + data: Data, + *, + target: ExplanationTarget, + index: int | None, +) -> RepresentationExplanation: + """Attribute ``target`` with PyG Captum ``IntegratedGradients``. + + Requires the optional ``[explain]`` extra (Captum). + + Parameters + ---------- + model : GraphKoopmanModel + Homogeneous model to explain. + data : Data + Snapshot with ``x`` and ``edge_index``. + target : {"latent", "one_step_forecast", "reconstruction"} + Quantity to attribute. + index : int or None + Node index for IG; defaults to ``0`` when ``None``. + + Returns + ------- + RepresentationExplanation + Attribute masks and metadata for ``algorithm='integrated_gradients'``. + """ + captum_explainer_cls, explainer_cls, model_config_cls = _import_captum_stack() + + node_index = 0 if index is None else int(index) + if node_index < 0: + raise ValueError(f"index must be >= 0, got {node_index}") + + wrapper = _HomogeneousExplainModule( + model, + target=target, + reduce_to_node_scalar=True, + ) + wrapper.eval() + explainer = explainer_cls( + model=wrapper, + algorithm=captum_explainer_cls("IntegratedGradients"), + explanation_type="model", + model_config=model_config_cls( + mode="regression", + task_level="node", + return_type="raw", + ), + node_mask_type="attributes", + edge_mask_type=None, + ) + explanation = explainer( + data.x, + data.edge_index, + index=node_index, + ) + node_mask = explanation.node_mask + feature_mask = node_mask if isinstance(node_mask, Tensor) else None + edge_mask = getattr(explanation, "edge_mask", None) + if edge_mask is not None and not isinstance(edge_mask, Tensor): + edge_mask = None + return RepresentationExplanation( + target=target, + node_mask=node_mask if isinstance(node_mask, Tensor) else None, + edge_mask=edge_mask, + feature_mask=feature_mask, + algorithm="integrated_gradients", + index=node_index, + ) + + +def explain_representation( + model: GraphKoopmanModel, + data: Data | HeteroData, + *, + target: ExplanationTarget = "one_step_forecast", + algorithm: ExplanationAlgorithm = "gnn_explainer", + index: int | None = None, + epochs: int = 50, + lr: float = 0.01, +) -> RepresentationExplanation: + """Attribute a homogeneous model quantity with node / edge masks. + + Honesty contract + ---------------- + Returned masks are **interpretive** under ``algorithm``. They are + **non-causal**, **not** ResDMD residual bounds, and **not** + :class:`~koopman_graph.analysis.ModeEnergyAttribution`. + + Parameters + ---------- + model : GraphKoopmanModel + Trained (or randomly initialized) homogeneous model. + data : Data + Homogeneous snapshot with ``x`` and ``edge_index``. + target : {"latent", "one_step_forecast", "reconstruction"}, optional + Quantity to attribute. ``latent`` uses the encoder output; + ``one_step_forecast`` uses encode → Koopman advance → decode; + ``reconstruction`` uses encode → decode. + algorithm : {"gnn_explainer", "integrated_gradients"}, optional + ``gnn_explainer`` wraps PyG :class:`~torch_geometric.explain.GNNExplainer` + (core). ``integrated_gradients`` wraps PyG ``CaptumExplainer`` and + requires ``pip install "koopman-graph[explain]"``. + index : int or None, optional + Optional node index forwarded to PyG ``Explainer``. For + ``integrated_gradients``, defaults to ``0`` when omitted. + epochs, lr : int / float, optional + GNNExplainer mask-optimization hyperparameters (ignored for IG). + + Returns + ------- + RepresentationExplanation + Masks and metadata. GNNExplainer object masks are typically + ``node_mask`` ``(num_nodes, 1)`` and ``edge_mask`` ``(num_edges,)``. + IG attribute masks are typically ``node_mask`` / + ``feature_mask`` ``(num_nodes, num_features)`` with ``edge_mask`` + often ``None``. + + Raises + ------ + TypeError, ValueError + Unsupported data / model modes, or unknown target / algorithm. + ImportError + If ``integrated_gradients`` is requested without Captum. + """ + if target not in _SUPPORTED_TARGETS: + allowed = ", ".join(sorted(_SUPPORTED_TARGETS)) + raise ValueError(f"Unsupported target {target!r}; allowed: {allowed}") + if algorithm not in {"gnn_explainer", "integrated_gradients"}: + raise ValueError( + f"Unsupported algorithm {algorithm!r}; " + "allowed: 'gnn_explainer', 'integrated_gradients'" + ) + + _reject_unsupported_model(model) + snapshot = _reject_unsupported_data(data) + was_training = model.training + try: + model.eval() + for parameter in model.parameters(): + parameter.requires_grad_(False) + if algorithm == "integrated_gradients": + return _run_integrated_gradients( + model, + snapshot, + target=target, # type: ignore[arg-type] + index=index, + ) + return _run_gnn_explainer( + model, + snapshot, + target=target, # type: ignore[arg-type] + index=index, + epochs=epochs, + lr=lr, + ) + finally: + for parameter in model.parameters(): + parameter.requires_grad_(True) + model.train(was_training) diff --git a/src/koopman_graph/cli/__init__.py b/src/koopman_graph/cli/__init__.py new file mode 100644 index 0000000..2e06ef9 --- /dev/null +++ b/src/koopman_graph/cli/__init__.py @@ -0,0 +1,22 @@ +"""Config-driven CLI for KoopmanGraph (``koopman-graph`` console script). + +Power-user package: import as ``koopman_graph.cli``. Intentionally omitted +from root ``koopman_graph.__all__``. Other library modules must not import +this package (CLI is a leaf façade). +""" + +from koopman_graph.cli.config import ( + ConfigError, + load_config, + load_train_config, + validate_train_config, +) +from koopman_graph.cli.main import main + +__all__ = [ + "ConfigError", + "load_config", + "load_train_config", + "main", + "validate_train_config", +] diff --git a/src/koopman_graph/cli/__main__.py b/src/koopman_graph/cli/__main__.py new file mode 100644 index 0000000..5085059 --- /dev/null +++ b/src/koopman_graph/cli/__main__.py @@ -0,0 +1,10 @@ +"""Module entry: ``python -m koopman_graph.cli``.""" + +from __future__ import annotations + +import sys + +from koopman_graph.cli.main import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/koopman_graph/cli/config.py b/src/koopman_graph/cli/config.py new file mode 100644 index 0000000..9d83248 --- /dev/null +++ b/src/koopman_graph/cli/config.py @@ -0,0 +1,261 @@ +"""Load and validate CLI train configs (JSON always; YAML via ``[cli]``).""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +# Top-level sections for train configs (Appendix B). +_TRAIN_TOP_LEVEL = frozenset({"model", "data", "fit", "checkpoint"}) + +# MVP allowlists — extend in code + docs when new CLI surfaces ship. +_MODEL_KEYS = frozenset( + { + "encoder", + "decoder", + "in_channels", + "hidden_channels", + "latent_dim", + "num_layers", + "time_step", + "dynamics_mode", + "koopman", + "koopman_parameterization", + "koopman_init_mode", + "koopman_init_scale", + "n_delays", + "control_dim", + "control_mode", + } +) + +_DATA_KINDS = frozenset({"synthetic_path", "cached_sequence"}) + +_DATA_KEYS = frozenset( + { + "kind", + "num_nodes", + "num_timesteps", + "seed", + "in_channels", + "path", + "feature_dim", + } +) + +_FIT_KEYS = frozenset( + { + "epochs", + "lr", + "device", + "batch_size", + "window_length", + "max_grad_norm", + "use_amp", + "early_stopping_patience", + "early_stopping_min_delta", + "restore_best_weights", + } +) + +_CHECKPOINT_KEYS = frozenset({"path", "format"}) + + +class ConfigError(ValueError): + """Invalid CLI config (unknown key, bad type, or missing required field). + + Notes + ----- + Raised by config loaders and train/predict handlers for actionable CLI + failures (exit code 1 at the boundary). + """ + + +def _require_mapping(value: object, *, path: str) -> dict[str, Any]: + """Return ``value`` as a ``dict`` or raise :class:`ConfigError`. + + Parameters + ---------- + value : object + Candidate mapping from a parsed config. + path : str + Dotted config path used in error messages (for example ``"model"``). + + Returns + ------- + dict of str to Any + The same mapping after key-type checks. + + Raises + ------ + ConfigError + If ``value`` is not a ``dict`` or contains non-string keys. + """ + if not isinstance(value, dict): + msg = f"{path} must be a mapping, got {type(value).__name__}" + raise ConfigError(msg) + # JSON object keys are str; reject non-str keys early. + for key in value: + if not isinstance(key, str): + msg = f"{path} keys must be strings, got {type(key).__name__}" + raise ConfigError(msg) + return value + + +def _reject_unknown_keys( + section: dict[str, Any], + *, + allowed: frozenset[str], + path: str, +) -> None: + """Raise if ``section`` contains keys outside ``allowed``. + + Parameters + ---------- + section : dict + Config subsection to check. + allowed : frozenset of str + Allowlisted keys for this subsection. + path : str + Dotted config path prefix for error messages. + + Raises + ------ + ConfigError + If any key is not in ``allowed``. + """ + unknown = sorted(set(section) - allowed) + if unknown: + dotted = ", ".join(f"{path}.{key}" for key in unknown) + msg = f"Unknown config key(s): {dotted}" + raise ConfigError(msg) + + +def load_config(path: str | Path) -> dict[str, Any]: + """Load a JSON or YAML config file into a plain ``dict``. + + Parameters + ---------- + path : str or Path + Config file path. Suffix ``.json`` uses the stdlib; ``.yaml`` / + ``.yml`` require PyYAML (``pip install 'koopman-graph[cli]'``). + + Returns + ------- + dict + Parsed mapping (not yet schema-validated). + + Raises + ------ + ConfigError + If the file is missing, has an unsupported suffix, fails to parse, + or does not contain a top-level mapping. + ImportError + If a YAML file is requested but PyYAML is not installed. + """ + destination = Path(path) + if not destination.is_file(): + msg = f"Config file not found: {destination}" + raise ConfigError(msg) + + suffix = destination.suffix.lower() + text = destination.read_text(encoding="utf-8") + if suffix == ".json": + try: + payload = json.loads(text) + except json.JSONDecodeError as exc: + msg = f"Invalid JSON in {destination}: {exc.msg}" + raise ConfigError(msg) from exc + elif suffix in {".yaml", ".yml"}: + try: + import yaml + except ImportError as exc: + msg = ( + "YAML configs require PyYAML. Install with: " + "pip install 'koopman-graph[cli]' (or: pip install 'pyyaml>=6')" + ) + raise ImportError(msg) from exc + try: + payload = yaml.safe_load(text) + except yaml.YAMLError as exc: + msg = f"Invalid YAML in {destination}: {exc}" + raise ConfigError(msg) from exc + else: + msg = ( + f"Unsupported config suffix {suffix!r} for {destination}; " + "use .json, .yaml, or .yml" + ) + raise ConfigError(msg) + + return _require_mapping(payload, path="config") + + +def validate_train_config(config: dict[str, Any]) -> dict[str, Any]: + """Validate a train config against the MVP allowlisted schema. + + Parameters + ---------- + config : dict + Parsed config mapping. + + Returns + ------- + dict + The same mapping after validation (mutated only by type checks; + callers may treat the return as the validated config). + + Raises + ------ + ConfigError + On unknown keys, missing required sections, or invalid ``data.kind``. + """ + _reject_unknown_keys(config, allowed=_TRAIN_TOP_LEVEL, path="config") + + if "model" not in config: + raise ConfigError("Missing required section: model") + if "data" not in config: + raise ConfigError("Missing required section: data") + + model = _require_mapping(config["model"], path="model") + _reject_unknown_keys(model, allowed=_MODEL_KEYS, path="model") + + data = _require_mapping(config["data"], path="data") + _reject_unknown_keys(data, allowed=_DATA_KEYS, path="data") + kind = data.get("kind") + if kind is None: + raise ConfigError("Missing required key: data.kind") + if kind not in _DATA_KINDS: + allowed = ", ".join(sorted(_DATA_KINDS)) + msg = f"Unsupported data.kind {kind!r}; allowed: {allowed}" + raise ConfigError(msg) + if kind == "cached_sequence" and "path" not in data: + raise ConfigError("data.kind='cached_sequence' requires data.path") + + if "fit" in config: + fit = _require_mapping(config["fit"], path="fit") + _reject_unknown_keys(fit, allowed=_FIT_KEYS, path="fit") + + if "checkpoint" in config: + checkpoint = _require_mapping(config["checkpoint"], path="checkpoint") + _reject_unknown_keys(checkpoint, allowed=_CHECKPOINT_KEYS, path="checkpoint") + if "path" not in checkpoint: + raise ConfigError("checkpoint section requires checkpoint.path") + + return config + + +def load_train_config(path: str | Path) -> dict[str, Any]: + """Load and validate a train config file. + + Parameters + ---------- + path : str or Path + JSON or YAML config path. + + Returns + ------- + dict + Validated train config. + """ + return validate_train_config(load_config(path)) diff --git a/src/koopman_graph/cli/main.py b/src/koopman_graph/cli/main.py new file mode 100644 index 0000000..8946698 --- /dev/null +++ b/src/koopman_graph/cli/main.py @@ -0,0 +1,177 @@ +"""Argparse root for the ``koopman-graph`` console script.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Sequence + +from koopman_graph import __version__ + +_PROG = "koopman-graph" + + +def _build_parser() -> argparse.ArgumentParser: + """Construct the root parser (subcommands added in later tasks). + + Returns + ------- + argparse.ArgumentParser + Root parser with ``version``, ``train``, and ``predict`` subcommands. + """ + parser = argparse.ArgumentParser( + prog=_PROG, + description=("Config-driven train / predict entry points for KoopmanGraph."), + ) + parser.add_argument( + "--version", + action="version", + version=f"{_PROG} {__version__}", + help="Print package version and exit.", + ) + subparsers = parser.add_subparsers(dest="command", metavar="COMMAND") + version_parser = subparsers.add_parser( + "version", + help="Print package version and exit.", + ) + version_parser.set_defaults(_handler=_handle_version) + + train_parser = subparsers.add_parser( + "train", + help="Train from a JSON/YAML config and save a checkpoint.", + ) + train_parser.add_argument( + "--config", + required=True, + help="Path to a train config (.json / .yaml / .yml).", + ) + train_parser.add_argument( + "--out", + default=None, + help=( + "Optional output directory; relative checkpoint.path values " + "are resolved under this directory." + ), + ) + train_parser.set_defaults(_handler=_handle_train) + + predict_parser = subparsers.add_parser( + "predict", + help=( + "Load a checkpoint, forecast from a data source, and write a " + ".pt forecast payload." + ), + description=( + "Load a safetensors_v1 / legacy checkpoint, take the first " + "snapshot from --data as the initial graph, run " + "GraphKoopmanModel.predict for --steps, and write a torch " + ".pt dict with keys 'steps', 'forecasts' (list of Data), and " + "'summary' to --out." + ), + ) + predict_parser.add_argument( + "--checkpoint", + required=True, + help="Checkpoint path (safetensors_v1 directory/.kgckpt or legacy .pt).", + ) + predict_parser.add_argument( + "--data", + required=True, + help=( + "Initial-sequence source: .pt GraphSnapshotSequence, or " + "JSON/YAML with a data section (e.g. synthetic_path config)." + ), + ) + predict_parser.add_argument( + "--steps", + type=int, + default=5, + help="Autoregressive forecast horizon (default: 5).", + ) + predict_parser.add_argument( + "--out", + required=True, + help="Output .pt path for the forecast payload.", + ) + predict_parser.set_defaults(_handler=_handle_predict) + return parser + + +def _handle_version(_args: argparse.Namespace) -> int: + """Print ``koopman-graph {version}`` to stdout. + + Parameters + ---------- + _args : argparse.Namespace + Unused parsed arguments (required by the handler protocol). + + Returns + ------- + int + Always ``0``. + """ + print(f"{_PROG} {__version__}") + return 0 + + +def _handle_train(args: argparse.Namespace) -> int: + """Delegate to :func:`koopman_graph.cli.train.handle_train`. + + Parameters + ---------- + args : argparse.Namespace + Parsed ``train`` arguments (``config``, optional ``out``). + + Returns + ------- + int + Process exit code from the train handler. + """ + from koopman_graph.cli.train import handle_train + + return handle_train(args) + + +def _handle_predict(args: argparse.Namespace) -> int: + """Delegate to :func:`koopman_graph.cli.predict.handle_predict`. + + Parameters + ---------- + args : argparse.Namespace + Parsed ``predict`` arguments (``checkpoint``, ``data``, ``steps``, + ``out``). + + Returns + ------- + int + Process exit code from the predict handler. + """ + from koopman_graph.cli.predict import handle_predict + + return handle_predict(args) + + +def main(argv: Sequence[str] | None = None) -> int: + """CLI entry point used by ``[project.scripts]`` and ``python -m``. + + Parameters + ---------- + argv : sequence of str or None, optional + Argument vector without the program name. Defaults to ``sys.argv[1:]``. + + Returns + ------- + int + Process exit code (``0`` success; argparse errors typically ``2``). + """ + parser = _build_parser() + args = parser.parse_args(list(argv) if argv is not None else None) + handler = getattr(args, "_handler", None) + if handler is None: + parser.print_help() + return 0 + return int(handler(args)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/koopman_graph/cli/predict.py b/src/koopman_graph/cli/predict.py new file mode 100644 index 0000000..b5d167d --- /dev/null +++ b/src/koopman_graph/cli/predict.py @@ -0,0 +1,161 @@ +"""``koopman-graph predict`` — load checkpoint, roll out, write forecasts.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Any + +import torch +from torch_geometric.data import Data + +from koopman_graph.cli.config import ConfigError, load_config +from koopman_graph.cli.train import ( + build_sequence_from_config, + load_cached_sequence, +) +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.serialization import load_checkpoint + + +def load_sequence_from_path(path: str | Path) -> GraphSnapshotSequence: + """Load a sequence from a ``.pt`` cache or a JSON/YAML data config. + + Parameters + ---------- + path : str or Path + Either a pickle ``GraphSnapshotSequence`` / list of ``Data``, or a + JSON/YAML file containing a ``data`` section (or a bare data mapping + with ``kind``). + + Returns + ------- + GraphSnapshotSequence + Loaded or synthesized trajectory. + """ + destination = Path(path) + if not destination.is_file(): + raise ConfigError(f"Data path not found: {destination}") + + suffix = destination.suffix.lower() + if suffix in {".json", ".yaml", ".yml"}: + payload = load_config(destination) + if "data" in payload: + data_cfg = payload["data"] + if not isinstance(data_cfg, dict): + raise ConfigError("config.data must be a mapping") + elif "kind" in payload: + data_cfg = payload + else: + raise ConfigError( + f"{destination} must contain a 'data' section or a top-level " + "'kind' field" + ) + return build_sequence_from_config(data_cfg) + + if suffix in {".pt", ".pth"}: + return load_cached_sequence( + {"kind": "cached_sequence", "path": str(destination)} + ) + + msg = ( + f"Unsupported data path suffix {suffix!r} for {destination}; " + "use .json/.yaml/.yml or .pt/.pth" + ) + raise ConfigError(msg) + + +def run_predict( + *, + checkpoint: str | Path, + data: str | Path, + steps: int, + out: str | Path, +) -> tuple[Path, dict[str, Any] | None]: + """Load a checkpoint, forecast ``steps``, and write a ``.pt`` payload. + + Parameters + ---------- + checkpoint : str or Path + Safe or legacy checkpoint path. + data : str or Path + Sequence source (see :func:`load_sequence_from_path`). + steps : int + Autoregressive forecast horizon (must be >= 1). + out : str or Path + Destination ``.pt`` file for the forecast payload. + + Returns + ------- + path : Path + Path written. + summary : dict or None + Optional shape summary for stdout. + """ + if steps < 1: + raise ConfigError(f"--steps must be >= 1, got {steps}") + + model = load_checkpoint(checkpoint) + sequence = load_sequence_from_path(data) + if len(sequence) < 1: + raise ConfigError("Data sequence is empty") + + initial = sequence[0] + forecasts = model.predict(initial, steps=steps) + if not forecasts: + raise ConfigError("predict returned an empty forecast list") + + out_path = Path(out) + out_path.parent.mkdir(parents=True, exist_ok=True) + first = forecasts[0] + summary: dict[str, Any] | None = None + if isinstance(first, Data) and first.x is not None: + summary = { + "num_nodes": int(first.x.shape[0]), + "feature_dim": int(first.x.shape[1]), + "num_forecasts": len(forecasts), + } + payload: dict[str, Any] = { + "steps": steps, + "forecasts": forecasts, + "summary": summary, + } + torch.save(payload, out_path) + return out_path, summary + + +def handle_predict(args: argparse.Namespace) -> int: + """Argparse handler for the ``predict`` subcommand. + + Parameters + ---------- + args : argparse.Namespace + Parsed CLI namespace with ``checkpoint``, ``data``, ``steps``, and + ``out``. + + Returns + ------- + int + ``0`` on success; ``1`` on config / I/O / type errors. + """ + try: + out_path, summary = run_predict( + checkpoint=args.checkpoint, + data=args.data, + steps=args.steps, + out=args.out, + ) + if summary is not None: + print( + "wrote forecast: " + f"{out_path} (steps={args.steps}, " + f"nodes={summary.get('num_nodes')}, " + f"feature_dim={summary.get('feature_dim')})" + ) + else: + print(f"wrote forecast: {out_path} (steps={args.steps})") + except (ConfigError, ValueError, FileNotFoundError, OSError, TypeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + return 0 diff --git a/src/koopman_graph/cli/train.py b/src/koopman_graph/cli/train.py new file mode 100644 index 0000000..51789d8 --- /dev/null +++ b/src/koopman_graph/cli/train.py @@ -0,0 +1,422 @@ +"""``koopman-graph train`` — config-driven fit + safe checkpoint save.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Any + +import torch +from torch_geometric.data import Data + +from koopman_graph import ( + GATDecoder, + GATEncoder, + GNNDecoder, + GNNEncoder, + GraphKoopmanModel, + SAGEDecoder, + SAGEEncoder, +) +from koopman_graph.cli.config import ConfigError, load_train_config +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.datasets.topology import path_edge_index + +_ENCODER_KINDS = frozenset({"gcn", "gat", "sage"}) +_MODEL_PASSTHROUGH = frozenset( + { + "time_step", + "dynamics_mode", + "koopman", + "koopman_parameterization", + "koopman_init_mode", + "koopman_init_scale", + "n_delays", + "control_dim", + "control_mode", + } +) +_FIT_PASSTHROUGH = frozenset( + { + "epochs", + "lr", + "device", + "batch_size", + "window_length", + "max_grad_norm", + "use_amp", + "early_stopping_patience", + "early_stopping_min_delta", + "restore_best_weights", + } +) + + +def _require_int(value: object, *, path: str) -> int: + """Coerce a config value to ``int`` or raise :class:`ConfigError`. + + Parameters + ---------- + value : object + Candidate integer from the config. + path : str + Dotted config path for error messages. + + Returns + ------- + int + Validated integer (bool rejected). + + Raises + ------ + ConfigError + If ``value`` is not an ``int`` (or is a ``bool``). + """ + if isinstance(value, bool) or not isinstance(value, int): + msg = f"{path} must be an int, got {type(value).__name__}" + raise ConfigError(msg) + return value + + +def _require_float(value: object, *, path: str) -> float: + """Coerce a config value to ``float`` or raise :class:`ConfigError`. + + Parameters + ---------- + value : object + Candidate number from the config. + path : str + Dotted config path for error messages. + + Returns + ------- + float + Validated float (bool rejected). + + Raises + ------ + ConfigError + If ``value`` is not an ``int`` or ``float`` (or is a ``bool``). + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + msg = f"{path} must be a number, got {type(value).__name__}" + raise ConfigError(msg) + return float(value) + + +def _build_encoder_decoder( + model_cfg: dict[str, Any], +) -> tuple[Any, Any, int]: + """Build matched encoder/decoder peers and return ``latent_dim``. + + Parameters + ---------- + model_cfg : dict + Validated ``model`` section from a train config. + + Returns + ------- + encoder, decoder, latent_dim + Constructed peers and latent width. + + Raises + ------ + ConfigError + If encoder/decoder kinds or channel sizes are invalid. + """ + encoder_kind = model_cfg.get("encoder") + if not isinstance(encoder_kind, str): + raise ConfigError("model.encoder is required (string kind)") + encoder_kind = encoder_kind.lower() + if encoder_kind not in _ENCODER_KINDS: + allowed = ", ".join(sorted(_ENCODER_KINDS)) + msg = f"Unsupported model.encoder {encoder_kind!r}; CLI MVP allows: {allowed}" + raise ConfigError(msg) + + decoder_kind = model_cfg.get("decoder", encoder_kind) + if not isinstance(decoder_kind, str): + raise ConfigError("model.decoder must be a string when set") + decoder_kind = decoder_kind.lower() + if decoder_kind != encoder_kind: + msg = ( + f"model.decoder {decoder_kind!r} must match model.encoder " + f"{encoder_kind!r} in the CLI MVP" + ) + raise ConfigError(msg) + + for key in ("in_channels", "hidden_channels", "latent_dim"): + if key not in model_cfg: + raise ConfigError(f"Missing required key: model.{key}") + + in_channels = _require_int(model_cfg["in_channels"], path="model.in_channels") + hidden_channels = _require_int( + model_cfg["hidden_channels"], path="model.hidden_channels" + ) + latent_dim = _require_int(model_cfg["latent_dim"], path="model.latent_dim") + num_layers = _require_int(model_cfg.get("num_layers", 2), path="model.num_layers") + if in_channels < 1 or hidden_channels < 1 or latent_dim < 1 or num_layers < 1: + raise ConfigError("model channel / layer sizes must be positive") + + peer_kwargs = { + "in_channels": in_channels, + "hidden_channels": hidden_channels, + "latent_dim": latent_dim, + "num_layers": num_layers, + } + if encoder_kind == "gcn": + encoder = GNNEncoder(**peer_kwargs) + decoder = GNNDecoder( + latent_dim=latent_dim, + hidden_channels=hidden_channels, + out_channels=in_channels, + num_layers=num_layers, + ) + elif encoder_kind == "gat": + encoder = GATEncoder(**peer_kwargs) + decoder = GATDecoder( + latent_dim=latent_dim, + hidden_channels=hidden_channels, + out_channels=in_channels, + num_layers=num_layers, + ) + else: + encoder = SAGEEncoder(**peer_kwargs) + decoder = SAGEDecoder( + latent_dim=latent_dim, + hidden_channels=hidden_channels, + out_channels=in_channels, + num_layers=num_layers, + ) + return encoder, decoder, latent_dim + + +def build_model_from_config(model_cfg: dict[str, Any]) -> GraphKoopmanModel: + """Construct :class:`~koopman_graph.GraphKoopmanModel` from CLI ``model``. + + Parameters + ---------- + model_cfg : dict + Validated ``model`` section. + + Returns + ------- + GraphKoopmanModel + Homogeneous model ready for ``fit``. + """ + encoder, decoder, latent_dim = _build_encoder_decoder(model_cfg) + kwargs: dict[str, Any] = {} + for key in _MODEL_PASSTHROUGH: + if key in model_cfg: + kwargs[key] = model_cfg[key] + if "time_step" not in kwargs: + kwargs["time_step"] = 1.0 + return GraphKoopmanModel( + encoder=encoder, + decoder=decoder, + latent_dim=latent_dim, + **kwargs, + ) + + +def _build_synthetic_path_sequence(data_cfg: dict[str, Any]) -> GraphSnapshotSequence: + """Build a seeded decay trajectory on a path graph. + + Parameters + ---------- + data_cfg : dict + ``data`` section with ``kind='synthetic_path'`` fields. + + Returns + ------- + GraphSnapshotSequence + Synthetic trajectory for smoke training. + """ + num_nodes = _require_int(data_cfg.get("num_nodes", 8), path="data.num_nodes") + num_timesteps = _require_int( + data_cfg.get("num_timesteps", 40), path="data.num_timesteps" + ) + seed = _require_int(data_cfg.get("seed", 0), path="data.seed") + feature_dim = _require_int( + data_cfg.get("feature_dim", data_cfg.get("in_channels", 3)), + path="data.feature_dim", + ) + if num_nodes < 2 or num_timesteps < 2 or feature_dim < 1: + raise ConfigError( + "synthetic_path requires num_nodes>=2, num_timesteps>=2, feature_dim>=1" + ) + + generator = torch.Generator().manual_seed(seed) + edge_index = path_edge_index(num_nodes) + x0 = torch.randn(num_nodes, feature_dim, generator=generator) + snapshots = [ + Data(x=x0 * (0.9**t), edge_index=edge_index) for t in range(num_timesteps) + ] + return GraphSnapshotSequence(snapshots) + + +def load_cached_sequence(data_cfg: dict[str, Any]) -> GraphSnapshotSequence: + """Load a trusted ``GraphSnapshotSequence`` (or list of ``Data``) from disk. + + Parameters + ---------- + data_cfg : dict + ``data`` section with ``kind='cached_sequence'`` and ``path``. + + Returns + ------- + GraphSnapshotSequence + Loaded trajectory. + + Notes + ----- + Uses ``torch.load`` pickle — trusted-source only (see ``SECURITY.md``). + """ + path = Path(str(data_cfg["path"])) + if not path.is_file(): + raise ConfigError(f"Cached sequence not found: {path}") + # Trust boundary: pickle load — same as teaching caches (see SECURITY.md). + payload = torch.load(path, map_location="cpu", weights_only=False) + if isinstance(payload, GraphSnapshotSequence): + return payload + if isinstance(payload, list) and payload and isinstance(payload[0], Data): + return GraphSnapshotSequence(payload) + msg = ( + f"Cached sequence at {path} must be a GraphSnapshotSequence " + "or a non-empty list of Data snapshots" + ) + raise ConfigError(msg) + + +def build_sequence_from_config(data_cfg: dict[str, Any]) -> GraphSnapshotSequence: + """Build the training sequence for an allowlisted ``data.kind``. + + Parameters + ---------- + data_cfg : dict + Validated ``data`` section. + + Returns + ------- + GraphSnapshotSequence + Synthetic or cached trajectory. + """ + kind = data_cfg["kind"] + if kind == "synthetic_path": + return _build_synthetic_path_sequence(data_cfg) + if kind == "cached_sequence": + return load_cached_sequence(data_cfg) + msg = f"Unsupported data.kind {kind!r}" + raise ConfigError(msg) + + +def _resolve_checkpoint_path( + config: dict[str, Any], + *, + out_dir: Path | None, +) -> Path: + """Resolve the checkpoint destination path (directory or ``.kgckpt``). + + Parameters + ---------- + config : dict + Validated train config (may omit ``checkpoint``). + out_dir : Path or None + Optional base directory for relative ``checkpoint.path`` values. + + Returns + ------- + Path + Absolute or resolved destination path (parent dirs created). + """ + checkpoint = config.get("checkpoint") or {} + raw = checkpoint.get("path", "artifacts/model.kgckpt") + path = Path(str(raw)) + if out_dir is not None and not path.is_absolute(): + path = out_dir / path + path.parent.mkdir(parents=True, exist_ok=True) + return path + + +def _fit_kwargs(fit_cfg: dict[str, Any]) -> dict[str, Any]: + """Filter allowlisted fit kwargs for :meth:`GraphKoopmanModel.fit`. + + Parameters + ---------- + fit_cfg : dict + ``fit`` section from the train config. + + Returns + ------- + dict + Keyword arguments forwarded to ``model.fit``. + """ + kwargs: dict[str, Any] = {} + for key in _FIT_PASSTHROUGH: + if key in fit_cfg: + kwargs[key] = fit_cfg[key] + return kwargs + + +def run_train( + config: dict[str, Any], + *, + out_dir: str | Path | None = None, +) -> Path: + """Train from a validated config and save a safetensors checkpoint. + + Parameters + ---------- + config : dict + Validated train config (see :func:`load_train_config`). + out_dir : str, Path, or None, optional + When set, relative ``checkpoint.path`` values are resolved under this + directory. + + Returns + ------- + Path + Checkpoint path written by :meth:`GraphKoopmanModel.save`. + """ + model = build_model_from_config(config["model"]) + sequence = build_sequence_from_config(config["data"]) + fit_cfg = config.get("fit") or {} + model.fit(sequence, **_fit_kwargs(fit_cfg)) + + checkpoint_cfg = config.get("checkpoint") or {} + fmt = checkpoint_cfg.get("format", "safetensors_v1") + if fmt not in {"safetensors_v1", "legacy_pt"}: + raise ConfigError(f"Unsupported checkpoint.format {fmt!r}") + + destination = _resolve_checkpoint_path( + config, + out_dir=None if out_dir is None else Path(out_dir), + ) + model.save(destination, format=fmt) # type: ignore[arg-type] + return destination + + +def handle_train(args: argparse.Namespace) -> int: + """Argparse handler for the ``train`` subcommand. + + Parameters + ---------- + args : argparse.Namespace + Parsed CLI namespace with ``config`` and optional ``out``. + + Returns + ------- + int + ``0`` on success; ``1`` on config / I/O / type errors. + """ + try: + config = load_train_config(args.config) + out_dir = Path(args.out) if args.out is not None else None + if out_dir is not None: + out_dir.mkdir(parents=True, exist_ok=True) + destination = run_train(config, out_dir=out_dir) + except (ConfigError, ValueError, FileNotFoundError, OSError, TypeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + print(f"wrote checkpoint: {destination}") + return 0 diff --git a/src/koopman_graph/distributed/_fit_epochs.py b/src/koopman_graph/distributed/_fit_epochs.py index 14715fe..b8afd8d 100644 --- a/src/koopman_graph/distributed/_fit_epochs.py +++ b/src/koopman_graph/distributed/_fit_epochs.py @@ -355,7 +355,8 @@ def fit_epochs_distributed( core.load_state_dict(best_state_dict) broadcast_state_fn(core) if checkpoint_path is not None and is_main_fn(): - save_checkpoint(cast(Any, core), checkpoint_path) + # File-path best-epoch checkpoints stay legacy .pt (not directories). + save_checkpoint(cast(Any, core), checkpoint_path, format="legacy_pt") if not restore_best_weights and last_state_dict is not None: core.load_state_dict(last_state_dict) broadcast_state_fn(core) diff --git a/src/koopman_graph/distributed/lightning_module.py b/src/koopman_graph/distributed/lightning_module.py index f931f00..da998ac 100644 --- a/src/koopman_graph/distributed/lightning_module.py +++ b/src/koopman_graph/distributed/lightning_module.py @@ -315,7 +315,8 @@ def export_format1_checkpoint(self, path: str | Path) -> None: """ from koopman_graph.serialization import save_checkpoint - save_checkpoint(self.model, path) + # Keep format-1 .pt export as a single pickle file for Trainer hooks. + save_checkpoint(self.model, path, format="legacy_pt") return KoopmanLightningModule diff --git a/src/koopman_graph/model/estimator.py b/src/koopman_graph/model/estimator.py index b344977..67cef80 100644 --- a/src/koopman_graph/model/estimator.py +++ b/src/koopman_graph/model/estimator.py @@ -71,6 +71,7 @@ from koopman_graph.training import ( EarlyStoppingMonitor, ExtraLosses, + FitCallback, FitHistory, LossWeights, LossWeightSchedule, @@ -1391,19 +1392,30 @@ def disable_online_adaptation(self, *, unfreeze: bool = True) -> None: unfreeze=unfreeze, ) - def save(self, path: str | Path) -> None: + def save( + self, + path: str | Path, + *, + format: Literal["safetensors_v1", "legacy_pt"] = "safetensors_v1", + ) -> None: """Persist model weights and architecture configuration to disk. Parameters ---------- path : str or Path - Destination checkpoint file (``.pt``). Parent directories are - created when missing. + Destination **directory** for the default ``safetensors_v1`` + container, a ``.kgckpt`` / ``.zip`` path for the zip bundle of + the same layout, or a ``.pt`` file path when + ``format="legacy_pt"``. Parent directories are created when + missing. + format : {"safetensors_v1", "legacy_pt"}, optional + On-disk container. Default ``safetensors_v1``; pass + ``legacy_pt`` for the pickle escape hatch. """ import importlib serialization = importlib.import_module("koopman_graph.serialization") - serialization.save_checkpoint(self, path) + serialization.save_checkpoint(self, path, format=format) @classmethod def load( @@ -1413,17 +1425,22 @@ def load( map_location: str | torch.device | None = None, physics_lifting_fn: PhysicsLiftingFn | None = None, ) -> GraphKoopmanModel: - """Load a trained model from a checkpoint file. + """Load a trained model from a checkpoint file or directory. Reconstructs encoder, decoder, and Koopman operator architecture from - the saved configuration and restores learned weights. + the saved configuration and restores learned weights. Auto-detects + ``safetensors_v1`` directories and ``.kgckpt`` / ``.zip`` bundles, then + legacy ``.pt`` pickle files. Parameters ---------- path : str or Path - Checkpoint file produced by :meth:`save`. + Checkpoint file or directory produced by :meth:`save` / + :func:`~koopman_graph.serialization.save_checkpoint`. map_location : str, torch.device, or None, optional - Device mapping forwarded to :func:`torch.load`. + Device mapping forwarded to the underlying loader + (:func:`torch.load` for ``.pt``; device string for + ``safetensors_v1``). physics_lifting_fn : callable or None, optional Custom physics lifting function required when the checkpoint stores hybrid observables without a registered preset. @@ -2090,6 +2107,7 @@ def fit( validation_sequence: ValidationInput = None, restore_best_weights: bool = False, checkpoint_path: str | Path | None = None, + callbacks: Sequence[FitCallback] | None = None, strategy: Literal["ddp"] | None = None, find_unused_parameters: bool | None = None, **optimizer_kwargs: Any, @@ -2223,6 +2241,13 @@ def fit( checkpoint_path : str, Path, or None, optional When set, write a checkpoint at the lowest-loss epoch using :meth:`save`. Default is ``None``. + callbacks : sequence of FitCallback or None, optional + Observe-only fit hooks forwarded to + :func:`~koopman_graph.training.run_fit_loop` when + ``strategy=None``. Default ``None`` skips hooks. Not supported + with ``strategy="ddp"`` yet (raises ``ValueError``); use + single-process fit or Lightning loggers for distributed + tracking. strategy : {"ddp"} or None, optional Training orchestration backend. ``None`` (default) uses :func:`~koopman_graph.training.run_fit_loop` (single-process, @@ -2254,7 +2279,8 @@ def fit( ``early_stopping_monitor="val"`` without ``validation_sequence``, validation list length mismatches training trajectories, fewer than two snapshots are provided for training or validation, or - ``strategy`` is not ``None`` / ``"ddp"``. + ``strategy`` is not ``None`` / ``"ddp"``, or ``callbacks`` is set + with ``strategy="ddp"``. """ uses_simplicial_modules(self.encoder, self.decoder) uses_sheaf_modules(self.encoder, self.decoder) @@ -2302,8 +2328,20 @@ def fit( **optimizer_kwargs, } if strategy is None: - return run_fit_loop(self, prepared.train_sequences, **loop_kwargs) + return run_fit_loop( + self, + prepared.train_sequences, + callbacks=callbacks, + **loop_kwargs, + ) if strategy == "ddp": + if callbacks is not None: + msg = ( + "fit(..., callbacks=...) is not supported with " + 'strategy="ddp" yet; use strategy=None (single-process) ' + "or Lightning loggers for distributed tracking" + ) + raise ValueError(msg) from koopman_graph.distributed import run_ddp_fit_loop return run_ddp_fit_loop( diff --git a/src/koopman_graph/serialization.py b/src/koopman_graph/serialization.py index 7f491b8..0e7aa22 100644 --- a/src/koopman_graph/serialization.py +++ b/src/koopman_graph/serialization.py @@ -62,6 +62,36 @@ multi-version checkpoint tracking begins at 1.0. Loaders accept only ``{1}``; ``format_version`` 2 and other lineages remain unsupported. +On-disk containers + ``safetensors_v1`` (default; prefer for sharing) + Same three members — ``meta.json`` (container marker + + ``format_version`` + ``package_version``), ``config.json`` (JSON + architecture config from :func:`build_model_config`), and + ``model.safetensors`` (weight tensors) — as either a **directory** + or a **zip bundle** (``.kgckpt`` / ``.zip``). Logical architecture + schema remains format-1; the container only changes how bytes are + stored. Default :func:`save_checkpoint` / + :meth:`GraphKoopmanModel.save` writes a directory unless ``path`` + ends in ``.kgckpt`` or ``.zip``, in which case it writes a zip with + those members only (no pickle). Weights load via safetensors rather + than pickle; see repository ``SECURITY.md`` for trust boundaries. + ``legacy_pt`` + Single ``torch.save`` pickle file (``.pt`` / ``.pth``) holding + ``format_version``, ``package_version``, ``config``, and + ``state_dict``. Loading executes pickle and remains a trusted-source + trust boundary (see ``SECURITY.md``). Pass ``format="legacy_pt"`` + for the pickle escape hatch. Training ``checkpoint_path`` writers + keep ``legacy_pt`` explicitly so file-path best-epoch checkpoints + remain single ``.pt`` files. + +Load detection + 1. Directory with ``meta.json`` → ``safetensors_v1`` directory. + 2. File that is a zip containing root members ``meta.json``, + ``config.json``, and ``model.safetensors`` → ``safetensors_v1`` zip + (``.kgckpt`` / ``.zip``; also matched when those markers are present + so pickle is not used when safe markers exist). + 3. Other files → legacy ``torch.load`` pickle path. + Custom injected operators (anything other than the built-in serializable operator classes registered in this module) are **not** round-trippable: :func:`build_model_config` / :meth:`GraphKoopmanModel.save` raise rather than @@ -70,13 +100,19 @@ from __future__ import annotations -from collections.abc import Sequence +import json +import zipfile +from collections.abc import Mapping, Sequence from copy import deepcopy from importlib.metadata import PackageNotFoundError, version from pathlib import Path -from typing import Any +from tempfile import TemporaryDirectory +from typing import Any, Literal import torch +from safetensors.torch import load_file as safetensors_load_file +from safetensors.torch import save as safetensors_save_bytes +from safetensors.torch import save_file as safetensors_save_file from torch import nn from koopman_graph.data.hetero_layout import validate_latent_dims @@ -127,6 +163,20 @@ FORMAT_VERSION = 1 SUPPORTED_FORMAT_VERSIONS = frozenset({1}) +CheckpointFormat = Literal["safetensors_v1", "legacy_pt"] +SAFE_CONTAINER = "safetensors_v1" +SAFE_META_FILENAME = "meta.json" +SAFE_CONFIG_FILENAME = "config.json" +SAFE_WEIGHTS_FILENAME = "model.safetensors" +SAFE_BUNDLE_SUFFIXES = frozenset({".kgckpt", ".zip"}) +SAFE_ZIP_MEMBER_NAMES = frozenset( + { + SAFE_META_FILENAME, + SAFE_CONFIG_FILENAME, + SAFE_WEIGHTS_FILENAME, + } +) + # Keys always written by :func:`build_model_config` for the current format-1 # baseline. Sparse historical payloads that omit these are rejected on load. _FORMAT_1_REQUIRED_KEYS = frozenset( @@ -1496,7 +1546,172 @@ def build_checkpoint(model: ModeShapeModel) -> dict[str, Any]: } -def save_checkpoint(model: ModeShapeModel, path: str | Path) -> None: +def _json_ready(value: Any) -> Any: + """Convert nested configs to JSON-serializable values. + + Parameters + ---------- + value : Any + Config fragment that may contain tuples. + + Returns + ------- + Any + Structure with tuples replaced by lists. + """ + if isinstance(value, dict): + return {str(key): _json_ready(item) for key, item in value.items()} + if isinstance(value, tuple): + return [_json_ready(item) for item in value] + if isinstance(value, list): + return [_json_ready(item) for item in value] + return value + + +def _state_dict_for_safetensors( + state_dict: Mapping[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Detach, clone, and contiguous-ize tensors for safetensors I/O. + + Parameters + ---------- + state_dict : mapping of str to Tensor + Model ``state_dict``. + + Returns + ------- + dict of str to Tensor + CPU contiguous copies suitable for :func:`safetensors.torch.save_file`. + """ + return { + key: tensor.detach().cpu().contiguous().clone() + for key, tensor in state_dict.items() + } + + +def _safetensors_v1_member_payload( + model: ModeShapeModel, +) -> tuple[str, str, dict[str, torch.Tensor]]: + """Build ``meta.json`` / ``config.json`` text and CPU weight tensors. + + Parameters + ---------- + model : ModeShapeModel + Model whose config and ``state_dict`` are serialized. + + Returns + ------- + meta_text, config_text, weights + JSON text for ``meta.json`` and ``config.json``, plus a CPU + contiguous weight mapping for safetensors. + """ + config = _json_ready(build_model_config(model)) + meta = { + "container": SAFE_CONTAINER, + "format_version": FORMAT_VERSION, + "package_version": _package_version(), + } + meta_text = json.dumps(meta, indent=2, sort_keys=True) + "\n" + config_text = json.dumps(config, indent=2, sort_keys=True) + "\n" + weights = _state_dict_for_safetensors(model.state_dict()) + return meta_text, config_text, weights + + +def _save_safetensors_v1_directory(model: ModeShapeModel, path: Path) -> None: + """Write a ``safetensors_v1`` directory checkpoint. + + Parameters + ---------- + model : ModeShapeModel + Model to serialize. + path : Path + Destination directory (created if missing). Must not be an existing + regular file. + + Raises + ------ + ValueError + If ``path`` exists as a file. + """ + if path.exists() and path.is_file(): + msg = ( + f"safetensors_v1 checkpoint path must be a directory; " + f"got existing file {path}" + ) + raise ValueError(msg) + + path.mkdir(parents=True, exist_ok=True) + meta_text, config_text, weights = _safetensors_v1_member_payload(model) + (path / SAFE_META_FILENAME).write_text(meta_text, encoding="utf-8") + (path / SAFE_CONFIG_FILENAME).write_text(config_text, encoding="utf-8") + safetensors_save_file(weights, path / SAFE_WEIGHTS_FILENAME) + + +def _save_safetensors_v1_zip(model: ModeShapeModel, path: Path) -> None: + """Write a ``safetensors_v1`` zip / ``.kgckpt`` checkpoint. + + Parameters + ---------- + model : ModeShapeModel + Model to serialize. + path : Path + Destination ``.kgckpt`` or ``.zip`` file path. Must not be an + existing directory. + + Raises + ------ + ValueError + If ``path`` exists as a directory. + """ + if path.exists() and path.is_dir(): + msg = ( + f"safetensors_v1 zip checkpoint path must be a file; " + f"got existing directory {path}" + ) + raise ValueError(msg) + + path.parent.mkdir(parents=True, exist_ok=True) + meta_text, config_text, weights = _safetensors_v1_member_payload(model) + weight_bytes = safetensors_save_bytes(weights) + with zipfile.ZipFile(path, mode="w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(SAFE_META_FILENAME, meta_text) + archive.writestr(SAFE_CONFIG_FILENAME, config_text) + archive.writestr(SAFE_WEIGHTS_FILENAME, weight_bytes) + + +def _is_safetensors_v1_zip(path: Path) -> bool: + """Return True when ``path`` is a zip with root ``safetensors_v1`` members. + + PyTorch ``.pt`` archives are often zip files themselves; this probe + requires the safe-container member names so legacy pickle archives are + not mistaken for ``safetensors_v1``. + + Parameters + ---------- + path : Path + Candidate file path. + + Returns + ------- + bool + ``True`` when the zip root contains the safe-container member set. + """ + if not path.is_file() or not zipfile.is_zipfile(path): + return False + try: + with zipfile.ZipFile(path, mode="r") as archive: + names = set(archive.namelist()) + except zipfile.BadZipFile: + return False + return SAFE_ZIP_MEMBER_NAMES.issubset(names) + + +def save_checkpoint( + model: ModeShapeModel, + path: str | Path, + *, + format: CheckpointFormat = "safetensors_v1", +) -> None: """Persist a trained model checkpoint to disk. Parameters @@ -1504,11 +1719,33 @@ def save_checkpoint(model: ModeShapeModel, path: str | Path) -> None: model : ModeShapeModel Model to serialize. path : str or Path - Destination ``.pt`` file path. + When ``format="safetensors_v1"`` (default): destination **directory**, + or a ``.kgckpt`` / ``.zip`` file for the zip bundle of the same three + members. When ``format="legacy_pt"``: destination ``.pt`` file path. + format : {"safetensors_v1", "legacy_pt"}, optional + On-disk container. Default ``safetensors_v1`` writes JSON config + + safetensors weights; pass ``legacy_pt`` for the pickle ``.pt`` escape + hatch (still used by ``fit(..., checkpoint_path=...)`` writers). + + Raises + ------ + ValueError + If ``format`` is unknown, or ``safetensors_v1`` path type conflicts + with an existing file/directory. """ destination = Path(path) - destination.parent.mkdir(parents=True, exist_ok=True) - torch.save(build_checkpoint(model), destination) + if format == "legacy_pt": + destination.parent.mkdir(parents=True, exist_ok=True) + torch.save(build_checkpoint(model), destination) + return + if format == "safetensors_v1": + if destination.suffix.lower() in SAFE_BUNDLE_SUFFIXES: + _save_safetensors_v1_zip(model, destination) + else: + _save_safetensors_v1_directory(model, destination) + return + msg = f"Unsupported checkpoint format {format!r}" + raise ValueError(msg) def _allocate_adaptive_topology_from_state( @@ -1545,42 +1782,193 @@ def _allocate_adaptive_topology_from_state( adaptive.set_num_nodes(int(source.shape[0]), device=source.device) -def load_checkpoint( - path: str | Path, +def _resolve_safetensors_device( + map_location: str | torch.device | None, +) -> str: + """Normalize ``map_location`` for :func:`safetensors.torch.load_file`. + + Parameters + ---------- + map_location : str, torch.device, or None + Device mapping from :func:`load_checkpoint`. Callables are not + supported on the safetensors path. + + Returns + ------- + str + Device string (``"cpu"`` when ``map_location`` is ``None``). + """ + if map_location is None: + return "cpu" + return str(map_location) + + +def _load_safetensors_v1_directory( + path: Path, *, - map_location: str | torch.device | None = None, - physics_lifting_fn: PhysicsLiftingFn | None = None, -) -> ModeShapeModel: - """Load a trained model from a checkpoint file. + map_location: str | torch.device | None, +) -> tuple[dict[str, Any], dict[str, torch.Tensor], int]: + """Load config and weights from a ``safetensors_v1`` directory. Parameters ---------- - path : str or Path - Checkpoint ``.pt`` file produced by :func:`save_checkpoint`. - map_location : str, torch.device, or None, optional - Device mapping forwarded to :func:`torch.load`. - physics_lifting_fn : callable or None, optional - Custom physics lifting function for hybrid checkpoints without a stored - preset. + path : Path + Checkpoint directory containing ``meta.json``, ``config.json``, and + ``model.safetensors``. + map_location : str, torch.device, or None + Device for weight tensors. Returns ------- - GraphKoopmanModel - Reconstructed model with restored weights in evaluation mode. + config : dict + Architecture configuration. + state_dict : dict of str to Tensor + Model weights. + format_version : int + Logical checkpoint schema version from ``meta.json``. Raises ------ ValueError - If the checkpoint format version is unsupported or the payload is invalid. - FileNotFoundError - If ``path`` does not exist. + If metadata, config, or weights are missing or invalid. """ - destination = Path(path) - if not destination.is_file(): - msg = f"Checkpoint file not found: {destination}" - raise FileNotFoundError(msg) + meta_path = path / SAFE_META_FILENAME + config_path = path / SAFE_CONFIG_FILENAME + weights_path = path / SAFE_WEIGHTS_FILENAME + if not meta_path.is_file(): + msg = f"safetensors_v1 checkpoint missing {SAFE_META_FILENAME}: {path}" + raise ValueError(msg) + if not config_path.is_file(): + msg = f"safetensors_v1 checkpoint missing {SAFE_CONFIG_FILENAME}: {path}" + raise ValueError(msg) + if not weights_path.is_file(): + msg = f"safetensors_v1 checkpoint missing {SAFE_WEIGHTS_FILENAME}: {path}" + raise ValueError(msg) + + try: + meta = json.loads(meta_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + msg = f"Invalid {SAFE_META_FILENAME} in safetensors_v1 checkpoint: {path}" + raise ValueError(msg) from exc + if not isinstance(meta, dict): + msg = f"{SAFE_META_FILENAME} must be a JSON object: {path}" + raise ValueError(msg) + container = meta.get("container") + if container != SAFE_CONTAINER: + msg = ( + f"Unsupported safetensors container {container!r} in {meta_path}; " + f"expected {SAFE_CONTAINER!r}" + ) + raise ValueError(msg) + format_version = meta.get("format_version", FORMAT_VERSION) + if format_version not in SUPPORTED_FORMAT_VERSIONS: + supported = ", ".join( + str(version) for version in sorted(SUPPORTED_FORMAT_VERSIONS) + ) + msg = ( + f"Unsupported checkpoint format_version {format_version!r}; " + f"supported versions: {supported}" + ) + raise ValueError(msg) + + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + msg = f"Invalid {SAFE_CONFIG_FILENAME} in safetensors_v1 checkpoint: {path}" + raise ValueError(msg) from exc + if not isinstance(config, dict): + msg = f"{SAFE_CONFIG_FILENAME} must be a JSON object: {path}" + raise ValueError(msg) + + state_dict = safetensors_load_file( + weights_path, + device=_resolve_safetensors_device(map_location), + ) + return config, state_dict, int(format_version) + + +def _load_safetensors_v1_zip( + path: Path, + *, + map_location: str | torch.device | None, +) -> tuple[dict[str, Any], dict[str, torch.Tensor], int]: + """Load config and weights from a ``safetensors_v1`` zip / ``.kgckpt``. - payload = torch.load(destination, map_location=map_location, weights_only=False) + Extracts the three root members into a temporary directory and reuses the + directory loader validation path. Only exact member names are extracted + (no nested paths). + + Parameters + ---------- + path : Path + Zip archive containing ``meta.json``, ``config.json``, and + ``model.safetensors`` at the archive root. + map_location : str, torch.device, or None + Device for weight tensors. + + Returns + ------- + config : dict + Architecture configuration. + state_dict : dict of str to Tensor + Model weights. + format_version : int + Logical checkpoint schema version from ``meta.json``. + + Raises + ------ + ValueError + If members are missing or the archive is not a valid zip. + """ + try: + with zipfile.ZipFile(path, mode="r") as archive: + names = set(archive.namelist()) + missing = sorted(SAFE_ZIP_MEMBER_NAMES - names) + if missing: + msg = f"safetensors_v1 zip missing members {missing} in {path}" + raise ValueError(msg) + with TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + for member in SAFE_ZIP_MEMBER_NAMES: + archive.extract(member, path=tmp_path) + return _load_safetensors_v1_directory( + tmp_path, + map_location=map_location, + ) + except zipfile.BadZipFile as exc: + msg = f"Invalid safetensors_v1 zip checkpoint: {path}" + raise ValueError(msg) from exc + + +def _load_legacy_pt_file( + path: Path, + *, + map_location: str | torch.device | None, +) -> tuple[dict[str, Any], dict[str, torch.Tensor], int]: + """Load config and weights from a legacy pickle ``.pt`` checkpoint. + + Parameters + ---------- + path : Path + Checkpoint file path. + map_location : str, torch.device, or None + Device mapping forwarded to :func:`torch.load`. + + Returns + ------- + config : dict + Architecture configuration. + state_dict : dict of str to Tensor + Model weights. + format_version : int + Logical checkpoint schema version. + + Raises + ------ + ValueError + If the payload shape or format version is invalid. + """ + payload = torch.load(path, map_location=map_location, weights_only=False) if not isinstance(payload, dict): msg = "Checkpoint must be a dictionary payload" raise ValueError(msg) @@ -1601,8 +1989,35 @@ def load_checkpoint( if not isinstance(config, dict) or not isinstance(state_dict, dict): msg = "Checkpoint must contain 'config' and 'state_dict' dictionaries" raise ValueError(msg) + return config, state_dict, int(format_version) + + +def _assemble_model_from_checkpoint( + config: dict[str, Any], + state_dict: dict[str, torch.Tensor], + *, + format_version: int, + physics_lifting_fn: PhysicsLiftingFn | None, +) -> ModeShapeModel: + """Reconstruct a model and load weights from parsed checkpoint pieces. - migrated_config = _migrate_config(config, format_version=int(format_version)) + Parameters + ---------- + config : dict + Checkpoint config mapping (possibly pre-migration). + state_dict : dict of str to Tensor + Weight tensors to load. + format_version : int + On-disk format version used for config migration. + physics_lifting_fn : callable or None + Optional custom physics lifting for hybrid reconstructions. + + Returns + ------- + ModeShapeModel + Reconstructed model in evaluation mode with loaded weights. + """ + migrated_config = _migrate_config(config, format_version=format_version) if migrated_config.get("koopman_kind") == "hetero_graph": _validate_hetero_latent_dims_vs_state(migrated_config, state_dict) model = reconstruct_model(migrated_config, physics_lifting_fn=physics_lifting_fn) @@ -1614,6 +2029,91 @@ def load_checkpoint( return model +def load_checkpoint( + path: str | Path, + *, + map_location: str | torch.device | None = None, + physics_lifting_fn: PhysicsLiftingFn | None = None, +) -> ModeShapeModel: + """Load a trained model from a checkpoint file or directory. + + Auto-detects on-disk containers: + + 1. **Directory** with ``meta.json`` → ``safetensors_v1`` directory + (JSON config + safetensors weights; no pickle). + 2. **Zip** whose root members include ``meta.json``, ``config.json``, + and ``model.safetensors`` → ``safetensors_v1`` zip / ``.kgckpt`` + (same three files; no pickle). Matched by member names so safe + markers win over legacy pickle even when the file is a zip archive. + 3. **Other files** (typically ``.pt`` / ``.pth``) → legacy + ``torch.load`` pickle payload (trusted-source trust boundary; see + ``SECURITY.md``). + + Parameters + ---------- + path : str or Path + Checkpoint ``.pt`` file, ``safetensors_v1`` directory, or + ``.kgckpt`` / ``.zip`` bundle produced by :func:`save_checkpoint`. + map_location : str, torch.device, or None, optional + Device mapping. For legacy ``.pt`` files, forwarded to + :func:`torch.load`. For ``safetensors_v1``, normalized to a device + string (``None`` → ``"cpu"``); callable ``map_location`` values are + not supported on that path. + physics_lifting_fn : callable or None, optional + Custom physics lifting function for hybrid checkpoints without a stored + preset. + + Returns + ------- + GraphKoopmanModel + Reconstructed model with restored weights in evaluation mode. + + Raises + ------ + ValueError + If the checkpoint layout, container, format version, or payload is + invalid. + FileNotFoundError + If ``path`` does not exist. + """ + destination = Path(path) + if not destination.exists(): + msg = f"Checkpoint path not found: {destination}" + raise FileNotFoundError(msg) + + if destination.is_dir(): + config, state_dict, format_version = _load_safetensors_v1_directory( + destination, + map_location=map_location, + ) + elif _is_safetensors_v1_zip(destination): + config, state_dict, format_version = _load_safetensors_v1_zip( + destination, + map_location=map_location, + ) + elif destination.is_file(): + if destination.suffix.lower() in SAFE_BUNDLE_SUFFIXES: + msg = ( + f"Path {destination} has a safetensors_v1 bundle suffix but " + f"is not a zip with {sorted(SAFE_ZIP_MEMBER_NAMES)}" + ) + raise ValueError(msg) + config, state_dict, format_version = _load_legacy_pt_file( + destination, + map_location=map_location, + ) + else: + msg = f"Unsupported checkpoint path type: {destination}" + raise ValueError(msg) + + return _assemble_model_from_checkpoint( + config, + state_dict, + format_version=format_version, + physics_lifting_fn=physics_lifting_fn, + ) + + def snapshot_state_dict(module: nn.Module) -> dict[str, torch.Tensor]: """Return a detached copy of a module's ``state_dict`` for checkpointing. diff --git a/src/koopman_graph/tracking/__init__.py b/src/koopman_graph/tracking/__init__.py new file mode 100644 index 0000000..93d8381 --- /dev/null +++ b/src/koopman_graph/tracking/__init__.py @@ -0,0 +1,26 @@ +"""Experiment-tracking adapters for fit callbacks. + +Capability layout +----------------- +``csv`` + :class:`~koopman_graph.tracking.CsvFitLogger` — stdlib CSV writer + implementing :class:`~koopman_graph.training.FitCallback`. +``tensorboard`` + :class:`~koopman_graph.tracking.TensorBoardFitLogger` — optional + ``SummaryWriter`` adapter (requires peer ``tensorboard`` install). + +Power-user module: import as ``koopman_graph.tracking``. Types are +intentionally omitted from root ``koopman_graph.__all__`` (callbacks are +root-exported; adapters stay here). This package may import +:mod:`koopman_graph.training` types; ``training`` must not import tracking +adapters. +""" + +from koopman_graph.tracking.csv import CSV_FIT_FIELDNAMES, CsvFitLogger +from koopman_graph.tracking.tensorboard import TensorBoardFitLogger + +__all__ = [ + "CSV_FIT_FIELDNAMES", + "CsvFitLogger", + "TensorBoardFitLogger", +] diff --git a/src/koopman_graph/tracking/csv.py b/src/koopman_graph/tracking/csv.py new file mode 100644 index 0000000..503becc --- /dev/null +++ b/src/koopman_graph/tracking/csv.py @@ -0,0 +1,162 @@ +"""CSV adapter for :class:`~koopman_graph.training.FitCallback`.""" + +from __future__ import annotations + +import csv +from collections.abc import Mapping +from pathlib import Path +from typing import Any, TextIO + +from koopman_graph.training.history import FitHistory, TrainingLossBreakdown + +# Keys from :meth:`TrainingLossBreakdown.to_floats` in stable column order. +_LOSS_TERM_KEYS: tuple[str, ...] = ( + "total", + "reconstruction", + "forward", + "backward", + "rollout", + "eigenvalue", + "lie", + "pde", + "sparsity", + "worst_case", + "vamp2", +) + +CSV_FIT_FIELDNAMES: tuple[str, ...] = ( + "epoch", + *(f"train_{key}" for key in _LOSS_TERM_KEYS), + *(f"val_{key}" for key in _LOSS_TERM_KEYS), +) + + +def _breakdown_row( + prefix: str, + breakdown: TrainingLossBreakdown | None, +) -> dict[str, str]: + """Map a loss breakdown to CSV cells under ``prefix_`` keys. + + Parameters + ---------- + prefix : str + Column prefix (``"train"`` or ``"val"``). + breakdown : TrainingLossBreakdown or None + Epoch loss terms; ``None`` yields empty string cells. + + Returns + ------- + dict of str to str + Mapping from ``{prefix}_{term}`` to stringified floats (or ``""``). + """ + if breakdown is None: + return {f"{prefix}_{key}": "" for key in _LOSS_TERM_KEYS} + floats = breakdown.to_floats() + return {f"{prefix}_{key}": f"{floats[key]!s}" for key in _LOSS_TERM_KEYS} + + +class CsvFitLogger: + """Write per-epoch train/val loss terms to a CSV file. + + Implements :class:`~koopman_graph.training.FitCallback`. Opens the file on + ``on_fit_start``, appends one row per ``on_epoch_end``, and closes on + ``on_fit_end``. Observe-only: does not mutate the model. + + Parameters + ---------- + path : str or Path + Destination CSV path. Parent directories are created when missing. + overwrite : bool, optional + When ``True`` (default), truncate an existing file. When ``False``, + append without rewriting the header (caller must ensure the file + already has a compatible header). + + Notes + ----- + Column names are :data:`CSV_FIT_FIELDNAMES`. Validation columns are left + empty when ``val_breakdown`` is ``None``. + """ + + def __init__(self, path: str | Path, *, overwrite: bool = True) -> None: + """Store destination path and overwrite policy. + + Parameters + ---------- + path : str or Path + Destination CSV path. + overwrite : bool, optional + Truncate on start when ``True`` (default). + """ + self.path = Path(path) + self.overwrite = overwrite + self._file: TextIO | None = None + self._writer: csv.DictWriter[str] | None = None + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + """Open the CSV and write the header when starting a new file. + + Parameters + ---------- + model : Any + Unused model reference (observe-only). + fit_kwargs : Mapping[str, Any] + Unused fit keyword arguments. + """ + del model, fit_kwargs + self.path.parent.mkdir(parents=True, exist_ok=True) + mode = "w" if self.overwrite else "a" + self._file = self.path.open(mode, newline="", encoding="utf-8") + self._writer = csv.DictWriter( + self._file, + fieldnames=list(CSV_FIT_FIELDNAMES), + extrasaction="ignore", + ) + if self.overwrite or self.path.stat().st_size == 0: + self._writer.writeheader() + self._file.flush() + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + """Append one epoch row of train/val loss terms. + + Parameters + ---------- + epoch : int + Zero-based epoch index. + train_breakdown : TrainingLossBreakdown + Mean training loss terms for the epoch. + val_breakdown : TrainingLossBreakdown or None + Mean validation loss terms when validation is enabled. + history_so_far : FitHistory or None + Unused partial history. + """ + del history_so_far + if self._writer is None or self._file is None: + msg = "CsvFitLogger.on_epoch_end called before on_fit_start" + raise RuntimeError(msg) + row: dict[str, str] = {"epoch": str(epoch)} + row.update(_breakdown_row("train", train_breakdown)) + row.update(_breakdown_row("val", val_breakdown)) + self._writer.writerow(row) + self._file.flush() + + def on_fit_end(self, *, history: FitHistory) -> None: + """Flush and close the CSV file. + + Parameters + ---------- + history : FitHistory + Unused final fit history. + """ + del history + if self._file is not None: + self._file.flush() + self._file.close() + self._file = None + self._writer = None diff --git a/src/koopman_graph/tracking/tensorboard.py b/src/koopman_graph/tracking/tensorboard.py new file mode 100644 index 0000000..9030edd --- /dev/null +++ b/src/koopman_graph/tracking/tensorboard.py @@ -0,0 +1,182 @@ +"""TensorBoard adapter for :class:`~koopman_graph.training.FitCallback`.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from koopman_graph.training.history import FitHistory, TrainingLossBreakdown + +# Match :mod:`koopman_graph.tracking.csv` term order for comparable dashboards. +_LOSS_TERM_KEYS: tuple[str, ...] = ( + "total", + "reconstruction", + "forward", + "backward", + "rollout", + "eigenvalue", + "lie", + "pde", + "sparsity", + "worst_case", + "vamp2", +) + + +def _require_summary_writer() -> Any: + """Import ``SummaryWriter`` or raise with an install hint. + + Returns + ------- + type + ``torch.utils.tensorboard.SummaryWriter`` class. + + Raises + ------ + ImportError + If the ``tensorboard`` peer package is not installed. + """ + try: + from torch.utils.tensorboard import SummaryWriter + except ImportError as exc: + msg = ( + "TensorBoardFitLogger requires the tensorboard package " + "(peer install for torch.utils.tensorboard). " + "Install with: pip install tensorboard" + ) + raise ImportError(msg) from exc + return SummaryWriter + + +def _log_breakdown( + writer: Any, + *, + prefix: str, + breakdown: TrainingLossBreakdown, + epoch: int, +) -> None: + """Write scalar tags ``{prefix}/{term}`` for one breakdown. + + Parameters + ---------- + writer : SummaryWriter + Open TensorBoard writer. + prefix : str + Tag prefix (``"train"`` or ``"val"``). + breakdown : TrainingLossBreakdown + Epoch loss terms to log. + epoch : int + Global step / epoch index. + """ + floats = breakdown.to_floats() + for key in _LOSS_TERM_KEYS: + writer.add_scalar(f"{prefix}/{key}", floats[key], global_step=epoch) + + +class TensorBoardFitLogger: + """Log per-epoch train/val loss terms with TensorBoard ``SummaryWriter``. + + Implements :class:`~koopman_graph.training.FitCallback`. Requires the + peer ``tensorboard`` package (not pinned in this library; Q10). Observe-only. + + Parameters + ---------- + log_dir : str or Path + Directory for TensorBoard event files (created if missing). + **writer_kwargs + Forwarded to ``torch.utils.tensorboard.SummaryWriter`` (for example + ``comment`` or ``max_queue``). + + Raises + ------ + ImportError + If ``tensorboard`` / ``SummaryWriter`` is not importable. The message + includes an install hint. + """ + + def __init__(self, log_dir: str | Path, **writer_kwargs: Any) -> None: + """Store log directory and SummaryWriter kwargs. + + Parameters + ---------- + log_dir : str or Path + Directory for TensorBoard event files. + **writer_kwargs + Forwarded to ``SummaryWriter`` on ``on_fit_start``. + """ + self.log_dir = Path(log_dir) + self._writer_kwargs = writer_kwargs + self._writer: Any | None = None + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + """Create a ``SummaryWriter`` under ``log_dir``. + + Parameters + ---------- + model : Any + Unused model reference (observe-only). + fit_kwargs : Mapping[str, Any] + Unused fit keyword arguments. + """ + del model, fit_kwargs + summary_writer_cls = _require_summary_writer() + self.log_dir.mkdir(parents=True, exist_ok=True) + self._writer = summary_writer_cls( + log_dir=str(self.log_dir), + **self._writer_kwargs, + ) + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + """Log train (and optional val) scalars for the completed epoch. + + Parameters + ---------- + epoch : int + Zero-based epoch index used as ``global_step``. + train_breakdown : TrainingLossBreakdown + Mean training loss terms for the epoch. + val_breakdown : TrainingLossBreakdown or None + Mean validation loss terms when validation is enabled. + history_so_far : FitHistory or None + Unused partial history. + """ + del history_so_far + if self._writer is None: + msg = "TensorBoardFitLogger.on_epoch_end called before on_fit_start" + raise RuntimeError(msg) + _log_breakdown( + self._writer, + prefix="train", + breakdown=train_breakdown, + epoch=epoch, + ) + if val_breakdown is not None: + _log_breakdown( + self._writer, + prefix="val", + breakdown=val_breakdown, + epoch=epoch, + ) + self._writer.flush() + + def on_fit_end(self, *, history: FitHistory) -> None: + """Flush and close the ``SummaryWriter``. + + Parameters + ---------- + history : FitHistory + Unused final fit history. + """ + del history + if self._writer is not None: + self._writer.flush() + self._writer.close() + self._writer = None diff --git a/src/koopman_graph/training/__init__.py b/src/koopman_graph/training/__init__.py index 3ffa4c3..8532e59 100644 --- a/src/koopman_graph/training/__init__.py +++ b/src/koopman_graph/training/__init__.py @@ -34,8 +34,13 @@ ``resolve_training_sequences`` / ``resolve_validation_sequences``. ``loop`` ``run_fit_loop`` plus early-stopping / LR-scheduler helpers. +``callbacks`` + Observe-only :class:`FitCallback` protocol and :class:`NoOpFitCallback` + (fit-loop wiring is separate; optional trackers live under + ``koopman_graph.tracking`` when added). """ +from koopman_graph.training.callbacks import FitCallback, NoOpFitCallback from koopman_graph.training.device import resolve_device, sequence_to_device from koopman_graph.training.epochs import ( eval_one_epoch, @@ -81,10 +86,12 @@ __all__ = [ "EarlyStoppingMonitor", "ExtraLosses", + "FitCallback", "FitHistory", "LRSchedulerFactory", "LossWeightSchedule", "LossWeights", + "NoOpFitCallback", "TrainingInput", "ValidationInput", "compute_backward_consistency_sequence_loss", diff --git a/src/koopman_graph/training/callbacks.py b/src/koopman_graph/training/callbacks.py new file mode 100644 index 0000000..0afbf72 --- /dev/null +++ b/src/koopman_graph/training/callbacks.py @@ -0,0 +1,127 @@ +"""Observe-only hooks for :func:`~koopman_graph.training.loop.run_fit_loop`. + +Callbacks must not mutate model parameters. Wiring into the fit loop is +added in a later task; this module defines the protocol and a no-op +implementation for typing and smoke tests. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, runtime_checkable + +from koopman_graph.training.history import FitHistory, TrainingLossBreakdown + + +@runtime_checkable +class FitCallback(Protocol): + """Protocol for epoch-level observers during ``fit`` / ``run_fit_loop``. + + Implementations should treat all arguments as read-only. Do not mutate + model parameters, optimizer state, or loss breakdown tensors in place. + + Notes + ----- + Concrete adapters (CSV, TensorBoard, custom loggers) implement this + protocol; :class:`NoOpFitCallback` is the observe-nothing default. + """ + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + """Called once before the first training epoch. + + Parameters + ---------- + model : Any + Trainable model about to be fit (typically + :class:`~koopman_graph.model.GraphKoopmanModel`). + fit_kwargs : Mapping[str, Any] + Keyword arguments supplied to the fit loop (epochs, device, + loss weights, etc.). Contents are informational and may omit + internal-only resolved objects. + """ + ... + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + """Called after each completed epoch. + + Parameters + ---------- + epoch : int + Zero-based epoch index that just finished. + train_breakdown : TrainingLossBreakdown + Mean training loss terms for the epoch. + val_breakdown : TrainingLossBreakdown or None + Mean validation loss terms when validation is enabled. + history_so_far : FitHistory or None + History accumulated through this epoch when available; may be + ``None`` until the loop builds a public history object. + """ + ... + + def on_fit_end(self, *, history: FitHistory) -> None: + """Called once after training finishes (including early stop). + + Parameters + ---------- + history : FitHistory + Final immutable training history returned to the caller. + """ + ... + + +class NoOpFitCallback: + """Concrete :class:`FitCallback` that records nothing and mutates nothing. + + Notes + ----- + Useful as a typed default when no logging adapter is attached. + """ + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + """No-op fit start hook. + + Parameters + ---------- + model : Any + Unused model reference. + fit_kwargs : Mapping[str, Any] + Unused fit keyword arguments. + """ + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + """No-op epoch-end hook. + + Parameters + ---------- + epoch : int + Unused epoch index. + train_breakdown : TrainingLossBreakdown + Unused training loss breakdown. + val_breakdown : TrainingLossBreakdown or None + Unused validation loss breakdown. + history_so_far : FitHistory or None + Unused partial history. + """ + + def on_fit_end(self, *, history: FitHistory) -> None: + """No-op fit end hook. + + Parameters + ---------- + history : FitHistory + Unused final fit history. + """ diff --git a/src/koopman_graph/training/loop.py b/src/koopman_graph/training/loop.py index c16340c..d52d036 100644 --- a/src/koopman_graph/training/loop.py +++ b/src/koopman_graph/training/loop.py @@ -28,6 +28,7 @@ from koopman_graph.graph_utils import snapshot_hyperedge_index from koopman_graph.nn.heterogeneous import resolve_multiplex_relation_inputs from koopman_graph.protocols import TrainableKoopmanModel +from koopman_graph.training.callbacks import FitCallback from koopman_graph.training.device import resolve_device, sequence_to_device from koopman_graph.training.epochs import ( eval_one_epoch, @@ -45,6 +46,107 @@ ) from koopman_graph.training.schedules import resolve_loss_weights_for_epoch + +def _build_fit_history( + *, + losses: list[float], + reconstruction_losses: list[float], + forward_losses: list[float], + backward_losses: list[float], + rollout_losses: list[float], + eigenvalue_losses: list[float], + lie_losses: list[float], + pde_losses: list[float], + sparsity_losses: list[float], + worst_case_losses: list[float], + vamp2_losses: list[float], + val_losses: list[float] | None, + val_reconstruction_losses: list[float] | None, + val_forward_losses: list[float] | None, + val_backward_losses: list[float] | None, + val_rollout_losses: list[float] | None, + val_eigenvalue_losses: list[float] | None, + val_lie_losses: list[float] | None, + val_pde_losses: list[float] | None, + val_sparsity_losses: list[float] | None, + val_worst_case_losses: list[float] | None, + val_vamp2_losses: list[float] | None, + stopped_early: bool, + best_epoch: int | None, + best_loss: float | None, +) -> FitHistory: + """Build an immutable :class:`FitHistory` from in-progress epoch lists. + + Parameters + ---------- + losses, reconstruction_losses, forward_losses, backward_losses, + rollout_losses, eigenvalue_losses, lie_losses, pde_losses, + sparsity_losses, worst_case_losses, vamp2_losses : list of float + Per-epoch training loss series (same length as ``losses``). + val_losses, val_reconstruction_losses, val_forward_losses, + val_backward_losses, val_rollout_losses, val_eigenvalue_losses, + val_lie_losses, val_pde_losses, val_sparsity_losses, + val_worst_case_losses, val_vamp2_losses : list of float or None + Optional per-epoch validation series; ``None`` when validation + was not run. + stopped_early : bool + Whether early stopping terminated training. + best_epoch : int or None + Best epoch index when tracked. + best_loss : float or None + Best monitored loss when tracked. + + Returns + ------- + FitHistory + Frozen history with tuple series. + """ + return FitHistory( + loss=tuple(losses), + epochs=len(losses), + reconstruction_loss=tuple(reconstruction_losses), + forward_loss=tuple(forward_losses), + backward_loss=tuple(backward_losses), + rollout_loss=tuple(rollout_losses), + eigenvalue_loss=tuple(eigenvalue_losses), + lie_loss=tuple(lie_losses), + pde_loss=tuple(pde_losses), + sparsity_loss=tuple(sparsity_losses), + worst_case_loss=tuple(worst_case_losses), + vamp2_loss=tuple(vamp2_losses), + val_loss=None if val_losses is None else tuple(val_losses), + val_reconstruction_loss=( + None + if val_reconstruction_losses is None + else tuple(val_reconstruction_losses) + ), + val_forward_loss=( + None if val_forward_losses is None else tuple(val_forward_losses) + ), + val_backward_loss=( + None if val_backward_losses is None else tuple(val_backward_losses) + ), + val_rollout_loss=( + None if val_rollout_losses is None else tuple(val_rollout_losses) + ), + val_eigenvalue_loss=( + None if val_eigenvalue_losses is None else tuple(val_eigenvalue_losses) + ), + val_lie_loss=None if val_lie_losses is None else tuple(val_lie_losses), + val_pde_loss=None if val_pde_losses is None else tuple(val_pde_losses), + val_sparsity_loss=( + None if val_sparsity_losses is None else tuple(val_sparsity_losses) + ), + val_worst_case_loss=( + None if val_worst_case_losses is None else tuple(val_worst_case_losses) + ), + val_vamp2_loss=(None if val_vamp2_losses is None else tuple(val_vamp2_losses)), + stopped_early=stopped_early, + best_epoch=best_epoch, + best_loss=best_loss, + ) + + __all__ = [ "bind_pending_orbit_ties", "eval_one_epoch", @@ -278,6 +380,7 @@ def run_fit_loop( val_sequences: Sequence[SnapshotSequence] | None = None, restore_best_weights: bool = False, checkpoint_path: str | Path | None = None, + callbacks: Sequence[FitCallback] | None = None, **optimizer_kwargs: Any, ) -> FitHistory: """Run the multi-epoch training loop for a trainable Koopman model. @@ -288,6 +391,20 @@ def run_fit_loop( ``early_stopping_monitor`` via :func:`resolve_early_stopping_monitor` before calling this helper (``"auto"`` is not accepted here). + Callback hooks + -------------- + When ``callbacks`` is set, each :class:`~koopman_graph.training.FitCallback` + is invoked in list order as: + + * ``on_fit_start`` once after device / optimizer setup and before epoch 0 + * ``on_epoch_end`` after each epoch's train (and optional val) losses are + recorded and best-epoch tracking is updated, before early-stop ``break`` + * ``on_fit_end`` once with the final :class:`FitHistory`, after optional + best-weight restore / checkpoint write + + Callbacks are observe-only: do not mutate model parameters or optimizer + state. Default ``None`` preserves prior behavior. + Parameters ---------- model : TrainableKoopmanModel @@ -357,6 +474,8 @@ def run_fit_loop( Reload in-memory weights from the lowest-loss epoch when ``True``. checkpoint_path : str, Path, or None, optional Write a checkpoint at the lowest-loss epoch when set. + callbacks : sequence of FitCallback or None, optional + Observe-only fit hooks. Default ``None`` skips all hook calls. **optimizer_kwargs Extra keyword arguments for the optimizer constructor. @@ -470,6 +589,28 @@ def run_fit_loop( track_best = restore_best_weights or checkpoint_path is not None epochs_without_improvement = 0 stopped_early = False + active_callbacks: Sequence[FitCallback] = () if callbacks is None else callbacks + fit_kwargs: dict[str, Any] = { + "epochs": epochs, + "lr": lr, + "device": train_device, + "loss_weights": loss_weights, + "loss_weight_schedule": loss_weight_schedule, + "extra_losses": extra_losses, + "rollout_horizon": rollout_horizon, + "window_length": window_length, + "batch_size": batch_size, + "max_grad_norm": max_grad_norm, + "use_amp": use_amp, + "early_stopping_patience": early_stopping_patience, + "early_stopping_min_delta": early_stopping_min_delta, + "early_stopping_monitor": early_stopping_monitor, + "restore_best_weights": restore_best_weights, + "checkpoint_path": checkpoint_path, + "has_val_sequences": val_sequences is not None, + } + for callback in active_callbacks: + callback.on_fit_start(model=model, fit_kwargs=fit_kwargs) for epoch in range(epochs): epoch_weights = resolve_loss_weights_for_epoch( @@ -535,8 +676,9 @@ def run_fit_loop( vamp2_losses.append(term_values["vamp2"]) monitored_loss = term_values["total"] + epoch_val_breakdown = None if val_sequences is not None: - val_breakdown = eval_one_epoch( + epoch_val_breakdown = eval_one_epoch( model, val_sequences, epoch_weights, @@ -544,7 +686,7 @@ def run_fit_loop( rollout_horizon=rollout_horizon, rollout_start_indices=epoch_rollout_starts, ) - val_terms = val_breakdown.to_floats() + val_terms = epoch_val_breakdown.to_floats() assert val_losses is not None assert val_reconstruction_losses is not None assert val_forward_losses is not None @@ -575,6 +717,42 @@ def run_fit_loop( best_epoch = epoch best_state_dict = snapshot_state_dict(module) + if active_callbacks: + history_so_far = _build_fit_history( + losses=losses, + reconstruction_losses=reconstruction_losses, + forward_losses=forward_losses, + backward_losses=backward_losses, + rollout_losses=rollout_losses, + eigenvalue_losses=eigenvalue_losses, + lie_losses=lie_losses, + pde_losses=pde_losses, + sparsity_losses=sparsity_losses, + worst_case_losses=worst_case_losses, + vamp2_losses=vamp2_losses, + val_losses=val_losses, + val_reconstruction_losses=val_reconstruction_losses, + val_forward_losses=val_forward_losses, + val_backward_losses=val_backward_losses, + val_rollout_losses=val_rollout_losses, + val_eigenvalue_losses=val_eigenvalue_losses, + val_lie_losses=val_lie_losses, + val_pde_losses=val_pde_losses, + val_sparsity_losses=val_sparsity_losses, + val_worst_case_losses=val_worst_case_losses, + val_vamp2_losses=val_vamp2_losses, + stopped_early=False, + best_epoch=best_epoch, + best_loss=best_loss, + ) + for callback in active_callbacks: + callback.on_epoch_end( + epoch=epoch, + train_breakdown=breakdown, + val_breakdown=epoch_val_breakdown, + history_so_far=history_so_far, + ) + if early_stopping_patience is not None: stop, best_loss_for_stop, epochs_without_improvement = should_stop_early( epoch_loss=monitored_loss, @@ -593,51 +771,42 @@ def run_fit_loop( last_state_dict = snapshot_state_dict(module) module.load_state_dict(best_state_dict) if checkpoint_path is not None: - save_checkpoint(model, checkpoint_path) # type: ignore[arg-type] + # File-path best-epoch checkpoints stay legacy .pt (not directories). + save_checkpoint( + model, # type: ignore[arg-type] + checkpoint_path, + format="legacy_pt", + ) if not restore_best_weights and last_state_dict is not None: module.load_state_dict(last_state_dict) - return FitHistory( - loss=tuple(losses), - epochs=len(losses), - reconstruction_loss=tuple(reconstruction_losses), - forward_loss=tuple(forward_losses), - backward_loss=tuple(backward_losses), - rollout_loss=tuple(rollout_losses), - eigenvalue_loss=tuple(eigenvalue_losses), - lie_loss=tuple(lie_losses), - pde_loss=tuple(pde_losses), - sparsity_loss=tuple(sparsity_losses), - worst_case_loss=tuple(worst_case_losses), - vamp2_loss=tuple(vamp2_losses), - val_loss=None if val_losses is None else tuple(val_losses), - val_reconstruction_loss=( - None - if val_reconstruction_losses is None - else tuple(val_reconstruction_losses) - ), - val_forward_loss=( - None if val_forward_losses is None else tuple(val_forward_losses) - ), - val_backward_loss=( - None if val_backward_losses is None else tuple(val_backward_losses) - ), - val_rollout_loss=( - None if val_rollout_losses is None else tuple(val_rollout_losses) - ), - val_eigenvalue_loss=( - None if val_eigenvalue_losses is None else tuple(val_eigenvalue_losses) - ), - val_lie_loss=None if val_lie_losses is None else tuple(val_lie_losses), - val_pde_loss=None if val_pde_losses is None else tuple(val_pde_losses), - val_sparsity_loss=( - None if val_sparsity_losses is None else tuple(val_sparsity_losses) - ), - val_worst_case_loss=( - None if val_worst_case_losses is None else tuple(val_worst_case_losses) - ), - val_vamp2_loss=(None if val_vamp2_losses is None else tuple(val_vamp2_losses)), + history = _build_fit_history( + losses=losses, + reconstruction_losses=reconstruction_losses, + forward_losses=forward_losses, + backward_losses=backward_losses, + rollout_losses=rollout_losses, + eigenvalue_losses=eigenvalue_losses, + lie_losses=lie_losses, + pde_losses=pde_losses, + sparsity_losses=sparsity_losses, + worst_case_losses=worst_case_losses, + vamp2_losses=vamp2_losses, + val_losses=val_losses, + val_reconstruction_losses=val_reconstruction_losses, + val_forward_losses=val_forward_losses, + val_backward_losses=val_backward_losses, + val_rollout_losses=val_rollout_losses, + val_eigenvalue_losses=val_eigenvalue_losses, + val_lie_losses=val_lie_losses, + val_pde_losses=val_pde_losses, + val_sparsity_losses=val_sparsity_losses, + val_worst_case_losses=val_worst_case_losses, + val_vamp2_losses=val_vamp2_losses, stopped_early=stopped_early, best_epoch=best_epoch, best_loss=best_loss, ) + for callback in active_callbacks: + callback.on_fit_end(history=history) + return history diff --git a/src/koopman_graph/tuning/__init__.py b/src/koopman_graph/tuning/__init__.py new file mode 100644 index 0000000..dfaccb6 --- /dev/null +++ b/src/koopman_graph/tuning/__init__.py @@ -0,0 +1,41 @@ +"""Hyperparameter-search helpers for user-defined Tune / Optuna trials. + +Capability layout +----------------- +``metrics`` + :func:`~koopman_graph.tuning.fit_history_metrics` — flatten a + :class:`~koopman_graph.training.FitHistory` into scalar floats for + trial reporters. +``ray_tune`` + :func:`~koopman_graph.tuning.run_ray_tune` — thin Ray Tune + :class:`~ray.tune.Tuner` façade (lazy import; requires ``[ray]``). +``search_spaces`` + :func:`~koopman_graph.tuning.example_lr_loguniform_space` and + :func:`~koopman_graph.tuning.example_lr_latent_dim_space` — *example* + smoke scaffolds only (not scientific defaults; requires ``[ray]``). + +This package is **not** an AutoML product. Search spaces, searchable +hyperparameters, and objective choice remain caller-owned. Helpers support +Ray Tune / Optuna (and similar) integrations; they do not invent +recommended scientific defaults. + +Power-user module: import as ``koopman_graph.tuning``. Intentionally +omitted from root ``koopman_graph.__all__``. May import +:mod:`koopman_graph.training` types; ``training`` must not import +``tuning``. Importing this package does not import Ray; Ray is loaded +lazily inside Ray-backed helpers. +""" + +from koopman_graph.tuning.metrics import fit_history_metrics +from koopman_graph.tuning.ray_tune import run_ray_tune +from koopman_graph.tuning.search_spaces import ( + example_lr_latent_dim_space, + example_lr_loguniform_space, +) + +__all__ = [ + "example_lr_latent_dim_space", + "example_lr_loguniform_space", + "fit_history_metrics", + "run_ray_tune", +] diff --git a/src/koopman_graph/tuning/metrics.py b/src/koopman_graph/tuning/metrics.py new file mode 100644 index 0000000..023f654 --- /dev/null +++ b/src/koopman_graph/tuning/metrics.py @@ -0,0 +1,81 @@ +"""Flatten :class:`~koopman_graph.training.FitHistory` for HPO reporters.""" + +from __future__ import annotations + +from koopman_graph.training.history import FitHistory + + +def _last(series: tuple[float, ...] | None) -> float | None: + """Return the last finite-friendly float from a series, or ``None``. + + Parameters + ---------- + series : tuple of float or None + Epoch series; empty or ``None`` yields ``None``. + + Returns + ------- + float or None + Last element as ``float``, or ``None`` when unavailable. + """ + if series is None or len(series) == 0: + return None + return float(series[-1]) + + +def fit_history_metrics(history: FitHistory) -> dict[str, float]: + """Extract scalar metrics from a fit history for Tune / Optuna reporters. + + Returns a flat ``str → float`` mapping. Optional series that are missing + or empty are omitted (no sentinel NaNs). Boolean flags are encoded as + ``0.0`` / ``1.0``. + + Parameters + ---------- + history : FitHistory + History returned by :meth:`~koopman_graph.GraphKoopmanModel.fit` or + :func:`~koopman_graph.training.run_fit_loop`. + + Returns + ------- + dict of str to float + Keys always include ``epochs`` and ``stopped_early``. When the + corresponding series or fields are present, also includes + ``loss`` / ``final_loss``, ``reconstruction_loss``, ``best_loss``, + ``best_epoch``, ``val_loss`` / ``final_val_loss``, and + ``val_reconstruction_loss``. + + Notes + ----- + This helper does not choose an optimization objective. Callers decide + which key (e.g. ``loss`` vs ``val_loss``) to minimize or maximize. + """ + metrics: dict[str, float] = { + "epochs": float(history.epochs), + "stopped_early": 1.0 if history.stopped_early else 0.0, + } + + final_loss = _last(history.loss) + if final_loss is not None: + metrics["loss"] = final_loss + metrics["final_loss"] = final_loss + + reconstruction = _last(history.reconstruction_loss) + if reconstruction is not None: + metrics["reconstruction_loss"] = reconstruction + + if history.best_loss is not None: + metrics["best_loss"] = float(history.best_loss) + if history.best_epoch is not None: + metrics["best_epoch"] = float(history.best_epoch) + + final_val = _last(history.val_loss) + if final_val is not None: + metrics["val_loss"] = final_val + metrics["final_val_loss"] = final_val + + val_reconstruction = _last(history.val_reconstruction_loss) + if val_reconstruction is not None: + metrics["val_reconstruction_loss"] = val_reconstruction + + return metrics diff --git a/src/koopman_graph/tuning/ray_tune.py b/src/koopman_graph/tuning/ray_tune.py new file mode 100644 index 0000000..965f3c0 --- /dev/null +++ b/src/koopman_graph/tuning/ray_tune.py @@ -0,0 +1,153 @@ +"""Thin Ray Tune wrapper (lazy import; requires ``[ray]``). + +Search spaces and trainables remain caller-owned — this module is not an +AutoML product. Prefer the examples-only script +``examples/scripts/ray_tune_koopman_example.py`` for a full smoke pattern +(including metrics from :func:`~koopman_graph.tuning.fit_history_metrics`). +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any + +_RAY_INSTALL_HINT = 'pip install "koopman-graph[ray]"' + + +def _prefer_driver_interpreter_for_uv_run() -> None: + """Disable Ray's ``uv run`` worker rewrite when unset. + + Must run before ``import ray`` so Ray reads the flag at import time. + + Notes + ----- + Sets ``RAY_ENABLE_UV_RUN_RUNTIME_ENV=0`` via ``setdefault`` so existing + user overrides are preserved. + """ + os.environ.setdefault("RAY_ENABLE_UV_RUN_RUNTIME_ENV", "0") + + +def _import_ray_tune() -> tuple[Any, Any, Any]: + """Import ``ray``, ``ray.tune``, and ``ray.train.RunConfig`` lazily. + + Returns + ------- + ray_module, tune_module, RunConfig + Imported symbols. + + Raises + ------ + ImportError + If Ray / Ray Tune is not installed (``[ray]`` extra missing). + """ + _prefer_driver_interpreter_for_uv_run() + try: + ray_module = importlib.import_module("ray") + tune_module = importlib.import_module("ray.tune") + run_config_cls = importlib.import_module("ray.train").RunConfig + except ImportError as exc: + msg = ( + "Ray Tune is required for koopman_graph.tuning.run_ray_tune; " + f"install with: {_RAY_INSTALL_HINT}" + ) + raise ImportError(msg) from exc + return ray_module, tune_module, run_config_cls + + +def run_ray_tune( + trainable: Callable[..., Any] | type, + param_space: Mapping[str, Any], + *, + num_samples: int = 1, + metric: str = "loss", + mode: str = "min", + storage_path: str | Path | None = None, + run_name: str | None = None, + ray_init_kwargs: Mapping[str, Any] | None = None, + **tuner_kwargs: Any, +) -> Any: + """Run a Ray Tune search with a caller-owned trainable and param space. + + This is a thin façade over :class:`ray.tune.Tuner`. It does not define + searchable hyperparameters or scientific defaults. Requires the + ``[ray]`` extra. + + Parameters + ---------- + trainable : callable or class + Ray Tune trainable (function or Trainable subclass). + param_space : mapping + Search space / fixed config owned by the caller. + num_samples : int, optional + Number of samples for :class:`~ray.tune.TuneConfig` (default 1). + metric : str, optional + Metric name reported by the trainable (default ``\"loss\"``). + mode : str, optional + Optimization mode ``\"min\"`` or ``\"max\"`` (default ``\"min\"``). + storage_path : str, Path, or None, optional + When set, builds a :class:`~ray.train.RunConfig` with this storage + directory. + run_name : str or None, optional + Optional Tune run name when ``storage_path`` is set (default + ``koopman_graph_tune``). + ray_init_kwargs : mapping or None, optional + Forwarded to ``ray.init`` when Ray is not already initialized. + **tuner_kwargs + Extra keyword arguments forwarded to :class:`ray.tune.Tuner` + (for example a pre-built ``tune_config`` overrides the default + constructed here only if you pass ``tune_config=`` explicitly — + prefer the dedicated parameters above for the common case). + + Returns + ------- + ResultGrid + Return value of :meth:`ray.tune.Tuner.fit`. + + Raises + ------ + ImportError + If Ray Tune is not installed. + ValueError + If ``num_samples < 1``. + + See Also + -------- + examples/scripts/ray_tune_koopman_example.py + Examples-only smoke that owns the search space and trainable body. + """ + if num_samples < 1: + msg = f"num_samples must be >= 1, got {num_samples}" + raise ValueError(msg) + + ray_module, tune_module, run_config_cls = _import_ray_tune() + + if not ray_module.is_initialized(): + init_kwargs = dict(ray_init_kwargs or {}) + ray_module.init(**init_kwargs) + + tune_config = tuner_kwargs.pop("tune_config", None) + if tune_config is None: + tune_config = tune_module.TuneConfig( + num_samples=num_samples, + metric=metric, + mode=mode, + ) + + run_config = tuner_kwargs.pop("run_config", None) + if run_config is None and storage_path is not None: + run_config = run_config_cls( + name=run_name or "koopman_graph_tune", + storage_path=str(Path(storage_path).resolve()), + ) + + tuner = tune_module.Tuner( + trainable, + param_space=dict(param_space), + tune_config=tune_config, + run_config=run_config, + **tuner_kwargs, + ) + return tuner.fit() diff --git a/src/koopman_graph/tuning/search_spaces.py b/src/koopman_graph/tuning/search_spaces.py new file mode 100644 index 0000000..7672a81 --- /dev/null +++ b/src/koopman_graph/tuning/search_spaces.py @@ -0,0 +1,127 @@ +"""Example-only Ray Tune param-space scaffolds (smoke tests, not science). + +These helpers build small Ray Tune–compatible ``param_space`` dicts for local +smoke runs. They are **not** recommended scientific defaults, domain oracles, +or an AutoML search-space DSL. Real studies must define their own ranges and +searchable keys. + +Requires the ``[ray]`` extra. Ray is imported lazily when a helper is called +so :mod:`koopman_graph.tuning` stays importable without Ray. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Sequence +from typing import Any + +_RAY_INSTALL_HINT = 'pip install "koopman-graph[ray]"' + + +def _import_tune() -> Any: + """Import ``ray.tune`` with an actionable install hint. + + Returns + ------- + module + The ``ray.tune`` module. + + Raises + ------ + ImportError + If Ray Tune is not installed (``[ray]`` extra missing). + """ + try: + return importlib.import_module("ray.tune") + except ImportError as exc: + msg = ( + "Ray Tune is required for koopman_graph.tuning search-space " + f"helpers; install with: {_RAY_INSTALL_HINT}" + ) + raise ImportError(msg) from exc + + +def example_lr_loguniform_space( + *, + lr_low: float = 1e-3, + lr_high: float = 1e-1, + **fixed: Any, +) -> dict[str, Any]: + """Build an *example* param space with log-uniform learning rate. + + For smoke tests only — not a recommended scientific default. + + Parameters + ---------- + lr_low, lr_high : float, optional + Bounds for ``tune.loguniform`` (defaults ``1e-3`` … ``1e-1``). + **fixed + Additional fixed config keys (for example ``epochs``, ``seed``). + + Returns + ------- + dict + Ray Tune–compatible ``param_space`` mapping. + + Raises + ------ + ImportError + If Ray Tune is not installed. + ValueError + If ``lr_low >= lr_high`` or bounds are non-positive. + """ + if lr_low <= 0 or lr_high <= 0: + raise ValueError("lr_low and lr_high must be positive") + if lr_low >= lr_high: + raise ValueError(f"lr_low must be < lr_high, got {lr_low} >= {lr_high}") + + tune = _import_tune() + space: dict[str, Any] = {"lr": tune.loguniform(lr_low, lr_high)} + space.update(fixed) + return space + + +def example_lr_latent_dim_space( + *, + lr_low: float = 1e-3, + lr_high: float = 1e-1, + latent_dims: Sequence[int] = (8, 16), + **fixed: Any, +) -> dict[str, Any]: + """Build an *example* param space over ``lr`` and ``latent_dim``. + + For smoke tests only — not a recommended scientific default. The + ``latent_dims`` values are illustrative integers for tiny models. + + Parameters + ---------- + lr_low, lr_high : float, optional + Bounds for ``tune.loguniform`` on ``lr``. + latent_dims : sequence of int, optional + Discrete choices for ``latent_dim`` (default ``(8, 16)``). + **fixed + Additional fixed config keys. + + Returns + ------- + dict + Ray Tune–compatible ``param_space`` mapping. + + Raises + ------ + ImportError + If Ray Tune is not installed. + ValueError + If ``latent_dims`` is empty or contains non-positive values, or if + learning-rate bounds are invalid. + """ + dims = tuple(int(d) for d in latent_dims) + if not dims: + raise ValueError("latent_dims must contain at least one positive int") + if any(d < 1 for d in dims): + raise ValueError("latent_dims entries must be >= 1") + + space = example_lr_loguniform_space(lr_low=lr_low, lr_high=lr_high, **fixed) + tune = _import_tune() + space["latent_dim"] = tune.choice(list(dims)) + return space diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py new file mode 100644 index 0000000..b2cdfa6 --- /dev/null +++ b/tests/test_callbacks.py @@ -0,0 +1,221 @@ +"""Smoke tests for the observe-only :class:`FitCallback` protocol.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import asdict +from typing import Any + +import pytest +import torch + +from koopman_graph import ( + FitCallback, + GNNDecoder, + GNNEncoder, + GraphKoopmanModel, + NoOpFitCallback, +) +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.training import FitCallback as TrainingFitCallback +from koopman_graph.training import NoOpFitCallback as TrainingNoOpFitCallback +from koopman_graph.training import run_fit_loop +from koopman_graph.training.history import FitHistory, TrainingLossBreakdown + + +def test_noop_is_runtime_checkable_fit_callback() -> None: + """``NoOpFitCallback`` satisfies ``FitCallback`` at runtime.""" + callback = NoOpFitCallback() + assert isinstance(callback, FitCallback) + + +def test_missing_hook_fails_runtime_checkable() -> None: + """Structural objects missing a required hook are not FitCallbacks.""" + + class _Incomplete: + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + return None + + def on_fit_end(self, *, history: FitHistory) -> None: + return None + + assert not isinstance(_Incomplete(), FitCallback) + + +def test_noop_hooks_accept_protocol_arguments() -> None: + """No-op hooks accept the designed keyword shapes without error.""" + callback = NoOpFitCallback() + breakdown = TrainingLossBreakdown.zeros(torch.device("cpu")) + history = FitHistory(loss=(1.0,), epochs=1) + + callback.on_fit_start(model=object(), fit_kwargs={"epochs": 1}) + callback.on_epoch_end( + epoch=0, + train_breakdown=breakdown, + val_breakdown=None, + history_so_far=None, + ) + callback.on_fit_end(history=history) + + +def test_fit_callback_exported_from_training_package() -> None: + """Protocol symbols are part of the curated ``training`` ``__all__``.""" + import koopman_graph.training as training + + assert "FitCallback" in training.__all__ + assert "NoOpFitCallback" in training.__all__ + assert training.FitCallback is TrainingFitCallback + assert training.NoOpFitCallback is TrainingNoOpFitCallback + + +def test_fit_callback_exported_from_root_package() -> None: + """Q11: ``FitCallback`` / ``NoOpFitCallback`` are root ``__all__`` exports.""" + import koopman_graph as kg + + assert "FitCallback" in kg.__all__ + assert "NoOpFitCallback" in kg.__all__ + assert kg.FitCallback is FitCallback + assert kg.NoOpFitCallback is NoOpFitCallback + assert FitCallback is TrainingFitCallback + assert NoOpFitCallback is TrainingNoOpFitCallback + + +class _RecordingFitCallback: + """Test double that records hook order and epoch indices.""" + + def __init__(self) -> None: + self.events: list[str] = [] + self.epochs: list[int] = [] + self.final_history: FitHistory | None = None + + def on_fit_start(self, *, model: Any, fit_kwargs: Mapping[str, Any]) -> None: + assert "epochs" in fit_kwargs + self.events.append("start") + + def on_epoch_end( + self, + *, + epoch: int, + train_breakdown: TrainingLossBreakdown, + val_breakdown: TrainingLossBreakdown | None, + history_so_far: FitHistory | None, + ) -> None: + assert history_so_far is not None + assert history_so_far.epochs == epoch + 1 + self.events.append(f"epoch:{epoch}") + self.epochs.append(epoch) + + def on_fit_end(self, *, history: FitHistory) -> None: + self.events.append("end") + self.final_history = history + + +def _tiny_model() -> GraphKoopmanModel: + """Build a small discrete GCN for callback loop tests.""" + return GraphKoopmanModel( + encoder=GNNEncoder( + in_channels=3, hidden_channels=8, latent_dim=4, num_layers=2 + ), + decoder=GNNDecoder( + latent_dim=4, hidden_channels=8, out_channels=3, num_layers=2 + ), + latent_dim=4, + time_step=0.1, + ) + + +def test_run_fit_loop_callback_order( + scaling_sequence: GraphSnapshotSequence, +) -> None: + """Recording callback sees start → epochs → end with zero-based indices.""" + model = _tiny_model() + recorder = _RecordingFitCallback() + torch.manual_seed(0) + history = run_fit_loop( + model, + [scaling_sequence], + epochs=3, + lr=1e-2, + device="cpu", + callbacks=[recorder], + ) + + assert recorder.events == ["start", "epoch:0", "epoch:1", "epoch:2", "end"] + assert recorder.epochs == [0, 1, 2] + assert recorder.final_history == history + assert history.epochs == 3 + + +def test_noop_callbacks_preserve_fit_history( + scaling_sequence: GraphSnapshotSequence, +) -> None: + """Same seed: NoOp callbacks leave FitHistory identical to no callbacks.""" + torch.manual_seed(1) + model_a = _tiny_model() + state = deepcopy(model_a.state_dict()) + model_b = _tiny_model() + model_b.load_state_dict(deepcopy(state)) + + torch.manual_seed(2) + history_a = run_fit_loop( + model_a, + [scaling_sequence], + epochs=2, + lr=1e-2, + device="cpu", + callbacks=None, + ) + torch.manual_seed(2) + history_b = run_fit_loop( + model_b, + [scaling_sequence], + epochs=2, + lr=1e-2, + device="cpu", + callbacks=[NoOpFitCallback()], + ) + + assert asdict(history_a) == asdict(history_b) + for key in state: + torch.testing.assert_close( + model_a.state_dict()[key], + model_b.state_dict()[key], + rtol=0.0, + atol=0.0, + ) + + +def test_model_fit_forwards_callbacks( + scaling_sequence: GraphSnapshotSequence, +) -> None: + """Public ``GraphKoopmanModel.fit`` forwards callbacks to the fit loop.""" + model = _tiny_model() + recorder = _RecordingFitCallback() + torch.manual_seed(3) + history = model.fit( + scaling_sequence, + epochs=2, + lr=1e-2, + device="cpu", + callbacks=[recorder], + ) + + assert recorder.events == ["start", "epoch:0", "epoch:1", "end"] + assert recorder.final_history == history + + +def test_model_fit_ddp_rejects_callbacks( + scaling_sequence: GraphSnapshotSequence, +) -> None: + """``strategy="ddp"`` rejects non-None callbacks until DDP wiring exists.""" + model = _tiny_model() + with pytest.raises(ValueError, match='strategy="ddp"'): + model.fit( + scaling_sequence, + epochs=1, + lr=1e-2, + device="cpu", + callbacks=[NoOpFitCallback()], + strategy="ddp", + ) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..0fc122e --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,350 @@ +"""Smoke tests for the ``koopman-graph`` CLI skeleton.""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import pytest + +from koopman_graph import __version__ +from koopman_graph.cli import ( + ConfigError, + load_config, + load_train_config, + main, + validate_train_config, +) + + +def test_main_version_flag_prints_package_version( + capsys: pytest.CaptureFixture[str], +) -> None: + """``--version`` prints ``koopman-graph {version}`` and exits 0.""" + with pytest.raises(SystemExit) as exc_info: + main(["--version"]) + assert exc_info.value.code == 0 + captured = capsys.readouterr() + assert captured.out.strip() == f"koopman-graph {__version__}" + + +def test_main_version_subcommand( + capsys: pytest.CaptureFixture[str], +) -> None: + """``version`` subcommand matches the ``--version`` string.""" + assert main(["version"]) == 0 + captured = capsys.readouterr() + assert captured.out.strip() == f"koopman-graph {__version__}" + + +def test_main_help_exits_zero(capsys: pytest.CaptureFixture[str]) -> None: + """``--help`` exits 0 and mentions the program name.""" + with pytest.raises(SystemExit) as exc_info: + main(["--help"]) + assert exc_info.value.code == 0 + captured = capsys.readouterr() + assert "koopman-graph" in captured.out + + +def test_module_invocation_version() -> None: + """``python -m koopman_graph.cli --version`` succeeds.""" + result = subprocess.run( + [sys.executable, "-m", "koopman_graph.cli", "--version"], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert result.stdout.strip() == f"koopman-graph {__version__}" + + +def test_console_script_version() -> None: + """Installed ``koopman-graph --version`` succeeds (editable install).""" + result = subprocess.run( + ["koopman-graph", "--version"], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert result.stdout.strip() == f"koopman-graph {__version__}" + + +def test_library_modules_do_not_import_cli() -> None: + """Layer rule: non-cli package sources must not import ``cli``.""" + package_root = Path(__file__).resolve().parents[1] / "src" / "koopman_graph" + cli_root = package_root / "cli" + offenders: list[str] = [] + for path in package_root.rglob("*.py"): + if cli_root in path.parents or path.parent == cli_root: + continue + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("#"): + continue + from_cli = stripped.startswith("from koopman_graph.cli") + import_cli = stripped.startswith("import koopman_graph.cli") + if from_cli or import_cli: + offenders.append(str(path.relative_to(package_root.parent))) + break + assert offenders == [] + + +def _minimal_train_config() -> dict[str, object]: + """Return a valid MVP train config matching the design sketch.""" + return { + "model": { + "encoder": "gcn", + "in_channels": 3, + "hidden_channels": 32, + "latent_dim": 16, + "num_layers": 2, + "dynamics_mode": "discrete", + }, + "data": { + "kind": "synthetic_path", + "num_nodes": 8, + "num_timesteps": 40, + "seed": 0, + }, + "fit": {"epochs": 20, "lr": 1e-3, "device": "cpu"}, + "checkpoint": {"path": "artifacts/model.kgckpt"}, + } + + +def test_load_train_config_json(tmp_path: Path) -> None: + """JSON train configs load without requiring PyYAML.""" + import json + + path = tmp_path / "train.json" + path.write_text(json.dumps(_minimal_train_config()), encoding="utf-8") + loaded = load_train_config(path) + assert loaded["data"]["kind"] == "synthetic_path" + assert loaded["model"]["latent_dim"] == 16 + + +def test_validate_rejects_unknown_model_key() -> None: + """Unknown keys report a dotted path.""" + config = _minimal_train_config() + assert isinstance(config["model"], dict) + config["model"]["not_a_real_kwarg"] = 1 + with pytest.raises(ConfigError, match=r"model\.not_a_real_kwarg"): + validate_train_config(config) # type: ignore[arg-type] + + +def test_validate_rejects_unknown_data_kind() -> None: + """``data.kind`` must be on the MVP allowlist.""" + config = _minimal_train_config() + assert isinstance(config["data"], dict) + config["data"]["kind"] = "arbitrary_import" + with pytest.raises(ConfigError, match="data.kind"): + validate_train_config(config) # type: ignore[arg-type] + + +def test_validate_requires_model_and_data() -> None: + """Train configs require ``model`` and ``data`` sections.""" + with pytest.raises(ConfigError, match="model"): + validate_train_config({"data": {"kind": "synthetic_path"}}) + + +def test_load_config_yaml(tmp_path: Path) -> None: + """YAML loads when PyYAML is installed (``[cli]`` extra).""" + pytest.importorskip("yaml") + path = tmp_path / "train.yaml" + path.write_text( + "\n".join( + [ + "model:", + " encoder: gcn", + " in_channels: 3", + " hidden_channels: 16", + " latent_dim: 8", + " num_layers: 2", + "data:", + " kind: synthetic_path", + " num_nodes: 5", + " num_timesteps: 10", + " seed: 1", + ] + ) + + "\n", + encoding="utf-8", + ) + loaded = load_train_config(path) + assert loaded["model"]["encoder"] == "gcn" + assert loaded["data"]["num_nodes"] == 5 + + +def test_load_config_unsupported_suffix(tmp_path: Path) -> None: + """Non JSON/YAML suffixes raise ConfigError.""" + path = tmp_path / "train.toml" + path.write_text("x = 1\n", encoding="utf-8") + with pytest.raises(ConfigError, match="Unsupported config suffix"): + load_config(path) + + +def test_cli_extra_declares_pyyaml() -> None: + """``[cli]`` optional dependency pins PyYAML.""" + root = Path(__file__).resolve().parents[1] + text = (root / "pyproject.toml").read_text(encoding="utf-8") + assert "\ncli = [\n" in text + assert '"pyyaml>=6"' in text + + +def test_train_synthetic_writes_safetensors_checkpoint( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + """``train`` on the bundled synthetic config writes a loadable checkpoint.""" + import zipfile + + from koopman_graph.serialization import ( + SAFE_ZIP_MEMBER_NAMES, + load_checkpoint, + ) + + root = Path(__file__).resolve().parents[1] + config_path = root / "examples" / "cli" / "synthetic_train.json" + assert config_path.is_file() + + code = main( + [ + "train", + "--config", + str(config_path), + "--out", + str(tmp_path), + ] + ) + assert code == 0 + captured = capsys.readouterr() + assert "wrote checkpoint:" in captured.out + + checkpoint = tmp_path / "model.kgckpt" + assert checkpoint.is_file() + assert zipfile.is_zipfile(checkpoint) + with zipfile.ZipFile(checkpoint, mode="r") as archive: + assert SAFE_ZIP_MEMBER_NAMES.issubset(set(archive.namelist())) + + loaded = load_checkpoint(checkpoint) + assert loaded.latent_dim == 8 + + +def test_train_bad_config_exits_one( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + """Unknown config keys exit with code 1 and a dotted-path message.""" + import json + + path = tmp_path / "bad.json" + path.write_text( + json.dumps( + { + "model": { + "encoder": "gcn", + "in_channels": 3, + "hidden_channels": 8, + "latent_dim": 4, + "not_allowed": True, + }, + "data": {"kind": "synthetic_path", "num_nodes": 5, "num_timesteps": 8}, + } + ), + encoding="utf-8", + ) + code = main(["train", "--config", str(path), "--out", str(tmp_path / "out")]) + assert code == 1 + err = capsys.readouterr().err + assert "error:" in err + assert "model.not_allowed" in err + + +def test_cli_e2e_train_predict_smoke( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + """CI smoke: train → predict on the bundled synthetic config (CPU, tiny).""" + import json + import time + + import torch + + root = Path(__file__).resolve().parents[1] + config_path = root / "examples" / "cli" / "synthetic_train.json" + assert config_path.is_file() + config = json.loads(config_path.read_text(encoding="utf-8")) + assert config["fit"]["device"] == "cpu" + assert int(config["fit"]["epochs"]) <= 5 + + train_out = tmp_path / "train" + forecast_path = tmp_path / "forecast.pt" + started = time.perf_counter() + + train_code = main(["train", "--config", str(config_path), "--out", str(train_out)]) + assert train_code == 0 + checkpoint = train_out / "model.kgckpt" + assert checkpoint.is_file() + + predict_code = main( + [ + "predict", + "--checkpoint", + str(checkpoint), + "--data", + str(config_path), + "--steps", + "4", + "--out", + str(forecast_path), + ] + ) + assert predict_code == 0 + elapsed = time.perf_counter() - started + assert elapsed < 60.0, f"e2e smoke too slow for default CI: {elapsed:.1f}s" + + out = capsys.readouterr().out + assert "wrote checkpoint:" in out + assert "wrote forecast:" in out + assert forecast_path.is_file() + + payload = torch.load(forecast_path, map_location="cpu", weights_only=False) + assert payload["steps"] == 4 + assert len(payload["forecasts"]) == 4 + assert payload["summary"]["num_nodes"] == 8 + assert payload["summary"]["feature_dim"] == 3 + + +def test_predict_missing_checkpoint_exits_one( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + """Missing checkpoint path exits with code 1.""" + root = Path(__file__).resolve().parents[1] + config_path = root / "examples" / "cli" / "synthetic_train.json" + code = main( + [ + "predict", + "--checkpoint", + str(tmp_path / "missing.kgckpt"), + "--data", + str(config_path), + "--out", + str(tmp_path / "out.pt"), + ] + ) + assert code == 1 + assert "error:" in capsys.readouterr().err + + +def test_predict_help_documents_io(capsys: pytest.CaptureFixture[str]) -> None: + """``predict --help`` documents checkpoint, data, and output paths.""" + with pytest.raises(SystemExit) as exc_info: + main(["predict", "--help"]) + assert exc_info.value.code == 0 + help_text = capsys.readouterr().out + assert "--checkpoint" in help_text + assert "--data" in help_text + assert "--out" in help_text + assert "forecast" in help_text.lower() diff --git a/tests/test_distributed_ray.py b/tests/test_distributed_ray.py index 609229a..f7c9a73 100644 --- a/tests/test_distributed_ray.py +++ b/tests/test_distributed_ray.py @@ -217,6 +217,7 @@ def fake_init(**kwargs: object) -> None: assert len(histories) == 1 +@pytest.mark.ray def test_fit_ensemble_with_ray_two_members_finite_histories() -> None: """Ray path fits two members with finite losses under a local runtime.""" ray = pytest.importorskip("ray") diff --git a/tests/test_distributed_ray_train.py b/tests/test_distributed_ray_train.py index 0664bdb..e8dd55f 100644 --- a/tests/test_distributed_ray_train.py +++ b/tests/test_distributed_ray_train.py @@ -135,6 +135,7 @@ def test_run_ray_train_fit_loop_rejects_bad_num_workers() -> None: ) +@pytest.mark.ray def test_run_ray_train_fit_loop_world_size_one_smoke() -> None: """World-size-1 TorchTrainer smoke reuses ``_fit_epochs`` successfully.""" try: @@ -176,6 +177,7 @@ def test_unwrapped_state_dict_strips_ddp_module_prefix() -> None: assert not any(key.startswith("module.") for key in state) +@pytest.mark.ray def test_ray_train_checkpoint_path_matches_single_process_keys( tmp_path: Path, ) -> None: @@ -190,7 +192,7 @@ def test_ray_train_checkpoint_path_matches_single_process_keys( sequence = _tiny_sequence() single = _tiny_model(seed=0) single_path = tmp_path / "single.pt" - save_checkpoint(single, single_path) + save_checkpoint(single, single_path, format="legacy_pt") single_keys = set( torch.load(single_path, map_location="cpu", weights_only=False)["state_dict"] ) @@ -215,6 +217,7 @@ def test_ray_train_checkpoint_path_matches_single_process_keys( assert ray_keys == single_keys +@pytest.mark.ray def test_world_size_one_ray_train_matches_native_ddp_loss() -> None: """Seeded world-size-1 Ray Train loss matches ``run_ddp_fit_loop``. @@ -263,6 +266,7 @@ def test_world_size_one_ray_train_matches_native_ddp_loss() -> None: ) +@pytest.mark.ray def test_ray_train_result_payload_has_no_module_prefix(tmp_path: Path) -> None: """Train result ``state_dict`` restored onto the driver has bare keys.""" try: @@ -283,13 +287,13 @@ def test_ray_train_result_payload_has_no_module_prefix(tmp_path: Path) -> None: seed=2, ) out = tmp_path / "after_ray.pt" - save_checkpoint(model, out) + save_checkpoint(model, out, format="legacy_pt") keys = set(torch.load(out, map_location="cpu", weights_only=False)["state_dict"]) assert not any(key.startswith("module.") for key in keys) # Same key set as a never-wrapped peer architecture. peer = _tiny_model(seed=9) peer_path = tmp_path / "peer.pt" - save_checkpoint(peer, peer_path) + save_checkpoint(peer, peer_path, format="legacy_pt") peer_keys = set( torch.load(peer_path, map_location="cpu", weights_only=False)["state_dict"] ) diff --git a/tests/test_explain_representation.py b/tests/test_explain_representation.py new file mode 100644 index 0000000..8a1ce7f --- /dev/null +++ b/tests/test_explain_representation.py @@ -0,0 +1,376 @@ +"""Tests for representation-explanation result types and GNNExplainer path.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from unittest.mock import MagicMock, patch + +import pytest +import torch +from torch_geometric.data import Data, HeteroData + +from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel +from koopman_graph import analysis as kg_analysis +from koopman_graph.analysis import RepresentationExplanation, explain_representation +from koopman_graph.analysis.explain import RepresentationExplanation as direct + + +def _path_data(*, num_nodes: int = 6, in_channels: int = 3) -> Data: + """Build a small path-graph snapshot.""" + sources = list(range(num_nodes - 1)) + list(range(1, num_nodes)) + targets = list(range(1, num_nodes)) + list(range(num_nodes - 1)) + edge_index = torch.tensor([sources, targets], dtype=torch.long) + x = torch.randn(num_nodes, in_channels) + return Data(x=x, edge_index=edge_index) + + +def _tiny_model(*, in_channels: int = 3, latent_dim: int = 4) -> GraphKoopmanModel: + """Build a tiny homogeneous GCN Koopman model.""" + return GraphKoopmanModel( + encoder=GNNEncoder(in_channels, 8, latent_dim, num_layers=2), + decoder=GNNDecoder(latent_dim, 8, in_channels, num_layers=2), + latent_dim=latent_dim, + time_step=0.1, + ) + + +def test_representation_explanation_exported_from_analysis() -> None: + """Result type is on the analysis public surface (not root ``__all__``).""" + import koopman_graph + + assert RepresentationExplanation is direct + assert "RepresentationExplanation" in kg_analysis.__all__ + assert "explain_representation" in kg_analysis.__all__ + assert "RepresentationExplanation" not in koopman_graph.__all__ + assert "explain_representation" not in koopman_graph.__all__ + assert not hasattr(koopman_graph, "RepresentationExplanation") + assert not hasattr(koopman_graph, "explain_representation") + with pytest.raises(ImportError): + exec("from koopman_graph import explain_representation") + with pytest.raises(ImportError): + exec("from koopman_graph import RepresentationExplanation") + + +def test_representation_explanation_fields_match_design() -> None: + """Frozen dataclass matches the design §3.5 field sketch.""" + node_mask = torch.ones(4) + edge_mask = torch.zeros(6) + explanation = RepresentationExplanation( + target="one_step_forecast", + node_mask=node_mask, + edge_mask=edge_mask, + feature_mask=None, + algorithm="gnn_explainer", + index=2, + ) + assert explanation.target == "one_step_forecast" + assert explanation.node_mask is node_mask + assert explanation.edge_mask is edge_mask + assert explanation.feature_mask is None + assert explanation.algorithm == "gnn_explainer" + assert explanation.index == 2 + with pytest.raises(FrozenInstanceError): + explanation.algorithm = "integrated_gradients" # type: ignore[misc] + + +def test_representation_explanation_allows_all_none_masks() -> None: + """Masks may be entirely absent before algorithm wrappers ship.""" + explanation = RepresentationExplanation( + target="latent", + node_mask=None, + edge_mask=None, + feature_mask=None, + algorithm="pending", + index=None, + ) + assert explanation.node_mask is None + assert explanation.edge_mask is None + assert explanation.feature_mask is None + assert explanation.index is None + + +def test_honesty_docstrings_non_causal_not_resdmd_not_mode_energy() -> None: + """Module and class docs must state interpretive / non-causal honesty.""" + import koopman_graph.analysis.explain as explain_mod + + module_doc = explain_mod.__doc__ or "" + class_doc = RepresentationExplanation.__doc__ or "" + fn_doc = explain_representation.__doc__ or "" + forbidden = ( + "causal discovery", + "guarantees causality", + "interventional ground truth", + "proves causality", + "causal topology recovery", + ) + for doc in (module_doc, class_doc, fn_doc): + lowered = doc.lower() + assert "interpretive" in lowered + assert "non-causal" in lowered or "not causal" in lowered + assert "not" in lowered + assert "ResDMD" in doc or "resdmd" in lowered + assert "ModeEnergyAttribution" in doc + for phrase in forbidden: + assert phrase not in lowered + + +def test_explain_representation_gnn_explainer_shape_contract_all_targets() -> None: + """GNNExplainer object masks match documented ``(N, 1)`` / ``(E,)`` shapes.""" + torch.manual_seed(0) + model = _tiny_model() + data = _path_data() + num_nodes = int(data.x.shape[0]) + num_edges = int(data.edge_index.shape[1]) + + for target in ("latent", "one_step_forecast", "reconstruction"): + result = explain_representation( + model, + data, + target=target, # type: ignore[arg-type] + algorithm="gnn_explainer", + epochs=5, + lr=0.05, + ) + assert result.node_mask is not None + assert result.edge_mask is not None + assert result.node_mask.shape[0] == num_nodes + assert result.node_mask.ndim == 2 + assert result.node_mask.shape[-1] == 1 + assert result.edge_mask.shape == (num_edges,) + assert torch.all(torch.isfinite(result.node_mask)) + assert torch.all(torch.isfinite(result.edge_mask)) + + +def test_explain_representation_rejects_unknown_algorithm_and_target() -> None: + """Unknown algorithm / target names raise actionable ValueError.""" + model = _tiny_model() + data = _path_data() + with pytest.raises(ValueError, match="Unsupported algorithm"): + explain_representation( + model, + data, + algorithm="not_a_real_algorithm", # type: ignore[arg-type] + epochs=1, + ) + with pytest.raises(ValueError, match="Unsupported target"): + explain_representation( + model, + data, + target="not_a_real_target", # type: ignore[arg-type] + epochs=1, + ) + + +def test_explain_representation_rejects_delay_models() -> None: + """Delay embeddings are outside the homogeneous explain MVP.""" + model = GraphKoopmanModel( + encoder=GNNEncoder(6, 8, 4, num_layers=2), + decoder=GNNDecoder(4, 8, 3, num_layers=2), + latent_dim=4, + time_step=0.1, + n_delays=2, + ) + data = _path_data(in_channels=3) + with pytest.raises(ValueError, match="n_delays"): + explain_representation(model, data, epochs=1) + + +def test_explain_representation_ig_attribute_shape_matches_features() -> None: + """IG attribute masks match ``Data.x`` shape when Captum is installed.""" + pytest.importorskip("captum") + torch.manual_seed(3) + model = _tiny_model(in_channels=3) + data = _path_data(in_channels=3) + result = explain_representation( + model, + data, + target="one_step_forecast", + algorithm="integrated_gradients", + index=0, + ) + assert result.node_mask is not None + assert result.node_mask.shape == tuple(data.x.shape) + assert result.feature_mask is not None + assert result.feature_mask.shape == tuple(data.x.shape) + + +def test_explain_representation_gnn_explainer_seed_stability() -> None: + """Same seed + model/data yields close GNNExplainer masks across runs.""" + + def _run() -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(42) + model = _tiny_model() + data = _path_data() + # Freeze data randomness by rebuilding under the same seed above. + result = explain_representation( + model, + data, + target="latent", + algorithm="gnn_explainer", + epochs=8, + lr=0.05, + index=0, + ) + assert result.node_mask is not None + assert result.edge_mask is not None + return result.node_mask.detach().clone(), result.edge_mask.detach().clone() + + node_a, edge_a = _run() + node_b, edge_b = _run() + assert torch.allclose(node_a, node_b, atol=1e-5, rtol=1e-4) + assert torch.allclose(edge_a, edge_b, atol=1e-5, rtol=1e-4) + + +def test_explain_representation_gnn_explainer_latent_and_forecast() -> None: + """GNNExplainer returns finite masks with documented shapes.""" + torch.manual_seed(0) + model = _tiny_model() + data = _path_data() + num_nodes = int(data.x.shape[0]) + num_edges = int(data.edge_index.shape[1]) + + latent = explain_representation( + model, + data, + target="latent", + algorithm="gnn_explainer", + epochs=5, + lr=0.05, + ) + assert latent.algorithm == "gnn_explainer" + assert latent.target == "latent" + assert latent.node_mask is not None + assert latent.edge_mask is not None + assert latent.node_mask.shape[0] == num_nodes + assert latent.edge_mask.shape == (num_edges,) + assert torch.all(torch.isfinite(latent.node_mask)) + assert torch.all(torch.isfinite(latent.edge_mask)) + + forecast = explain_representation( + model, + data, + target="one_step_forecast", + algorithm="gnn_explainer", + epochs=5, + lr=0.05, + index=0, + ) + assert forecast.target == "one_step_forecast" + assert forecast.index == 0 + assert forecast.node_mask is not None + assert forecast.edge_mask is not None + assert forecast.node_mask.shape[0] == num_nodes + assert forecast.edge_mask.shape == (num_edges,) + + +def test_explain_representation_reconstruction_target() -> None: + """Reconstruction target uses encode → decode via PyG Explainer.""" + torch.manual_seed(1) + model = _tiny_model() + data = _path_data() + result = explain_representation( + model, + data, + target="reconstruction", + algorithm="gnn_explainer", + epochs=5, + lr=0.05, + ) + assert result.target == "reconstruction" + assert result.node_mask is not None + assert result.edge_mask is not None + + +def test_explain_representation_rejects_hetero_data() -> None: + """HeteroData is rejected with an actionable message.""" + model = _tiny_model() + hetero = HeteroData() + hetero["n"].x = torch.randn(3, 3) + hetero["n", "to", "n"].edge_index = torch.tensor([[0, 1], [1, 2]], dtype=torch.long) + with pytest.raises(TypeError, match="homogeneous"): + explain_representation(model, hetero, epochs=2) + + +def test_explain_representation_rejects_integrated_gradients_without_captum() -> None: + """Missing Captum raises ImportError with the ``[explain]`` install hint.""" + from koopman_graph.analysis import explain as explain_mod + + model = _tiny_model() + data = _path_data() + with ( + patch.object( + explain_mod, + "_import_captum_stack", + side_effect=ImportError( + "algorithm='integrated_gradients' requires Captum; " + 'install with: pip install "koopman-graph[explain]"' + ), + ), + pytest.raises(ImportError, match=r"koopman-graph\[explain\]"), + ): + explain_representation( + model, + data, + algorithm="integrated_gradients", + ) + + +def test_explain_representation_integrated_gradients_when_captum_installed() -> None: + """IG path returns attribute masks when Captum is installed.""" + pytest.importorskip("captum") + torch.manual_seed(2) + model = _tiny_model() + data = _path_data() + result = explain_representation( + model, + data, + target="latent", + algorithm="integrated_gradients", + index=1, + ) + assert result.algorithm == "integrated_gradients" + assert result.target == "latent" + assert result.index == 1 + assert result.node_mask is not None + assert result.node_mask.ndim == 2 + assert result.node_mask.shape[0] == data.num_nodes + assert result.feature_mask is not None + assert torch.equal(result.feature_mask, result.node_mask) + + +def test_explain_extra_declares_captum() -> None: + """``[explain]`` optional dependency pins Captum.""" + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + text = (root / "pyproject.toml").read_text(encoding="utf-8") + assert "\nexplain = [\n" in text + assert '"captum>=0.7"' in text + + +def test_explain_representation_uses_torch_geometric_explainer() -> None: + """Implementation routes through PyG Explainer rather than a custom loop.""" + model = _tiny_model() + data = _path_data() + fake_explanation = MagicMock() + fake_explanation.node_mask = torch.ones(data.num_nodes, 1) + fake_explanation.edge_mask = torch.ones(data.num_edges) + fake_explanation.feature_mask = None + + with patch("torch_geometric.explain.Explainer") as explainer_cls: + instance = MagicMock(return_value=fake_explanation) + explainer_cls.return_value = instance + with patch("torch_geometric.explain.GNNExplainer") as gnn_cls: + gnn_cls.return_value = MagicMock() + result = explain_representation( + model, + data, + target="latent", + algorithm="gnn_explainer", + epochs=3, + ) + explainer_cls.assert_called_once() + gnn_cls.assert_called_once() + instance.assert_called_once() + assert result.node_mask is not None + assert torch.equal(result.node_mask, fake_explanation.node_mask) diff --git a/tests/test_hetero_fabric_lightning_ray.py b/tests/test_hetero_fabric_lightning_ray.py index fad62ad..9f5ec6e 100644 --- a/tests/test_hetero_fabric_lightning_ray.py +++ b/tests/test_hetero_fabric_lightning_ray.py @@ -196,6 +196,7 @@ def test_koopman_lightning_module_hetero_parity_vs_run_fit_loop() -> None: assert torch.allclose(param_a, param_b, rtol=0, atol=1e-5), name_a +@pytest.mark.ray def test_fit_ensemble_with_ray_hetero_two_members() -> None: """Ray ensemble members accept multiplex inputs once single-process fit does.""" ray = pytest.importorskip("ray") diff --git a/tests/test_hetero_reverse_relations.py b/tests/test_hetero_reverse_relations.py index 37a65c7..a9f25be 100644 --- a/tests/test_hetero_reverse_relations.py +++ b/tests/test_hetero_reverse_relations.py @@ -231,7 +231,7 @@ def test_checkpoint_round_trip_preserves_synthesize_flag(tmp_path: Path) -> None koopman_synthesize_reverse_relations=True, ) path = tmp_path / "rev.pt" - model.save(path) + model.save(path, format="legacy_pt") loaded = GraphKoopmanModel.load(path) assert loaded.synthesize_reverse_relations is True assert loaded.koopman.num_relations == 4 diff --git a/tests/test_hetero_serialization_dims.py b/tests/test_hetero_serialization_dims.py index eae835e..b0047c8 100644 --- a/tests/test_hetero_serialization_dims.py +++ b/tests/test_hetero_serialization_dims.py @@ -161,7 +161,7 @@ def test_incomplete_latent_dims_rejected(tmp_path: Path) -> None: """Missing type keys in latent_dims raise a clear re-save error.""" model = _rectangular_model() path = tmp_path / "rect.pt" - model.save(path) + model.save(path, format="legacy_pt") payload = torch.load(path, map_location="cpu", weights_only=False) payload["config"]["latent_dims"] = {"a": 2} broken = tmp_path / "incomplete.pt" @@ -176,7 +176,7 @@ def test_stripped_latent_dims_with_rectangular_weights_rejected( """Rectangular weights without config.latent_dims are rejected.""" model = _rectangular_model() path = tmp_path / "rect.pt" - model.save(path) + model.save(path, format="legacy_pt") payload = torch.load(path, map_location="cpu", weights_only=False) del payload["config"]["latent_dims"] broken = tmp_path / "broken.pt" @@ -189,7 +189,7 @@ def test_mismatched_self_factor_shape_rejected(tmp_path: Path) -> None: """latent_dims that disagree with self-factor shapes raise clearly.""" model = _rectangular_model() path = tmp_path / "rect.pt" - model.save(path) + model.save(path, format="legacy_pt") payload = torch.load(path, map_location="cpu", weights_only=False) payload["config"]["latent_dims"] = {"a": 3, "b": 3} broken = tmp_path / "shape_mismatch.pt" diff --git a/tests/test_package.py b/tests/test_package.py index 8374c61..7bb23a5 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -24,6 +24,7 @@ "DiffConvDecoder", "DiffConvEncoder", "E3EquivariantEncoder", + "FitCallback", "GATDecoder", "GATEncoder", "GNNDecoder", @@ -41,6 +42,7 @@ "KoopmanOperator", "KoopmanSpectrum", "MultiTrajectory", + "NoOpFitCallback", "RelGraphDecoder", "RelGraphEncoder", "SAGEDecoder", @@ -295,8 +297,10 @@ def test_analysis_secondaries_demoted_from_root() -> None: "detect_anomaly", "discrete_spectrum_at_delta_t", "dynamical_similarity", + "explain_representation", "koopman_std", "plot_spectrum", + "RepresentationExplanation", "spectral_residuals", "SpectralResidualReport", "spectrum_distance", @@ -308,6 +312,7 @@ def test_analysis_secondaries_demoted_from_root() -> None: from koopman_graph.analysis import ( AnomalyDetectionResult, + RepresentationExplanation, SpectralResidualReport, calibrate_anomaly_threshold, compute_generator_spectrum, @@ -315,6 +320,7 @@ def test_analysis_secondaries_demoted_from_root() -> None: detect_anomaly, discrete_spectrum_at_delta_t, dynamical_similarity, + explain_representation, koopman_std, plot_spectrum, spectral_residuals, @@ -322,6 +328,7 @@ def test_analysis_secondaries_demoted_from_root() -> None: ) assert AnomalyDetectionResult is not None + assert RepresentationExplanation is not None assert SpectralResidualReport is not None assert callable(calibrate_anomaly_threshold) assert callable(compute_generator_spectrum) @@ -329,6 +336,7 @@ def test_analysis_secondaries_demoted_from_root() -> None: assert callable(detect_anomaly) assert callable(discrete_spectrum_at_delta_t) assert callable(dynamical_similarity) + assert callable(explain_representation) assert callable(koopman_std) assert callable(plot_spectrum) assert callable(spectral_residuals) diff --git a/tests/test_patch_coverage_013.py b/tests/test_patch_coverage_013.py new file mode 100644 index 0000000..5356b64 --- /dev/null +++ b/tests/test_patch_coverage_013.py @@ -0,0 +1,998 @@ +"""Patch-coverage guards for the 0.13.0 Codecov patch gate (target >= 90%).""" + +from __future__ import annotations + +import builtins +import importlib +import json +import runpy +import sys +import zipfile +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from torch import nn +from torch_geometric.data import Data + +import koopman_graph.analysis.explain as explain_mod +import koopman_graph.cli.config as config_mod +import koopman_graph.cli.predict as predict_mod +import koopman_graph.cli.train as train_mod +import koopman_graph.distributed.lightning_module as lightning_mod +import koopman_graph.serialization as serialization_mod +import koopman_graph.tracking.tensorboard as tensorboard_mod +import koopman_graph.tuning.ray_tune as ray_tune_mod +import koopman_graph.tuning.search_spaces as search_spaces_mod +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.tracking.csv import CsvFitLogger +from koopman_graph.training.history import TrainingLossBreakdown + +cli_main_mod = importlib.import_module("koopman_graph.cli.main") + + +def _minimal_model_config(*, encoder: str = "gcn") -> dict[str, object]: + """Return the smallest valid CLI model config.""" + return { + "encoder": encoder, + "in_channels": 2, + "hidden_channels": 3, + "latent_dim": 2, + } + + +def _minimal_train_config() -> dict[str, object]: + """Return a valid train config for validation branch tests.""" + return { + "model": _minimal_model_config(), + "data": {"kind": "synthetic_path"}, + } + + +def _snapshot_sequence() -> GraphSnapshotSequence: + """Return a deterministic two-snapshot homogeneous sequence.""" + edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + snapshots = [ + Data(x=torch.full((2, 2), float(step)), edge_index=edge_index) + for step in range(2) + ] + return GraphSnapshotSequence(snapshots) + + +def _safe_meta(*, format_version: int = 1) -> dict[str, object]: + """Return valid safetensors metadata.""" + return { + "container": serialization_mod.SAFE_CONTAINER, + "format_version": format_version, + "package_version": "0.13.0", + } + + +def _write_safe_directory( + path: Path, + *, + meta_text: str | None = None, + config_text: str | None = None, + include_config: bool = True, +) -> None: + """Write enough safe-checkpoint members to reach parser validation.""" + path.mkdir() + (path / serialization_mod.SAFE_META_FILENAME).write_text( + meta_text if meta_text is not None else json.dumps(_safe_meta()), + encoding="utf-8", + ) + if include_config: + (path / serialization_mod.SAFE_CONFIG_FILENAME).write_text( + config_text if config_text is not None else "{}", + encoding="utf-8", + ) + (path / serialization_mod.SAFE_WEIGHTS_FILENAME).write_bytes(b"placeholder") + + +def _breakdown() -> TrainingLossBreakdown: + """Return a deterministic all-zero loss breakdown.""" + return TrainingLossBreakdown.zeros(torch.device("cpu")) + + +def _unsupported_explain_model(**overrides: object) -> SimpleNamespace: + """Return a model-shaped object for explain-MVP validation.""" + attributes: dict[str, object] = { + "n_delays": 1, + "control_dim": 0, + "uses_hetero_koopman": False, + "_uses_relgraph_encode": lambda: False, + "uses_hypergraph_koopman": False, + "_uses_hypergraph_encode": lambda: False, + "learns_pairwise_topology": False, + "adaptive_topology": None, + } + attributes.update(overrides) + return SimpleNamespace(**attributes) + + +@pytest.mark.parametrize("value", [True, 1.2, "1"]) +def test_train_require_int_rejects_non_integer_values(value: object) -> None: + """The CLI integer validator rejects bool and non-integer values.""" + with pytest.raises(config_mod.ConfigError, match="model.value must be an int"): + train_mod._require_int(value, path="model.value") + + +@pytest.mark.parametrize("value", [False, "1.0", None]) +def test_train_require_float_rejects_non_numeric_values(value: object) -> None: + """The CLI float validator rejects bool and non-numeric values.""" + with pytest.raises(config_mod.ConfigError, match="must be a number"): + train_mod._require_float(value, path="fit.lr") + + +def test_train_require_float_accepts_integer() -> None: + """The CLI float validator converts an integer to float.""" + assert train_mod._require_float(2, path="fit.lr") == pytest.approx( + 2.0, rel=0.0, abs=0.0 + ) + + +@pytest.mark.parametrize( + ("model_config", "message"), + [ + ({}, "model.encoder is required"), + ({"encoder": "transformer"}, "Unsupported model.encoder"), + ({"encoder": "gcn", "decoder": 1}, "model.decoder must be a string"), + ( + {"encoder": "gcn", "decoder": "gat"}, + "must match model.encoder", + ), + ({"encoder": "gcn"}, "Missing required key"), + ( + { + "encoder": "gcn", + "in_channels": 0, + "hidden_channels": 2, + "latent_dim": 2, + }, + "must be positive", + ), + ], +) +def test_train_encoder_decoder_rejects_invalid_configs( + model_config: dict[str, object], + message: str, +) -> None: + """Encoder/decoder config guards raise actionable errors.""" + with pytest.raises(config_mod.ConfigError, match=message): + train_mod._build_encoder_decoder(model_config) + + +@pytest.mark.parametrize("kind", ["gat", "sage"]) +def test_train_builds_non_gcn_encoder_decoder_branches(kind: str) -> None: + """GAT and GraphSAGE CLI model branches construct matched peers.""" + encoder, decoder, latent_dim = train_mod._build_encoder_decoder( + _minimal_model_config(encoder=kind) + ) + assert encoder is not None + assert decoder is not None + assert latent_dim == 2 + + +def test_train_model_defaults_time_step() -> None: + """An omitted time step is forwarded to the model as 1.0.""" + sentinel = object() + with patch.object( + train_mod, + "GraphKoopmanModel", + return_value=sentinel, + ) as model_cls: + result = train_mod.build_model_from_config(_minimal_model_config()) + assert result is sentinel + assert model_cls.call_args.kwargs["time_step"] == pytest.approx( + 1.0, rel=0.0, abs=0.0 + ) + + +def test_train_synthetic_sequence_rejects_invalid_dimensions() -> None: + """Synthetic path data enforces its node, time, and feature minima.""" + with pytest.raises(config_mod.ConfigError, match="num_nodes>=2"): + train_mod._build_synthetic_path_sequence( + {"num_nodes": 1, "num_timesteps": 2, "feature_dim": 1} + ) + + +def test_train_cached_sequence_loads_supported_payloads(tmp_path: Path) -> None: + """Cached loaders accept sequence objects and non-empty Data lists.""" + path = tmp_path / "sequence.pt" + sequence = _snapshot_sequence() + torch.save(sequence, path) + loaded_sequence = train_mod.load_cached_sequence({"path": path}) + assert len(loaded_sequence) == 2 + + torch.save(list(sequence), path) + loaded_list = train_mod.load_cached_sequence({"path": path}) + assert len(loaded_list) == 2 + + +def test_train_cached_sequence_rejects_missing_and_invalid_payloads( + tmp_path: Path, +) -> None: + """Cached loaders reject missing paths and unsupported payload shapes.""" + missing = tmp_path / "missing.pt" + with pytest.raises(config_mod.ConfigError, match="not found"): + train_mod.load_cached_sequence({"path": missing}) + + invalid = tmp_path / "invalid.pt" + torch.save({"not": "a sequence"}, invalid) + with pytest.raises(config_mod.ConfigError, match="must be a GraphSnapshotSequence"): + train_mod.load_cached_sequence({"path": invalid}) + + +def test_train_sequence_dispatches_cached_and_rejects_unknown_kind() -> None: + """Data-kind dispatch covers the cached and unknown branches.""" + sentinel = _snapshot_sequence() + with patch.object(train_mod, "load_cached_sequence", return_value=sentinel): + result = train_mod.build_sequence_from_config( + {"kind": "cached_sequence", "path": "ignored.pt"} + ) + assert result is sentinel + + with pytest.raises(config_mod.ConfigError, match="Unsupported data.kind"): + train_mod.build_sequence_from_config({"kind": "remote_code"}) + + +def test_train_rejects_unknown_checkpoint_format() -> None: + """Training refuses checkpoint formats outside the explicit allowlist.""" + fake_model = MagicMock() + with ( + patch.object(train_mod, "build_model_from_config", return_value=fake_model), + patch.object( + train_mod, + "build_sequence_from_config", + return_value=_snapshot_sequence(), + ), + pytest.raises(config_mod.ConfigError, match="checkpoint.format"), + ): + train_mod.run_train( + { + "model": {}, + "data": {}, + "checkpoint": {"format": "unsafe_pickle"}, + } + ) + + +def test_serialization_json_and_device_helpers_cover_tuple_paths() -> None: + """Tuple configs become lists and explicit devices become strings.""" + assert serialization_mod._json_ready({"shape": (1, (2, 3))}) == { + "shape": [1, [2, 3]] + } + assert serialization_mod._json_ready([{"nested": (4,)}]) == [{"nested": [4]}] + assert serialization_mod._resolve_safetensors_device(torch.device("cpu")) == "cpu" + + +def test_estimator_ddp_strategy_rejects_callbacks() -> None: + """Distributed ``strategy='ddp'`` refuses ``callbacks=`` until supported.""" + from koopman_graph.model.estimator import GraphKoopmanModel + from koopman_graph.nn import GNNDecoder, GNNEncoder + from koopman_graph.training.callbacks import NoOpFitCallback + + model = GraphKoopmanModel( + encoder=GNNEncoder(2, 4, 3, num_layers=1), + decoder=GNNDecoder(3, 4, 2, num_layers=1), + latent_dim=3, + time_step=1.0, + ) + edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long) + sequence = GraphSnapshotSequence( + [Data(x=torch.randn(2, 2), edge_index=edge_index) for _ in range(4)] + ) + with pytest.raises(ValueError, match="callbacks"): + model.fit( + sequence, + epochs=1, + strategy="ddp", + callbacks=[NoOpFitCallback()], + ) + + +def test_serialization_zip_writer_rejects_existing_directory( + tmp_path: Path, +) -> None: + """A zip checkpoint cannot overwrite an existing directory.""" + destination = tmp_path / "model.kgckpt" + destination.mkdir() + with pytest.raises(ValueError, match="must be a file"): + serialization_mod._save_safetensors_v1_zip(object(), destination) + + +def test_serialization_zip_probe_covers_false_and_bad_zip(tmp_path: Path) -> None: + """The safe-zip probe rejects absent files and malformed archives.""" + missing = tmp_path / "missing.kgckpt" + assert not serialization_mod._is_safetensors_v1_zip(missing) + + candidate = tmp_path / "candidate.kgckpt" + candidate.write_bytes(b"not a zip") + with ( + patch.object(serialization_mod.zipfile, "is_zipfile", return_value=True), + patch.object( + serialization_mod.zipfile, + "ZipFile", + side_effect=zipfile.BadZipFile("bad"), + ), + ): + assert not serialization_mod._is_safetensors_v1_zip(candidate) + + +def test_serialization_save_rejects_unknown_format(tmp_path: Path) -> None: + """The checkpoint writer rejects unknown format tokens.""" + with pytest.raises(ValueError, match="Unsupported checkpoint format"): + serialization_mod.save_checkpoint( + object(), + tmp_path / "model", + format="unknown", # type: ignore[arg-type] + ) + + +def test_serialization_directory_loader_rejects_missing_config( + tmp_path: Path, +) -> None: + """A safe directory must contain config.json.""" + checkpoint = tmp_path / "missing_config" + _write_safe_directory(checkpoint, include_config=False) + with pytest.raises(ValueError, match=serialization_mod.SAFE_CONFIG_FILENAME): + serialization_mod._load_safetensors_v1_directory( + checkpoint, + map_location=None, + ) + + +@pytest.mark.parametrize( + ("meta_text", "message"), + [ + ("{", "Invalid meta.json"), + ("[]", "meta.json must be a JSON object"), + (json.dumps(_safe_meta(format_version=999)), "format_version"), + ], +) +def test_serialization_directory_loader_rejects_invalid_metadata( + tmp_path: Path, + meta_text: str, + message: str, +) -> None: + """Malformed, non-object, and unsupported metadata are rejected.""" + checkpoint = tmp_path / "bad_meta" + _write_safe_directory(checkpoint, meta_text=meta_text) + with pytest.raises(ValueError, match=message): + serialization_mod._load_safetensors_v1_directory( + checkpoint, + map_location=None, + ) + + +def test_serialization_directory_loader_rejects_non_object_config( + tmp_path: Path, +) -> None: + """Safe config JSON must decode to an object.""" + checkpoint = tmp_path / "bad_config" + _write_safe_directory(checkpoint, config_text="[]") + with pytest.raises(ValueError, match="config.json must be a JSON object"): + serialization_mod._load_safetensors_v1_directory( + checkpoint, + map_location=None, + ) + + +def test_serialization_zip_loader_rejects_missing_members_and_bad_zip( + tmp_path: Path, +) -> None: + """The direct zip loader reports missing members and malformed archives.""" + incomplete = tmp_path / "incomplete.kgckpt" + with zipfile.ZipFile(incomplete, mode="w") as archive: + archive.writestr(serialization_mod.SAFE_META_FILENAME, "{}") + with pytest.raises(ValueError, match="missing members"): + serialization_mod._load_safetensors_v1_zip( + incomplete, + map_location=None, + ) + + malformed = tmp_path / "malformed.kgckpt" + malformed.write_bytes(b"not a zip") + with pytest.raises(ValueError, match="Invalid safetensors_v1 zip"): + serialization_mod._load_safetensors_v1_zip( + malformed, + map_location=None, + ) + + +def test_serialization_load_rejects_unsupported_path_type() -> None: + """An existing path that is neither file nor directory is rejected.""" + destination = MagicMock() + destination.exists.return_value = True + destination.is_dir.return_value = False + destination.is_file.return_value = False + with ( + patch.object(serialization_mod, "Path", return_value=destination), + pytest.raises(ValueError, match="Unsupported checkpoint path type"), + ): + serialization_mod.load_checkpoint("special-device") + + +def test_tensorboard_summary_writer_import_success() -> None: + """The TensorBoard peer import returns SummaryWriter when available.""" + writer_cls = object() + fake_module = ModuleType("torch.utils.tensorboard") + fake_module.SummaryWriter = writer_cls # type: ignore[attr-defined] + with patch.dict(sys.modules, {"torch.utils.tensorboard": fake_module}): + assert tensorboard_mod._require_summary_writer() is writer_cls + + +def test_tensorboard_logger_lifecycle_with_mocked_writer(tmp_path: Path) -> None: + """Mocked TensorBoard logging covers start, train/val, flush, and close.""" + writer = MagicMock() + writer_cls = MagicMock(return_value=writer) + logger = tensorboard_mod.TensorBoardFitLogger( + tmp_path / "events", + comment="coverage", + ) + + with pytest.raises(RuntimeError, match="before on_fit_start"): + logger.on_epoch_end( + epoch=0, + train_breakdown=_breakdown(), + val_breakdown=None, + history_so_far=None, + ) + + with patch.object( + tensorboard_mod, + "_require_summary_writer", + return_value=writer_cls, + ): + logger.on_fit_start(model=object(), fit_kwargs={}) + logger.on_epoch_end( + epoch=3, + train_breakdown=_breakdown(), + val_breakdown=_breakdown(), + history_so_far=None, + ) + logger.on_fit_end(history=MagicMock()) + + writer_cls.assert_called_once_with( + log_dir=str(tmp_path / "events"), + comment="coverage", + ) + assert writer.add_scalar.call_count == 22 + assert writer.flush.call_count == 2 + writer.close.assert_called_once_with() + assert logger._writer is None + + +def test_config_mapping_validation_errors() -> None: + """Mapping validation rejects non-mappings and non-string keys.""" + with pytest.raises(config_mod.ConfigError, match="must be a mapping"): + config_mod._require_mapping([], path="config") + with pytest.raises(config_mod.ConfigError, match="keys must be strings"): + config_mod._require_mapping({1: "bad"}, path="config") + + +def test_config_load_rejects_missing_and_invalid_json(tmp_path: Path) -> None: + """Missing files and malformed JSON produce ConfigError.""" + with pytest.raises(config_mod.ConfigError, match="Config file not found"): + config_mod.load_config(tmp_path / "missing.json") + + invalid = tmp_path / "invalid.json" + invalid.write_text("{", encoding="utf-8") + with pytest.raises(config_mod.ConfigError, match="Invalid JSON"): + config_mod.load_config(invalid) + + +def test_config_load_yaml_import_and_parse_errors(tmp_path: Path) -> None: + """YAML loading reports missing PyYAML and parser failures.""" + path = tmp_path / "invalid.yaml" + path.write_text("invalid: [", encoding="utf-8") + real_import = builtins.__import__ + + def blocked_import( + name: str, + globals: object = None, + locals: object = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> object: + if name == "yaml": + raise ImportError("simulated missing yaml") + return real_import(name, globals, locals, fromlist, level) + + with ( + patch.object(builtins, "__import__", side_effect=blocked_import), + pytest.raises(ImportError, match="YAML configs require PyYAML"), + ): + config_mod.load_config(path) + + class FakeYamlError(Exception): + """Synthetic YAML parser error.""" + + fake_yaml = ModuleType("yaml") + fake_yaml.YAMLError = FakeYamlError # type: ignore[attr-defined] + fake_yaml.safe_load = MagicMock(side_effect=FakeYamlError("bad yaml")) # type: ignore[attr-defined] + with ( + patch.dict(sys.modules, {"yaml": fake_yaml}), + pytest.raises(config_mod.ConfigError, match="Invalid YAML"), + ): + config_mod.load_config(path) + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ({"model": {}}, "Missing required section: data"), + ( + {"model": {}, "data": {}}, + "Missing required key: data.kind", + ), + ( + {"model": {}, "data": {"kind": "cached_sequence"}}, + "requires data.path", + ), + ( + { + "model": {}, + "data": {"kind": "synthetic_path"}, + "checkpoint": {}, + }, + "checkpoint.path", + ), + ], +) +def test_config_validate_required_data_and_checkpoint_fields( + config: dict[str, object], + message: str, +) -> None: + """Train config validation covers missing required fields.""" + with pytest.raises(config_mod.ConfigError, match=message): + config_mod.validate_train_config(config) + + +class _NonTensorForecast: + """Model-shaped callable that returns a non-tensor forecast.""" + + def __call__(self, data: Data) -> dict[str, Data]: + return {"decoded": data} + + +def test_explain_wrapper_rejects_non_tensor_forecast() -> None: + """One-step explanation requires a tensor model decode.""" + wrapper = explain_mod._HomogeneousExplainModule( + _NonTensorForecast(), # type: ignore[arg-type] + target="one_step_forecast", + ) + with pytest.raises(TypeError, match="requires a tensor decode"): + wrapper( + torch.ones(2, 2), + torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + ) + + +@pytest.mark.parametrize( + ("model", "message"), + [ + (_unsupported_explain_model(control_dim=1), "control_dim"), + ( + _unsupported_explain_model(uses_hetero_koopman=True), + "RelGraph / hetero_graph", + ), + ( + _unsupported_explain_model(uses_hypergraph_koopman=True), + "hypergraph", + ), + ( + _unsupported_explain_model(adaptive_topology=object()), + "adaptive / learned", + ), + ], +) +def test_explain_rejects_unsupported_model_modes( + model: SimpleNamespace, + message: str, +) -> None: + """Explain-MVP validation covers control, hetero, hypergraph, and adaptive.""" + with pytest.raises(ValueError, match=message): + explain_mod._reject_unsupported_model(model) # type: ignore[arg-type] + + +def test_explain_rejects_invalid_data_inputs() -> None: + """Explain data validation rejects non-Data and missing graph fields.""" + with pytest.raises(TypeError, match="torch_geometric.data.Data"): + explain_mod._reject_unsupported_data(object()) # type: ignore[arg-type] + with pytest.raises(ValueError, match="Data.x and Data.edge_index"): + explain_mod._reject_unsupported_data(Data()) + + +@pytest.mark.parametrize( + ("epochs", "lr", "message"), + [(0, 0.01, "epochs"), (1, 0.0, "lr")], +) +def test_explain_gnn_explainer_validates_hyperparameters( + epochs: int, + lr: float, + message: str, +) -> None: + """GNNExplainer rejects invalid epochs and learning rates.""" + data = Data( + x=torch.ones(2, 2), + edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + ) + with pytest.raises(ValueError, match=message): + explain_mod._run_gnn_explainer( + MagicMock(), + data, + target="latent", + index=None, + epochs=epochs, + lr=lr, + ) + + +def test_explain_captum_import_error_hint() -> None: + """The direct Captum importer adds the optional-extra install hint.""" + real_import = builtins.__import__ + + def blocked_import( + name: str, + globals: object = None, + locals: object = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> object: + if name == "captum": + raise ImportError("simulated missing captum") + return real_import(name, globals, locals, fromlist, level) + + with ( + patch.object(builtins, "__import__", side_effect=blocked_import), + pytest.raises(ImportError, match=r"koopman-graph\[explain\]"), + ): + explain_mod._import_captum_stack() + + +def test_explain_integrated_gradients_rejects_negative_index() -> None: + """Integrated gradients validates the selected node index.""" + data = Data( + x=torch.ones(2, 2), + edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + ) + with ( + patch.object( + explain_mod, + "_import_captum_stack", + return_value=(MagicMock(), MagicMock(), MagicMock()), + ), + pytest.raises(ValueError, match="index must be >= 0"), + ): + explain_mod._run_integrated_gradients( + nn.Identity(), # type: ignore[arg-type] + data, + target="latent", + index=-1, + ) + + +def test_explain_integrated_gradients_discards_non_tensor_edge_mask() -> None: + """Integrated gradients normalizes a non-tensor edge mask to None.""" + + class FakeExplainer: + """Return a deterministic explanation without invoking the wrapper.""" + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + + def __call__( + self, + x: torch.Tensor, + edge_index: torch.Tensor, + *, + index: int, + ) -> SimpleNamespace: + del edge_index, index + return SimpleNamespace(node_mask=torch.ones_like(x), edge_mask=object()) + + data = Data( + x=torch.ones(2, 2), + edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.long), + ) + with patch.object( + explain_mod, + "_import_captum_stack", + return_value=(MagicMock(), FakeExplainer, MagicMock()), + ): + result = explain_mod._run_integrated_gradients( + nn.Identity(), # type: ignore[arg-type] + data, + target="latent", + index=1, + ) + assert result.edge_mask is None + assert result.node_mask is not None + + +def test_predict_sequence_loader_covers_config_and_file_branches( + tmp_path: Path, +) -> None: + """Prediction data loading covers config, cache, and suffix validation.""" + config_path = tmp_path / "data.json" + config_path.write_text("{}", encoding="utf-8") + sentinel = _snapshot_sequence() + + with ( + patch.object(predict_mod, "load_config", return_value={"data": []}), + pytest.raises(config_mod.ConfigError, match="config.data must be a mapping"), + ): + predict_mod.load_sequence_from_path(config_path) + + with ( + patch.object( + predict_mod, + "load_config", + return_value={"kind": "synthetic_path"}, + ), + patch.object( + predict_mod, + "build_sequence_from_config", + return_value=sentinel, + ), + ): + assert predict_mod.load_sequence_from_path(config_path) is sentinel + + with ( + patch.object(predict_mod, "load_config", return_value={"other": "value"}), + pytest.raises(config_mod.ConfigError, match="must contain a 'data' section"), + ): + predict_mod.load_sequence_from_path(config_path) + + cache_path = tmp_path / "data.pt" + cache_path.write_bytes(b"placeholder") + with patch.object( + predict_mod, + "load_cached_sequence", + return_value=sentinel, + ) as cached_loader: + assert predict_mod.load_sequence_from_path(cache_path) is sentinel + cached_loader.assert_called_once() + + unsupported = tmp_path / "data.csv" + unsupported.write_text("x\n", encoding="utf-8") + with pytest.raises(config_mod.ConfigError, match="Unsupported data path suffix"): + predict_mod.load_sequence_from_path(unsupported) + + with pytest.raises(config_mod.ConfigError, match="Data path not found"): + predict_mod.load_sequence_from_path(tmp_path / "missing.pt") + + +def test_predict_run_rejects_invalid_steps_empty_data_and_empty_forecast( + tmp_path: Path, +) -> None: + """Prediction validates horizon, input sequence, and forecast output.""" + with pytest.raises(config_mod.ConfigError, match="--steps must be >= 1"): + predict_mod.run_predict( + checkpoint="model", + data="data", + steps=0, + out=tmp_path / "out.pt", + ) + + with ( + patch.object(predict_mod, "load_checkpoint", return_value=MagicMock()), + patch.object(predict_mod, "load_sequence_from_path", return_value=[]), + pytest.raises(config_mod.ConfigError, match="Data sequence is empty"), + ): + predict_mod.run_predict( + checkpoint="model", + data="data", + steps=1, + out=tmp_path / "out.pt", + ) + + model = MagicMock() + model.predict.return_value = [] + with ( + patch.object(predict_mod, "load_checkpoint", return_value=model), + patch.object( + predict_mod, + "load_sequence_from_path", + return_value=[Data(x=torch.ones(1, 1))], + ), + pytest.raises(config_mod.ConfigError, match="empty forecast"), + ): + predict_mod.run_predict( + checkpoint="model", + data="data", + steps=1, + out=tmp_path / "out.pt", + ) + + +def test_predict_handler_prints_path_without_summary( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + """The predict handler reports successful forecasts without shape metadata.""" + output = tmp_path / "forecast.pt" + args = SimpleNamespace( + checkpoint="model", + data="data", + steps=2, + out=output, + ) + with patch.object(predict_mod, "run_predict", return_value=(output, None)): + assert predict_mod.handle_predict(args) == 0 # type: ignore[arg-type] + assert f"wrote forecast: {output} (steps=2)" in capsys.readouterr().out + + +def test_ray_import_helpers_cover_success_failure_and_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ray helpers set the uv flag and wrap lazy-import failures.""" + monkeypatch.delenv("RAY_ENABLE_UV_RUN_RUNTIME_ENV", raising=False) + ray_tune_mod._prefer_driver_interpreter_for_uv_run() + assert sys.modules["os"].environ["RAY_ENABLE_UV_RUN_RUNTIME_ENV"] == "0" + + fake_ray = ModuleType("ray") + fake_tune = ModuleType("ray.tune") + fake_run_config = object() + fake_train = SimpleNamespace(RunConfig=fake_run_config) + + def import_ray(name: str) -> object: + return { + "ray": fake_ray, + "ray.tune": fake_tune, + "ray.train": fake_train, + }[name] + + with patch.object( + ray_tune_mod.importlib, + "import_module", + side_effect=import_ray, + ): + assert ray_tune_mod._import_ray_tune() == ( + fake_ray, + fake_tune, + fake_run_config, + ) + + with ( + patch.object( + ray_tune_mod.importlib, + "import_module", + side_effect=ImportError("missing ray"), + ), + pytest.raises(ImportError, match=r"koopman-graph\[ray\]"), + ): + ray_tune_mod._import_ray_tune() + + +def test_ray_tune_initializes_uninitialized_runtime() -> None: + """The wrapper initializes Ray with caller-provided kwargs when needed.""" + fake_ray = ModuleType("ray") + fake_ray.is_initialized = MagicMock(return_value=False) # type: ignore[attr-defined] + fake_ray.init = MagicMock() # type: ignore[attr-defined] + fake_tuner = MagicMock() + fake_tuner.fit.return_value = "results" + fake_tune = ModuleType("ray.tune") + fake_tune.TuneConfig = MagicMock(return_value=object()) # type: ignore[attr-defined] + fake_tune.Tuner = MagicMock(return_value=fake_tuner) # type: ignore[attr-defined] + + with patch.object( + ray_tune_mod, + "_import_ray_tune", + return_value=(fake_ray, fake_tune, MagicMock()), + ): + result = ray_tune_mod.run_ray_tune( + lambda config: config, + {"lr": 0.01}, + ray_init_kwargs={"local_mode": True}, + ) + assert result == "results" + fake_ray.init.assert_called_once_with(local_mode=True) + + +def test_search_space_import_helper_success_and_failure() -> None: + """The Ray Tune search-space importer handles both lazy-import outcomes.""" + fake_tune = ModuleType("ray.tune") + with patch.object( + search_spaces_mod.importlib, + "import_module", + return_value=fake_tune, + ): + assert search_spaces_mod._import_tune() is fake_tune + + with ( + patch.object( + search_spaces_mod.importlib, + "import_module", + side_effect=ImportError("missing ray"), + ), + pytest.raises(ImportError, match=r"koopman-graph\[ray\]"), + ): + search_spaces_mod._import_tune() + + +@pytest.mark.parametrize( + ("function", "kwargs", "message"), + [ + ( + search_spaces_mod.example_lr_loguniform_space, + {"lr_low": 0.0}, + "must be positive", + ), + ( + search_spaces_mod.example_lr_latent_dim_space, + {"latent_dims": ()}, + "at least one", + ), + ( + search_spaces_mod.example_lr_latent_dim_space, + {"latent_dims": (2, 0)}, + "entries must be >= 1", + ), + ], +) +def test_search_space_validates_positive_domains( + function: object, + kwargs: dict[str, object], + message: str, +) -> None: + """Example search spaces reject non-positive and empty domains.""" + with pytest.raises(ValueError, match=message): + function(**kwargs) # type: ignore[operator] + + +def test_cli_entry_points_cover_help_and_main_guards( + capsys: pytest.CaptureFixture[str], +) -> None: + """Direct and module entry points print help and exit successfully.""" + assert cli_main_mod.main([]) == 0 + assert "usage: koopman-graph" in capsys.readouterr().out + + main_path = Path(cli_main_mod.__file__) + with ( + patch.object(sys, "argv", [str(main_path)]), + pytest.raises(SystemExit) as main_exit, + ): + runpy.run_path(str(main_path), run_name="__main__") + assert main_exit.value.code == 0 + + module_path = main_path.with_name("__main__.py") + with ( + patch.object(sys, "argv", [str(module_path)]), + pytest.raises(SystemExit) as module_exit, + ): + runpy.run_module("koopman_graph.cli", run_name="__main__", alter_sys=True) + assert module_exit.value.code == 0 + + +def test_csv_logger_rejects_epoch_before_start(tmp_path: Path) -> None: + """CSV logging requires on_fit_start before the first epoch.""" + logger = CsvFitLogger(tmp_path / "fit.csv") + with pytest.raises(RuntimeError, match="before on_fit_start"): + logger.on_epoch_end( + epoch=0, + train_breakdown=_breakdown(), + val_breakdown=None, + history_so_far=None, + ) + + +def test_lightning_export_forces_legacy_checkpoint_format(tmp_path: Path) -> None: + """Lightning's Trainer hook exports an unwrapped legacy checkpoint.""" + fake_lightning = SimpleNamespace(LightningModule=nn.Module) + with patch.object(lightning_mod, "_import_lightning", return_value=fake_lightning): + module_cls = lightning_mod._build_koopman_lightning_module() + model = nn.Linear(2, 2) + module = module_cls(model) + destination = tmp_path / "lightning.pt" + + with patch.object(serialization_mod, "save_checkpoint") as save_checkpoint: + module.export_format1_checkpoint(destination) + save_checkpoint.assert_called_once_with( + model, + destination, + format="legacy_pt", + ) diff --git a/tests/test_safe_checkpoint.py b/tests/test_safe_checkpoint.py new file mode 100644 index 0000000..be2b824 --- /dev/null +++ b/tests/test_safe_checkpoint.py @@ -0,0 +1,278 @@ +"""Parity tests: safetensors_v1 directory / .kgckpt ≡ legacy_pt pickle.""" + +from __future__ import annotations + +import zipfile +from collections.abc import Callable +from pathlib import Path + +import pytest +import torch +from torch_geometric.data import Data + +from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.serialization import ( + SAFE_CONFIG_FILENAME, + SAFE_META_FILENAME, + SAFE_WEIGHTS_FILENAME, + SAFE_ZIP_MEMBER_NAMES, + build_model_config, + load_checkpoint, + save_checkpoint, +) + +# Safe path detaches/clones to contiguous CPU tensors; legacy pickle preserves +# the same float32 values. Bit-identical agreement is required — any drift is a +# serialization bug, not float noise (rtol=0, atol=0). +_PARITY_RTOL = 0.0 +_PARITY_ATOL = 0.0 + + +def _json_ready(value: object) -> object: + """Mirror serialization's tuple→list coercion for config comparison.""" + if isinstance(value, dict): + return {str(key): _json_ready(item) for key, item in value.items()} + if isinstance(value, tuple): + return [_json_ready(item) for item in value] + if isinstance(value, list): + return [_json_ready(item) for item in value] + return value + + +def _assert_state_dict_close( + left: dict[str, torch.Tensor], + right: dict[str, torch.Tensor], +) -> None: + """Assert two state dicts match exactly on CPU float tensors.""" + assert set(left) == set(right) + for key in left: + torch.testing.assert_close( + left[key].detach().cpu(), + right[key].detach().cpu(), + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + +def _assert_safe_legacy_parity( + model: GraphKoopmanModel, + tmp_path: Path, + *, + predict: Callable[[GraphKoopmanModel], object], +) -> None: + """Save both containers, reload, and compare config / weights / predict.""" + model.eval() + with torch.no_grad(): + before = predict(model) + + safe_dir = tmp_path / "safe_ckpt" + legacy_path = tmp_path / "legacy.pt" + save_checkpoint(model, safe_dir, format="safetensors_v1") + save_checkpoint(model, legacy_path, format="legacy_pt") + + assert safe_dir.is_dir() + assert (safe_dir / SAFE_META_FILENAME).is_file() + assert (safe_dir / SAFE_CONFIG_FILENAME).is_file() + assert (safe_dir / SAFE_WEIGHTS_FILENAME).is_file() + assert legacy_path.is_file() + + safe_loaded = load_checkpoint(safe_dir) + legacy_loaded = load_checkpoint(legacy_path) + assert not safe_loaded.training + assert not legacy_loaded.training + + original_config = _json_ready(build_model_config(model)) + assert _json_ready(build_model_config(safe_loaded)) == original_config + assert _json_ready(build_model_config(legacy_loaded)) == original_config + + _assert_state_dict_close(model.state_dict(), safe_loaded.state_dict()) + _assert_state_dict_close(model.state_dict(), legacy_loaded.state_dict()) + _assert_state_dict_close(safe_loaded.state_dict(), legacy_loaded.state_dict()) + + with torch.no_grad(): + safe_pred = predict(safe_loaded) + legacy_pred = predict(legacy_loaded) + + _assert_predictions_close(before, safe_pred) + _assert_predictions_close(before, legacy_pred) + _assert_predictions_close(safe_pred, legacy_pred) + + +def _assert_predictions_close(left: object, right: object) -> None: + """Compare nested prediction tensors / Data.x values exactly.""" + if isinstance(left, list) and isinstance(right, list): + assert len(left) == len(right) + for left_item, right_item in zip(left, right, strict=True): + _assert_predictions_close(left_item, right_item) + return + if isinstance(left, Data) and isinstance(right, Data): + torch.testing.assert_close( + left.x, + right.x, + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + return + if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor): + torch.testing.assert_close( + left, + right, + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + return + msg = f"Unsupported prediction types: {type(left)!r} vs {type(right)!r}" + raise TypeError(msg) + + +def test_safe_vs_legacy_parity_discrete_gcn( + scaling_sequence: GraphSnapshotSequence, + tmp_path: Path, +) -> None: + """Discrete GCN: safetensors_v1 and legacy_pt agree after short fit.""" + encoder = GNNEncoder(in_channels=3, hidden_channels=16, latent_dim=8) + decoder = GNNDecoder(latent_dim=8, hidden_channels=16, out_channels=3) + model = GraphKoopmanModel( + encoder=encoder, + decoder=decoder, + latent_dim=8, + time_step=0.1, + ) + torch.manual_seed(0) + model.fit(scaling_sequence, epochs=2, lr=1e-2) + + def predict(loaded: GraphKoopmanModel) -> list[torch.Tensor]: + return [ + graph.x.detach().clone() + for graph in loaded.predict(scaling_sequence[0], steps=3) + ] + + _assert_safe_legacy_parity(model, tmp_path, predict=predict) + + +def test_safe_vs_legacy_parity_continuous( + synthetic_edge_index: torch.Tensor, + tmp_path: Path, +) -> None: + """Continuous dissipative model: both containers preserve predict_at.""" + torch.manual_seed(1) + encoder = GNNEncoder(in_channels=3, hidden_channels=8, latent_dim=4, num_layers=2) + decoder = GNNDecoder(latent_dim=4, hidden_channels=8, out_channels=3, num_layers=2) + model = GraphKoopmanModel( + encoder, + decoder, + latent_dim=4, + time_step=0.1, + dynamics_mode="continuous", + koopman_parameterization="dissipative", + ) + graph = Data(x=torch.randn(5, 3), edge_index=synthetic_edge_index) + + def predict(loaded: GraphKoopmanModel) -> list[Data]: + return list(loaded.predict_at(graph, step_deltas=[0.1, 0.1])) + + _assert_safe_legacy_parity(model, tmp_path, predict=predict) + + +def test_safe_vs_legacy_parity_delay_embedding( + synthetic_edge_index: torch.Tensor, + tmp_path: Path, +) -> None: + """Delay-embedding GCN: both containers preserve n_delays and predict.""" + torch.manual_seed(2) + # in_channels = n_delays * feature_dim = 2 * 3 + model = GraphKoopmanModel( + encoder=GNNEncoder(in_channels=6, hidden_channels=8, latent_dim=4), + decoder=GNNDecoder(latent_dim=4, hidden_channels=8, out_channels=3), + latent_dim=4, + time_step=1.0, + n_delays=2, + ) + x0 = torch.randn(5, 3) + snapshots = [ + Data(x=x0 * (0.95**t), edge_index=synthetic_edge_index) for t in range(6) + ] + sequence = GraphSnapshotSequence(snapshots) + model.fit(sequence, epochs=1, lr=1e-2) + assert model.n_delays == 2 + + def predict(loaded: GraphKoopmanModel) -> list[torch.Tensor]: + assert loaded.n_delays == 2 + return [ + graph.x.detach().clone() for graph in loaded.predict(sequence[0], steps=2) + ] + + _assert_safe_legacy_parity(model, tmp_path, predict=predict) + + +def _small_discrete_model() -> GraphKoopmanModel: + """Untrained tiny GCN for zip-container unit tests.""" + torch.manual_seed(3) + return GraphKoopmanModel( + encoder=GNNEncoder( + in_channels=3, hidden_channels=8, latent_dim=4, num_layers=2 + ), + decoder=GNNDecoder( + latent_dim=4, hidden_channels=8, out_channels=3, num_layers=2 + ), + latent_dim=4, + time_step=0.1, + ) + + +def test_kgckpt_zip_round_trip_and_directory_parity( + synthetic_edge_index: torch.Tensor, + tmp_path: Path, +) -> None: + """``.kgckpt`` zip matches directory safetensors_v1 (exact tensors).""" + model = _small_discrete_model() + model.eval() + graph = Data(x=torch.randn(5, 3), edge_index=synthetic_edge_index) + with torch.no_grad(): + before = [g.x.detach().clone() for g in model.predict(graph, steps=2)] + + safe_dir = tmp_path / "safe_dir" + kgckpt_path = tmp_path / "model.kgckpt" + save_checkpoint(model, safe_dir, format="safetensors_v1") + save_checkpoint(model, kgckpt_path, format="safetensors_v1") + + assert kgckpt_path.is_file() + assert zipfile.is_zipfile(kgckpt_path) + with zipfile.ZipFile(kgckpt_path, mode="r") as archive: + names = set(archive.namelist()) + assert SAFE_ZIP_MEMBER_NAMES.issubset(names) + assert not any(name.endswith(".pt") for name in names) + + dir_loaded = load_checkpoint(safe_dir) + zip_loaded = load_checkpoint(kgckpt_path) + assert not zip_loaded.training + + original_config = _json_ready(build_model_config(model)) + assert _json_ready(build_model_config(dir_loaded)) == original_config + assert _json_ready(build_model_config(zip_loaded)) == original_config + _assert_state_dict_close(model.state_dict(), zip_loaded.state_dict()) + _assert_state_dict_close(dir_loaded.state_dict(), zip_loaded.state_dict()) + + with torch.no_grad(): + zip_pred = [g.x.detach().clone() for g in zip_loaded.predict(graph, steps=2)] + _assert_predictions_close(before, zip_pred) + + +def test_kgckpt_suffix_rejects_non_safe_zip(tmp_path: Path) -> None: + """``.kgckpt`` without safetensors members must not fall through to pickle.""" + bad_path = tmp_path / "broken.kgckpt" + with zipfile.ZipFile(bad_path, mode="w") as archive: + archive.writestr("readme.txt", "not a checkpoint\n") + + with pytest.raises(ValueError, match="bundle suffix"): + load_checkpoint(bad_path) + + +def test_zip_suffix_round_trip(tmp_path: Path) -> None: + """Explicit ``.zip`` suffix also selects the safetensors_v1 bundle writer.""" + model = _small_discrete_model() + zip_path = tmp_path / "model.zip" + save_checkpoint(model, zip_path, format="safetensors_v1") + loaded = load_checkpoint(zip_path) + _assert_state_dict_close(model.state_dict(), loaded.state_dict()) diff --git a/tests/test_serialization.py b/tests/test_serialization.py index c318ed9..073f186 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from pathlib import Path import pytest @@ -24,6 +25,10 @@ from koopman_graph.data import GraphSnapshotSequence from koopman_graph.serialization import ( FORMAT_VERSION, + SAFE_CONFIG_FILENAME, + SAFE_CONTAINER, + SAFE_META_FILENAME, + SAFE_WEIGHTS_FILENAME, build_checkpoint, build_model_config, load_checkpoint, @@ -292,6 +297,177 @@ def test_save_checkpoint_uses_current_format_version( } +def test_save_checkpoint_safetensors_v1_directory_layout( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify safetensors_v1 writes meta/config/weights as a directory.""" + assert FORMAT_VERSION == 1 + destination = tmp_path / "safe_ckpt" + save_checkpoint(graph_koopman_model, destination, format="safetensors_v1") + + meta_path = destination / SAFE_META_FILENAME + config_path = destination / SAFE_CONFIG_FILENAME + weights_path = destination / SAFE_WEIGHTS_FILENAME + assert meta_path.is_file() + assert config_path.is_file() + assert weights_path.is_file() + + meta = json.loads(meta_path.read_text(encoding="utf-8")) + assert meta["container"] == SAFE_CONTAINER + assert meta["format_version"] == FORMAT_VERSION + assert isinstance(meta["package_version"], str) + + config = json.loads(config_path.read_text(encoding="utf-8")) + assert config == json.loads(json.dumps(config)) + assert config["latent_dim"] == graph_koopman_model.latent_dim + assert config["encoder"]["type"] == "gcn" + assert config["decoder"]["type"] == "gcn" + + +def test_save_checkpoint_safetensors_v1_rejects_existing_file( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify safetensors_v1 refuses to treat an existing file as a directory.""" + file_path = tmp_path / "not_a_dir.pt" + file_path.write_bytes(b"not-a-checkpoint") + with pytest.raises(ValueError, match="must be a directory"): + save_checkpoint(graph_koopman_model, file_path, format="safetensors_v1") + + +def test_save_checkpoint_default_is_safetensors_v1( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify default save_checkpoint writes a safetensors_v1 directory.""" + destination = tmp_path / "default_safe" + save_checkpoint(graph_koopman_model, destination) + assert destination.is_dir() + assert (destination / SAFE_META_FILENAME).is_file() + assert (destination / SAFE_CONFIG_FILENAME).is_file() + assert (destination / SAFE_WEIGHTS_FILENAME).is_file() + loaded = load_checkpoint(destination) + assert not loaded.training + + +def test_save_checkpoint_legacy_pt_escape_hatch( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify format='legacy_pt' still writes a pickle .pt file.""" + checkpoint = tmp_path / "legacy.pt" + save_checkpoint(graph_koopman_model, checkpoint, format="legacy_pt") + assert checkpoint.is_file() + payload = torch.load(checkpoint, map_location="cpu", weights_only=False) + assert payload["format_version"] == FORMAT_VERSION + assert isinstance(payload["config"], dict) + assert isinstance(payload["state_dict"], dict) + + +def test_model_save_default_is_safetensors_v1( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify GraphKoopmanModel.save defaults to safetensors_v1.""" + destination = tmp_path / "model_safe" + graph_koopman_model.save(destination) + assert destination.is_dir() + assert (destination / SAFE_META_FILENAME).is_file() + loaded = GraphKoopmanModel.load(destination) + assert isinstance(loaded, GraphKoopmanModel) + assert not loaded.training + + +def test_model_save_legacy_pt_escape_and_load( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify Model.save(..., format='legacy_pt') round-trips via load.""" + checkpoint = tmp_path / "model_legacy.pt" + graph_koopman_model.save(checkpoint, format="legacy_pt") + assert checkpoint.is_file() + loaded = GraphKoopmanModel.load(checkpoint) + assert isinstance(loaded, GraphKoopmanModel) + + +def test_load_checkpoint_safetensors_v1_round_trip( + trained_gcn_model: GraphKoopmanModel, + scaling_sequence: GraphSnapshotSequence, + tmp_path: Path, +) -> None: + """Verify safetensors_v1 save → load preserves weights and predictions.""" + destination = tmp_path / "safe_ckpt" + save_checkpoint(trained_gcn_model, destination, format="safetensors_v1") + loaded = load_checkpoint(destination) + assert not loaded.training + + original = trained_gcn_model.state_dict() + restored = loaded.state_dict() + assert set(original) == set(restored) + for key in original: + torch.testing.assert_close(restored[key], original[key]) + + initial_graph = scaling_sequence[0] + for original_pred, loaded_pred in zip( + _predictions(trained_gcn_model, initial_graph), + _predictions(loaded, initial_graph), + strict=True, + ): + torch.testing.assert_close(original_pred, loaded_pred) + + +def test_load_checkpoint_safetensors_v1_missing_meta_raises( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify directories without meta.json are rejected.""" + destination = tmp_path / "broken" + save_checkpoint(graph_koopman_model, destination, format="safetensors_v1") + (destination / SAFE_META_FILENAME).unlink() + with pytest.raises(ValueError, match="missing meta.json"): + load_checkpoint(destination) + + +def test_load_checkpoint_safetensors_v1_missing_weights_raises( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify directories without model.safetensors are rejected.""" + destination = tmp_path / "broken_weights" + save_checkpoint(graph_koopman_model, destination, format="safetensors_v1") + (destination / SAFE_WEIGHTS_FILENAME).unlink() + with pytest.raises(ValueError, match="missing model.safetensors"): + load_checkpoint(destination) + + +def test_load_checkpoint_safetensors_v1_bad_container_raises( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify unsupported container markers are rejected.""" + destination = tmp_path / "bad_container" + save_checkpoint(graph_koopman_model, destination, format="safetensors_v1") + meta_path = destination / SAFE_META_FILENAME + meta = json.loads(meta_path.read_text(encoding="utf-8")) + meta["container"] = "not_a_real_container" + meta_path.write_text(json.dumps(meta), encoding="utf-8") + with pytest.raises(ValueError, match="Unsupported safetensors container"): + load_checkpoint(destination) + + +def test_load_checkpoint_safetensors_v1_invalid_config_raises( + graph_koopman_model: GraphKoopmanModel, + tmp_path: Path, +) -> None: + """Verify truncated config.json raises an actionable ValueError.""" + destination = tmp_path / "bad_config" + save_checkpoint(graph_koopman_model, destination, format="safetensors_v1") + (destination / SAFE_CONFIG_FILENAME).write_text("{not-json", encoding="utf-8") + with pytest.raises(ValueError, match="Invalid config.json"): + load_checkpoint(destination) + + def test_load_retired_format2_checkpoint_rejected( graph_koopman_model: GraphKoopmanModel, tmp_path: Path, @@ -353,7 +529,7 @@ def test_load_sparse_historical_format1_checkpoint_rejected( def test_load_checkpoint_missing_file_raises(tmp_path: Path) -> None: """Verify missing checkpoint paths raise FileNotFoundError.""" - with pytest.raises(FileNotFoundError, match="Checkpoint file not found"): + with pytest.raises(FileNotFoundError, match="Checkpoint path not found"): load_checkpoint(tmp_path / "missing.pt") @@ -537,9 +713,10 @@ def test_save_checkpoint_creates_parent_directories( tmp_path: Path, ) -> None: """Verify save creates nested parent directories.""" - path = tmp_path / "nested" / "dir" / "model.pt" + path = tmp_path / "nested" / "dir" / "model_ckpt" save_checkpoint(graph_koopman_model, path) - assert path.is_file() + assert path.is_dir() + assert (path / SAFE_WEIGHTS_FILENAME).is_file() def test_odo_model_round_trip_preserves_predictions( diff --git a/tests/test_tracking.py b/tests/test_tracking.py new file mode 100644 index 0000000..b744244 --- /dev/null +++ b/tests/test_tracking.py @@ -0,0 +1,155 @@ +"""Tests for ``koopman_graph.tracking`` fit adapters.""" + +from __future__ import annotations + +import csv +from pathlib import Path + +import pytest +import torch + +from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel +from koopman_graph.data import GraphSnapshotSequence +from koopman_graph.tracking import ( + CSV_FIT_FIELDNAMES, + CsvFitLogger, + TensorBoardFitLogger, +) + + +def _tiny_model() -> GraphKoopmanModel: + """Build a small discrete GCN for tracking integration tests.""" + return GraphKoopmanModel( + encoder=GNNEncoder( + in_channels=3, hidden_channels=8, latent_dim=4, num_layers=2 + ), + decoder=GNNDecoder( + latent_dim=4, hidden_channels=8, out_channels=3, num_layers=2 + ), + latent_dim=4, + time_step=0.1, + ) + + +def test_csv_fit_logger_writes_expected_rows( + scaling_sequence: GraphSnapshotSequence, + tmp_path: Path, +) -> None: + """Tiny fit with CsvFitLogger writes header + one row per epoch.""" + csv_path = tmp_path / "fit_log.csv" + logger = CsvFitLogger(csv_path) + model = _tiny_model() + torch.manual_seed(0) + history = model.fit( + scaling_sequence, + epochs=3, + lr=1e-2, + device="cpu", + callbacks=[logger], + ) + + assert csv_path.is_file() + with csv_path.open(newline="", encoding="utf-8") as handle: + reader = csv.DictReader(handle) + assert reader.fieldnames == list(CSV_FIT_FIELDNAMES) + rows = list(reader) + + assert len(rows) == history.epochs == 3 + assert [int(row["epoch"]) for row in rows] == [0, 1, 2] + for row, expected_total in zip(rows, history.loss, strict=True): + assert float(row["train_total"]) == expected_total + assert row["train_reconstruction"] != "" + assert row["val_total"] == "" + + +def test_csv_fit_logger_writes_val_columns( + scaling_sequence: GraphSnapshotSequence, + tmp_path: Path, +) -> None: + """Validation breakdown populates ``val_*`` columns when provided.""" + csv_path = tmp_path / "fit_val.csv" + model = _tiny_model() + torch.manual_seed(1) + history = model.fit( + scaling_sequence, + epochs=2, + lr=1e-2, + device="cpu", + validation_sequence=scaling_sequence, + callbacks=[CsvFitLogger(csv_path)], + ) + + with csv_path.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + + assert len(rows) == 2 + assert history.val_loss is not None + for row, expected_val in zip(rows, history.val_loss, strict=True): + assert float(row["val_total"]) == expected_val + assert row["val_reconstruction"] != "" + + +def test_training_does_not_import_tracking() -> None: + """Layer rule: training sources must not import tracking adapters.""" + package_root = Path(__file__).resolve().parents[1] / "src" / "koopman_graph" + training_root = package_root / "training" + offenders: list[str] = [] + for path in training_root.rglob("*.py"): + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("#"): + continue + from_tracking = stripped.startswith("from koopman_graph.tracking") + import_tracking = stripped.startswith("import koopman_graph.tracking") + if from_tracking or import_tracking: + offenders.append(str(path.relative_to(package_root.parent))) + break + assert offenders == [] + + +def test_tensorboard_fit_logger_writes_event_files( + scaling_sequence: GraphSnapshotSequence, + tmp_path: Path, +) -> None: + """Tiny fit writes TensorBoard event files when tensorboard is installed.""" + pytest.importorskip("tensorboard") + log_dir = tmp_path / "tb_runs" + model = _tiny_model() + torch.manual_seed(2) + model.fit( + scaling_sequence, + epochs=2, + lr=1e-2, + device="cpu", + validation_sequence=scaling_sequence, + callbacks=[TensorBoardFitLogger(log_dir)], + ) + + event_files = list(log_dir.rglob("events.out.tfevents.*")) + assert event_files, f"expected TensorBoard events under {log_dir}" + + +def test_tensorboard_fit_logger_import_error_hint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Missing tensorboard raises ImportError with an install hint.""" + import builtins + + import koopman_graph.tracking.tensorboard as tb_mod + + real_import = builtins.__import__ + + def blocked( + name: str, + globals: object = None, + locals: object = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ) -> object: + if name == "torch.utils.tensorboard": + raise ImportError("simulated missing tensorboard") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", blocked) + with pytest.raises(ImportError, match="pip install tensorboard"): + tb_mod._require_summary_writer() diff --git a/tests/test_tuning.py b/tests/test_tuning.py new file mode 100644 index 0000000..d1c8414 --- /dev/null +++ b/tests/test_tuning.py @@ -0,0 +1,384 @@ +"""Unit tests for ``koopman_graph.tuning`` (no Ray required for unit suite).""" + +from __future__ import annotations + +import ast +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from koopman_graph.training.history import FitHistory +from koopman_graph.tuning import ( + example_lr_latent_dim_space, + example_lr_loguniform_space, + fit_history_metrics, + run_ray_tune, +) + + +def test_fit_history_metrics_train_only() -> None: + """Final train loss and reconstruction are reported without val keys.""" + history = FitHistory( + loss=(1.0, 0.5, 0.25), + epochs=3, + reconstruction_loss=(0.9, 0.4, 0.2), + stopped_early=False, + ) + metrics = fit_history_metrics(history) + assert metrics["epochs"] == 3.0 + assert metrics["stopped_early"] == 0.0 + assert metrics["loss"] == 0.25 + assert metrics["final_loss"] == 0.25 + assert metrics["reconstruction_loss"] == 0.2 + assert "val_loss" not in metrics + assert "best_loss" not in metrics + + +def test_fit_history_metrics_with_val_and_best() -> None: + """Validation and best-epoch fields flatten when present.""" + history = FitHistory( + loss=(1.0, 0.8), + epochs=2, + reconstruction_loss=(0.7, 0.6), + val_loss=(1.2, 0.9), + val_reconstruction_loss=(1.1, 0.85), + stopped_early=True, + best_epoch=1, + best_loss=0.8, + ) + metrics = fit_history_metrics(history) + assert metrics["stopped_early"] == 1.0 + assert metrics["val_loss"] == 0.9 + assert metrics["final_val_loss"] == 0.9 + assert metrics["val_reconstruction_loss"] == 0.85 + assert metrics["best_epoch"] == 1.0 + assert metrics["best_loss"] == 0.8 + + +def test_fit_history_metrics_empty_loss_omits_loss_keys() -> None: + """Empty loss series omit ``loss`` rather than inventing NaNs.""" + history = FitHistory(loss=(), epochs=0) + metrics = fit_history_metrics(history) + assert metrics == {"epochs": 0.0, "stopped_early": 0.0} + + +def test_fit_history_metrics_early_stop_flag() -> None: + """Early stopping encodes ``stopped_early`` as ``1.0`` with short epochs.""" + history = FitHistory( + loss=(1.0, 0.5), + epochs=2, + reconstruction_loss=(0.8, 0.4), + stopped_early=True, + best_epoch=1, + best_loss=0.5, + ) + metrics = fit_history_metrics(history) + assert metrics["stopped_early"] == 1.0 + assert metrics["epochs"] == 2.0 + assert metrics["loss"] == 0.5 + assert metrics["best_loss"] == 0.5 + assert metrics["best_epoch"] == 1.0 + + +def test_fit_history_metrics_missing_val_series_omitted() -> None: + """``val_loss=None`` and empty val tuples omit validation metric keys.""" + none_history = FitHistory( + loss=(0.5,), + epochs=1, + reconstruction_loss=(0.4,), + val_loss=None, + val_reconstruction_loss=None, + ) + none_metrics = fit_history_metrics(none_history) + assert "val_loss" not in none_metrics + assert "final_val_loss" not in none_metrics + assert "val_reconstruction_loss" not in none_metrics + + empty_history = FitHistory( + loss=(0.5,), + epochs=1, + reconstruction_loss=(0.4,), + val_loss=(), + val_reconstruction_loss=(), + ) + empty_metrics = fit_history_metrics(empty_history) + assert "val_loss" not in empty_metrics + assert "final_val_loss" not in empty_metrics + assert "val_reconstruction_loss" not in empty_metrics + + +def test_fit_history_metrics_partial_optionals() -> None: + """Empty reconstruction and one-sided best fields stay omitted.""" + no_recon = FitHistory( + loss=(1.0, 0.2), + epochs=2, + reconstruction_loss=(), + ) + metrics = fit_history_metrics(no_recon) + assert metrics["loss"] == 0.2 + assert "reconstruction_loss" not in metrics + + best_loss_only = FitHistory(loss=(0.3,), epochs=1, best_loss=0.3) + metrics_loss = fit_history_metrics(best_loss_only) + assert metrics_loss["best_loss"] == 0.3 + assert "best_epoch" not in metrics_loss + + best_epoch_only = FitHistory(loss=(0.3,), epochs=1, best_epoch=0) + metrics_epoch = fit_history_metrics(best_epoch_only) + assert metrics_epoch["best_epoch"] == 0.0 + assert "best_loss" not in metrics_epoch + + +def test_fit_history_metrics_values_are_floats() -> None: + """Returned mapping values are plain floats (Tune/Optuna-friendly).""" + history = FitHistory( + loss=(1.0, 0.5), + epochs=2, + reconstruction_loss=(0.9, 0.4), + val_loss=(1.1, 0.6), + val_reconstruction_loss=(1.0, 0.55), + stopped_early=True, + best_epoch=1, + best_loss=0.5, + ) + metrics = fit_history_metrics(history) + assert metrics + assert all(type(value) is float for value in metrics.values()) + assert all(value == value for value in metrics.values()) # no NaN + + +def test_tuning_not_on_root_all() -> None: + """``tuning`` stays a power-user import (off root ``__all__``).""" + import koopman_graph + + assert "fit_history_metrics" not in koopman_graph.__all__ + assert "run_ray_tune" not in koopman_graph.__all__ + + +def _top_level_ray_import_lines(path: Path) -> list[str]: + """Return source lines that import Ray at module scope.""" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + offenders: list[str] = [] + for node in tree.body: + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "ray" or alias.name.startswith("ray."): + offenders.append(f"{path.name}:{node.lineno}") + elif ( + isinstance(node, ast.ImportFrom) + and node.module is not None + and (node.module == "ray" or node.module.startswith("ray.")) + ): + offenders.append(f"{path.name}:{node.lineno}") + return offenders + + +def test_tuning_has_no_toplevel_ray_imports() -> None: + """Importing ``tuning`` must not execute a top-level ``import ray``.""" + package_root = Path(__file__).resolve().parents[1] / "src" / "koopman_graph" + tuning_root = package_root / "tuning" + offenders: list[str] = [] + for path in tuning_root.rglob("*.py"): + offenders.extend(_top_level_ray_import_lines(path)) + assert offenders == [] + + +def test_importing_tuning_does_not_load_ray() -> None: + """``import koopman_graph.tuning`` must not import Ray by itself.""" + import subprocess + + script = ( + "import sys\n" + "import koopman_graph.tuning\n" + "assert 'ray' not in sys.modules, sorted(sys.modules)\n" + ) + completed = subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + ) + assert completed.returncode == 0, completed.stderr + + +def test_training_does_not_import_tuning() -> None: + """Layer rule: ``training`` must not import ``tuning``.""" + package_root = Path(__file__).resolve().parents[1] / "src" / "koopman_graph" + training_root = package_root / "training" + offenders: list[str] = [] + for path in training_root.rglob("*.py"): + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("#"): + continue + from_tuning = stripped.startswith("from koopman_graph.tuning") + import_tuning = stripped.startswith("import koopman_graph.tuning") + if from_tuning or import_tuning: + offenders.append(str(path.relative_to(package_root.parent))) + break + assert offenders == [] + + +def test_run_ray_tune_rejects_nonpositive_samples() -> None: + """``num_samples`` must be at least 1.""" + with pytest.raises(ValueError, match="num_samples"): + run_ray_tune(lambda config: None, {"lr": 1e-3}, num_samples=0) + + +def test_run_ray_tune_missing_ray_message() -> None: + """Missing Ray raises ImportError with the ``[ray]`` install hint.""" + from koopman_graph.tuning import ray_tune as ray_tune_mod + + with ( + patch.object( + ray_tune_mod, + "_import_ray_tune", + side_effect=ImportError( + "Ray Tune is required for koopman_graph.tuning.run_ray_tune; " + 'install with: pip install "koopman-graph[ray]"' + ), + ), + pytest.raises(ImportError, match=r"koopman-graph\[ray\]"), + ): + ray_tune_mod.run_ray_tune(lambda config: None, {"lr": 1e-3}, num_samples=1) + + +def test_run_ray_tune_builds_tuner_with_mocks(tmp_path: Path) -> None: + """Wrapper constructs TuneConfig / Tuner and returns ``fit()`` result.""" + from koopman_graph.tuning import ray_tune as ray_tune_mod + + fake_results = SimpleNamespace(ok=True) + fake_tune_config = object() + fake_run_config = object() + fake_tuner = MagicMock() + fake_tuner.fit.return_value = fake_results + + fake_ray = ModuleType("ray") + fake_ray.is_initialized = MagicMock(return_value=True) # type: ignore[attr-defined] + fake_ray.init = MagicMock() # type: ignore[attr-defined] + + fake_tune = ModuleType("ray.tune") + fake_tune.TuneConfig = MagicMock(return_value=fake_tune_config) # type: ignore[attr-defined] + fake_tune.Tuner = MagicMock(return_value=fake_tuner) # type: ignore[attr-defined] + + fake_run_config_cls = MagicMock(return_value=fake_run_config) + + def _trainable(config: dict[str, object]) -> None: + return None + + with patch.object( + ray_tune_mod, + "_import_ray_tune", + return_value=(fake_ray, fake_tune, fake_run_config_cls), + ): + result = ray_tune_mod.run_ray_tune( + _trainable, + {"lr": 0.01, "epochs": 1}, + num_samples=2, + metric="loss", + mode="min", + storage_path=tmp_path / "tune", + run_name="unit_test_tune", + ) + + assert result is fake_results + fake_ray.init.assert_not_called() + fake_tune.TuneConfig.assert_called_once_with( + num_samples=2, + metric="loss", + mode="min", + ) + fake_run_config_cls.assert_called_once() + run_kwargs = fake_run_config_cls.call_args.kwargs + assert run_kwargs["name"] == "unit_test_tune" + assert run_kwargs["storage_path"] == str((tmp_path / "tune").resolve()) + fake_tune.Tuner.assert_called_once() + tuner_kwargs = fake_tune.Tuner.call_args + assert tuner_kwargs.args[0] is _trainable + assert tuner_kwargs.kwargs["param_space"] == {"lr": 0.01, "epochs": 1} + assert tuner_kwargs.kwargs["tune_config"] is fake_tune_config + assert tuner_kwargs.kwargs["run_config"] is fake_run_config + fake_tuner.fit.assert_called_once_with() + + +def test_example_search_space_names_and_docs_are_honest() -> None: + """Helpers are explicitly example-only (not scientific defaults).""" + from koopman_graph.tuning import search_spaces + + assert example_lr_loguniform_space.__name__.startswith("example_") + assert example_lr_latent_dim_space.__name__.startswith("example_") + module_doc = search_spaces.__doc__ or "" + assert "not" in module_doc.lower() + assert "scientific" in module_doc.lower() or "defaults" in module_doc.lower() + assert "smoke" in module_doc.lower() or "example" in module_doc.lower() + + +def test_example_lr_loguniform_space_with_mocked_tune() -> None: + """LR helper returns a Tune-compatible dict via lazy ``ray.tune``.""" + from koopman_graph.tuning import search_spaces as spaces_mod + + fake_sampler = object() + fake_tune = ModuleType("ray.tune") + fake_tune.loguniform = MagicMock(return_value=fake_sampler) # type: ignore[attr-defined] + + with patch.object(spaces_mod, "_import_tune", return_value=fake_tune): + space = spaces_mod.example_lr_loguniform_space( + lr_low=1e-4, + lr_high=1e-2, + epochs=2, + seed=0, + ) + + assert space["lr"] is fake_sampler + assert space["epochs"] == 2 + assert space["seed"] == 0 + fake_tune.loguniform.assert_called_once_with(1e-4, 1e-2) + + +def test_example_lr_latent_dim_space_with_mocked_tune() -> None: + """Joint lr / latent_dim example space includes both samplers.""" + from koopman_graph.tuning import search_spaces as spaces_mod + + fake_lr = object() + fake_dim = object() + fake_tune = ModuleType("ray.tune") + fake_tune.loguniform = MagicMock(return_value=fake_lr) # type: ignore[attr-defined] + fake_tune.choice = MagicMock(return_value=fake_dim) # type: ignore[attr-defined] + + with patch.object(spaces_mod, "_import_tune", return_value=fake_tune): + space = spaces_mod.example_lr_latent_dim_space( + latent_dims=(4, 8), + epochs=1, + ) + + assert space["lr"] is fake_lr + assert space["latent_dim"] is fake_dim + assert space["epochs"] == 1 + fake_tune.choice.assert_called_once_with([4, 8]) + + +def test_example_search_space_missing_ray_message() -> None: + """Missing Ray raises ImportError with the ``[ray]`` install hint.""" + from koopman_graph.tuning import search_spaces as spaces_mod + + with ( + patch.object( + spaces_mod, + "_import_tune", + side_effect=ImportError( + "Ray Tune is required for koopman_graph.tuning search-space " + 'helpers; install with: pip install "koopman-graph[ray]"' + ), + ), + pytest.raises(ImportError, match=r"koopman-graph\[ray\]"), + ): + spaces_mod.example_lr_loguniform_space() + + +def test_example_search_space_validates_bounds() -> None: + """Invalid learning-rate bounds raise ``ValueError`` before Ray use.""" + with pytest.raises(ValueError, match="lr_low"): + example_lr_loguniform_space(lr_low=0.1, lr_high=0.01) diff --git a/uv.lock b/uv.lock index cae27b4..e3f992e 100644 --- a/uv.lock +++ b/uv.lock @@ -282,6 +282,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/fe/6bea5c9162869c5beba5d9c8abbed835ec85bf1ec1fba05a3822325c45f3/build-1.5.0-py3-none-any.whl", hash = "sha256:13f3eecb844759ab66efec90ca17639bbf14dc06cb2fdf37a9010322d9c50a6f", size = 26018, upload-time = "2026-04-30T03:18:23.644Z" }, ] +[[package]] +name = "captum" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "matplotlib", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "packaging" }, + { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, + { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f7/9e/f31758e20c3d2b42b05420936b75f663e7aeae9ba63b84aee5cad1b044a3/captum-0.9.0.tar.gz", hash = "sha256:e4868e22b68224f1a5355d2b7966803c021cf9fe5d893f51f587b4002a666b54", size = 363381, upload-time = "2026-04-17T06:11:39.577Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/75/c99571fee171bb9d34d3e1df5fdb16d57e2c4c15339eb220e6e6e4539f1c/captum-0.9.0-py3-none-any.whl", hash = "sha256:cda38e1d42c37591d71560bb2f5813ac4cae8d8543afac45279c7efd5945997e", size = 455196, upload-time = "2026-04-17T06:11:38.229Z" }, +] + [[package]] name = "certifi" version = "2026.6.17" @@ -1762,6 +1781,7 @@ wheels = [ name = "koopman-graph" source = { editable = "." } dependencies = [ + { name = "safetensors" }, { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.13.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" }, { name = "torch-geometric" }, @@ -1771,6 +1791,9 @@ dependencies = [ baselines-ode = [ { name = "torchdiffeq" }, ] +cli = [ + { name = "pyyaml" }, +] dask = [ { name = "dask", extra = ["distributed"] }, ] @@ -1808,6 +1831,9 @@ docs = [ equivariance = [ { name = "e3nn" }, ] +explain = [ + { name = "captum" }, +] lightning = [ { name = "lightning" }, ] @@ -1839,6 +1865,7 @@ symmetry = [ [package.metadata] requires-dist = [ { name = "build", marker = "extra == 'dev'", specifier = ">=1.2.0" }, + { name = "captum", marker = "extra == 'explain'", specifier = ">=0.7" }, { name = "dask", extras = ["distributed"], marker = "extra == 'dask'", specifier = ">=2024.1.0" }, { name = "dask", extras = ["distributed"], marker = "extra == 'distributed'", specifier = ">=2024.1.0" }, { name = "deeptime", marker = "extra == 'msm'", specifier = ">=0.4" }, @@ -1860,9 +1887,11 @@ requires-dist = [ { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.4.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=6.0.0" }, { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.6.0" }, + { name = "pyyaml", marker = "extra == 'cli'", specifier = ">=6" }, { name = "ray", extras = ["train"], marker = "extra == 'distributed'", specifier = ">=2.9.0" }, { name = "ray", extras = ["train"], marker = "extra == 'ray'", specifier = ">=2.9.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.15.0" }, + { name = "safetensors", specifier = ">=0.4" }, { name = "scipy", marker = "extra == 'dev'", specifier = ">=1.11.0" }, { name = "scipy", marker = "extra == 'mpc'", specifier = ">=1.11.0" }, { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.0.0" }, @@ -1872,7 +1901,7 @@ requires-dist = [ { name = "torch-geometric", specifier = ">=2.6.1" }, { name = "torchdiffeq", marker = "extra == 'baselines-ode'", specifier = ">=0.2.3" }, ] -provides-extras = ["dev", "docs", "rl", "symmetry", "mpc", "lightning", "ray", "dask", "md", "msm", "baselines-ode", "baselines-graphcast", "equivariance", "distributed"] +provides-extras = ["cli", "explain", "dev", "docs", "rl", "symmetry", "mpc", "lightning", "ray", "dask", "md", "msm", "baselines-ode", "baselines-graphcast", "equivariance", "distributed"] [[package]] name = "lightning" @@ -4205,6 +4234,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/57/c9/e69b1ff4c8b69093ef08b8919ab767af0569666865b39c30a8795d88d3c6/ruff-0.15.22-py3-none-win_arm64.whl", hash = "sha256:e1168075b72158510839f250027659cdd78476f40507dd517892304c41318661", size = 11298172, upload-time = "2026-07-16T15:14:10.51Z" }, ] +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + [[package]] name = "scikit-learn" version = "1.7.2"