Add anonymous reproduction code

code/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+.pytest_cache/
+.venv/
+*.egg-info/
+/runs/
+/outputs/

code/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 The Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

code/README.md

@@ -0,0 +1,93 @@
+# Stabilizing Probabilistic Forecasts via Stochastic Interpolation Cascades
+
+Code for reproducing the experiments accompanying an anonymous research submission on conditional generative forecasting.
+
+We compare Direct Stochastic Interpolation (Direct SI), Cascaded Stochastic Interpolation (Cascaded SI), and Elucidating Diffusion Models (EDM) on five dynamical systems: 2-D Euler, surface quasi-geostrophic flow, heterogeneous waves, nonlinear Klein–Gordon, and passive scalar advection–diffusion.
+
+The repository contains the model implementations (`bifurcation/`), configurations (`configs/`), recorded experimental results (`records/`), and reproduction scripts (`reproduce/`).
+
+## Installation
+
+From the extracted code directory:
+
+```bash
+uv sync
+export CLEAN_FID_CACHE=$HOME/.cache/clean-fid
+uv run python -c "from bifurcation.fid import backend as B; B.build_extractor('cpu')"
+```
+
+The dependencies are pinned in `uv.lock`. The original experiments used different PyTorch builds for some datasets, so minor numerical differences across environments are possible.
+
+## Data and checkpoints
+
+The datasets and trained models are available on Hugging Face:
+
+- [Euler](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/euler)
+- [SQG](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/sqg)
+- [Heterogeneous waves](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/het-waves)
+- [Nonlinear Klein–Gordon](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/nonlinear-kg)
+- [Passive scalar](https://huggingface.co/datasets/coarse-grained-forecasting-benchmark/passive-scalar)
+- [Model checkpoints](https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models)
+
+For example, download the Euler dataset and the checkpoints:
+
+```bash
+hf download coarse-grained-forecasting-benchmark/euler \
+    --repo-type dataset --local-dir data/euler
+
+hf download coarse-grained-forecasting-benchmark/forecasting-models \
+    --local-dir models
+
+export EULER_DATA="$PWD/data/euler/data"
+export EULER_TRUTH_BANK="$PWD/data/euler/truth_bank"
+```
+
+The other systems follow the same directory structure.
+
+## Reproducing the results
+
+To display the recorded results of Tables 2, 5, and 6:
+
+```bash
+uv run python -m reproduce.tables
+```
+
+To generate 16 conditional forecasts for one Euler test condition using Cascaded SI:
+
+```bash
+uv run python -m reproduce.table2 generate \
+    --system euler --method cascaded_si \
+    --checkpoint models/euler/cascaded-si-S10/model.pt \
+    --out runs/example --gids 36 --n-samples 16 --device cpu
+```
+
+To reproduce a Table 2 entry, generating 10,000 forecasts per test condition and computing csFID-10k:
+
+```bash
+uv run python -m reproduce.table2 run \
+    --system euler --method cascaded_si \
+    --checkpoint models/euler/cascaded-si-S10/model.pt \
+    --out runs/euler_cascaded_si --device cuda
+```
+
+Use `--method direct_si` or `--method edm` for the baselines. For Cascaded SI, `--S` selects a specific cascade depth.
+
+The scripts in `reproduce/figures/` and `reproduce/appendix/` generate the paper's figures and additional analyses from the recorded results, datasets, or forecast samples, as applicable.
+
+## Training
+
+Models can also be trained using the configurations of the selected experiments:
+
+```bash
+uv run python -m bifurcation.train \
+    --config records/euler/runs/cascaded_si_S10/config.yaml \
+    --data "$EULER_DATA" --out runs/euler_S10
+```
+
+The training configurations, model hyperparameters, and evaluation settings are provided in `records/` and `configs/`.
+
+## License
+
+MIT License.
+
+The EDM implementation follows [Karras et al. (2022)](https://github.com/NVlabs/edm). Image evaluation uses [Clean-FID](https://github.com/GaParmar/clean-fid).

code/bifurcation/__init__.py

@@ -0,0 +1,3 @@
+"""bifurcation — models, methods, data loading, sampling and the csFID metric for conditional forecasting of PDE ensembles."""
+
+__version__ = "0.1.0"

code/bifurcation/config.py

@@ -0,0 +1,100 @@
+"""Configuration files of the repository: the dataset cards, the evaluation settings and the recorded runs.
+
+* ``configs/<system>/dataset.yaml`` names the data of one system; ``${VAR}`` entries are read from the environment.
+* ``configs/evaluation.yaml`` holds the csFID settings and the sampling seeds shared by all systems.
+* ``records/<system>/units.json`` lists the models of the paper and their recorded scores.
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+from pathlib import Path
+
+import yaml
+
+REPO = Path(__file__).resolve().parents[1]
+CONFIGS = REPO / "configs"
+RECORDS = REPO / "records"
+SYSTEMS = ("euler", "sqg", "waves", "nlkg", "passive_scalar")
+
+_ENV_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}")
+
+
+def expand_env(value):
+    """``${VAR}`` / ``${VAR:-default}`` expansion in strings (recursively in dicts and lists)."""
+    if isinstance(value, str):
+        def sub(m):
+            var, default = m.group(1), m.group(2)
+            v = os.environ.get(var)
+            if v is None or v == "":
+                if default is None:
+                    raise KeyError(f"environment variable {var} is not set (used in a config)")
+                return default
+            return v
+        return _ENV_RE.sub(sub, value)
+    if isinstance(value, dict):
+        return {k: expand_env(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [expand_env(v) for v in value]
+    return value
+
+
+def load_yaml(path: str | os.PathLike) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f) or {}
+
+
+def dataset_card(system: str, expand: bool = False) -> dict:
+    """The dataset card of ``system``; with ``expand`` the ``${VAR}`` paths are resolved from the environment."""
+    if system not in SYSTEMS:
+        raise KeyError(f"unknown system {system!r} (one of {SYSTEMS})")
+    card = load_yaml(CONFIGS / system / "dataset.yaml")
+    return expand_env(card) if expand else card
+
+
+def data_paths(system: str, data: str | None = None, truth_bank: str | None = None) -> tuple[Path, Path]:
+    """The package and truth-bank directories: the given values, else the environment variables of the card."""
+    card = dataset_card(system)
+    if data is None:
+        data = expand_env(card["data"])
+    if truth_bank is None:
+        truth_bank = expand_env(card["truth_bank"])
+    return Path(data).resolve(), Path(truth_bank).resolve()
+
+
+def evaluation_settings() -> dict:
+    return load_yaml(CONFIGS / "evaluation.yaml")
+
+
+def records(system: str) -> dict:
+    return json.loads((RECORDS / system / "units.json").read_text())
+
+
+def unit(system: str, method: str, S: int | None = None) -> tuple[dict, dict]:
+    """The recorded model of ``system`` for a paper method (``direct_si``, ``edm``, ``cascaded_si``); with ``S`` the
+    Cascaded SI model of that depth (Table 6). Returns (records, unit)."""
+    rec = records(system)
+    if method == "edm":
+        if S is not None:
+            raise SystemExit("--S applies to cascaded_si only")
+        key = "edm"
+    elif method == "direct_si":
+        if S not in (None, 1):
+            raise SystemExit("direct_si is S = 1; use --method cascaded_si --S <S> for other depths")
+        key = rec["table2"]["direct_si"]
+    elif method == "cascaded_si":
+        key = rec["table2"]["cascaded_si"] if S is None else f"cascaded_si_S{int(S):02d}"
+    else:
+        raise SystemExit(f"unknown method {method!r}")
+    u = next((u for u in rec["units"] if u["key"] == key), None)
+    if u is None:
+        raise SystemExit(f"{system}: no recorded model {key} (depths: {[u['S'] for u in rec['units'] if u['key'] != 'edm']})")
+    return rec, u
+
+
+def run_config(system: str, u: dict, data: Path) -> dict:
+    """The recorded run configuration of a model with the dataset path filled in."""
+    cfg = load_yaml(REPO / u["config"])
+    cfg["data"]["packaged_dataset_path"] = str(Path(data).resolve()) + "/"
+    return cfg

code/bifurcation/data/fine.py

@@ -0,0 +1,548 @@
+"""Loaders of a dataset package: ``trajectories_fine.npz`` holds the blocks ``vecs_train`` (conditions, realizations,
+snapshots, D), ``vecs_eval`` and optionally ``vecs_new``, keyed by ``gids_train`` / ``gids_eval`` / ``gids_new``, and
+the snapshot times ``t_values``; ``split.json`` names the training, validation and test conditions;
+``normalization.npz`` the per-coordinate mean and std.
+
+  * training pairs are read from ``vecs_train`` only, restricted to the (condition, realization) rows of a pair
+    manifest (``bifurcation.data.manifests``); segments are ``boundary_idx`` snapshot pairs and the conditioning
+    times are the physical times of the stored grid;
+  * validation and test inputs are ``vecs_eval`` at the snapshot ``fine_horizon_idx``;
+  * x_0 of a condition is member 0 at t = 0 (all members share it; asserted at 1e-9).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import zipfile
+from typing import List, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+FINE_DT = 0.025
+FINE_FILE = "trajectories_fine.npz"
+ENDPOINT_SCHEMA = "endpoint_pairs_v1"      # the endpoint manifests of bifurcation.data.manifests
+
+
+# --------------------------------------------------------------------------- #
+# npz member access (uncompressed members -> memmap; else plain load)
+# --------------------------------------------------------------------------- #
+
+def _npz_member_memmap(npz_path: str, name: str) -> np.ndarray:
+    """Memory-map one array stored in an UNCOMPRESSED ``.npz``.
+
+    The packages store every member with ``compress_type == 0``, so the raw ``.npy`` bytes sit contiguously
+    inside the zip and can be memmapped: fancy-indexing then reads only the pages of the selected trajectories.
+
+    Falls back to ``np.load`` (full read) if the member is compressed.
+    """
+    member = name if name.endswith(".npy") else name + ".npy"
+    with zipfile.ZipFile(npz_path) as zf:
+        info = zf.getinfo(member)
+        if info.compress_type != zipfile.ZIP_STORED:
+            return np.load(npz_path)[name]
+        with open(npz_path, "rb") as f:
+            f.seek(info.header_offset)
+            local = f.read(30)
+            if local[:4] != b"PK\x03\x04":
+                raise IOError(f"bad local zip header for {member} in {npz_path}")
+            name_len  = int.from_bytes(local[26:28], "little")
+            extra_len = int.from_bytes(local[28:30], "little")
+            f.seek(info.header_offset + 30 + name_len + extra_len)
+            version = np.lib.format.read_magic(f)
+            np.lib.format._check_version(version)
+            shape, fortran, dtype = np.lib.format._read_array_header(f, version)
+            offset = f.tell()
+    order = "F" if fortran else "C"
+    return np.memmap(npz_path, dtype=dtype, mode="r", offset=offset,
+                     shape=shape, order=order)
+
+
+def _sha256_file(path: str) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _resolve_path(path: str) -> str:
+    """Absolute paths pass through; relative paths resolve against the repository root."""
+    if os.path.isabs(path):
+        return path
+    return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), path)
+
+
+def _load_manifest(slice_manifest: str) -> dict:
+    path = _resolve_path(slice_manifest)
+    with open(path) as f:
+        man = json.load(f)
+    man["_manifest_path"] = path
+    # Provenance guard: the manifest must certify the master ordering it was
+    # cut from, and that file must still hash to the recorded value.
+    master = os.path.join(os.path.dirname(path), "master_order.npz")
+    if not os.path.isfile(master):
+        raise FileNotFoundError(
+            f"master_order.npz not found next to slice manifest {path} — "
+            "slice provenance cannot be verified")
+    got = _sha256_file(master)
+    want = man["master_order_sha256"]
+    if got != want:
+        raise AssertionError(
+            f"master_order.npz hash mismatch for {path}: {got} != {want}")
+    return man
+
+
+# --------------------------------------------------------------------------- #
+# training dataset
+# --------------------------------------------------------------------------- #
+
+class SegmentPairDataset(Dataset):
+    """``(x_tk, x_tk1, t_k, t_k1)`` pairs from ``vecs_train`` per a pair manifest.
+
+    Checks: exactly ``S`` pairs per selected trajectory, exactly ``B == N_traj * S`` pairs in total,
+    boundaries on stored snapshots, physical times consistent with the manifest, no validation or test gid.
+    """
+
+    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
+                 seed: int = 0,
+                 expected_t_values: list[float] | None = None):
+        man = _load_manifest(slice_manifest)
+        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
+
+        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
+        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
+        S       = int(man["S"])
+        n_traj  = int(man["N_traj"])
+        B       = int(man["B"])
+        rows    = np.asarray(man["trajectories"], dtype=np.int64)  # (N,3)
+
+        # --- manifest self-consistency ---
+        if rows.shape != (n_traj, 3):
+            raise AssertionError(
+                f"manifest trajectories shape {rows.shape} != ({n_traj}, 3)")
+        if n_traj * S != B:
+            raise AssertionError(f"N_traj*S = {n_traj * S} != B = {B}")
+        if len(boundary_idx) != S + 1 or boundary_idx[0] != 0:
+            raise AssertionError(f"boundary_idx must be S+1 long, start at 0; "
+                                 f"got {boundary_idx.tolist()}")
+        if np.any(np.diff(boundary_idx) != boundary_idx[1] - boundary_idx[0]):
+            raise AssertionError("boundary_idx must be uniformly strided")
+        uniq = np.unique(rows[:, :2], axis=0)
+        if uniq.shape[0] != n_traj:
+            raise AssertionError("duplicate (group_row, member_col) in slice")
+
+        gids_train = _npz_member_memmap(fine_path, "gids_train")
+        t_values   = np.asarray(_npz_member_memmap(fine_path, "t_values"))
+        if boundary_idx[-1] > t_values.shape[0] - 1:
+            raise AssertionError("boundary_idx exceeds fine snapshot grid")
+        if not np.allclose(t_values[boundary_idx], t_vals_man, atol=1e-12):
+            raise AssertionError(
+                f"manifest t_values {t_vals_man.tolist()} do not match the "
+                f"stored grid at boundary_idx {boundary_idx.tolist()}")
+        if expected_t_values is not None:
+            exp = np.asarray([float(t) for t in expected_t_values])
+            if len(exp) != len(t_vals_man) or not np.allclose(exp, t_vals_man,
+                                                              atol=1e-9):
+                raise AssertionError(
+                    f"sampling.t_values {exp.tolist()} != slice-manifest "
+                    f"t_values {t_vals_man.tolist()} — the sampler chain and "
+                    "the training segments must describe the same physics")
+        # gid column must be exactly gids_train[group_row] and never an
+        # eval/new gid (train block only contains canonical train gids, but
+        # assert the mapping so a corrupted manifest fails loudly).
+        if not np.array_equal(np.asarray(gids_train)[rows[:, 0]], rows[:, 2]):
+            raise AssertionError("manifest gid column inconsistent with "
+                                 "gids_train[group_row]")
+        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
+            split = json.load(f)
+        forbidden = set(split["val"]) | set(split["test"])
+        if forbidden & set(rows[:, 2].tolist()):
+            raise AssertionError("slice manifest contains val/test gids")
+
+        # --- gather (N_traj, S+1, D) raw states, normalize, build pairs ---
+        vecs = _npz_member_memmap(fine_path, "vecs_train")   # (80,200,41,D)
+        seg = np.asarray(vecs[rows[:, 0], rows[:, 1]])[:, boundary_idx, :]
+        if not np.all(np.isfinite(seg)):
+            raise AssertionError("non-finite values in selected trajectories")
+        D = seg.shape[-1]
+
+        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
+        mean = norm["mean"].astype(np.float32)
+        std  = norm["std"].astype(np.float32)
+
+        starts = seg[:, :-1, :].reshape(-1, D)               # (N*S, D)
+        ends   = seg[:,  1:, :].reshape(-1, D)
+        # float32 normalization arithmetic
+        x_tk  = ((starts.astype(np.float32) - mean) / std).astype(np.float32)
+        x_tk1 = ((ends.astype(np.float32)   - mean) / std).astype(np.float32)
+
+        t_k_grid  = np.broadcast_to(t_vals_man[:-1], (n_traj, S)).reshape(-1)
+        t_k1_grid = np.broadcast_to(t_vals_man[1:],  (n_traj, S)).reshape(-1)
+
+        self.x_tk  = torch.from_numpy(x_tk)
+        self.x_tk1 = torch.from_numpy(x_tk1)
+        self.t_k   = torch.from_numpy(t_k_grid.astype(np.float32))
+        self.t_k1  = torch.from_numpy(t_k1_grid.astype(np.float32))
+        if self.x_tk.shape[0] != B:
+            raise AssertionError(f"built {self.x_tk.shape[0]} pairs, expected {B}")
+
+        # Provenance of each pair (gid, member, segment).
+        self.pair_gid     = np.repeat(rows[:, 2], S)
+        self.pair_member  = np.repeat(rows[:, 1], S)
+        self.pair_segment = np.tile(np.arange(S), n_traj)
+        self.n_segments = S
+        self.manifest = man
+        self.seed = int(seed)
+
+    def __len__(self) -> int:
+        return self.x_tk.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.x_tk[idx], self.x_tk1[idx], self.t_k[idx], self.t_k1[idx]
+
+
+class ValSegmentPairDataset(Dataset):
+    """Validation pairs of the training objective: the run's segmentation applied to the validation gids of
+    ``vecs_eval``, first ``n_members`` members per gid (the ``val_loss`` column of the log)."""
+
+    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
+                 n_members: int = 20):
+        man = _load_manifest(slice_manifest)
+        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
+        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
+        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
+        S = int(man["S"])
+
+        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
+            split = json.load(f)
+        val_gids = sorted(split["val"])
+
+        gids_eval = np.asarray(_npz_member_memmap(fine_path, "gids_eval"))
+        vecs_eval = _npz_member_memmap(fine_path, "vecs_eval")
+        rows = []
+        for gid in val_gids:
+            where = np.nonzero(gids_eval == gid)[0]
+            if len(where) != 1:
+                raise AssertionError(f"val gid {gid} not found once in gids_eval")
+            rows.append(int(where[0]))
+        n_members = int(n_members)
+        if not (1 <= n_members <= vecs_eval.shape[1]):
+            raise ValueError(f"n_members {n_members} out of range")
+
+        seg = np.asarray(vecs_eval[rows][:, :n_members, boundary_idx, :])
+        if not np.all(np.isfinite(seg)):
+            raise AssertionError("non-finite values in vecs_eval slice")
+        n_traj = seg.shape[0] * seg.shape[1]
+        D = seg.shape[-1]
+        seg = seg.reshape(n_traj, S + 1, D)
+
+        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
+        mean = norm["mean"].astype(np.float32)
+        std  = norm["std"].astype(np.float32)
+        starts = seg[:, :-1, :].reshape(-1, D)
+        ends   = seg[:,  1:, :].reshape(-1, D)
+        self.x_tk  = torch.from_numpy(((starts.astype(np.float32) - mean) / std)
+                                      .astype(np.float32))
+        self.x_tk1 = torch.from_numpy(((ends.astype(np.float32) - mean) / std)
+                                      .astype(np.float32))
+        self.t_k   = torch.from_numpy(np.broadcast_to(t_vals_man[:-1],
+                                      (n_traj, S)).reshape(-1).astype(np.float32))
+        self.t_k1  = torch.from_numpy(np.broadcast_to(t_vals_man[1:],
+                                      (n_traj, S)).reshape(-1).astype(np.float32))
+        self.n_segments = S
+
+    def __len__(self) -> int:
+        return self.x_tk.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.x_tk[idx], self.x_tk1[idx], self.t_k[idx], self.t_k1[idx]
+
+
+# --------------------------------------------------------------------------- #
+# endpoint datasets
+# --------------------------------------------------------------------------- #
+
+class EndpointPairDataset(Dataset):
+    """``(x_0, x_T)`` endpoint pairs from ``vecs_train`` (block 0) and/or ``vecs_new`` (block 1) per an endpoint
+    manifest (rows ``[block, group_row, member_col, gid]``), what the Direct SI and EDM losses consume.
+    Normalization: float64 arithmetic, then cast to float32 (as ``get_fine_split_set``).
+
+    Checks: exactly ``B == N_traj`` pairs (S = 1), boundaries exactly ``[0, idx_T]`` on stored snapshots, the gid
+    mapping against the package's own key arrays, every gid in the training split, no duplicate row.
+    """
+
+    _BLOCKS = {0: ("vecs_train", "gids_train"), 1: ("vecs_new", "gids_new")}
+
+    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
+                 seed: int = 0,
+                 expected_t_values: list[float] | None = None):
+        man = _load_manifest(slice_manifest)
+        if man.get("schema") != ENDPOINT_SCHEMA:
+            raise AssertionError(
+                f"expected schema {ENDPOINT_SCHEMA}, got "
+                f"{man.get('schema')!r} in {slice_manifest}")
+        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
+
+        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
+        t_vals_man   = np.asarray(man["t_values"], dtype=np.float64)
+        S      = int(man["S"])
+        n_traj = int(man["N_traj"])
+        B      = int(man["B"])
+        rows   = np.asarray(man["trajectories"], dtype=np.int64)   # (N,4)
+
+        # --- manifest self-consistency ---
+        if S != 1 or n_traj != B:
+            raise AssertionError(f"endpoint slice must have S=1, N_traj=B; "
+                                 f"got S={S} N_traj={n_traj} B={B}")
+        if rows.shape != (n_traj, 4):
+            raise AssertionError(
+                f"manifest trajectories shape {rows.shape} != ({n_traj}, 4)")
+        if len(boundary_idx) != 2 or boundary_idx[0] != 0:
+            raise AssertionError(
+                f"endpoint boundary_idx must be [0, idx_T]; got "
+                f"{boundary_idx.tolist()}")
+        if not set(np.unique(rows[:, 0]).tolist()) <= set(self._BLOCKS):
+            raise AssertionError(
+                f"unknown block ids {np.unique(rows[:, 0]).tolist()}")
+        if np.unique(rows[:, :3], axis=0).shape[0] != n_traj:
+            raise AssertionError(
+                "duplicate (block, group_row, member_col) in slice")
+
+        t_values = np.asarray(_npz_member_memmap(fine_path, "t_values"))
+        if boundary_idx[-1] > t_values.shape[0] - 1:
+            raise AssertionError("boundary_idx exceeds fine snapshot grid")
+        if not np.allclose(t_values[boundary_idx], t_vals_man, atol=1e-12):
+            raise AssertionError(
+                f"manifest t_values {t_vals_man.tolist()} do not match the "
+                f"stored grid at boundary_idx {boundary_idx.tolist()}")
+        if expected_t_values is not None:
+            exp = np.asarray([float(t) for t in expected_t_values])
+            if len(exp) != len(t_vals_man) or not np.allclose(exp, t_vals_man,
+                                                              atol=1e-9):
+                raise AssertionError(
+                    f"sampling.t_values {exp.tolist()} != slice-manifest "
+                    f"t_values {t_vals_man.tolist()}")
+
+        # --- gid guards against the dataset's own key arrays + split.json ---
+        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
+            split = json.load(f)
+        train_set = set(split["train"])
+        forbidden = set(split["val"]) | set(split["test"])
+        for blk, (_, gid_key) in self._BLOCKS.items():
+            sub = rows[rows[:, 0] == blk]
+            if sub.shape[0] == 0:
+                continue
+            block_gids = np.asarray(_npz_member_memmap(fine_path, gid_key))
+            if not np.array_equal(block_gids[sub[:, 1]], sub[:, 3]):
+                raise AssertionError(
+                    f"manifest gid column inconsistent with "
+                    f"{gid_key}[group_row] (block {blk})")
+        man_gids = set(rows[:, 3].tolist())
+        if man_gids & forbidden:
+            raise AssertionError("slice manifest contains val/test gids")
+        if not man_gids <= train_set:
+            raise AssertionError("slice manifest contains gids outside "
+                                 "split.train")
+
+        # --- gather (N, 2, D) raw endpoints in manifest row order ---
+        first_blk = int(rows[0, 0])
+        probe = _npz_member_memmap(fine_path, self._BLOCKS[first_blk][0])
+        D = probe.shape[-1]
+        seg = np.empty((n_traj, 2, D), dtype=np.float64)
+        for blk, (vec_key, _) in self._BLOCKS.items():
+            sel = np.nonzero(rows[:, 0] == blk)[0]
+            if sel.size == 0:
+                continue
+            vecs = _npz_member_memmap(fine_path, vec_key)
+            seg[sel] = np.asarray(
+                vecs[rows[sel, 1], rows[sel, 2]])[:, boundary_idx, :]
+        if not np.all(np.isfinite(seg)):
+            raise AssertionError("non-finite values in selected trajectories")
+
+        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
+        mean, std = norm["mean"], norm["std"]
+        # Shared endpoint convention: float64 arithmetic, then float32 cast.
+        self.x0 = torch.from_numpy(((seg[:, 0] - mean) / std).astype(np.float32))
+        self.xT = torch.from_numpy(((seg[:, 1] - mean) / std).astype(np.float32))
+        if self.x0.shape[0] != B:
+            raise AssertionError(f"built {self.x0.shape[0]} pairs, expected {B}")
+
+        # Provenance of each pair (block, member, gid).
+        self.pair_block  = rows[:, 0].copy()
+        self.pair_member = rows[:, 2].copy()
+        self.pair_gid    = rows[:, 3].copy()
+        self.n_segments = 1
+        self.manifest = man
+        self.seed = int(seed)
+
+    def __len__(self) -> int:
+        return self.x0.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.x0[idx], self.xT[idx]
+
+
+class ValEndpointPairDataset(Dataset):
+    """Validation endpoint pairs ``(x_0, x_T)`` at the manifest's ``[0, idx_T]`` boundaries from the validation gids
+    of ``vecs_eval``, first ``n_members`` members per gid (the ``val_loss`` column of the log)."""
+
+    def __init__(self, packaged_dataset_path: str, slice_manifest: str,
+                 n_members: int = 20):
+        man = _load_manifest(slice_manifest)
+        if man.get("schema") != ENDPOINT_SCHEMA:
+            raise AssertionError(
+                f"expected schema {ENDPOINT_SCHEMA}, got "
+                f"{man.get('schema')!r} in {slice_manifest}")
+        fine_path = os.path.join(packaged_dataset_path, FINE_FILE)
+        boundary_idx = np.asarray(man["boundary_idx"], dtype=np.int64)
+
+        with open(os.path.join(packaged_dataset_path, "split.json")) as f:
+            split = json.load(f)
+        val_gids = sorted(split["val"])
+
+        gids_eval = np.asarray(_npz_member_memmap(fine_path, "gids_eval"))
+        vecs_eval = _npz_member_memmap(fine_path, "vecs_eval")
+        rows = []
+        for gid in val_gids:
+            where = np.nonzero(gids_eval == gid)[0]
+            if len(where) != 1:
+                raise AssertionError(f"val gid {gid} not found once in gids_eval")
+            rows.append(int(where[0]))
+        n_members = int(n_members)
+        if not (1 <= n_members <= vecs_eval.shape[1]):
+            raise ValueError(f"n_members {n_members} out of range")
+
+        seg = np.asarray(vecs_eval[rows][:, :n_members][:, :, boundary_idx, :])
+        if not np.all(np.isfinite(seg)):
+            raise AssertionError("non-finite values in vecs_eval slice")
+        D = seg.shape[-1]
+        seg = seg.reshape(-1, 2, D).astype(np.float64)
+
+        norm = np.load(os.path.join(packaged_dataset_path, "normalization.npz"))
+        mean, std = norm["mean"], norm["std"]
+        self.x0 = torch.from_numpy(((seg[:, 0] - mean) / std).astype(np.float32))
+        self.xT = torch.from_numpy(((seg[:, 1] - mean) / std).astype(np.float32))
+        self.n_segments = 1
+
+    def __len__(self) -> int:
+        return self.x0.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.x0[idx], self.xT[idx]
+
+
+# --------------------------------------------------------------------------- #
+# sampling and validation inputs at the forecast horizon
+# --------------------------------------------------------------------------- #
+
+def _x0_and_cloud_from_bank(bank_dir: str, gid: int, idx_T: int, cloud_n: int = 200):
+    """x_0 and a truth cloud of a test condition that the package does not carry (waves, NLKG), read from the
+    truth bank (group_<gid>.npz); used only for the test split when ``x0_bank`` is given. Fails loudly on any
+    inconsistency."""
+    p = os.path.join(bank_dir, f"group_{int(gid):04d}.npz")
+    if not os.path.isfile(p):
+        raise FileNotFoundError(
+            f"test-x0-bank fallback: no group_{int(gid):04d}.npz under {bank_dir} for gid {gid}")
+    with np.load(p) as z:
+        if int(z["gid"]) != int(gid):
+            raise AssertionError(
+                f"test-x0-bank fallback: {p} declares gid {int(z['gid'])}, expected {gid} "
+                "(bank file must be looked up and verified by gid, never by row position)")
+        vecs = np.asarray(z["vecs"])                        # (N, n_snapshots, D)
+    if not (0 <= idx_T < vecs.shape[1]):
+        raise ValueError(f"test-x0-bank fallback: fine_horizon_idx {idx_T} outside bank grid for gid {gid}")
+    x0_members = vecs[:, 0, :]
+    spread = float(np.max(np.abs(x0_members - x0_members[0])))
+    if spread > 1e-9:
+        raise AssertionError(
+            f"test-x0-bank fallback: gid {gid} x0 differs across bank members by {spread:.3e} — "
+            "shared-IC assumption violated")
+    n = min(int(cloud_n), vecs.shape[0])
+    return x0_members[0], np.asarray(vecs[:n, idx_T, :])
+
+
+def get_fine_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None
+                       ) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
+    """Per-gid x0 + truth cloud at the snapshot ``cfg.method.fine_horizon_idx``:
+    ``(x0_per_gid (n,D) f32, truth_clouds (n,V,D) f32, gids list)``.
+
+    x0 convention: member 0 at t=0. Cross-member spread at t=0 is asserted < 1e-9.
+
+    ``x0_bank`` (test split only): a truth-bank directory (group_<gid>.npz) used for a test gid absent from the
+    package's ``vecs_eval`` block; it changes nothing for the other splits or for a gid the package carries.
+    """
+    packaged_dir = cfg["data"]["packaged_dataset_path"]
+    method_cfg = cfg.get("method", {})
+    idx_T = int(method_cfg["fine_horizon_idx"])
+    fine_path = os.path.join(packaged_dir, FINE_FILE)
+
+    t_values = np.asarray(_npz_member_memmap(fine_path, "t_values"))
+    if not (1 <= idx_T <= t_values.shape[0] - 1):
+        raise ValueError(f"fine_horizon_idx {idx_T} outside grid")
+    # The horizon index must agree with the sampler chain's terminal time.
+    samp_t = method_cfg.get("sampling", {}).get("t_values")
+    if samp_t is not None:
+        if abs(float(samp_t[-1]) - float(t_values[idx_T])) > 1e-9:
+            raise AssertionError(
+                f"fine_horizon_idx={idx_T} (t={t_values[idx_T]}) != "
+                f"sampling.t_values[-1]={samp_t[-1]}")
+
+    with open(os.path.join(packaged_dir, "split.json")) as f:
+        split = json.load(f)
+    gids_override = cfg["data"].get("gids_override")
+    if gids_override is not None:
+        gids = sorted(int(g) for g in gids_override)
+    else:
+        if split_name not in split:
+            raise KeyError(f"split '{split_name}' not in split.json")
+        gids = sorted(split[split_name])
+        if split_name == "train":
+            # the training block carries the conditions of vecs_train only (not vecs_new)
+            canonical = set(np.asarray(
+                _npz_member_memmap(fine_path, "gids_train")).tolist())
+            gids = [g for g in gids if g in canonical]
+
+    if split_name in ("val", "test") and gids_override is None:
+        block, gid_key = "vecs_eval", "gids_eval"
+    else:
+        # the training conditions of vecs_train; any other gid fails loudly below
+        block, gid_key = "vecs_train", "gids_train"
+    block_gids = np.asarray(_npz_member_memmap(fine_path, gid_key))
+    vecs = _npz_member_memmap(fine_path, block)
+
+    norm = np.load(os.path.join(packaged_dir, "normalization.npz"))
+    mean, std = norm["mean"], norm["std"]
+
+    x0_list, cloud_list = [], []
+    for gid in gids:
+        where = np.nonzero(block_gids == gid)[0]
+        if len(where) != 1:
+            if split_name == "test" and x0_bank is not None:
+                x0_g, cloud_g = _x0_and_cloud_from_bank(x0_bank, gid, idx_T)
+                x0_list.append(x0_g); cloud_list.append(cloud_g)
+                continue
+            raise KeyError(
+                f"gid {gid} not found in {gid_key}: the package serves the conditions of "
+                "vecs_train and the validation and test conditions of vecs_eval")
+        r = int(where[0])
+        x0_members = np.asarray(vecs[r, :, 0, :])          # (V, D)
+        spread = float(np.max(np.abs(x0_members - x0_members[0])))
+        if spread > 1e-9:
+            raise AssertionError(
+                f"gid {gid}: x0 differs across members by {spread:.3e} — "
+                "shared-IC assumption violated")
+        x0_list.append(x0_members[0])
+        cloud_list.append(np.asarray(vecs[r, :, idx_T, :]))
+    x0_raw = np.stack(x0_list, axis=0)
+    clouds_raw = np.stack(cloud_list, axis=0)
+    if not (np.all(np.isfinite(x0_raw)) and np.all(np.isfinite(clouds_raw))):
+        raise AssertionError("non-finite values in fine split set")
+
+    # Same normalization arithmetic as get_split_set (float64 then cast).
+    x0 = torch.from_numpy(((x0_raw - mean) / std).astype(np.float32))
+    clouds = torch.from_numpy(((clouds_raw - mean) / std).astype(np.float32))
+    return x0, clouds, gids

code/bifurcation/data/loaders.py

@@ -0,0 +1,51 @@
+"""Training pairs and evaluation inputs of a dataset package (``trajectories_fine.npz``, ``split.json``,
+``normalization.npz``).
+
+* ``get_dataloaders(cfg, seed)`` -> ``(train_loader, val_loader)``: segment pairs ``(x_tk, x_tk1, t_k, t_k1)`` for
+  Cascaded SI (``method.data_mode: fine_segment_pairs``) or endpoint pairs ``(x_0, x_T)`` for Direct SI and EDM
+  (``fine_endpoint_pairs``), from the pair manifest named by ``method.fine_slice_manifest``.
+* ``get_split_set(cfg, split, seed, x0_bank)`` -> ``(x0_per_gid, truth_clouds, gids)`` at the forecast horizon
+  ``method.fine_horizon_idx``, the inputs of sampling and of validation.
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+from torch.utils.data import DataLoader
+
+from .fine import (EndpointPairDataset, SegmentPairDataset, ValEndpointPairDataset,  # noqa: F401
+                   ValSegmentPairDataset, get_fine_split_set)
+
+DATA_MODES = ("fine_segment_pairs", "fine_endpoint_pairs")
+
+
+def get_dataloaders(cfg: dict, seed: int = 0) -> Tuple[DataLoader, DataLoader]:
+    packaged_dir = cfg["data"]["packaged_dataset_path"]
+    batch_size = int(cfg["training"]["batch_size"])
+    method_cfg = cfg.get("method", {})
+    data_mode = method_cfg.get("data_mode")
+    num_workers = int(cfg["training"].get("num_workers", 0))
+    if data_mode == "fine_segment_pairs":
+        train_ds = SegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
+                                      expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
+        val_ds = ValSegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
+                                       n_members=int(method_cfg.get("fine_val_members", 20)))
+    elif data_mode == "fine_endpoint_pairs":
+        train_ds = EndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed,
+                                       expected_t_values=method_cfg.get("sampling", {}).get("t_values"))
+        val_ds = ValEndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"],
+                                        n_members=int(method_cfg.get("fine_val_members", 20)))
+    else:
+        raise ValueError(f"method.data_mode must be one of {DATA_MODES}, got {data_mode!r}")
+    g = torch.Generator()
+    g.manual_seed(seed)
+    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, generator=g,
+                              pin_memory=num_workers > 0)
+    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers,
+                            pin_memory=num_workers > 0)
+    return train_loader, val_loader
+
+
+def get_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None):
+    return get_fine_split_set(cfg, split_name, seed=seed, x0_bank=x0_bank)

code/bifurcation/data/manifests.py

@@ -0,0 +1,124 @@
+"""Pair manifests: which trajectories of ``trajectories_fine.npz`` a training run reads, and at which snapshots.
+
+A manifest directory holds ``slice_<name>.json`` and ``master_order.npz`` (the same rows; the loaders re-hash it).
+Rows are ``[group_row, member_col, gid]`` into ``vecs_train`` (segment pairs) or ``[block, group_row, member_col, gid]``
+with block 0 = ``vecs_train`` (endpoint pairs). ``bifurcation.train`` writes the manifest of a run before training.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+
+from .fine import ENDPOINT_SCHEMA as _fine_endpoint_schema
+from .fine import FINE_FILE, _npz_member_memmap, _sha256_file
+
+SEGMENT_SCHEMA = "segment_pairs_v1"
+ENDPOINT_SCHEMA = _fine_endpoint_schema
+
+
+def block_rows_segment(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
+    """Training conditions 0..n_ic-1 x realizations 0..members-1; rows [group_row, member_col, gid]."""
+    rows = np.empty((n_ic * members, 3), dtype=np.int64)
+    i = 0
+    for gi in range(n_ic):
+        for mj in range(members):
+            rows[i] = (gi, mj, int(gids_train[gi]))
+            i += 1
+    return rows
+
+
+def block_rows_endpoint(gids_train: np.ndarray, n_ic: int, members: int) -> np.ndarray:
+    """The same enumeration for endpoint pairs; rows [block, group_row, member_col, gid] with block 0 = vecs_train."""
+    rows = np.empty((n_ic * members, 4), dtype=np.int64)
+    i = 0
+    for gi in range(n_ic):
+        for mj in range(members):
+            rows[i] = (0, gi, mj, int(gids_train[gi]))
+            i += 1
+    return rows
+
+
+def segment_geometry(T: float, S: int, t_values: np.ndarray) -> dict:
+    """Uniform segmentation of [0, T] on the stored snapshot grid: n_T = index of T (an exact snapshot), S must
+    divide n_T; boundary_idx = [0, stride, ..., n_T]; the segment times are read from the stored grid."""
+    t_values = np.asarray(t_values, dtype=np.float64)
+    hits = np.nonzero(np.isclose(t_values, T, atol=1e-9))[0]
+    if len(hits) != 1:
+        raise ValueError(f"T={T} is not a stored snapshot time (grid {t_values[:3]}...)")
+    n_T = int(hits[0])
+    if n_T % S != 0:
+        raise ValueError(f"S={S} does not divide n_T={n_T} (T={T}); admissible S: "
+                         f"{[s for s in range(1, n_T + 1) if n_T % s == 0]}")
+    stride = n_T // S
+    boundary_idx = [q * stride for q in range(S + 1)]
+    return {"T": float(T), "S": int(S), "n_T": n_T, "stride": stride, "delta": float(T) / S,
+            "boundary_idx": boundary_idx, "t_values": [float(t_values[b]) for b in boundary_idx]}
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")
+
+
+def _write_master(slice_dir: Path, rows: np.ndarray, note: str) -> str:
+    slice_dir.mkdir(parents=True, exist_ok=True)
+    path = slice_dir / "master_order.npz"
+    if rows.shape[1] == 3:
+        payload = dict(group_row=rows[:, 0].astype(np.int32), member_col=rows[:, 1].astype(np.int32),
+                       gid=rows[:, 2].astype(np.int32))
+    else:
+        payload = dict(block=rows[:, 0].astype(np.int32), group_row=rows[:, 1].astype(np.int32),
+                       member_col=rows[:, 2].astype(np.int32), gid=rows[:, 3].astype(np.int32))
+    payload["algorithm"] = np.bytes_(note.encode())
+    np.savez(path, **payload)
+    return _sha256_file(str(path))
+
+
+def dataset_provenance(packaged_dir: str | os.PathLike) -> dict:
+    packaged_dir = Path(packaged_dir)
+    return {"dataset_path": str(packaged_dir / FINE_FILE),
+            "split_sha256": _sha256_file(str(packaged_dir / "split.json")),
+            "normalization_sha256": _sha256_file(str(packaged_dir / "normalization.npz"))}
+
+
+def write_segment_manifest(slice_dir: Path, rows3: np.ndarray, geometry: dict, name: str,
+                           packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
+    S = int(geometry["S"])
+    master_sha = _write_master(slice_dir, rows3, f"complete block ({name}): canonical enumeration")
+    man = {"schema": SEGMENT_SCHEMA, "arm": name, "cell": name, **geometry,
+           "N_traj": int(rows3.shape[0]), "B": int(rows3.shape[0]) * S,
+           "trajectories": rows3.tolist(), "master_order_sha256": master_sha,
+           **dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
+    out = slice_dir / f"slice_{name}.json"
+    out.write_text(json.dumps(man))
+    return out
+
+
+def write_endpoint_manifest(slice_dir: Path, rows4: np.ndarray, idx_T: int, t_values: np.ndarray, name: str,
+                            packaged_dir: str | os.PathLike, extra: dict | None = None) -> Path:
+    master_sha = _write_master(slice_dir, rows4, f"complete block ({name}): canonical enumeration, block 0 = vecs_train")
+    man = {"schema": ENDPOINT_SCHEMA, "arm": name, "S": 1,
+           "T": float(t_values[idx_T]), "n_T": int(idx_T),
+           "N_traj": int(rows4.shape[0]), "B": int(rows4.shape[0]),
+           "boundary_idx": [0, int(idx_T)], "t_values": [0.0, float(t_values[idx_T])],
+           "blocks": {"0": "vecs_train"},
+           "trajectories": rows4.tolist(), "master_order_sha256": master_sha,
+           **dataset_provenance(packaged_dir), "built_at": _now(), **(extra or {})}
+    out = slice_dir / f"slice_{name}.json"
+    out.write_text(json.dumps(man))
+    return out
+
+
+def fine_grid(packaged_dir: str | os.PathLike) -> dict:
+    """Key arrays of the fine file (memmapped; cheap)."""
+    fine = str(Path(packaged_dir) / FINE_FILE)
+    out = {"t_values": np.asarray(_npz_member_memmap(fine, "t_values")),
+           "gids_train": np.asarray(_npz_member_memmap(fine, "gids_train")),
+           "gids_eval": np.asarray(_npz_member_memmap(fine, "gids_eval"))}
+    for k in ("vecs_train", "vecs_eval"):
+        out[k + "_shape"] = tuple(_npz_member_memmap(fine, k).shape)
+    return out

code/bifurcation/fid/__init__.py

@@ -0,0 +1,15 @@
+"""The image metrics: csFID (the paper's metric), cFID and cKID.
+
+  backend.py   the pinned Clean-FID pipeline: one resize (PIL bicubic, 299x299), the TF-Inception-v3 weights
+               (2015-12-05), the pool3 feature (cFID, cKID) and the spatial mixed_6/conv feature (csFID)
+  adapters.py  the dataset adapter psi_D: physical state -> uint8 (H, W, 3) image, with constants fitted on
+               training data only (records/<system>/fid_adapter.json)
+  euler.py, passive_scalar.py, heterogeneous_waves(_augmented).py, nlkg.py   the adapters of the five systems
+  evaluate.py  the metrics of all test conditions of a model against the truth bank
+"""
+from .adapters import ADAPTERS, FIDAdapter, get_adapter, register_adapter  # noqa: F401
+from . import euler as _euler  # noqa: F401  (registers euler_coarse_vorticity)
+from . import passive_scalar as _passive_scalar  # noqa: F401  (registers passive_scalar_coarse)
+from . import heterogeneous_waves as _heterogeneous_waves  # noqa: F401  (base class of the waves adapter)
+from . import heterogeneous_waves_augmented as _heterogeneous_waves_augmented  # noqa: F401  (registers heterogeneous_waves_augmented_coarse: waves)
+from . import nlkg as _nlkg  # noqa: F401  (registers nlkg_coarse_u: NLKG; reads the displacement u)

code/bifurcation/fid/adapters.py

@@ -0,0 +1,153 @@
+"""Dataset adapters psi_D: physical state -> three-channel uint8 image.
+
+Contract (enforced by the base class):
+  * ``fit(train_states_raw)`` computes the adapter's numerical constants from training data only and fixes
+    them (``params``; records/<system>/fid_adapter.json);
+  * ``to_images(states_raw)`` is deterministic given ``params``; no per-image normalization;
+  * signed fields with a meaningful zero: symmetric clip [-c, c], c = Q_q(|x_train|); nonnegative fields:
+    [0, c], c = Q_q(x_train); the map to uint8 is the fixed linear map of the range onto [0, 255];
+  * scalar fields -> three identical channels unless the dataset card defines a physical 3-channel construction.
+
+States enter raw (physical units, un-standardized). configs/<system>/dataset.yaml names the adapter and its parameters.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+ADAPTERS: dict[str, type] = {}
+
+
+def register_adapter(name: str):
+    def deco(cls):
+        cls.name = name
+        ADAPTERS[name] = cls
+        return cls
+    return deco
+
+
+def get_adapter(name: str, **kwargs) -> "FIDAdapter":
+    if name not in ADAPTERS:
+        raise KeyError(f"unknown FID adapter {name!r} (registered: {sorted(ADAPTERS)})")
+    return ADAPTERS[name](**kwargs)
+
+
+def quantize_symmetric(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]:
+    """Signed field -> uint8 grayscale: rint((clip(f, -c, c) + c) / (2c) * 255); zero -> 128. Returns the
+    clipped-pixel fraction too."""
+    f = np.asarray(fields, dtype=np.float64)
+    clipped = float(np.mean((f > c) | (f < -c))) if f.size else 0.0
+    g = np.rint((np.clip(f, -c, c) + c) / (2.0 * c) * 255.0)
+    return np.clip(g, 0, 255).astype(np.uint8), clipped
+
+
+def quantize_nonnegative(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]:
+    f = np.asarray(fields, dtype=np.float64)
+    clipped = float(np.mean(f > c)) if f.size else 0.0
+    g = np.rint(np.clip(f, 0.0, c) / c * 255.0)
+    return np.clip(g, 0, 255).astype(np.uint8), clipped
+
+
+def gray_to_rgb(g: np.ndarray) -> np.ndarray:
+    return np.repeat(g[..., None], 3, axis=-1)
+
+
+class FIDAdapter:
+    name = "abstract"
+    field_kind = "signed"           # 'signed' (symmetric clip) | 'nonnegative' ([0, c])
+
+    def __init__(self, quantile: float = 0.9999, **kwargs):
+        self.quantile = float(quantile)
+        self.params: dict | None = None
+        self.kwargs = dict(kwargs)
+
+    # dataset-specific: raw state vectors (N, D) -> scalar fields (N, H, W)
+    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
+        raise NotImplementedError
+
+    def fit(self, train_states_raw: np.ndarray, source: dict | None = None, pixel_stride: int = 1,
+            chunk: int = 4000) -> dict:
+        """Constants from TRAINING data only. Freezes ``self.params``. ``pixel_stride`` subsamples the field
+        pixels (every k-th in each axis);
+        states are converted in chunks to bound memory."""
+        k = int(pixel_stride)
+        vals, shape = [], None
+        for i in range(0, len(train_states_raw), int(chunk)):
+            f = self.to_fields(train_states_raw[i:i + int(chunk)])
+            shape = list(f.shape[1:])
+            f = f[:, ::k, ::k] if f.ndim == 3 else f
+            vals.append((np.abs(f) if self.field_kind == "signed" else f).astype(np.float32).ravel())
+        vals = np.concatenate(vals)
+        c = float(np.quantile(vals, self.quantile))
+        rng = [-c, c] if self.field_kind == "signed" else [0.0, c]
+        self.params = {"adapter": self.name, "field_kind": self.field_kind, "quantile": self.quantile, "clip_c": c,
+                       "range": rng, "n_train_states": int(len(train_states_raw)), "field_shape": shape,
+                       "pixel_stride": k, "n_pixel_values": int(vals.size),
+                       "map": ("uint8 = rint((clip(w,-c,c)+c)/(2c)*255)" if self.field_kind == "signed"
+                               else "uint8 = rint(clip(w,0,c)/c*255)"),
+                       "channels": "3 identical (grayscale scalar field), no colormap",
+                       "source": source or {}, **self.card()}
+        return self.params
+
+    def card(self) -> dict:
+        """Dataset-card text of the adapter (channel construction, reconstruction) — override."""
+        return {}
+
+    def load_params(self, params: dict) -> "FIDAdapter":
+        if params.get("adapter") != self.name:
+            raise ValueError(f"params belong to adapter {params.get('adapter')!r}, not {self.name!r}")
+        self.params = dict(params)
+        self.quantile = float(params.get("quantile", self.quantile))
+        return self
+
+    def to_images(self, states_raw: np.ndarray) -> tuple[np.ndarray, float]:
+        """RAW states -> uint8 (N, H, W, 3) and the clipped-pixel fraction. Requires fitted params."""
+        if self.params is None:
+            raise RuntimeError("adapter not fitted / loaded (constants come from training data only)")
+        fields = self.to_fields(states_raw)
+        c = float(self.params["clip_c"])
+        g, frac = (quantize_symmetric(fields, c) if self.field_kind == "signed" else quantize_nonnegative(fields, c))
+        return gray_to_rgb(g), frac
+
+    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
+        """Optional physical diagnostics of a generated ensemble against the truth. Default: none."""
+        return {}
+
+    def save(self, path: Path) -> Path:
+        Path(path).write_text(json.dumps(self.params, indent=1) + "\n")
+        return Path(path)
+
+    @classmethod
+    def load(cls, path: Path, **kwargs) -> "FIDAdapter":
+        params = json.loads(Path(path).read_text())
+        return get_adapter(params["adapter"], quantile=params.get("quantile", 0.9999), **kwargs).load_params(params)
+
+
+@register_adapter("vector_grid")
+class VectorGridAdapter(FIDAdapter):
+    """Generic adapter for vector states with no field structure: the D coordinates are laid out row-major on
+    the smallest square grid (zero padding), each cell drawn as a ``cell`` x ``cell`` block (nearest
+    upsampling, no interpolation), symmetric clip from the training quantile. Used for datasets whose state is
+    a plain vector and for the synthetic test dataset."""
+    field_kind = "signed"
+
+    def __init__(self, quantile: float = 0.9999, cell: int = 8, **kwargs):
+        super().__init__(quantile, **kwargs)
+        self.cell = int(cell)
+
+    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
+        x = np.asarray(states_raw, dtype=np.float64)
+        if x.ndim == 1:
+            x = x[None]
+        N, D = x.shape
+        side = int(np.ceil(np.sqrt(D)))
+        grid = np.zeros((N, side * side), dtype=np.float64)
+        grid[:, :D] = x
+        grid = grid.reshape(N, side, side)
+        return np.repeat(np.repeat(grid, self.cell, axis=1), self.cell, axis=2)
+
+    def card(self) -> dict:
+        return {"reconstruction": f"row-major {int(np.ceil(np.sqrt(1)))}x? square grid of the state coordinates, zero-padded, "
+                                  f"each coordinate a {self.cell}x{self.cell} block (nearest)", "cell": self.cell}

code/bifurcation/fid/backend.py

@@ -0,0 +1,206 @@
+"""The Clean-FID pipeline, pinned and checked.
+
+Everything after the dataset adapter is fixed here for every dataset:
+  * images arrive as uint8 (N, H, W, 3) in [0, 255] (one fixed image range);
+  * a single resize to 299x299 with clean-fid's "clean" resizer (PIL bicubic, float output, antialiasing is
+    PIL's default for its resampling filters); the TorchScript Inception is built with resize_inside=False
+    and asserts 299x299 input, so a second resize cannot happen silently;
+  * the pinned Inception-v3 weights (inception-2015-12-05.pt, sha256 checked), pool3 features (2048-D);
+  * FID = clean-fid's frechet_distance formula; KID = clean-fid's kernel_distance (unbiased polynomial-kernel
+    MMD^2, num_subsets / max_subset_size from the evaluation settings).
+
+``pins()`` returns the record that every FID manifest carries; ``check_pins(expected)`` raises on any drift.
+"""
+from __future__ import annotations
+
+import hashlib
+import importlib.metadata as _md
+import os
+from pathlib import Path
+
+import numpy as np
+
+INCEPTION_DIR = Path(os.environ.get("CLEAN_FID_CACHE", Path.home() / ".cache" / "clean-fid"))
+INCEPTION_FILE = "inception-2015-12-05.pt"
+IMAGE_SIZE = 299
+
+# ── the spatial Inception feature of csFID (Nash et al., ICML 2021) ─────────────────────
+# The Inception-v3 activation "mixed_6/conv" -- the FIRST branch of the mixed_6 Inception block, BEFORE its
+# concatenation with tower/tower_1/tower_2 (openai/guided-diffusion evaluator.py FID_SPATIAL_NAME) -- is a
+# 192 x 17 x 17 tensor; its first 7 channels are kept and flattened NHWC to 17*17*7 = 2023 components per image,
+# without spatial average pooling (that pooling is what makes the 2048-D pool3 cFID feature layout-tolerant).
+# In this TorchScript port the tensor is the submodule layers.mixed_6.conv; the module names come from the TF op
+# names. The stepwise traversal used to read it reproduces layers.forward exactly on CPU (verify_spatial_tap).
+SPATIAL_LAYER = "mixed_6/conv"
+SPATIAL_CHANNELS = 7
+SPATIAL_GRID = 17
+SPATIAL_DIM = SPATIAL_GRID * SPATIAL_GRID * SPATIAL_CHANNELS          # 2023
+POOL3_DIM = 2048
+_MIXED5_SHAPE = (768, SPATIAL_GRID, SPATIAL_GRID)                     # canonical input of the mixed_6 block
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 22), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def pins() -> dict:
+    import PIL
+    import torch
+    p = INCEPTION_DIR / INCEPTION_FILE
+    return {"backend": "clean-fid", "clean_fid_version": _md.version("clean-fid"), "mode": "clean",
+            "resize": {"backend": "PIL", "filter": "bicubic", "size": IMAGE_SIZE, "once": True,
+                       "implementation": "cleanfid.resize.build_resizer('clean') = make_resizer('PIL', False, 'bicubic', (299, 299))"},
+            "inception": {"file": INCEPTION_FILE, "path": str(p), "sha256": _sha256(p) if p.is_file() else None,
+                          "feature_layer": "pool3", "feature_dim": POOL3_DIM, "resize_inside": False},
+            "spatial": {"feature_layer": SPATIAL_LAYER, "submodule": "layers.mixed_6.conv",
+                        "tap": "first branch of the mixed_6 block, before concatenation with tower/tower_1/tower_2",
+                        "channels_kept": SPATIAL_CHANNELS, "grid": [SPATIAL_GRID, SPATIAL_GRID], "feature_dim": SPATIAL_DIM,
+                        "flatten": "NHWC", "spatial_pooling": "none",
+                        "reference": "Nash et al., ICML 2021; openai/guided-diffusion evaluator.py FID_SPATIAL_NAME='mixed_6/conv:0'"},
+            "fid": "cleanfid.fid.frechet_distance (scipy sqrtm, eps 1e-6 fallback)",
+            "kid": "cleanfid.fid.kernel_distance (unbiased polynomial-kernel MMD^2)",
+            "torch": torch.__version__, "pillow": PIL.__version__, "numpy": np.__version__}
+
+
+def check_pins(expected: dict) -> dict:
+    """``expected`` = the fid settings (clean_fid_version, inception.sha256, inception.file). Raises on drift."""
+    got = pins()
+    probs = []
+    if str(expected.get("clean_fid_version")) != got["clean_fid_version"]:
+        probs.append(f"clean-fid {got['clean_fid_version']} != pinned {expected.get('clean_fid_version')}")
+    exp_inc = expected.get("inception") or {}
+    if exp_inc.get("file") and exp_inc["file"] != got["inception"]["file"]:
+        probs.append("inception file name differs")
+    if got["inception"]["sha256"] is None:
+        probs.append(f"Inception weights not found at {got['inception']['path']} (download once with cleanfid on a node with network)")
+    elif exp_inc.get("sha256") and exp_inc["sha256"] != got["inception"]["sha256"]:
+        probs.append(f"Inception weights sha256 {got['inception']['sha256'][:12]} != pinned {exp_inc['sha256'][:12]}")
+    if probs:
+        raise RuntimeError("FID pipeline pin violated: " + "; ".join(probs))
+    return got
+
+
+def build_extractor(device):
+    """clean-fid 'clean' extractor: TorchScript TF-Inception, weights from INCEPTION_DIR, no internal resize."""
+    from cleanfid.inception_torchscript import InceptionV3W
+    INCEPTION_DIR.mkdir(parents=True, exist_ok=True)
+    model = InceptionV3W(str(INCEPTION_DIR), download=not (INCEPTION_DIR / INCEPTION_FILE).is_file(),
+                         resize_inside=False).to(device).eval()
+    return model
+
+
+def inception_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
+    """uint8 (N, H, W, 3) -> float32 (N, 2048): a single clean resize per image, then the pinned network."""
+    import torch
+    from cleanfid.resize import build_resizer
+    if imgs_u8.dtype != np.uint8 or imgs_u8.ndim != 4 or imgs_u8.shape[-1] != 3:
+        raise ValueError(f"adapter output must be uint8 (N, H, W, 3); got {imgs_u8.dtype} {imgs_u8.shape}")
+    resizer = build_resizer("clean")
+    out = []
+    with torch.no_grad():
+        for i in range(0, len(imgs_u8), int(batch)):
+            xs = [resizer(im) for im in imgs_u8[i:i + int(batch)]]           # float32 (299, 299, 3) in [0, 255]
+            x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
+            if tuple(x.shape[-2:]) != (IMAGE_SIZE, IMAGE_SIZE):
+                raise AssertionError("resizer did not produce 299x299")
+            out.append(extractor(x).cpu().numpy())
+    return np.concatenate(out).astype(np.float32)
+
+
+def fid_from_features(f1: np.ndarray, f2: np.ndarray) -> float:
+    from cleanfid.fid import frechet_distance
+    f1 = np.asarray(f1, np.float64); f2 = np.asarray(f2, np.float64)
+    return float(frechet_distance(f1.mean(0), np.cov(f1, rowvar=False), f2.mean(0), np.cov(f2, rowvar=False)))
+
+
+def kid_from_features(f1: np.ndarray, f2: np.ndarray, num_subsets: int = 100, max_subset_size: int = 1000) -> float:
+    from cleanfid.fid import kernel_distance
+    return float(kernel_distance(np.asarray(f1, np.float64), np.asarray(f2, np.float64),
+                                 num_subsets=int(num_subsets), max_subset_size=int(max_subset_size)))
+
+
+# ── csFID feature extraction ────────────────────────────────────────────────────────────────────────────────
+
+def _pre_mixed6_stages(layers):
+    """The layer sequence of layers.forward up to (not including) mixed_6 -- mirrors its TorchScript source."""
+    return [layers.conv, layers.conv_1, layers.conv_2, layers.pool0, layers.conv_3, layers.conv_4, layers.pool1,
+            layers.mixed, layers.mixed_1, layers.mixed_2, layers.mixed_3, layers.mixed_4, layers.mixed_5]
+
+
+def _post_mixed6_stages(layers):
+    return [layers.mixed_6, layers.mixed_7, layers.mixed_8, layers.mixed_9, layers.mixed_10, layers.pool2]
+
+
+def spatial_tap(extractor, x_normalized):
+    """x already resized to 299x299 and normalized as InceptionV3W does ((x - 128) / 128) -> the mixed_6/conv
+    activation (N, 192, 17, 17)."""
+    L = extractor.layers
+    h = x_normalized
+    for st in _pre_mixed6_stages(L):
+        h = st.forward(h)
+    return L.mixed_6.conv.forward(h)
+
+
+def spatial_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
+    """uint8 (N, H, W, 3) -> float32 (N, 2023): the canonical sFID feature. Same clean resize and normalization
+    as inception_features; only the tapped tensor differs."""
+    import torch
+    from cleanfid.resize import build_resizer
+    if imgs_u8.dtype != np.uint8 or imgs_u8.ndim != 4 or imgs_u8.shape[-1] != 3:
+        raise ValueError(f"adapter output must be uint8 (N, H, W, 3); got {imgs_u8.dtype} {imgs_u8.shape}")
+    resizer = build_resizer("clean")
+    out = []
+    with torch.no_grad():
+        for i in range(0, len(imgs_u8), int(batch)):
+            xs = [resizer(im) for im in imgs_u8[i:i + int(batch)]]
+            x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
+            if tuple(x.shape[-2:]) != (IMAGE_SIZE, IMAGE_SIZE):
+                raise AssertionError("resizer did not produce 299x299")
+            sp = spatial_tap(extractor, (x - 128) / 128)[:, :SPATIAL_CHANNELS]
+            out.append(sp.permute(0, 2, 3, 1).reshape(sp.shape[0], -1).cpu().numpy())
+    feats = np.concatenate(out).astype(np.float32)
+    if feats.shape[1] != SPATIAL_DIM:
+        raise AssertionError(f"spatial feature dim {feats.shape[1]} != {SPATIAL_DIM}")
+    return feats
+
+
+def verify_spatial_tap(extractor=None) -> dict:
+    """Check on CPU that (a) the stepwise traversal reproduces layers.forward exactly (so the tap sits on the
+    cFID path) and (b) the tapped tensor is the canonical 17x17 mixed_6/conv representation. Raises on failure;
+    returns the record that MANIFEST_CSFID.json carries."""
+    import torch
+    ext = extractor if extractor is not None else build_extractor(torch.device("cpu"))
+    L = ext.layers
+    if not (hasattr(L, "mixed_5") and hasattr(L, "mixed_6") and hasattr(L.mixed_6, "conv")):
+        raise RuntimeError("this Inception build does not expose layers.mixed_6.conv -- canonical sFID unavailable")
+    x = (torch.rand(3, 3, IMAGE_SIZE, IMAGE_SIZE) * 255.0 - 128) / 128
+    with torch.no_grad():
+        h = x
+        for st in _pre_mixed6_stages(L):
+            h = st.forward(h)
+        mixed5 = h
+        sp = L.mixed_6.conv.forward(mixed5)
+        for st in _post_mixed6_stages(L):
+            h = st.forward(h)
+        diff = float((h.view(3, POOL3_DIM) - L.forward(x).view(3, POOL3_DIM)).abs().max())
+    flat_dim = int(sp[:, :SPATIAL_CHANNELS].permute(0, 2, 3, 1).reshape(3, -1).shape[1])
+    rec = {"mixed_5_output_shape_CHW": list(sp.shape[1:]) and list(mixed5.shape[1:]),
+           "mixed_6_conv_output_shape_CHW": list(sp.shape[1:]), "channels_total": int(sp.shape[1]),
+           "channels_kept": SPATIAL_CHANNELS, "feature_dim": flat_dim, "flatten": "NHWC",
+           "stepwise_vs_layers_forward_max_abs_diff_cpu": diff, "bitexact_cpu": diff == 0.0}
+    probs = []
+    if tuple(mixed5.shape[1:]) != _MIXED5_SHAPE:
+        probs.append(f"mixed_5 output {tuple(mixed5.shape[1:])} != canonical {_MIXED5_SHAPE}")
+    if tuple(sp.shape[2:]) != (SPATIAL_GRID, SPATIAL_GRID):
+        probs.append(f"mixed_6/conv spatial {tuple(sp.shape[2:])} != {SPATIAL_GRID}x{SPATIAL_GRID}")
+    if flat_dim != SPATIAL_DIM:
+        probs.append(f"feature dim {flat_dim} != {SPATIAL_DIM}")
+    if diff != 0.0:
+        probs.append(f"stepwise traversal differs from layers.forward on CPU by {diff}")
+    if probs:
+        raise RuntimeError("spatial tap verification failed: " + "; ".join(probs))
+    return rec

code/bifurcation/fid/euler.py

@@ -0,0 +1,58 @@
+"""Euler FID adapter: 196-D coarse spectral vorticity (|k| <= 8, Hermitian half, raw coefficients
+omega_hat / N_sim^2) -> the EXACT low-pass vorticity field on a 128x128 grid -> symmetric clip -> uint8 -> 3
+identical channels.
+
+The reconstruction is the same inverse FFT as ``bifurcation.models.spectral.FieldReconstructor`` (the mode
+index map is imported from there — one source of truth), vectorized over N states; equality with the
+reference per-vector reconstruction.
+
+Constant: c = Q_0.9999(|omega|) over the training pool (records/euler/fid_adapter.json). Physical diagnostic:
+the enstrophy ratio (mean enstrophy of the generated ensemble / mean enstrophy of the truth, from the coefficients).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from ..models.spectral import K_C, mode_int_wavenumbers
+from .adapters import FIDAdapter, register_adapter
+
+
+@register_adapter("euler_coarse_vorticity")
+class EulerVorticityAdapter(FIDAdapter):
+    field_kind = "signed"
+
+    def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs):
+        super().__init__(quantile, **kwargs)
+        self.n_disp, self.k_c = int(n_disp), int(k_c)
+        kx, ky = mode_int_wavenumbers(self.k_c)
+        self.n_modes = len(kx)
+        self.ix, self.iy = kx % self.n_disp, ky % self.n_disp
+        self.ix_c, self.iy_c = (-kx) % self.n_disp, (-ky) % self.n_disp
+        if self.k_c * 2 >= self.n_disp:
+            raise ValueError("n_disp must exceed 2 k_c (alias-free reconstruction)")
+
+    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
+        x = np.asarray(states_raw, dtype=np.float64)
+        if x.ndim == 1:
+            x = x[None]
+        if x.shape[1] != 2 * self.n_modes:
+            raise ValueError(f"expected {2 * self.n_modes}-D coarse spectral states, got D={x.shape[1]}")
+        N = x.shape[0]
+        coeff = x[:, :self.n_modes] + 1j * x[:, self.n_modes:]
+        OH = np.zeros((N, self.n_disp, self.n_disp), dtype=np.complex128)
+        OH[:, self.ix, self.iy] += coeff
+        OH[:, self.ix_c, self.iy_c] += np.conj(coeff)
+        return (self.n_disp ** 2) * np.fft.ifft2(OH, axes=(1, 2)).real
+
+    def card(self) -> dict:
+        return {"state": "196-D coarse spectral vorticity (98 complex modes |k| <= 8, Hermitian half, omega_hat / N_sim^2)",
+                "reconstruction": f"exact inverse FFT of the |k| <= {self.k_c} modes on a {self.n_disp}x{self.n_disp} grid (alias-free)",
+                "n_disp": self.n_disp, "k_c": self.k_c, "clip": "symmetric [-c, c], c = Q_q(|omega_T|) over the training pool's final-time fields",
+                "physical_diagnostic": "enstrophy_ratio = mean_gen sum|omega_hat_k|^2 / mean_truth sum|omega_hat_k|^2"}
+
+    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
+        def ens(v):
+            v = np.asarray(v, np.float64)
+            return float(np.mean(np.sum(v ** 2, axis=1)))       # sum over Re/Im of all Hermitian-half modes
+        et = ens(true_raw)
+        return {"enstrophy_ratio": ens(gen_raw) / et if et > 0 else float("nan")}

code/bifurcation/fid/evaluate.py

@@ -0,0 +1,224 @@
+"""Conditional image metrics per test condition: csFID (the paper's metric), cFID and cKID.
+
+  generated_states(samples, variant, packaged, n)   RAW generated states of every condition from a samples directory
+  truth_features / truth_spatial_features           cached features of the first n_truth bank members of one condition
+  evaluate_csfid_run(...)                           csFID of all conditions -> csfid_per_ic.csv, csfid_summary.json
+  evaluate_run(...)                                 cFID and cKID (pool3 features) -> per_ic_fid.csv, summary.json
+
+The truth bank holds ``<bank_dir>/group_<gid:04d>.npz`` with ``vecs`` (n_members, n_snapshots, D) in raw units; the
+first n_truth members are the reference ensemble of every method.
+"""
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from . import backend as BE
+from .adapters import FIDAdapter
+
+
+def generated_states(samples_dir: Path, variant: str, packaged: Path, n: int) -> dict:
+    """{gid: RAW generated states (n, D)} of ``samples_dir/<variant>/group_<gid>.npz`` (``gen_xT_norm``, un-normalized
+    with the package's normalization.npz)."""
+    z = np.load(Path(packaged) / "normalization.npz")
+    mean, std = np.asarray(z["mean"], np.float64), np.asarray(z["std"], np.float64)
+    out = {}
+    for p in sorted((Path(samples_dir) / variant).glob("group_[0-9][0-9][0-9][0-9].npz")):
+        with np.load(p) as f:
+            g = np.asarray(f["gen_xT_norm"], np.float64)[: int(n)]
+            out[int(f["gid"])] = g * std + mean
+    return out
+
+
+def bank_members(bank_dir: Path, gid: int, t_idx: int, n: int | None = None) -> np.ndarray:
+    """RAW states (n, D) of the first n members of test condition ``gid`` at snapshot ``t_idx``."""
+    p = Path(bank_dir) / f"group_{int(gid):04d}.npz"
+    if not p.is_file():
+        raise FileNotFoundError(f"truth bank has no group_{int(gid):04d}.npz under {bank_dir}")
+    with np.load(p) as z:
+        vecs = z["vecs"]
+        n_have = int(vecs.shape[0])
+        if n is not None and n_have < int(n):
+            raise ValueError(f"bank group {gid}: {n_have} members < required {n}")
+        return np.asarray(vecs[: (int(n) if n is not None else n_have), int(t_idx), :], dtype=np.float64)
+
+
+def bank_inventory(bank_dir: Path, gids: list[int]) -> dict:
+    out = {}
+    for g in gids:
+        p = Path(bank_dir) / f"group_{int(g):04d}.npz"
+        if not p.is_file():
+            out[int(g)] = 0
+            continue
+        with np.load(p) as z:
+            out[int(g)] = int(z["vecs"].shape[0])
+    return out
+
+
+def features_for(states_raw: np.ndarray, adapter: FIDAdapter, extractor, device, batch: int = 200) -> tuple[np.ndarray, float]:
+    imgs, clipped = adapter.to_images(states_raw)
+    return BE.inception_features(imgs, extractor, device, batch=batch), clipped
+
+
+def truth_features(bank_dir: Path, gid: int, t_idx: int, n_truth: int, adapter: FIDAdapter, extractor, device,
+                   cache_dir: Path, batch: int = 200) -> tuple[np.ndarray, dict]:
+    """Cached per (gid, n_truth, adapter clip_c): the truth side does not depend on the model."""
+    cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True)
+    tag = f"gid{int(gid):04d}_t{int(t_idx)}_n{int(n_truth)}_c{adapter.params['clip_c']:.6g}"
+    p = cache_dir / f"truth_{tag}.npz"
+    if p.is_file():
+        with np.load(p) as z:
+            return np.asarray(z["feats"]), {"cache": str(p), "clipped_frac": float(z["clipped_frac"]), "reused": True}
+    states = bank_members(bank_dir, gid, t_idx, n_truth)
+    feats, clipped = features_for(states, adapter, extractor, device, batch)
+    np.savez(p, feats=feats, clipped_frac=np.float64(clipped), gid=np.int64(gid), n=np.int64(n_truth))
+    return feats, {"cache": str(p), "clipped_frac": clipped, "reused": False}
+
+
+def spatial_features_for(states_raw: np.ndarray, adapter: FIDAdapter, extractor, device, batch: int = 200) -> tuple[np.ndarray, float]:
+    imgs, clipped = adapter.to_images(states_raw)
+    return BE.spatial_features(imgs, extractor, device, batch=batch), clipped
+
+
+def truth_spatial_features(bank_dir: Path, gid: int, t_idx: int, n_truth: int, adapter: FIDAdapter, extractor, device,
+                           cache_dir: Path, batch: int = 200) -> tuple[np.ndarray, dict]:
+    """Like truth_features, for the spatial feature (a separate cache file)."""
+    cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True)
+    tag = f"gid{int(gid):04d}_t{int(t_idx)}_n{int(n_truth)}_c{adapter.params['clip_c']:.6g}"
+    p = cache_dir / f"truth_spatial_{tag}.npz"
+    if p.is_file():
+        with np.load(p) as z:
+            return np.asarray(z["feats"]), {"cache": str(p), "clipped_frac": float(z["clipped_frac"]), "reused": True}
+    states = bank_members(bank_dir, gid, t_idx, n_truth)
+    feats, clipped = spatial_features_for(states, adapter, extractor, device, batch)
+    np.savez(p, feats=feats, clipped_frac=np.float64(clipped), gid=np.int64(gid), n=np.int64(n_truth))
+    return feats, {"cache": str(p), "clipped_frac": clipped, "reused": False}
+
+
+def fid_kid(f_truth: np.ndarray, f_gen: np.ndarray, kid_cfg: dict) -> dict:
+    return {"fid": BE.fid_from_features(f_truth, f_gen),
+            "kid": BE.kid_from_features(f_truth, f_gen, int(kid_cfg.get("num_subsets", 100)),
+                                        int(kid_cfg.get("max_subset_size", 1000)))}
+
+
+def summarize_per_ic(df: pd.DataFrame, cols=("fid", "kid")) -> dict:
+    out = {"n_ic": int(len(df))}
+    for c in cols:
+        v = df[c].to_numpy(dtype=float); v = v[np.isfinite(v)]
+        out[c] = {"mean": float(v.mean()) if len(v) else float("nan"),
+                  "sem": float(v.std(ddof=1) / np.sqrt(len(v))) if len(v) > 1 else float("nan"),
+                  "sd": float(v.std(ddof=1)) if len(v) > 1 else float("nan"), "n": int(len(v))}
+    return out
+
+
+def evaluate_run(gen_by_gid: dict, bank_dir: Path, t_idx: int, n_truth: int, adapter: FIDAdapter, out_dir: Path,
+                 cache_dir: Path, fid_cfg: dict, device, provenance: dict | None = None) -> dict:
+    """cFID and cKID. ``gen_by_gid``: {gid: RAW generated states (n_gen, D)}; writes out_dir/{per_ic_fid.csv,
+    summary.json, MANIFEST_FID.json, features/gen_gid<gid>.npz} and returns the summary (mean ± SEM over conditions)."""
+    out_dir = Path(out_dir); (out_dir / "features").mkdir(parents=True, exist_ok=True)
+    pins = BE.check_pins(fid_cfg)
+    extractor = BE.build_extractor(device)
+    batch = int(fid_cfg.get("feature_batch", 200))
+    rows, prov_ics = [], {}
+    for gid in sorted(gen_by_gid):
+        gen = np.asarray(gen_by_gid[gid], dtype=np.float64)
+        finite = np.isfinite(gen).all(axis=1)
+        n_bad = int((~finite).sum())
+        ft, tinfo = truth_features(bank_dir, gid, t_idx, n_truth, adapter, extractor, device, cache_dir, batch)
+        if n_bad == len(gen):
+            rows.append({"gid": int(gid), "fid": float("nan"), "kid": float("nan"), "n_truth": int(len(ft)),
+                         "n_gen": 0, "n_gen_nonfinite": n_bad, "clipped_frac_gen": float("nan"), "clipped_frac_truth": tinfo["clipped_frac"]})
+            continue
+        fg, clipped_g = features_for(gen[finite], adapter, extractor, device, batch)
+        np.savez(out_dir / "features" / f"gen_gid{int(gid):04d}.npz", feats=fg, gid=np.int64(gid))
+        m = fid_kid(ft, fg, fid_cfg.get("kid", {}))
+        rows.append({"gid": int(gid), **m, "n_truth": int(len(ft)), "n_gen": int(len(fg)), "n_gen_nonfinite": n_bad,
+                     "clipped_frac_gen": clipped_g, "clipped_frac_truth": tinfo["clipped_frac"]})
+        prov_ics[int(gid)] = tinfo
+    df = pd.DataFrame(rows).sort_values("gid")
+    df.to_csv(out_dir / "per_ic_fid.csv", index=False)
+    summary = summarize_per_ic(df)
+    summary["cFID"] = summary["fid"]; summary["cKID"] = summary["kid"]
+    summary["n_truth_per_ic"] = int(n_truth)
+    (out_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n")
+    man = {"schema": "fid_v1", "pins": pins, "adapter": adapter.params, "fid_settings": fid_cfg, "bank_dir": str(bank_dir),
+           "t_idx": int(t_idx), "n_truth": int(n_truth), "truth_features": prov_ics, "provenance": provenance or {},
+           "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
+    (out_dir / "MANIFEST_FID.json").write_text(json.dumps(man, indent=1, default=str) + "\n")
+    return summary
+
+
+def evaluate_csfid_run(gen_by_gid: dict, bank_dir: Path, t_idx: int, n_truth: int, adapter: FIDAdapter, out_dir: Path,
+                       cache_dir: Path, fid_cfg: dict, device, provenance: dict | None = None) -> dict:
+    """csFID: the Frechet distance of the spatial Inception feature (Nash et al., 2021) between the generated states
+    and the truth bank, per condition; writes out_dir/{csfid_per_ic.csv, csfid_summary.json, MANIFEST_CSFID.json} and
+    returns the summary (mean ± SEM over conditions)."""
+    out_dir = Path(out_dir); out_dir.mkdir(parents=True, exist_ok=True)
+    pins = BE.check_pins(fid_cfg)
+    sp = fid_cfg.get("spatial") or {}
+    want = {"layer": BE.SPATIAL_LAYER, "channels": BE.SPATIAL_CHANNELS, "grid": BE.SPATIAL_GRID, "dim": BE.SPATIAL_DIM,
+            "flatten": "NHWC", "pooling": "none"}
+    bad = {k: (sp[k], v) for k, v in want.items() if k in sp and sp[k] != v}
+    if bad:
+        raise RuntimeError(f"the fid.spatial settings disagree with the csFID implementation: {bad}")
+    pins["spatial_tap_check"] = BE.verify_spatial_tap()
+    extractor = BE.build_extractor(device)
+    batch = int(fid_cfg.get("feature_batch", 200))
+    rows, prov_ics = [], {}
+    for gid in sorted(gen_by_gid):
+        gen = np.asarray(gen_by_gid[gid], dtype=np.float64)
+        finite = np.isfinite(gen).all(axis=1)
+        n_bad = int((~finite).sum())
+        ft, tinfo = truth_spatial_features(bank_dir, gid, t_idx, n_truth, adapter, extractor, device, cache_dir, batch)
+        if n_bad == len(gen):
+            rows.append({"gid": int(gid), "sfid": float("nan"), "n_truth": int(len(ft)), "n_gen": 0, "n_gen_nonfinite": n_bad})
+            continue
+        fg, _ = spatial_features_for(gen[finite], adapter, extractor, device, batch)
+        rows.append({"gid": int(gid), "sfid": BE.fid_from_features(ft, fg), "n_truth": int(len(ft)), "n_gen": int(len(fg)),
+                     "n_gen_nonfinite": n_bad})
+        prov_ics[int(gid)] = tinfo
+    df = pd.DataFrame(rows).sort_values("gid")
+    df.to_csv(out_dir / "csfid_per_ic.csv", index=False)
+    summary = summarize_per_ic(df, cols=("sfid",))
+    summary["csFID"] = summary["sfid"]
+    summary["n_truth_per_ic"] = int(n_truth)
+    (out_dir / "csfid_summary.json").write_text(json.dumps(summary, indent=1) + "\n")
+    man = {"schema": "csfid_v1", "pins": pins, "adapter": adapter.params, "fid_settings": fid_cfg,
+           "bank_dir": str(bank_dir), "t_idx": int(t_idx), "n_truth": int(n_truth), "truth_features": prov_ics,
+           "provenance": provenance or {}, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
+    (out_dir / "MANIFEST_CSFID.json").write_text(json.dumps(man, indent=1, default=str) + "\n")
+    return summary
+
+
+def training_states_at(packaged_dir: Path, n_ic: int, members: int, snapshots) -> np.ndarray:
+    """RAW states of the training pool (the first ``n_ic`` training conditions x the first ``members`` realizations) at
+    the requested snapshots ('all' or a list of indices): the only data the adapter constants are fitted on."""
+    from ..data.fine import FINE_FILE, _npz_member_memmap
+    v = _npz_member_memmap(str(Path(packaged_dir) / FINE_FILE), "vecs_train")
+    if v.shape[0] < n_ic or v.shape[1] < members:
+        raise ValueError(f"vecs_train {tuple(v.shape)} smaller than the requested {n_ic} x {members} pool")
+    if snapshots == "all":
+        block = np.asarray(v[:n_ic, :members, :, :], dtype=np.float64)
+    else:
+        idx = [int(t) for t in (snapshots if isinstance(snapshots, (list, tuple)) else [snapshots])]
+        block = np.asarray(v[:n_ic, :members, :, :][:, :, idx, :], dtype=np.float64)
+    return block.reshape(-1, v.shape[-1])
+
+
+def fit_adapter_from_training(adapter: FIDAdapter, packaged_dir: Path, n_ic: int, pool: dict, t_idx_final: int) -> dict:
+    """``pool`` (dataset card): {members: int | 'all', snapshots: 'all' | 'final' | [indices], pixel_stride: int}."""
+    from ..data.fine import FINE_FILE, _npz_member_memmap
+    v = _npz_member_memmap(str(Path(packaged_dir) / FINE_FILE), "vecs_train")
+    members = int(v.shape[1]) if str(pool.get("members", "all")) == "all" else int(pool["members"])
+    snaps = pool.get("snapshots", "all")
+    snaps = [int(t_idx_final)] if snaps == "final" else snaps
+    states = training_states_at(packaged_dir, n_ic, members, snaps)
+    return adapter.fit(states, pixel_stride=int(pool.get("pixel_stride", 1)),
+                       source={"packaged_dataset_path": str(packaged_dir), "block": "vecs_train", "ics": f"canonical rows 0..{n_ic - 1}",
+                               "members": f"j < {members}", "snapshots": ("all" if snaps == "all" else list(snaps)),
+                               "pixel_stride": int(pool.get("pixel_stride", 1)), "n_states": int(len(states)),
+                               "rule": "training data only; validation, test and generated data never enter"})

code/bifurcation/fid/heterogeneous_waves.py

@@ -0,0 +1,59 @@
+"""Heterogeneous-wave FID adapter: the SAME 196-D coarse spectral observable as Euler and the passive scalar (98 complex modes |k| <= 8, Hermitian half, coefficients / N^2, shared mode order from
+``bifurcation.models.spectral``) -> the exact low-pass DISPLACEMENT field u on a 128x128 grid -> symmetric clip ->
+uint8 -> 3 identical channels.
+
+The reconstruction is byte-identical to the Euler adapter's (one inverse FFT, one mode-index map); only the physical
+reading differs. The observed field is the wave displacement u of
+
+    u_t = v,    v_t = div(c^2 grad u) - gamma v,    c^2(x, y) = 1 + 0.5 sin(2 pi x) sin(2 pi y),
+
+which is signed and zero-mean (the DC mode is not retained by the observable), so the symmetric clip about zero is
+the right quantization (a nonnegative [0, c] range would throw away half of every wave crest-trough pair).
+
+The velocity v and the unresolved band 8 < |k| <= 16 are hidden, which is what makes p(x_T | x_0) non-degenerate;
+neither enters the image, which is a function of the observed state alone.
+
+Physical diagnostics, both computed from the coefficient vectors:
+  * ``displacement_variance_ratio`` — mean sum |u_hat_k|^2, the direct analogue of Euler's enstrophy ratio and the
+    passive scalar's scalar-variance ratio;
+  * ``gradient_energy_ratio`` — mean sum |k|^2 |u_hat_k|^2, the resolved part of the potential energy density
+    c^2 |grad u|^2 up to the medium factor; a wave field can match the first moment while carrying its energy at
+    the wrong wavenumber, and this separates the two.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from ..models.spectral import K_C, mode_int_wavenumbers
+from .adapters import register_adapter
+from .euler import EulerVorticityAdapter
+
+
+@register_adapter("heterogeneous_waves_coarse")
+class HeterogeneousWavesAdapter(EulerVorticityAdapter):
+    field_kind = "signed"
+
+    def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs):
+        super().__init__(quantile, n_disp=n_disp, k_c=k_c, **kwargs)
+        kx, ky = mode_int_wavenumbers(self.k_c)
+        self._k2 = kx.astype(np.float64) ** 2 + ky.astype(np.float64) ** 2
+
+    def card(self) -> dict:
+        c = super().card()
+        c.update({"state": "196-D coarse spectral wave displacement (98 complex modes |k| <= 8, Hermitian half, u_hat / N^2)",
+                  "clip": "symmetric [-c, c], c = Q_q(|u|) over the training pool",
+                  "physical_diagnostic": "displacement_variance_ratio = mean_gen sum|u_hat_k|^2 / mean_truth sum|u_hat_k|^2; "
+                                         "gradient_energy_ratio = mean_gen sum|k|^2|u_hat_k|^2 / mean_truth sum|k|^2|u_hat_k|^2"})
+        return c
+
+    def _moments(self, v: np.ndarray) -> tuple[float, float]:
+        v = np.asarray(v, np.float64)
+        n = self.n_modes
+        power = v[:, :n] ** 2 + v[:, n:] ** 2          # |u_hat_k|^2 per retained mode
+        return float(np.mean(np.sum(power, axis=1))), float(np.mean(np.sum(self._k2 * power, axis=1)))
+
+    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
+        dg, gg = self._moments(gen_raw)
+        dt, gt = self._moments(true_raw)
+        return {"displacement_variance_ratio": dg / dt if dt > 0 else float("nan"),
+                "gradient_energy_ratio": gg / gt if gt > 0 else float("nan")}

code/bifurcation/fid/heterogeneous_waves_augmented.py

@@ -0,0 +1,39 @@
+"""FID adapter for the augmented state of the heterogeneous waves dataset.
+
+The raw state is 392-D = concat(u_coarse[196], v_hat_coarse[196]). The image and the physical diagnostics read
+the observed displacement field u only (the velocity is a model input, not the observable). This adapter slices the first 196
+coordinates and otherwise delegates to ``HeterogeneousWavesAdapter`` (identical reconstruction,
+clip rule, physical diagnostics) -- no numerical difference from the plain-u adapter once sliced.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .adapters import register_adapter
+from .heterogeneous_waves import HeterogeneousWavesAdapter
+
+_D_U = 196
+
+
+@register_adapter("heterogeneous_waves_augmented_coarse")
+class HeterogeneousWavesAugmentedAdapter(HeterogeneousWavesAdapter):
+    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
+        x = np.asarray(states_raw, dtype=np.float64)
+        if x.ndim == 1:
+            x = x[None]
+        if x.shape[1] not in (_D_U, 2 * _D_U):
+            raise ValueError(f"expected {_D_U}-D or {2 * _D_U}-D (augmented) states, got D={x.shape[1]}")
+        return super().to_fields(x[:, :_D_U])
+
+    def _moments(self, v: np.ndarray) -> tuple[float, float]:
+        v = np.asarray(v, np.float64)
+        if v.shape[1] == 2 * _D_U:
+            v = v[:, :_D_U]
+        return super()._moments(v)
+
+    def card(self) -> dict:
+        c = super().card()
+        c.update({"state": "392-D augmented state (196-D coarse displacement u + 196-D finite-difference "
+                           "coarse velocity v_hat); this adapter reads only the first 196 (displacement) "
+                           "coordinates -- identical reconstruction to heterogeneous_waves_coarse"})
+        return c

code/bifurcation/fid/nlkg.py

@@ -0,0 +1,65 @@
+"""NLKG image adapter: score displacement while tracking both displacement and velocity diagnostics.
+
+The 392-dimensional state contains 196 displacement coordinates followed by 196 velocity
+coordinates scaled by 1/4. Images reconstruct the displacement on a 128 x 128 grid.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from ..models.spectral import K_C, mode_int_wavenumbers
+from .adapters import register_adapter
+from .euler import EulerVorticityAdapter
+
+_D_C = 196
+
+
+@register_adapter("nlkg_coarse_u")
+class NLKGCoarseUAdapter(EulerVorticityAdapter):
+    field_kind = "signed"
+
+    def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs):
+        super().__init__(quantile, n_disp=n_disp, k_c=k_c, **kwargs)
+        kx, ky = mode_int_wavenumbers(self.k_c)
+        self._k2 = kx.astype(np.float64) ** 2 + ky.astype(np.float64) ** 2
+        if 2 * self.n_modes != _D_C:
+            raise ValueError("nlkg_coarse_u assumes the 196-D shared coarse operator per component")
+
+    @staticmethod
+    def _split(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        x = np.asarray(x, dtype=np.float64)
+        if x.ndim == 1:
+            x = x[None]
+        if x.shape[1] == 2 * _D_C:
+            return x[:, :_D_C], x[:, _D_C:]
+        if x.shape[1] == _D_C:
+            return x, None
+        raise ValueError(f"expected {2 * _D_C}-D two-component (u, v/4) states or {_D_C}-D displacement-only states, got D={x.shape[1]}")
+
+    def to_fields(self, states_raw: np.ndarray) -> np.ndarray:
+        re, _ = self._split(states_raw)
+        return super().to_fields(re)
+
+    def card(self) -> dict:
+        c = super().card()
+        c.update({"state": "392-D coarse state (196-D u + 196-D v/4, shared coarse operator |k| <= 8 on each); "
+                           "the image reads only the first component u -- identical reconstruction to euler_coarse_vorticity",
+                  "clip": "symmetric [-c, c], c = Q_q(|u_coarse|) over the training pool's final-time fields",
+                  "physical_diagnostic": "coarse_u_variance_ratio, gradient_energy_ratio (on u), coarse_v_variance_ratio (on v/4)"})
+        return c
+
+
+    def _power(self, v: np.ndarray) -> tuple[float, float]:
+        n = self.n_modes
+        power = v[:, :n] ** 2 + v[:, n:] ** 2
+        return float(np.mean(np.sum(power, axis=1))), float(np.mean(np.sum(self._k2 * power, axis=1)))
+
+    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
+        gr, gi = self._split(gen_raw); tr, ti = self._split(true_raw)
+        vg, eg = self._power(gr); vt, et = self._power(tr)
+        out = {"coarse_u_variance_ratio": vg / vt if vt > 0 else float("nan"),
+               "gradient_energy_ratio": eg / et if et > 0 else float("nan")}
+        if gi is not None and ti is not None:
+            vgi, _ = self._power(gi); vti, _ = self._power(ti)
+            out["coarse_v_variance_ratio"] = vgi / vti if vti > 0 else float("nan")
+        return out

code/bifurcation/fid/passive_scalar.py

@@ -0,0 +1,34 @@
+"""Passive-scalar FID adapter: the SAME 196-D coarse spectral observable as Euler (98 complex modes
+|k| <= 8, Hermitian half, coefficients / N^2, shared mode order from ``bifurcation.models.spectral``) -> the exact
+low-pass scalar field theta on a 128x128 grid -> symmetric clip -> uint8 -> 3 identical channels.
+
+The reconstruction is byte-identical to the Euler adapter's (one inverse FFT, one mode-index map); only the
+physical reading differs: the field is a passive scalar concentration anomaly, signed and zero-mean, so the
+symmetric clip about zero is the right quantization. Physical diagnostic: the scalar-variance ratio (mean sum|theta_hat_k|^2 of the generated cloud over the truth's) — the passive scalar's analogue of
+Euler's enstrophy ratio.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .adapters import register_adapter
+from .euler import EulerVorticityAdapter
+
+
+@register_adapter("passive_scalar_coarse")
+class PassiveScalarAdapter(EulerVorticityAdapter):
+    field_kind = "signed"
+
+    def card(self) -> dict:
+        c = super().card()
+        c.update({"state": "196-D coarse spectral passive scalar (98 complex modes |k| <= 8, Hermitian half, theta_hat / N^2)",
+                  "clip": "symmetric [-c, c], c = Q_q(|theta|) over the training pool",
+                  "physical_diagnostic": "scalar_variance_ratio = mean_gen sum|theta_hat_k|^2 / mean_truth sum|theta_hat_k|^2"})
+        return c
+
+    def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict:
+        def var(v):
+            v = np.asarray(v, np.float64)
+            return float(np.mean(np.sum(v ** 2, axis=1)))
+        vt = var(true_raw)
+        return {"scalar_variance_ratio": var(gen_raw) / vt if vt > 0 else float("nan")}

code/bifurcation/methods/__init__.py

@@ -0,0 +1,3 @@
+from .base import Method
+
+__all__ = ["Method"]

code/bifurcation/methods/base.py

@@ -0,0 +1,69 @@
+"""Method interface. EDM, Cascaded SI and Direct SI subclass this."""
+
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class Method(ABC):
+    """Training-method interface.
+
+    A method bundles together (a) the training-time loss, (b) the sampling
+    procedure used at eval time, and (c) the optimizer for the model. The
+    trainer is method-agnostic — it calls these three hooks and nothing else.
+    """
+
+    @abstractmethod
+    def loss(self, model: torch.nn.Module, batch: tuple[torch.Tensor, ...]) -> torch.Tensor:
+        """Per-batch training loss. batch = (x_0, x_T) tensors."""
+
+    @abstractmethod
+    def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int,
+               **variant_kwargs) -> torch.Tensor:
+        """Generate n_samples per x_0 conditioning vector. Returns [B, n_samples, D].
+
+        Methods that support multiple eval variants accept variant-specific
+        overrides as keyword arguments (e.g. ``noise_scale=0.25`` for SI). The
+        legal kwargs are whatever ``variants_for_eval()`` produces.
+        """
+
+    def variants_for_eval(self) -> dict[str, dict]:
+        """Eval-time sampling variants this method exposes.
+
+        Returns a dict ``{variant_label: sample_kwargs}``; ``bifurcation.sample`` writes
+        one directory per variant.
+
+        Default: one variant called ``default`` with no kwarg overrides — used
+        by EDM (the Heun sampler is deterministic; there's no analogous
+        noise-scale knob).
+        """
+        return {"default": {}}
+
+    def configure_optimizer(self, model: torch.nn.Module, cfg: dict) -> torch.optim.Optimizer:
+        """Build the optimizer from ``self.optimizer_cfg`` (set by the method's __init__).
+
+        Supported types:
+          - ``adam``  (default)  → ``torch.optim.Adam`` ; weight_decay ignored
+          - ``adamw``            → ``torch.optim.AdamW``; ``weight_decay`` honoured
+                                   (true decoupled decay — the right way to regularise
+                                   adaptive-LR optimisers)
+
+        Default (no ``optimizer`` block in the method section, or
+        ``optimizer.type == "adam"`` with ``weight_decay`` 0): plain Adam.
+
+        A non-None ``cfg["weight_decay"]`` (the training section, passed by
+        the trainer) overrides the method's block and selects AdamW with that decay.
+        """
+        opt_cfg = getattr(self, "optimizer_cfg", None) or {}
+        opt_type = str(opt_cfg.get("type", "adam")).lower()
+        wd = float(opt_cfg.get("weight_decay", 0.0))
+        params = {"lr": cfg["lr"], "betas": (0.9, 0.999), "eps": 1e-8}
+        wd_override = cfg.get("weight_decay")
+        if wd_override is not None:
+            return torch.optim.AdamW(model.parameters(),
+                                     weight_decay=float(wd_override), **params)
+        if opt_type == "adam":
+            return torch.optim.Adam(model.parameters(), **params)
+        if opt_type == "adamw":
+            return torch.optim.AdamW(model.parameters(), weight_decay=wd, **params)
+        raise ValueError(f"unknown optimizer.type: {opt_type!r}")

code/bifurcation/methods/cascaded_si/__init__.py

@@ -0,0 +1,3 @@
+from .method import CascadedSIMethod
+
+__all__ = ["CascadedSIMethod"]

code/bifurcation/methods/cascaded_si/core.py

@@ -0,0 +1,133 @@
+"""
+Cascaded SI: the Wiener-path (Foellmer) stochastic interpolant per segment, the segment-wise mirror of
+``methods.direct_si.core``. The trained drift b is the matched SDE drift (b = b_ODE + 0.5 sigma_tau^2 * score,
+baked in), so the sampler needs no score term and no second network.
+
+Identical Wiener coefficients to the direct construction; the differences:
+  - the anchor is the segment start x_start (= x_tk), not x0;
+  - the network takes (x, tau, x_start, t_k, t_k1), the five-channel segment head;
+  - sampling chains n_segments per-segment SDE integrations.
+
+sigma = amplitude; beta_power = 1 -> beta=tau (linear), 2 -> beta=tau^2 (b'(0)=0).
+Schedule: alpha=1-tau, sigma_tau = sigma*(1-tau). Realize W_tau = sqrt(tau)*z.
+Three distinct noise coefficients (not a derivative pair):
+    interpolant noise (coeff on z in x_tau) :  sigma_tau*sqrt(tau) = sigma(1-tau)sqrt(tau)
+    target      noise (coeff on z in R)     :  sigma_tau_dot*sqrt(tau) = -sigma*sqrt(tau)
+    sampler diffusion (coeff on dW)         :  sigma_tau = sigma(1-tau)
+beta_power only changes the DETERMINISTIC interpolant/target:
+    beta=tau   : I=(1-tau)x_tk+tau x_tk1,    dtIt=-x_tk+x_tk1
+    beta=tau^2 : I=(1-tau)x_tk+tau^2 x_tk1,  dtIt=-x_tk+2 tau x_tk1
+The Wiener target coeff -sigma*sqrt(tau) is bounded, so t_eps may be 0.0 (no
+endpoint singularity) and the sampler grid starts at tau=0.
+"""
+from __future__ import annotations
+import math
+import torch
+from torch import vmap
+
+def sigma_tau(tau, sigma):
+    return sigma * (1.0 - tau)
+
+def sigma_tau_dot(tau, sigma):
+    return -sigma * torch.ones_like(tau)
+
+def interp_It(tau, x0, x1, beta_power=1):
+    beta = tau if beta_power == 1 else tau * tau
+    return (1.0 - tau) * x0 + beta * x1
+
+def interp_dtIt(tau, x0, x1, beta_power=1):
+    beta_dot = 1.0 if beta_power == 1 else 2.0 * tau
+    return -x0 + beta_dot * x1
+
+def b_loss_per_sample(b_net, x_tk, x_tk1, tau, t_k, t_k1, sigma, beta_power=1):
+    """Wiener antithetic b-loss for one segment example.
+
+    Shapes:
+        x_tk, x_tk1 : (D,)
+        tau         : scalar
+        t_k, t_k1   : scalars (segment physical start/end)
+    Returns: scalar loss.
+    """
+    z = torch.randn_like(x_tk)
+    sqrt_t = torch.sqrt(tau)
+    It  = interp_It(tau, x_tk, x_tk1, beta_power)
+    env = sigma_tau(tau, sigma) * sqrt_t                 # interpolant noise coeff
+    xtp = It + env * z
+    xtm = It - env * z
+    tau_s  = tau.unsqueeze(0)
+    dtIt   = interp_dtIt(tau, x_tk, x_tk1, beta_power)
+    tcoeff = sigma_tau_dot(tau, sigma) * sqrt_t          # target noise coeff = -sigma*sqrt(tau)
+    tgt_p  = dtIt + tcoeff * z
+    tgt_m  = dtIt - tcoeff * z
+    btp = b_net(xtp.unsqueeze(0), tau_s,
+                x_tk.unsqueeze(0), t_k.view(1, 1), t_k1.view(1, 1)).squeeze(0)
+    btm = b_net(xtm.unsqueeze(0), tau_s,
+                x_tk.unsqueeze(0), t_k.view(1, 1), t_k1.view(1, 1)).squeeze(0)
+    return (0.5 * (btp ** 2).sum() - (btp * tgt_p).sum()
+          + 0.5 * (btm ** 2).sum() - (btm * tgt_m).sum())
+
+_batched_b_loss = vmap(
+    b_loss_per_sample,
+    in_dims=(None, 0, 0, 0, 0, 0, None, None),
+    randomness="different",
+)
+
+def local_si_loss(b_net, x_tk_b, x_tk1_b, t_k_b, t_k1_b, sigma,
+                  t_eps=0.0, beta_power=1):
+    # tau ~ U(t_eps, 1-t_eps). t_eps may be 0.0: the Wiener target coeff
+    # -sigma*sqrt(tau) is bounded (no endpoint singularity).
+    B = x_tk_b.shape[0]
+    device = x_tk_b.device
+    tau_batch = torch.empty(B, device=device).uniform_(t_eps, 1.0 - t_eps)
+    losses = _batched_b_loss(b_net, x_tk_b, x_tk1_b, tau_batch, t_k_b, t_k1_b,
+                             sigma, beta_power)
+    return losses.mean()
+
+def sde_segment(b_net, x_start_batch, t_k_val, t_k1_val,
+                n_steps, sigma, noise_scale=1.0):
+    # dY = b dtau + sigma_tau dW over one segment, anchor x_start. Trained drift
+    # already contains the score, so noise_scale = 1 is the matched sampler;
+    # diffusion = sigma_tau (not the envelope sigma_tau*sqrt(tau)). For Wiener,
+    # noise_scale != 1 is not marginal-preserving (no recovered score) -> use 1.
+    # beta_power does not enter here.
+    N = x_start_batch.shape[0]
+    device = x_start_batch.device
+    dt = 1.0 / n_steps
+    ts = torch.linspace(0.0, 1.0 - dt, n_steps, device=device)
+
+    t_k_t  = torch.tensor([[t_k_val]],  dtype=torch.float32, device=device).expand(N, 1)
+    t_k1_t = torch.tensor([[t_k1_val]], dtype=torch.float32, device=device).expand(N, 1)
+
+    x_start = x_start_batch.clone()
+    x = x_start.clone()
+
+    with torch.no_grad():
+        for t_val in ts:
+            tau = t_val.view(1, 1).expand(N, 1)
+            drift = b_net(x, tau, x_start, t_k_t, t_k1_t)
+            g = noise_scale * sigma_tau(t_val, sigma)
+            x = x + dt * drift + g * math.sqrt(dt) * torch.randn_like(x)
+    return x
+
+def rollout(b_net, x0, n_samples, t_values, n_sde_steps_per_seg, sigma,
+            noise_scale=1.0):
+    """Roll out n_samples trajectories from x0 across all segments.
+
+    Args
+    ----
+    x0        : (D,) or (1, D) — conditioning state.
+    t_values  : 1-D iterable of n_segments+1 physical times (e.g. 0.0..1.0).
+    Returns
+    -------
+    x_T : (n_samples, D) — endpoint cloud.
+    """
+    if x0.ndim == 1:
+        x = x0.unsqueeze(0).expand(n_samples, -1).clone()
+    else:
+        x = x0.expand(n_samples, -1).clone()
+    n_segments = len(t_values) - 1
+    for k in range(n_segments):
+        x = sde_segment(b_net, x,
+                        float(t_values[k]), float(t_values[k + 1]),
+                        n_sde_steps_per_seg, sigma, noise_scale)
+    return x

code/bifurcation/methods/cascaded_si/method.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+import torch
+from ..base import Method
+from . import core
+
+class CascadedSIMethod(Method):
+    def __init__(self, sigma_train, t_eps, sampling, beta_power=1,
+                 eval_noise_scales=None, optimizer_cfg=None):
+        self.sigma_train = float(sigma_train)
+        self.t_eps = float(t_eps)
+        self.beta_power = int(beta_power)
+        self.n_sde_steps_per_seg = int(sampling["n_sde_steps_per_seg"])
+        self.n_segments = int(sampling.get("n_segments", 10))
+        self.noise_scale = float(sampling.get("noise_scale", 1.0))
+        # Segment boundaries — uniform partition of [0, 1] by default, matching
+        # the normalized physical forecast horizon.
+        if "t_values" in sampling:
+            self.t_values = list(sampling["t_values"])
+        else:
+            self.t_values = [k / self.n_segments for k in range(self.n_segments + 1)]
+        self.eval_noise_scales = ([float(s) for s in eval_noise_scales]
+                                  if eval_noise_scales is not None else [self.noise_scale])
+        self.optimizer_cfg = optimizer_cfg or {}
+
+    @classmethod
+    def from_config(cls, cfg):
+        if cfg.get("anchor_ic", False):
+            raise ValueError("This release supports the paper's SI models without retained initial-state inputs")
+        return cls(
+            sigma_train=cfg["sigma_train"],
+            t_eps=cfg.get("t_eps", 0.0),
+            sampling=cfg["sampling"],
+            beta_power=cfg.get("beta_power", 1),
+            eval_noise_scales=cfg.get("eval_noise_scales"),
+            optimizer_cfg=cfg.get("optimizer"),
+        )
+
+    def variants_for_eval(self):
+        return {f"noise_scale_{_fmt(s)}": {"noise_scale": s} for s in self.eval_noise_scales}
+
+    def loss(self, model, batch):
+        x_tk, x_tk1, t_k, t_k1 = batch
+        return core.local_si_loss(model, x_tk, x_tk1, t_k, t_k1,
+                                  sigma=self.sigma_train, t_eps=self.t_eps,
+                                  beta_power=self.beta_power)
+
+    def sample(self, model, x_0, n_samples, noise_scale=None):
+        if noise_scale is None:
+            noise_scale = self.noise_scale
+        if x_0.ndim == 1:
+            x_0 = x_0.unsqueeze(0)
+        B = x_0.shape[0]
+        out = []
+        for i in range(B):
+            cloud = core.rollout(model, x_0[i], n_samples=n_samples,
+                                 t_values=self.t_values,
+                                 n_sde_steps_per_seg=self.n_sde_steps_per_seg,
+                                 sigma=self.sigma_train,
+                                 noise_scale=noise_scale)
+            out.append(cloud)
+        return torch.stack(out, dim=0)
+
+def _fmt(s):
+    if s == int(s):
+        return f"{int(s)}p0"
+    return f"{s}".replace(".", "p")

code/bifurcation/methods/direct_si/__init__.py

@@ -0,0 +1,3 @@
+from .method import DirectSIMethod
+
+__all__ = ["DirectSIMethod"]

code/bifurcation/methods/direct_si/core.py

@@ -0,0 +1,86 @@
+"""
+Direct SI: the Wiener-path (Foellmer) stochastic interpolant from x_0 to x_T. The trained drift b is the
+matched SDE drift (b = b_ODE + 0.5 sigma_tau^2 * score, baked in), so the sampler
+needs no score term and no second network.
+
+sigma = amplitude; beta_power = 1 -> beta=tau (linear), 2 -> beta=tau^2 (b'(0)=0).
+Schedule: alpha=1-tau, sigma_tau = sigma*(1-tau). Realize W_tau = sqrt(tau)*z.
+Three distinct noise coefficients (not a derivative pair):
+    interpolant noise (coeff on z in x_tau) :  sigma_tau*sqrt(tau) = sigma(1-tau)sqrt(tau)
+    target      noise (coeff on z in R)     :  sigma_tau_dot*sqrt(tau) = -sigma*sqrt(tau)
+    sampler diffusion (coeff on dW)         :  sigma_tau = sigma(1-tau)
+beta_power only changes the DETERMINISTIC interpolant/target:
+    beta=tau   : I=(1-tau)x0+tau x1,    dtIt=-x0+x1 = x1-x0
+    beta=tau^2 : I=(1-tau)x0+tau^2 x1,  dtIt=-x0+2 tau x1
+"""
+from __future__ import annotations
+import math
+import torch
+from torch import vmap
+
+def sigma_tau(tau, sigma):
+    return sigma * (1.0 - tau)
+
+def sigma_tau_dot(tau, sigma):
+    return -sigma * torch.ones_like(tau)
+
+def interp_It(tau, x0, x1, beta_power=1):
+    beta = tau if beta_power == 1 else tau * tau
+    return (1.0 - tau) * x0 + beta * x1
+
+def interp_dtIt(tau, x0, x1, beta_power=1):
+    beta_dot = 1.0 if beta_power == 1 else 2.0 * tau
+    return -x0 + beta_dot * x1
+
+def b_loss_per_sample(b_net, x0, x1, tau, sigma, beta_power=1):
+    z = torch.randn_like(x0)
+    sqrt_t = torch.sqrt(tau)
+    It  = interp_It(tau, x0, x1, beta_power)
+    env = sigma_tau(tau, sigma) * sqrt_t                 # interpolant noise coeff
+    xtp = It + env * z
+    xtm = It - env * z
+    tau_s  = tau.unsqueeze(0)
+    dtIt   = interp_dtIt(tau, x0, x1, beta_power)
+    tcoeff = sigma_tau_dot(tau, sigma) * sqrt_t          # target noise coeff = -sigma*sqrt(tau)
+    tgt_p  = dtIt + tcoeff * z
+    tgt_m  = dtIt - tcoeff * z
+    btp = b_net(xtp.unsqueeze(0), tau_s, x0.unsqueeze(0)).squeeze(0)
+    btm = b_net(xtm.unsqueeze(0), tau_s, x0.unsqueeze(0)).squeeze(0)
+    return (0.5 * (btp ** 2).sum() - (btp * tgt_p).sum()
+          + 0.5 * (btm ** 2).sum() - (btm * tgt_m).sum())
+
+_batched_b_loss = vmap(
+    b_loss_per_sample,
+    in_dims=(None, 0, 0, 0, None, None),
+    randomness="different",
+)
+
+def global_si_loss(b_net, x_0_batch, x_T_batch, sigma, t_eps=0.0, beta_power=1):
+    # tau ~ U(t_eps, 1-t_eps). t_eps may be 0.0: the Wiener target coeff
+    # -sigma*sqrt(tau) is bounded (no endpoint singularity).
+    B = x_0_batch.shape[0]
+    device = x_0_batch.device
+    tau_batch = torch.empty(B, device=device).uniform_(t_eps, 1.0 - t_eps)
+    losses = _batched_b_loss(b_net, x_0_batch, x_T_batch, tau_batch, sigma, beta_power)
+    return losses.mean()
+
+def sde_sample(b_net, x_0, n_samples, n_steps, sigma, noise_scale=1.0):
+    # dY = b dtau + sigma_tau dW, Y_0 = x0. Trained drift already contains the score,
+    # so noise_scale = 1 is the matched sampler; diffusion = sigma_tau (not the
+    # envelope sigma_tau*sqrt(tau)). For Wiener, noise_scale != 1 is not
+    # marginal-preserving (no recovered score) -> use 1. beta_power does not enter here.
+    device = x_0.device
+    dt = 1.0 / n_steps
+    ts = torch.linspace(0.0, 1.0 - dt, n_steps, device=device)
+    if x_0.ndim == 1:
+        x0_e = x_0.unsqueeze(0).expand(n_samples, -1)
+    else:
+        x0_e = x_0.expand(n_samples, -1)
+    x = x0_e.clone()
+    with torch.no_grad():
+        for t_val in ts:
+            tau = t_val.expand(n_samples)
+            drift = b_net(x, tau, x0_e)
+            g = noise_scale * sigma_tau(t_val, sigma)
+            x = x + dt * drift + g * math.sqrt(dt) * torch.randn_like(x)
+    return x

code/bifurcation/methods/direct_si/method.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+import torch
+from ..base import Method
+from . import core
+
+class DirectSIMethod(Method):
+    def __init__(self, sigma_train, t_eps, sampling, beta_power=1,
+                 eval_noise_scales=None, optimizer_cfg=None):
+        self.sigma_train = float(sigma_train)
+        self.t_eps = float(t_eps)
+        self.beta_power = int(beta_power)
+        self.n_sde_steps = int(sampling["n_sde_steps"])
+        self.noise_scale = float(sampling.get("noise_scale", 1.0))
+        self.eval_noise_scales = ([float(s) for s in eval_noise_scales]
+                                  if eval_noise_scales is not None else [self.noise_scale])
+        self.optimizer_cfg = optimizer_cfg or {}
+
+    @classmethod
+    def from_config(cls, cfg):
+        return cls(
+            sigma_train=cfg["sigma_train"],
+            t_eps=cfg.get("t_eps", 0.0),
+            sampling=cfg["sampling"],
+            beta_power=cfg.get("beta_power", 1),
+            eval_noise_scales=cfg.get("eval_noise_scales"),
+            optimizer_cfg=cfg.get("optimizer"),
+        )
+
+    def variants_for_eval(self):
+        return {f"noise_scale_{_fmt(s)}": {"noise_scale": s} for s in self.eval_noise_scales}
+
+    def loss(self, model, batch):
+        x_0, x_T = batch
+        return core.global_si_loss(model, x_0, x_T, sigma=self.sigma_train,
+                                   t_eps=self.t_eps, beta_power=self.beta_power)
+
+    def sample(self, model, x_0, n_samples, noise_scale=None):
+        if noise_scale is None:
+            noise_scale = self.noise_scale
+        if x_0.ndim == 1:
+            x_0 = x_0.unsqueeze(0)
+        B = x_0.shape[0]
+        out = []
+        for i in range(B):
+            cloud = core.sde_sample(model, x_0[i], n_samples=n_samples,
+                                    n_steps=self.n_sde_steps, sigma=self.sigma_train,
+                                    noise_scale=noise_scale)
+            out.append(cloud)
+        return torch.stack(out, dim=0)
+
+def _fmt(s):
+    if s == int(s):
+        return f"{int(s)}p0"
+    return f"{s}".replace(".", "p")

code/bifurcation/methods/edm/__init__.py

@@ -0,0 +1,3 @@
+from .method import EDMMethod
+
+__all__ = ["EDMMethod"]

code/bifurcation/methods/edm/core.py

@@ -0,0 +1,205 @@
+"""
+core.py — EDM core math for conditional diffusion.
+
+Reimplements the design of NVlabs/edm
+(https://github.com/NVlabs/edm; specifically ``training/training_loop.py``)
+following Karras et al. 2022, "Elucidating the Design Space of
+Diffusion-Based Generative Models", NeurIPS 2022.
+
+Method-only. No architecture, no data loading, no training loop.
+
+Six components:
+  - σ sampling at training time
+  - preconditioning (c_skip, c_out, c_in, c_noise)
+  - denoiser wrapper D_θ
+  - training loss with λ(σ) weighting
+  - σ schedule for sampling (ρ-parameterized)
+  - Heun's 2nd-order sampler with Euler fallback at σ=0
+"""
+
+import torch
+import math
+
+
+# --- 1. Training-time σ sampling ---
+
+def sample_training_sigma(batch_size, P_mean, P_std, device):
+    """ see Table 1 in EDM paper. Formula for noise distribution under Training section.
+    Returns: tensor of shape (batch_size,).
+    """
+    ln_sigma = torch.empty(batch_size, device=device).normal_(mean=P_mean, std=P_std)
+    sigma = torch.exp(ln_sigma)
+    return sigma
+
+
+# --- 2. Preconditioning (EDM Table 1) ---
+
+def c_skip(sigma, sigma_data):
+    """Skip-connection scaling."""
+    c_skip = sigma_data**2 / (sigma**2 +sigma_data**2)
+    return c_skip
+
+def c_out(sigma, sigma_data):
+    """Output scaling."""
+    c_out = sigma * sigma_data / (torch.sqrt(sigma_data**2+sigma**2))
+    return c_out
+
+def c_in(sigma, sigma_data):
+    """Input scaling."""
+    c_in=1 / (torch.sqrt(sigma**2+sigma_data**2))
+    return c_in
+
+def c_noise(sigma):
+    """Time-like scalar fed into the network."""
+    c_noise=0.25 * torch.log(sigma)
+    return c_noise
+
+
+
+
+# --- 3. Denoiser wrapper ---
+def denoise(F_theta, x_noisy, sigma, x_0, sigma_data):
+    """
+    Wraps the raw network F_theta into the preconditioned denoiser D_theta. This corresponds to equation 7 in the EDM paper.
+    D_θ(x, σ, x_0) = c_skip(σ)·x + c_out(σ) · F_θ(c_in(σ)·x, c_noise(σ), x_0)
+    Used identically in both training (target reconstruction) and sampling (Heuns method)
+    Returns predicted clean x_T (same shape as x_noisy).
+    """
+    # Step 1: normalize the noisy input
+    x_normalized = c_in(sigma, sigma_data) * x_noisy
+
+    # Step 2: embed the noise level
+    sigma_embedded = c_noise(sigma)
+
+    # Step 3: feed normalized inputs + conditioning through the network
+    F_out = F_theta(x_normalized, sigma_embedded, x_0)
+
+    # Step 4: combine via skip connection + scaled output
+    D_theta = c_skip(sigma, sigma_data) * x_noisy + c_out(sigma, sigma_data) * F_out
+
+    return D_theta
+
+
+# --- 4. Loss ---
+
+def lambda_weight(sigma, sigma_data):
+    """Per-σ loss weighting (EDM Table 1, section training)."""
+    lambda_sigma = (sigma**2 + sigma_data**2)/(sigma*sigma_data)**2
+    return lambda_sigma
+
+def edm_loss(F_theta, x_T_batch, x_0_batch, sigma_data, P_mean, P_std):
+    """
+    Compute the EDM training loss for one batch.
+
+    Args:
+        F_theta: the raw network
+        x_T_batch: clean data, shape [B, D]
+        x_0_batch: conditioning, shape [B, D]
+        sigma_data: fixed scalar
+        P_mean, P_std: log-normal sampling hyperparameters
+
+    Returns:
+        scalar loss
+    """
+    batch_size = x_T_batch.shape[0]
+    device = x_T_batch.device
+
+    # Step 1: sample sigma per example
+    sigma = sample_training_sigma(batch_size, P_mean, P_std, device)
+    # reshape for broadcasting: [B] -> [B, 1] for vector data
+    sigma = sigma.view(-1, 1)
+
+    # Step 2: sample noise
+    epsilon = torch.randn_like(x_T_batch)
+
+    # Step 3: add noise (EDM's simple forward process)
+    x_noisy = x_T_batch + sigma * epsilon
+
+    # Step 4: run through denoiser
+    D_theta = denoise(F_theta, x_noisy, sigma, x_0_batch, sigma_data)
+
+    # Step 5: compute weighted squared error
+    squared_error = (D_theta - x_T_batch).pow(2).sum(dim=-1)  # [B]
+    weighted_loss = lambda_weight(sigma.squeeze(-1), sigma_data) * squared_error  # [B]
+
+    # Step 6: average over batch
+    return weighted_loss.mean()
+
+
+# --- 5. Sampling-time σ schedule ---
+
+def sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device):
+    """ρ-parameterized power-law schedule.
+
+    Returns tensor of length N+1 with σ_0 = σ_max, σ_N = 0,
+    and σ_i for 0 < i < N as in EDM Eq. 5.
+    """
+    # i goes from 0 to N-1
+    i = torch.arange(N, device=device, dtype=torch.float32)
+
+    # Apply the power-law formula
+    inv_rho = 1.0 / rho
+    sigma_max_inv = sigma_max ** inv_rho
+    sigma_min_inv = sigma_min ** inv_rho
+    sigmas = (sigma_max_inv + (i / (N - 1)) * (sigma_min_inv - sigma_max_inv)) ** rho
+
+    # Append the final σ = 0 (we want to reach truly clean data)
+    sigmas = torch.cat([sigmas, torch.zeros(1, device=device)])
+    return sigmas
+
+
+# --- 6. Heun sampler (EDM Algorithm 1) ---
+
+def heun_sample(network, x_0, n_samples, N, sigma_min, sigma_max, rho, sigma_data, device):
+    """Sample from p(x_T | x_0) via Heun's 2nd-order method.
+
+    Initializes at x ~ N(0, σ_max² · I), iterates the σ schedule with
+    Heun steps, falls back to Euler at the final step where σ_next = 0.
+
+    Returns: tensor of shape (n_samples, x_0.shape[-1]).
+    """
+    # 1. Build the σ schedule (length N+1, ending at 0)
+    sigmas = sampling_sigma_schedule(N, sigma_min, sigma_max, rho, device)
+
+    # 2. Broadcast x_0 to (n_samples, D) so every generated sample uses
+    #    the same conditioning but different starting noise.
+    if x_0.ndim == 1:
+        x_0_batch = x_0.unsqueeze(0).expand(n_samples, -1)
+    else:
+        x_0_batch = x_0
+    D_dim = x_0_batch.shape[-1]
+
+    # 3. Initialize at pure noise scaled to σ_max.
+    #    y₀ ~ N(0, σ_max²·I) means y₀ = σ_max · ε for ε ~ N(0, I).
+    y = sigma_max * torch.randn(n_samples, D_dim, device=device)
+
+    # 4. Iterate Heun steps along the schedule.
+    for i in range(N):
+        sigma_cur = sigmas[i]       # current noise level
+        sigma_next = sigmas[i + 1]  # target noise level for this step
+
+        # Reshape σ to match batch shape for broadcasting
+        sigma_cur_b = sigma_cur.expand(n_samples, 1)
+
+        # --- Predictor: Euler step ---
+        # Evaluate the denoiser at the current point
+        D_cur = denoise(network, y, sigma_cur_b, x_0_batch, sigma_data)
+        # Compute the slope dy/dσ at the current point
+        d_cur = (y - D_cur) / sigma_cur
+        # Euler step toward σ_next (note: σ_next − σ_cur is negative since σ decreases)
+        y_pred = y + (sigma_next - sigma_cur) * d_cur
+
+        # --- Corrector: Heun's average slope ---
+        if sigma_next > 0:
+            # Evaluate the denoiser at the predicted next point
+            sigma_next_b = sigma_next.expand(n_samples, 1)
+            D_next = denoise(network, y_pred, sigma_next_b, x_0_batch, sigma_data)
+            # Slope at the predicted point
+            d_next = (y_pred - D_next) / sigma_next
+            # Take the average slope and redo the step
+            y = y + (sigma_next - sigma_cur) * 0.5 * (d_cur + d_next)
+        else:
+            # Final step: σ_next = 0 would divide by zero in d_next.
+            # Just use the Euler estimate. This is the only Euler-style step.
+            y = y_pred
+    return y

code/bifurcation/methods/edm/method.py

@@ -0,0 +1,74 @@
+"""EDMMethod — wraps the math in core.py."""
+
+import torch
+
+from ..base import Method
+from . import core
+
+
+class EDMMethod(Method):
+    """EDM (Karras et al. 2022) for conditional endpoint prediction.
+
+    Hyperparameters
+    ---------------
+    sigma_data : float
+        EDM σ_data; matches the empirical per-feature std of the (normalized)
+        training x_T. For data normalized to unit variance per feature this is
+        1.0.
+    P_mean, P_std : float
+        Log-normal training-σ sampler (Table 1 of the EDM paper).
+    sampling : dict
+        Keys: n_steps, sigma_min, sigma_max, rho. Used by Heun's sampler.
+    """
+
+    def __init__(self, sigma_data: float, P_mean: float, P_std: float, sampling: dict,
+                 optimizer_cfg: dict | None = None):
+        self.sigma_data = float(sigma_data)
+        self.P_mean = float(P_mean)
+        self.P_std = float(P_std)
+        self.n_steps = int(sampling["n_steps"])
+        self.sigma_min = float(sampling["sigma_min"])
+        self.sigma_max = float(sampling["sigma_max"])
+        self.rho = float(sampling["rho"])
+        # Picked up by base Method.configure_optimizer; defaults to plain Adam.
+        self.optimizer_cfg = optimizer_cfg or {}
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "EDMMethod":
+        return cls(
+            sigma_data=cfg["sigma_data"],
+            P_mean=cfg["P_mean"],
+            P_std=cfg["P_std"],
+            sampling=cfg["sampling"],
+            optimizer_cfg=cfg.get("optimizer"),
+        )
+
+    def loss(self, model: torch.nn.Module, batch) -> torch.Tensor:
+        x_0, x_T = batch
+        return core.edm_loss(model, x_T, x_0, self.sigma_data, self.P_mean, self.P_std)
+
+    def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int) -> torch.Tensor:
+        """Return Heun-sampled clouds, shape [B, n_samples, D].
+
+        x_0 may be either a single conditioning vector of shape (D,) or a
+        batch of shape (B, D). The output is always (B, n_samples, D).
+        """
+        device = next(model.parameters()).device
+        if x_0.ndim == 1:
+            x_0 = x_0.unsqueeze(0)
+        B = x_0.shape[0]
+        clouds = []
+        for i in range(B):
+            cloud = core.heun_sample(
+                model,
+                x_0[i],
+                n_samples=n_samples,
+                N=self.n_steps,
+                sigma_min=self.sigma_min,
+                sigma_max=self.sigma_max,
+                rho=self.rho,
+                sigma_data=self.sigma_data,
+                device=device,
+            )
+            clouds.append(cloud)
+        return torch.stack(clouds, dim=0)

code/bifurcation/metrics.py

@@ -0,0 +1,41 @@
+"""Distributional metrics in state space.
+
+``energy_distance`` is the validation metric of the paper's model selection (a V-statistic: ``md(X, X)`` and
+``md(Y, Y)`` average over all pairs including self-pairs; ``2 md(X, Y) - md(X, X) - md(Y, Y)``).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def _mean_pairwise_dist(A: np.ndarray, B: np.ndarray) -> float:
+    return np.mean(np.linalg.norm(A[:, None] - B[None, :], axis=-1))
+
+
+def energy_distance(X: np.ndarray, Y: np.ndarray) -> float:
+    md = _mean_pairwise_dist
+    return float(2 * md(X, Y) - md(X, X) - md(Y, Y))
+
+
+def energy_distance_cached_self(X: np.ndarray, Y: np.ndarray, md_yy: float) -> float:
+    """``energy_distance`` with the truth-self term ``md(Y, Y)`` supplied (it does not depend on X)."""
+    md = _mean_pairwise_dist
+    return float(2 * md(X, Y) - md(X, X) - float(md_yy))
+
+
+def energy_distance_u(X: np.ndarray, Y: np.ndarray) -> float:
+    """The unbiased U-statistic: the self terms exclude the diagonal, ``ED_U = ED_V - md(X,X)/(M-1) - md(Y,Y)/(N-1)``."""
+    md = _mean_pairwise_dist
+    mxx, myy = md(X, X), md(Y, Y)
+    M, N = X.shape[0], Y.shape[0]
+    return float(2 * md(X, Y) - mxx - myy - (mxx / (M - 1) if M > 1 else 0.0) - (myy / (N - 1) if N > 1 else 0.0))
+
+
+def effective_rank(X: np.ndarray) -> float:
+    """Effective rank of the sample covariance, tr(S)^2 / ||S||_F^2; NaN for fewer than 2 samples."""
+    X = np.asarray(X, dtype=np.float64)
+    if X.shape[0] < 2:
+        return float("nan")
+    S = np.cov(X.T)
+    fro2 = float((S * S).sum())
+    return float(S.trace() ** 2 / fro2) if fro2 > 0 else float("nan")

code/bifurcation/models/_common.py

@@ -0,0 +1,16 @@
+"""Shared helpers for backbones."""
+
+from __future__ import annotations
+
+import torch
+
+
+def _pack_scalar(s: torch.Tensor, B: int) -> torch.Tensor:
+    """Normalize a scalar / per-batch scalar to shape (B, 1)."""
+    if s.ndim == 0 or (s.ndim == 1 and s.numel() == 1):
+        return s.view(1, 1).expand(B, 1)
+    return s.reshape(B, 1)
+
+
+def count_parameters(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)

code/bifurcation/models/spectral.py

@@ -0,0 +1,97 @@
+"""Coarse spectral <-> field maps for the 196-dim state (numpy, no plotting).
+
+The (kx, ky) <-> index map is the single source of truth for the 196-dim coarse spectral vector shared by the
+five datasets.
+
+Reconstruction (exact inverse of extract_group.get_coarse_vec):
+  vec = [Re(oh_norm[mask]), Im(oh_norm[mask])], oh_norm = omega_hat / N_sim^2,
+  mask = (|k| <= k_c) & Hermitian-half. Un-normalize (clouds are per-feature
+  standardized), place each complex coefficient and its conjugate partner on an
+  N_disp grid, field = N_disp^2 * ifft2(OH).real.
+
+There is deliberately no default dataset path in this module: the normalization
+statistics that the wrapper un/re-normalizes with must come from the dataset
+the model is trained on (the registry injects ``data.packaged_dataset_path``).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+K_C = 8
+N_SIM = 1024  # the resolution get_coarse_vec normalized by (omega_hat / N_sim^2)
+
+
+def mode_int_wavenumbers(k_c: int = K_C, M: int = N_SIM):
+    """Integer (kx, ky) for each of the 98 complex modes, in get_coarse_vec
+    order. axis0<->kx, axis1<->ky (matches extract_group.make_kgrids)."""
+    k1d = np.fft.fftfreq(M, d=1.0 / M).astype(int)      # 0,1,..,-1
+    ky_grid, kx_grid = np.meshgrid(k1d, k1d)
+    k_mag = np.sqrt(kx_grid**2 + ky_grid**2)
+    hermitian = (kx_grid > 0) | ((kx_grid == 0) & (ky_grid > 0))
+    mask = (k_mag <= k_c) & hermitian
+    return kx_grid[mask], ky_grid[mask]                 # each (n_modes,)
+
+
+class FieldReconstructor:
+    def __init__(self, packaged: str | Path,
+                 n_disp: int = 128, k_c: int = K_C):
+        norm = np.load(Path(packaged) / "normalization.npz")
+        self.mean, self.std = norm["mean"], norm["std"]
+        self.n_disp = n_disp
+        self.kx, self.ky = mode_int_wavenumbers(k_c)
+        self.n_modes = len(self.kx)
+        # precompute grid placement indices (mod N for negative wavenumbers)
+        self.ix = self.kx % n_disp
+        self.iy = self.ky % n_disp
+        self.ix_c = (-self.kx) % n_disp
+        self.iy_c = (-self.ky) % n_disp
+
+    def _spectral_grid(self, vec_norm: np.ndarray) -> np.ndarray:
+        """(196,) standardized coarse vector -> full Hermitian ω̂ grid (n_disp²)."""
+        vec = np.asarray(vec_norm, float) * self.std + self.mean   # un-normalize
+        coeff = vec[:self.n_modes] + 1j * vec[self.n_modes:]       # oh_norm[mask]
+        OH = np.zeros((self.n_disp, self.n_disp), dtype=complex)
+        OH[self.ix,   self.iy]   += coeff
+        OH[self.ix_c, self.iy_c] += np.conj(coeff)
+        return OH
+
+    def reconstruct_field(self, vec_norm: np.ndarray,
+                          field: str = "vorticity") -> np.ndarray:
+        """(196,) standardized coarse vector -> (N_disp, N_disp) real field.
+
+        Only the 'vorticity' branch (the observed scalar itself) exists.
+        """
+        OH = self._spectral_grid(vec_norm)
+        n2 = self.n_disp ** 2
+        if field == "vorticity":
+            return n2 * np.fft.ifft2(OH).real
+        raise ValueError(f"unknown field: {field!r} (only 'vorticity' is supported here)")
+
+    def reconstruct_many(self, vecs_norm: np.ndarray,
+                         field: str = "vorticity") -> np.ndarray:
+        """(N, 196) -> (N, N_disp, N_disp)."""
+        return np.stack([self.reconstruct_field(v, field) for v in vecs_norm])
+
+    def extract_vec(self, field: np.ndarray) -> np.ndarray:
+        """(N_disp, N_disp) real field -> (2*n_modes,) standardized coarse vector.
+
+        Exact left-inverse of reconstruct_field when the field is band-limited
+        to |k| <= k_c on this grid (no aliasing): OH = fft2(field)/N_disp^2 undoes
+        the N_disp^2 scaling in reconstruct_field, the placed coefficients are
+        read back at the mode indices, and the per-feature standardization is
+        re-applied. Used to verify the FFT wrapper round-trips and as the forward
+        leg of the unet2 spectral<->field map."""
+        OH = np.fft.fft2(np.asarray(field, float)) / (self.n_disp ** 2)
+        coeff = OH[self.ix, self.iy]
+        vec = np.concatenate([coeff.real, coeff.imag])
+        return (vec - self.mean) / self.std
+
+    def roundtrip_maxerr(self, vecs_norm: np.ndarray) -> float:
+        """max |extract_vec(reconstruct_field(v)) - v| over the given vectors."""
+        vecs_norm = np.atleast_2d(np.asarray(vecs_norm, float))
+        errs = [np.abs(self.extract_vec(self.reconstruct_field(v)) - v).max()
+                for v in vecs_norm]
+        return float(max(errs))

code/bifurcation/models/unet2.py

@@ -0,0 +1,248 @@
+"""FFT-wrapped 2-D circular-convolution U-Net backbones: the endpoint head (3 channels) and the segment head (5 channels).
+
+The 196-dim coarse spectral state is not a periodic spatial signal, so the network reconstructs the field on a torus
+from the spectral vector, runs a 2-D circular-conv U-Net on it (spatial locality and translation equivariance) and
+maps the output back; the operator is wrapped in an exact FFT/iFFT pair, so input and output are the same 196-dim
+normalized spectral vector::
+
+    v_norm -> un-normalize -> Hermitian spectrum on N_g^2 -> ifft2 -> field
+           -> 2-D circular-conv U-Net -> fft2 -> extract |k|<=k_c -> re-normalize
+
+Noise and all metrics stay in the 196-dim spectral space. Grid N_g = 32 is alias-free for |k| <= k_c = 8
+(Nyquist 16 > 8), so the mode placement has no self-conjugate in-band mode and the wrapper is an exact identity
+on R^196 (round-trip at machine precision in float64, float32-eps in float32). The (kx, ky) <-> index map comes
+from ``bifurcation.models.spectral``.
+
+  - ``Conditional2D``: 3 field channels [x, x_0, sigma_broadcast]; EDM and Direct SI.
+  - ``LocalDrift2D``:  5 field channels [x_tau, x_tk, tau, t_k, t_k1]; Cascaded SI.
+
+The scalar conditioning is injected as a broadcast channel.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from ._common import _pack_scalar
+from .spectral import FieldReconstructor, K_C
+
+
+class PeriodicConv2d(nn.Module):
+    """Conv2d with cat-based circular padding on both spatial dims.
+
+    Cat-based (not ``padding_mode='circular'``) to stay vmap-safe under PyTorch 2.1. The manual wrap makes
+    the convolution genuinely periodic on the torus."""
+
+    def __init__(self, in_c: int, out_c: int, kernel_size: int, bias: bool = True):
+        super().__init__()
+        assert kernel_size % 2 == 1, "Use odd kernel_size for symmetric circular padding"
+        self.pad = kernel_size // 2
+        self.conv = nn.Conv2d(in_c, out_c, kernel_size, padding=0, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        p = self.pad
+        x = torch.cat([x[..., -p:, :], x, x[..., :p, :]], dim=-2)   # wrap rows
+        x = torch.cat([x[..., :, -p:], x, x[..., :, :p]], dim=-1)   # wrap cols
+        return self.conv(x)
+
+
+def _upsample2(x: torch.Tensor) -> torch.Tensor:
+    """Nearest-neighbour 2x upsample on both spatial dims."""
+    return x.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+
+
+class _SpectralFieldWrapper(nn.Module):
+    """Fixed, invertible reparameterization: normalized spectral vec <-> field.
+
+    The spectral->field map (un-normalize -> Hermitian placement -> ifft2 -> real)
+    and its inverse are both **exactly affine** in their input (the FFT and mode
+    placement are linear; un/re-normalization is affine). We therefore precompute
+    the affine maps as dense matrices once at construction — by passing the
+    standard basis through the verified-exact numpy ``FieldReconstructor`` (single
+    source of truth for the (kx,ky)<->index map) — and apply them as plain
+    matmuls at runtime. Matmul is **vmap-safe** (the SI losses vmap per-sample)
+    and differentiable; the maps are registered buffers so they move with
+    ``.to(device)`` / ``.double()``. Round-trip is exact to machine precision."""
+
+    def __init__(self, packaged: str, n_grid: int = 32,
+                 k_c: int = K_C, component: int | None = None):
+        super().__init__()
+        rec = FieldReconstructor(packaged=packaged, n_disp=n_grid, k_c=k_c)
+        ng = int(n_grid)
+        D = 2 * rec.n_modes
+        if component is not None:
+            # multi-component state (e.g. (u, v)): this wrapper owns one component, so it takes
+            # that component's own slice of the package's frozen normalization. component=None is the
+            # single-field behaviour and is untouched.
+            sl = slice(int(component) * D, (int(component) + 1) * D)
+            if len(rec.mean) < sl.stop:
+                raise ValueError(f"normalization has {len(rec.mean)} entries, too few for component {component}")
+            rec.mean, rec.std = rec.mean[sl], rec.std[sl]
+        # inv (affine, exact):  field_flat = A_inv @ v_norm + c_inv
+        c_inv = rec.reconstruct_field(np.zeros(D)).reshape(-1)
+        A_inv = np.empty((ng * ng, D))
+        e = np.zeros(D)
+        for j in range(D):
+            e[j] = 1.0
+            A_inv[:, j] = rec.reconstruct_field(e).reshape(-1) - c_inv
+            e[j] = 0.0
+        # fwd (affine, exact):  v_norm = A_fwd @ field_flat + c_fwd
+        c_fwd = rec.extract_vec(np.zeros((ng, ng)))
+        A_fwd = np.empty((D, ng * ng))
+        fz = np.zeros(ng * ng)
+        for kk in range(ng * ng):
+            fz[kk] = 1.0
+            A_fwd[:, kk] = rec.extract_vec(fz.reshape(ng, ng)) - c_fwd
+            fz[kk] = 0.0
+        # exact-identity sanity check (catches any self-conjugate-mode projection)
+        zc = np.random.default_rng(0).standard_normal(D)
+        rt = np.abs(A_fwd @ (A_inv @ zc + c_inv) + c_fwd - zc).max()
+        assert rt < 1e-9, f"wrapper is not an identity on R^{D} (round-trip {rt:.2e})"
+        self.n_grid, self.D, self.n_modes = ng, D, rec.n_modes
+        # store at full (float64) precision; cast to the working dtype per call so
+        # the round-trip is exact in float64 and float32-eps in float32 training.
+        self.register_buffer("A_inv", torch.from_numpy(A_inv))
+        self.register_buffer("c_inv", torch.from_numpy(c_inv))
+        self.register_buffer("A_fwd", torch.from_numpy(A_fwd))
+        self.register_buffer("c_fwd", torch.from_numpy(c_fwd))
+
+    def inv(self, v_norm: torch.Tensor) -> torch.Tensor:
+        """(B, D) normalized spectral -> (B, N_g, N_g) real vorticity field."""
+        field_flat = v_norm @ self.A_inv.t().to(v_norm) + self.c_inv.to(v_norm)
+        return field_flat.view(v_norm.shape[0], self.n_grid, self.n_grid)
+
+    def fwd(self, field: torch.Tensor) -> torch.Tensor:
+        """(B, N_g, N_g) real field -> (B, D) normalized spectral vector.
+
+        Exact left-inverse of ``inv`` for band-limited fields."""
+        flat = field.reshape(field.shape[0], -1)
+        return flat @ self.A_fwd.t().to(flat) + self.c_fwd.to(flat)
+
+
+class _PeriodicUNet2Core(nn.Module):
+    """N_g -> N_g/2 -> N_g/4 (bottleneck) -> N_g/2 -> N_g, periodic padding, skip connections."""
+
+    def __init__(self, in_channels: int, base_channels: int, dropout: float = 0.0,
+                 out_channels: int = 1):
+        super().__init__()
+        c = base_channels
+
+        # channel dropout (Dropout2d) after each block's SiLU, inserted only when dropout > 0
+        def _act():
+            return [nn.SiLU()] + ([nn.Dropout2d(float(dropout))] if dropout > 0 else [])
+
+        self.enc0 = nn.Sequential(PeriodicConv2d(in_channels, c,     3), *_act())
+        self.enc1 = nn.Sequential(PeriodicConv2d(c,           c * 2, 3), *_act())
+        self.enc2 = nn.Sequential(
+            PeriodicConv2d(c * 2, c * 4, 3), *_act(),
+            PeriodicConv2d(c * 4, c * 4, 3), *_act(),
+        )
+        self.pool = nn.AvgPool2d(2)
+        self.dec1 = nn.Sequential(PeriodicConv2d(c * 4 + c * 2, c * 2, 3), *_act())
+        self.dec0 = nn.Sequential(PeriodicConv2d(c * 2 + c,     c,     3), *_act())
+        self.out  = nn.Conv2d(c, int(out_channels), kernel_size=1)
+
+    def _forward_core(self, x: torch.Tensor) -> torch.Tensor:
+        s0 = self.enc0(x)
+        s1 = self.enc1(self.pool(s0))
+        s2 = self.enc2(self.pool(s1))
+        u1 = self.dec1(torch.cat([_upsample2(s2), s1], dim=1))
+        u0 = self.dec0(torch.cat([_upsample2(u1), s0], dim=1))
+        y = self.out(u0)
+        return y.squeeze(1) if y.shape[1] == 1 else y               # (B, N_g, N_g) or (B, C, N_g, N_g)
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+def _broadcast_scalar(s: torch.Tensor, B: int, ng: int) -> torch.Tensor:
+    """scalar / per-batch scalar -> (B, 1, N_g, N_g) constant field channel."""
+    return _pack_scalar(s, B).view(B, 1, 1, 1).expand(B, 1, ng, ng)
+
+
+def _require_packaged(cfg: dict) -> str:
+    """The dataset whose normalization.npz the FFT wrapper un/re-normalizes with. Injected by
+    ``bifurcation.registry.build_model`` from the data config; deliberately no default."""
+    if "packaged_dataset_path" not in cfg:
+        raise KeyError("model config lacks 'packaged_dataset_path' — build the model through "
+                       "bifurcation.registry.build_model(model_cfg, method_cfg, data_cfg)")
+    return str(cfg["packaged_dataset_path"])
+
+
+class Conditional2D(_PeriodicUNet2Core):
+    """3-channel-input FFT-wrapped 2-D U-Net for endpoint conditioning.
+
+    forward(x, sigma_embedded, x_0) -> (B, D). EDM passes sigma_embedded = c_noise(sigma), Direct SI
+    passes tau. Field channels: [x_field, x0_field, sigma_broadcast]."""
+
+    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
+                 k_c: int = K_C, packaged: str | None = None,
+                 dropout: float = 0.0):
+        super().__init__(in_channels=3, base_channels=base_channels, dropout=dropout)
+        self.D = D
+        if packaged is None:
+            raise ValueError(
+                "unet2 needs the training dataset's normalization: pass packaged (the "
+                "registry injects data.packaged_dataset_path)")
+        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
+        assert self.wrap.D == D, f"wrapper D {self.wrap.D} != model D {D}"
+
+    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor,
+                x_0: torch.Tensor) -> torch.Tensor:
+        B, ng = x.shape[0], self.wrap.n_grid
+        x_f = self.wrap.inv(x).unsqueeze(1)
+        x0_f = self.wrap.inv(x_0).unsqueeze(1)
+        sig_bc = _broadcast_scalar(sigma_embedded, B, ng)
+        field = self._forward_core(torch.cat([x_f, x0_f, sig_bc], dim=1))
+        return self.wrap.fwd(field)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2D":
+        return cls(D=int(cfg.get("D", 196)),
+                   base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)),
+                   k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg),
+                   dropout=float(cfg.get("dropout", 0.0)))
+
+
+class LocalDrift2D(_PeriodicUNet2Core):
+    """5-channel-input FFT-wrapped 2-D U-Net for the segment drift of Cascaded SI.
+
+    forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D). Field channels: [x_tau_field, x_tk_field, tau_bc,
+    t_k_bc, t_k1_bc]; the segment is identified by the two physical-time scalars broadcast spatially."""
+
+    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
+                 k_c: int = K_C, packaged: str | None = None,
+                 dropout: float = 0.0):
+        super().__init__(in_channels=5, base_channels=base_channels, dropout=dropout)
+        self.D = D
+        if packaged is None:
+            raise ValueError(
+                "unet2 needs the training dataset's normalization: pass packaged (the "
+                "registry injects data.packaged_dataset_path)")
+        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
+        assert self.wrap.D == D, f"wrapper D {self.wrap.D} != model D {D}"
+
+    def forward(self, x_tau: torch.Tensor, tau: torch.Tensor, x_tk: torch.Tensor,
+                t_k: torch.Tensor, t_k1: torch.Tensor) -> torch.Tensor:
+        B, ng = x_tau.shape[0], self.wrap.n_grid
+        xtau_f = self.wrap.inv(x_tau).unsqueeze(1)
+        xtk_f = self.wrap.inv(x_tk).unsqueeze(1)
+        tau_bc = _broadcast_scalar(tau, B, ng)
+        tk_bc = t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng)
+        tk1_bc = t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng)
+        field = self._forward_core(
+            torch.cat([xtau_f, xtk_f, tau_bc, tk_bc, tk1_bc], dim=1))
+        return self.wrap.fwd(field)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "LocalDrift2D":
+        return cls(D=int(cfg.get("D", 196)),
+                   base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)),
+                   k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg),
+                   dropout=float(cfg.get("dropout", 0.0)))

code/bifurcation/models/unet2_2comp.py

@@ -0,0 +1,87 @@
+"""Two-component U-Net for nonlinear Klein-Gordon forecasts.
+
+The state contains 196 coarse displacement coordinates and 196 scaled velocity
+coordinates. The segment head receives both components of the noisy state and
+conditioning state, together with interpolation time and segment boundaries.
+It emits both components through the shared periodic U-Net and frozen spectral
+normalization wrappers. Internal parameter names are retained for checkpoint loading."""
+from __future__ import annotations
+
+import torch
+
+from ..registry import register_model
+from .unet2 import K_C, _broadcast_scalar, _PeriodicUNet2Core, _require_packaged, _SpectralFieldWrapper
+
+
+class _TwoComponent(_PeriodicUNet2Core):
+    """Shared plumbing: one frozen wrapper applied to each half of the 2*D_c state vector."""
+
+    def __init__(self, D: int = 392, base_channels: int = 64, n_grid: int = 32, k_c: int = K_C,
+                 packaged: str | None = None, dropout: float = 0.0, in_channels: int = 7):
+        super().__init__(in_channels=in_channels, base_channels=base_channels, dropout=dropout,
+                         out_channels=2)
+        self.D = int(D)
+        if packaged is None:
+            raise ValueError("unet2_2comp needs the training dataset's normalization: pass packaged")
+        # one wrapper per component, each with that component's slice of the frozen normalization
+        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c, component=0)
+        self.wrap_im = _SpectralFieldWrapper(packaged, n_grid, k_c, component=1)
+        if self.wrap.D * 2 != self.D:
+            raise ValueError(f"two-component state needs D = 2 * {self.wrap.D}, got {self.D}")
+        self.Dc = self.wrap.D
+
+    def _fields(self, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """(B, 2*Dc) -> two (B, 1, Ng, Ng) field channels."""
+        return (self.wrap.inv(v[:, :self.Dc]).unsqueeze(1),
+                self.wrap_im.inv(v[:, self.Dc:]).unsqueeze(1))
+
+    def _vec(self, field2: torch.Tensor) -> torch.Tensor:
+        """(B, 2, Ng, Ng) -> (B, 2*Dc)."""
+        return torch.cat([self.wrap.fwd(field2[:, 0]), self.wrap_im.fwd(field2[:, 1])], dim=1)
+
+
+class LocalDrift2Comp(_TwoComponent):
+    """forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D); channels [Re/Im x_tau, Re/Im x_tk, tau, t_k, t_k1]."""
+
+    def forward(self, x_tau: torch.Tensor, tau: torch.Tensor, x_tk: torch.Tensor,
+                t_k: torch.Tensor, t_k1: torch.Tensor) -> torch.Tensor:
+        B, ng = x_tau.shape[0], self.wrap.n_grid
+        a_re, a_im = self._fields(x_tau)
+        b_re, b_im = self._fields(x_tk)
+        tau_bc = _broadcast_scalar(tau, B, ng)
+        tk_bc = t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng)
+        tk1_bc = t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng)
+        out = self._forward_core(torch.cat([a_re, a_im, b_re, b_im, tau_bc, tk_bc, tk1_bc], dim=1))
+        return self._vec(out)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "LocalDrift2Comp":
+        return cls(D=int(cfg.get("D", 392)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))
+
+
+class Conditional2Comp(LocalDrift2Comp):
+    """Endpoint head with the segment head's interface (t_k = t_start, t_k1 = t_end), as unet2_5ch."""
+
+    def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.t_start, self.t_end = float(t_start), float(t_end)
+
+    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        B = x.shape[0]
+        tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device)
+        tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device)
+        return super().forward(x, sigma_embedded, x_0, tk, tk1)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2Comp":
+        return cls(D=int(cfg.get("D", 392)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)),
+                   t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0)))
+
+
+register_model("unet2_2comp", endpoint_cls=Conditional2Comp, segment_cls=LocalDrift2Comp,
+               needs_dataset_norm=True,
+               description="two-component (Re,Im) coarse complex state; unet2 core, per-component FFT wrapper")

code/bifurcation/models/unet2_5ch.py

@@ -0,0 +1,50 @@
+"""``unet2_5ch``: the backbone of the Euler and SQG models.
+
+Channel parity between the endpoint and the local arms: the endpoint head IS the local head
+(``LocalDrift2D``, 5 input channels [x_τ, x_0, τ, t_start, t_end]) with the segment's start and end
+time fixed to the forecast horizon (t_start = 0, t_end = T). Both heads are therefore the identical
+module class with the identical parameter count (1,515,265 at base_channels 64); the two extra
+planes carry no information for an endpoint arm (constants) — the fix removes an *interface*
+asymmetry, not a capacity difference.
+
+``t_end`` must equal the forecast horizon: ``bifurcation.registry.build_model`` checks the model
+configuration's ``t_start`` / ``t_end`` against ``method.sampling.t_values`` and raises on disagreement.
+"""
+from __future__ import annotations
+
+import torch
+
+from ..registry import register_model
+from .unet2 import K_C, LocalDrift2D, _require_packaged
+
+
+class Conditional2D5ch(LocalDrift2D):
+    """Endpoint head with the segment head's exact interface: t_k = t_start, t_{k+1} = t_end.
+
+    forward(x, sigma_embedded, x_0) -> (B, D): global SI passes sigma_embedded = tau."""
+
+    def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.t_start, self.t_end = float(t_start), float(t_end)
+
+    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        B = x.shape[0]
+        tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device)
+        tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device)
+        return super().forward(x, sigma_embedded, x_0, tk, tk1)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2D5ch":
+        return cls(D=int(cfg.get("D", 196)),
+                   base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)),
+                   k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg),
+                   dropout=float(cfg.get("dropout", 0.0)),
+                   t_start=float(cfg.get("t_start", 0.0)),
+                   t_end=float(cfg.get("t_end", 1.0)))
+
+
+register_model("unet2_5ch", Conditional2D5ch, LocalDrift2D, needs_dataset_norm=True,
+               description="unet2 with channel parity: both heads are LocalDrift2D (5 ch); the endpoint head "
+                           "fixes (t_start, t_end) = (0, T). The backbone of the Euler and SQG models.")

code/bifurcation/models/unet2_5ch_coord.py

@@ -0,0 +1,45 @@
+"""``unet2_5ch_coord``: the backbone for datasets without translation symmetry (the passive scalar).
+
+The passive scalar's fixed cellular flow is not translation invariant, so the equivariant ``unet2`` cannot represent
+its conditional-mean map. This variant is ``unet2_5ch`` built on the ``unet2_coord`` core, which adds two coordinate
+input planes: both heads are ``LocalDrift2DCoord`` with 7 input channels [x_tau, x_tk, tau, t_start, t_end, X, Y],
+and the endpoint head fixes (t_start, t_end) = (0, T) as ``Conditional2D5ch`` does. The parameter count is that of
+unet2_5ch plus the two extra 3x3 input filters (1,152 at base_channels 64), identical for both heads, so EDM, Direct
+SI and Cascaded SI share one module. ``t_end`` must equal the forecast horizon (checked by
+``bifurcation.registry.build_model``).
+"""
+from __future__ import annotations
+
+import torch
+
+from ..registry import register_model
+from .spectral import K_C
+from .unet2 import _require_packaged
+from .unet2_coord import LocalDrift2DCoord
+
+
+class Conditional2D5chCoord(LocalDrift2DCoord):
+    """Endpoint head with the segment head's exact interface: t_k = t_start, t_{k+1} = t_end."""
+
+    def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.t_start, self.t_end = float(t_start), float(t_end)
+
+    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        B = x.shape[0]
+        tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device)
+        tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device)
+        return super().forward(x, sigma_embedded, x_0, tk, tk1)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2D5chCoord":
+        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)),
+                   t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0)))
+
+
+register_model("unet2_5ch_coord", Conditional2D5chCoord, LocalDrift2DCoord, needs_dataset_norm=True,
+               description="unet2_coord with channel parity: both heads are LocalDrift2DCoord (7 ch, not translation-"
+                           "equivariant); the endpoint head fixes (t_start, t_end) = (0, T). The backbone for "
+                           "datasets without translation symmetry (passive scalar).")

code/bifurcation/models/unet2_5ch_coord_aug.py

@@ -0,0 +1,110 @@
+"""Augmented-state variant of unet2_5ch_coord: the backbone of the heterogeneous-waves models.
+
+State x = (u, v_hat) in R^392: two independent 196-D spectral fields (coarse displacement, finite-
+difference coarse velocity), each with its OWN _SpectralFieldWrapper / normalization (velocity has a
+different physical scale than displacement). Everything else -- the periodic-conv trunk, coordinate
+planes, channel structure -- is IDENTICAL to unet2_5ch_coord; only the input/output channel counts
+differ (9 in / 2 out instead of 7 in / 1 out) to carry both fields through together.
+
+Standard single-path dataset-norm injection (needs_dataset_norm=True, like every other backbone): the
+u/v wrappers are built from TWO SUBDIRECTORIES of the one injected ``packaged_dataset_path``,
+
+    <packaged_dataset_path>/u_norm/normalization.npz   (196-D, displacement)
+    <packaged_dataset_path>/v_norm/normalization.npz   (196-D, velocity)
+
+so the run configuration names one dataset path, as for every other model.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .spectral import K_C
+from .unet2 import _PeriodicUNet2Core, _SpectralFieldWrapper, _broadcast_scalar, _upsample2, _require_packaged
+from .unet2_coord import _coord_planes
+
+
+class _DualHeadCore(_PeriodicUNet2Core):
+    """Same trunk as _PeriodicUNet2Core; final 1x1 conv widened from 1 to 2 output channels (predict
+    u_next and v_next together), and _forward_core no longer squeezes the channel axis away."""
+
+    def __init__(self, in_channels: int, base_channels: int, dropout: float = 0.0):
+        super().__init__(in_channels=in_channels, base_channels=base_channels, dropout=dropout)
+        self.out = nn.Conv2d(base_channels, 2, kernel_size=1)
+
+    def _forward_core(self, x: torch.Tensor) -> torch.Tensor:
+        s0 = self.enc0(x)
+        s1 = self.enc1(self.pool(s0))
+        s2 = self.enc2(self.pool(s1))
+        u1 = self.dec1(torch.cat([_upsample2(s2), s1], dim=1))
+        u0 = self.dec0(torch.cat([_upsample2(u1), s0], dim=1))
+        return self.out(u0)                                        # (B, 2, ng, ng) -- no squeeze
+
+
+class LocalDrift2DCoordAug(nn.Module):
+    """forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, 392); x_tau/x_tk are (B, 392) = concat(u, v)."""
+
+    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32, k_c: int = K_C,
+                 packaged: str | None = None, dropout: float = 0.0):
+        super().__init__()
+        if packaged is None:
+            raise ValueError("unet2_5ch_coord_aug needs the training dataset's normalization: pass packaged")
+        self.core = _DualHeadCore(in_channels=9, base_channels=base_channels, dropout=dropout)
+        self.wrap_u = _SpectralFieldWrapper(str(packaged) + "/u_norm", n_grid, k_c)
+        self.wrap_v = _SpectralFieldWrapper(str(packaged) + "/v_norm", n_grid, k_c)
+        assert self.wrap_u.D == D and self.wrap_v.D == D
+        self.D = D
+
+    def forward(self, x_tau, tau, x_tk, t_k, t_k1):
+        B, ng, D = x_tau.shape[0], self.wrap_u.n_grid, self.D
+        u_tau, v_tau = x_tau[:, :D], x_tau[:, D:]
+        u_tk, v_tk = x_tk[:, :D], x_tk[:, D:]
+        xt_u = self.wrap_u.inv(u_tau).unsqueeze(1); xt_v = self.wrap_v.inv(v_tau).unsqueeze(1)
+        xk_u = self.wrap_u.inv(u_tk).unsqueeze(1); xk_v = self.wrap_v.inv(v_tk).unsqueeze(1)
+        inp = torch.cat([xt_u, xt_v, xk_u, xk_v,
+                         _broadcast_scalar(tau, B, ng).to(xt_u),
+                         t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt_u),
+                         t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt_u),
+                         _coord_planes(B, ng, xt_u)], dim=1)
+        field = self.core._forward_core(inp)                        # (B, 2, ng, ng)
+        u_next = self.wrap_u.fwd(field[:, 0]); v_next = self.wrap_v.fwd(field[:, 1])
+        return torch.cat([u_next, v_next], dim=1)
+
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "LocalDrift2DCoordAug":
+        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))
+
+
+class Conditional2D5chCoordAug(LocalDrift2DCoordAug):
+    """Endpoint head with the segment head's exact interface (mirrors Conditional2D5chCoord):
+    forward(x, sigma_embedded, x_0) fixes (t_start, t_end) = (0, T)."""
+
+    def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.t_start, self.t_end = float(t_start), float(t_end)
+
+    def forward(self, x, sigma_embedded, x_0):
+        B = x.shape[0]
+        tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device)
+        tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device)
+        return super().forward(x, sigma_embedded, x_0, tk, tk1)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2D5chCoordAug":
+        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)),
+                   t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0)))
+
+
+from ..registry import register_model  # noqa: E402
+
+register_model("unet2_5ch_coord_aug", Conditional2D5chCoordAug, LocalDrift2DCoordAug, needs_dataset_norm=True,
+               description="Augmented-state (u, v_hat) backbone of the heterogeneous-waves models: two "
+                           "_SpectralFieldWrapper instances read from <packaged>/u_norm and <packaged>/v_norm; "
+                           "otherwise identical trunk and parameter count to unet2_5ch_coord.")

code/bifurcation/models/unet2_coord.py

@@ -0,0 +1,87 @@
+"""unet2 with coordinate channels — a NON-equivariant variant of the FFT-wrapped 2-D U-Net.
+
+The passive scalar's fixed cellular flow
+psi_K is not translation-invariant, so its conditional-mean map x_0 -> M_T x_0 is not
+translation-equivariant; ``unet2`` (exactly equivariant to shifts by its pooling stride)
+provably cannot represent it. This variant keeps unet2's field reconstruction, circular
+convolutions and heads but appends two constant coordinate planes (x, y in [-1, 1)) to the
+input channels, which breaks translation equivariance and lets position-dependent filters
+form.
+
+  - ``Conditional2DCoord``  forward(x, sigma_embedded, x_0) -> (B, D)       (5 input channels)
+  - ``LocalDrift2DCoord``   forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D)   (7 input channels)
+
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .spectral import K_C
+from .unet2 import _PeriodicUNet2Core, _SpectralFieldWrapper, _broadcast_scalar, _require_packaged
+
+
+def _coord_planes(B: int, ng: int, like: torch.Tensor) -> torch.Tensor:
+    """(B, 2, ng, ng) constant planes with x and y coordinates in [-1, 1)."""
+    c = torch.arange(ng, device=like.device, dtype=like.dtype) * (2.0 / ng) - 1.0
+    X = c.view(1, 1, ng, 1).expand(B, 1, ng, ng)
+    Y = c.view(1, 1, 1, ng).expand(B, 1, ng, ng)
+    return torch.cat([X, Y], dim=1)
+
+
+class Conditional2DCoord(_PeriodicUNet2Core):
+    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
+                 k_c: int = K_C, packaged: str | None = None, dropout: float = 0.0):
+        super().__init__(in_channels=5, base_channels=base_channels, dropout=dropout)
+        self.D = D
+        if packaged is None:
+            raise ValueError("unet2_coord needs the training dataset's normalization: pass packaged")
+        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
+        assert self.wrap.D == D
+
+    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor:
+        B, ng = x.shape[0], self.wrap.n_grid
+        x_f = self.wrap.inv(x).unsqueeze(1)
+        field = self._forward_core(torch.cat([x_f, self.wrap.inv(x_0).unsqueeze(1),
+                                              _broadcast_scalar(sigma_embedded, B, ng).to(x_f),
+                                              _coord_planes(B, ng, x_f)], dim=1))
+        return self.wrap.fwd(field)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "Conditional2DCoord":
+        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))
+
+
+class LocalDrift2DCoord(_PeriodicUNet2Core):
+    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
+                 k_c: int = K_C, packaged: str | None = None, dropout: float = 0.0):
+        super().__init__(in_channels=7, base_channels=base_channels, dropout=dropout)
+        self.D = D
+        if packaged is None:
+            raise ValueError("unet2_coord needs the training dataset's normalization: pass packaged")
+        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
+        assert self.wrap.D == D
+
+    def forward(self, x_tau: torch.Tensor, tau: torch.Tensor, x_tk: torch.Tensor,
+                t_k: torch.Tensor, t_k1: torch.Tensor) -> torch.Tensor:
+        B, ng = x_tau.shape[0], self.wrap.n_grid
+        xt = self.wrap.inv(x_tau).unsqueeze(1)
+        field = self._forward_core(torch.cat([
+            xt, self.wrap.inv(x_tk).unsqueeze(1), _broadcast_scalar(tau, B, ng).to(xt),
+            t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt), t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt),
+            _coord_planes(B, ng, xt)], dim=1))
+        return self.wrap.fwd(field)
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> "LocalDrift2DCoord":
+        return cls(D=int(cfg.get("D", 196)), base_channels=int(cfg.get("base_channels", 64)),
+                   n_grid=int(cfg.get("n_grid", 32)), k_c=int(cfg.get("k_c", K_C)),
+                   packaged=_require_packaged(cfg), dropout=float(cfg.get("dropout", 0.0)))
+
+
+from ..registry import register_model  # noqa: E402
+
+register_model("unet2_coord", Conditional2DCoord, LocalDrift2DCoord, needs_dataset_norm=True,
+               description="unet2 + two coordinate input planes (not translation-equivariant; passive scalar)")

code/bifurcation/registry.py

@@ -0,0 +1,127 @@
+"""Model and method factories shared by training, sampling and evaluation.
+
+Methods: ``direct_si``, ``cascaded_si`` and ``edm``. Models: ``unet2_5ch`` (Euler, SQG), ``unet2_5ch_coord``
+(passive scalar), ``unet2_5ch_coord_aug`` (heterogeneous waves) and ``unet2_2comp`` (NLKG); each run configuration
+under records/<system>/runs/ names its model and method."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+from .methods.edm import EDMMethod
+from .methods.direct_si import DirectSIMethod
+from .methods.cascaded_si import CascadedSIMethod
+from .models.unet2 import Conditional2D, LocalDrift2D
+
+SEGMENT_DATA_MODES = ("segment_pairs", "fine_segment_pairs")
+
+# reserved transient key: the method's segment boundaries in physical time, injected by
+# build_model for backbones that declare physical times in their own YAML (see module docstring)
+T_VALUES_KEY = "_method_t_values"
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    name: str
+    endpoint_cls: type
+    segment_cls: type
+    needs_dataset_norm: bool = False
+    description: str = ""
+
+
+MODELS: dict[str, ModelSpec] = {}
+METHODS: dict[str, Callable[[dict], object]] = {}
+
+
+def register_model(name: str, endpoint_cls: type, segment_cls: type,
+                   needs_dataset_norm: bool = False, description: str = "",
+                   overwrite: bool = False) -> ModelSpec:
+    if name in MODELS and not overwrite:
+        raise KeyError(f"model {name!r} already registered")
+    spec = ModelSpec(name, endpoint_cls, segment_cls, needs_dataset_norm, description)
+    MODELS[name] = spec
+    return spec
+
+
+def register_method(name: str, factory: Callable[[dict], object],
+                    overwrite: bool = False) -> None:
+    if name in METHODS and not overwrite:
+        raise KeyError(f"method {name!r} already registered")
+    METHODS[name] = factory
+
+
+def head_for_data_mode(data_mode: str | None) -> str:
+    return "segment" if data_mode in SEGMENT_DATA_MODES else "endpoint"
+
+
+def build_model(model_cfg: dict, method_cfg: dict | None = None,
+                data_cfg: dict | None = None):
+    """Construct the backbone named by ``model_cfg['name']``.
+
+    ``method_cfg['data_mode']`` selects the head. For models with ``needs_dataset_norm`` the dataset path is
+    taken from ``data_cfg`` (``packaged_dataset_path``); a model configuration that carries its own copy of the
+    key is accepted only if it agrees with the data config.
+    """
+    name = model_cfg["name"]
+    if name not in MODELS:
+        raise ValueError(f"unknown model: {name!r} (registered: {sorted(MODELS)})")
+    spec = MODELS[name]
+    cfg = dict(model_cfg)
+    if spec.needs_dataset_norm:
+        data_path = (data_cfg or {}).get("packaged_dataset_path")
+        model_path = cfg.get("packaged_dataset_path")
+        if data_path is None and model_path is None:
+            raise KeyError(
+                f"model {name!r} needs the dataset normalization: pass data_cfg with "
+                "'packaged_dataset_path' (model configurations do not carry dataset paths)")
+        if data_path is not None and model_path is not None:
+            if _norm_path(data_path) != _norm_path(model_path):
+                raise ValueError(
+                    f"model.packaged_dataset_path ({model_path}) != data.packaged_dataset_path "
+                    f"({data_path}) — the FFT wrapper must use the training dataset's statistics")
+        cfg["packaged_dataset_path"] = str(data_path if data_path is not None else model_path)
+    t_values = ((method_cfg or {}).get("sampling") or {}).get("t_values")
+    if t_values:
+        cfg[T_VALUES_KEY] = list(t_values)
+    data_mode = (method_cfg or {}).get("data_mode")
+    # a model configuration that declares the horizon planes must agree with the method's time grid:
+    # t_start = t_values[0], t_end = t_values[-1]
+    tv = ((method_cfg or {}).get("sampling") or {}).get("t_values")
+    if tv is not None and ("t_start" in cfg or "t_end" in cfg):
+        t0, t1 = float(tv[0]), float(tv[-1])
+        if "t_start" in cfg and abs(float(cfg["t_start"]) - t0) > 1e-9:
+            raise ValueError(f"model {name!r}: t_start {cfg['t_start']} != method t_values[0] = {t0}")
+        if "t_end" in cfg and abs(float(cfg["t_end"]) - t1) > 1e-9:
+            raise ValueError(f"model {name!r}: t_end {cfg['t_end']} != method horizon t_values[-1] = {t1} "
+                             f"(the model constant must equal the forecast horizon)")
+    cls = spec.segment_cls if head_for_data_mode(data_mode) == "segment" else spec.endpoint_cls
+    return cls.from_config(cfg)
+
+
+def build_method(method_cfg: dict):
+    name = method_cfg["name"]
+    if name not in METHODS:
+        raise ValueError(f"unknown method: {name!r} (registered: {sorted(METHODS)})")
+    return METHODS[name](method_cfg)
+
+
+def _norm_path(p: str) -> str:
+    import os
+    return os.path.normpath(os.path.abspath(str(p)))
+
+
+# ── built-ins ────────────────────────────────────────────────────────────────
+register_model("unet2", Conditional2D, LocalDrift2D, needs_dataset_norm=True,
+               description="FFT-wrapped 2-D circular-conv U-Net on the 32x32 coarse field "
+                           "(Conditional2D 3-ch / LocalDrift2D 5-ch)")
+register_method("edm", EDMMethod.from_config)
+register_method("direct_si", DirectSIMethod.from_config)
+register_method("cascaded_si", CascadedSIMethod.from_config)
+
+# the backbones register themselves on import; unet2_coord provides the coordinate planes of unet2_5ch_coord(_aug)
+from .models import unet2_coord as _unet2_coord  # noqa: F401,E402
+from .models import unet2_5ch as _unet2_5ch  # noqa: F401,E402
+from .models import unet2_5ch_coord as _unet2_5ch_coord  # noqa: F401,E402
+from .models import unet2_5ch_coord_aug as _unet2_5ch_coord_aug  # noqa: F401,E402
+from .models import unet2_2comp as _unet2_2comp  # noqa: F401,E402

code/bifurcation/sample.py

@@ -0,0 +1,142 @@
+"""Draw forecasts from a trained model: for every initial condition of a split, ``n_samples`` samples of x_T.
+
+    python -m bifurcation.sample --config CONFIG --data DIR --checkpoint FILE --out DIR [--split test] [--gids 5,16]
+        [--n-samples 10000] [--n-steps N] [--seed-base 777] [--chunk 2000] [--truth-bank DIR] [--device cuda]
+
+CONFIG is a run configuration (records/<system>/runs/<model>/config.yaml, or the config_resolved.yaml of a run trained
+here). The EMA weights of the checkpoint are used when the method keeps an EMA. ``--n-steps`` overrides the sampler
+resolution (steps per segment for Cascaded SI, SDE steps for Direct SI, Heun steps for EDM). The noise of condition
+``gid`` comes from the stream seeded with ``seed_base * 100003 + gid`` and is drawn in chunks of ``--chunk`` samples.
+``--truth-bank`` supplies x_0 of test conditions that are not in the package (waves, NLKG).
+
+Output: ``OUT/<variant>/group_<gid>.npz`` with ``gen_xT_norm`` (n_samples, D), ``x0_norm`` (D,), ``true_xT_norm``
+(the package's truth members at the horizon) and ``gid``; states are in the package's normalized units.
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .data.loaders import get_split_set
+from .registry import build_method, build_model
+from .utils import load_yaml
+
+
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--config", required=True, help="run configuration (YAML)")
+    p.add_argument("--data", required=True, help="dataset package directory")
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--out", required=True, help="output directory")
+    p.add_argument("--split", default="test", choices=("train", "val", "test"))
+    p.add_argument("--gids", default=None, help="comma-separated subset of the split's conditions")
+    p.add_argument("--n-samples", type=int, default=10000)
+    p.add_argument("--n-steps", type=int, default=None, help="sampler resolution override")
+    p.add_argument("--seed-base", type=int, default=777)
+    p.add_argument("--chunk", type=int, default=None, help="draw the samples of a condition in chunks of this size")
+    p.add_argument("--truth-bank", default=None, help="truth bank with the x_0 of test conditions absent from the package")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    return p.parse_args(argv)
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 24), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def main(argv=None) -> int:
+    args = parse_args(argv)
+    cfg = load_yaml(args.config)
+    cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/"
+    device = str(args.device)
+
+    allow_tf32 = cfg.get("training", {}).get("allow_tf32")
+    if allow_tf32 is not None:                        # the arithmetic class the model was trained under
+        torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
+        torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
+    seed = int(cfg.get("seed", cfg["data"].get("seed", 0)))
+
+    ckpt_path = Path(args.checkpoint).resolve()
+    ckpt = torch.load(ckpt_path, map_location=device)
+    model = build_model(cfg["model"], cfg["method"], cfg["data"])
+    ema_halflife = cfg.get("training", {}).get("ema_halflife_kimg", cfg["method"].get("ema_halflife_kimg"))
+    method_uses_ema = ema_halflife is not None and float(ema_halflife) > 0
+    if "ema_state_dict" in ckpt and method_uses_ema:
+        model.load_state_dict(ckpt["ema_state_dict"])
+        weights = "ema"
+    else:
+        model.load_state_dict(ckpt["model_state_dict"])
+        weights = "raw"
+    model.to(device).eval()
+    print(f"checkpoint: {ckpt_path.name} (step {ckpt.get('step')}, {weights} weights)")
+
+    if args.n_steps is not None:
+        sampling = cfg["method"]["sampling"]
+        for key in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps"):
+            if key in sampling:
+                break
+        else:
+            raise KeyError("--n-steps: method.sampling has none of n_sde_steps / n_sde_steps_per_seg / n_steps")
+        print(f"sampler resolution: method.sampling.{key} {sampling[key]} -> {args.n_steps}")
+        sampling[key] = args.n_steps
+    method = build_method(cfg["method"])
+
+    split_x0, truth_clouds, gids = get_split_set(cfg, args.split, seed=seed, x0_bank=args.truth_bank)
+    if args.gids is not None:
+        wanted = sorted(int(g) for g in args.gids.split(","))
+        missing = [g for g in wanted if g not in gids]
+        if missing:
+            raise ValueError(f"--gids not in split '{args.split}': {missing}")
+        sel = [gids.index(g) for g in wanted]
+        split_x0, truth_clouds, gids = split_x0[sel], truth_clouds[sel], wanted
+    split_x0, truth_clouds = split_x0.to(device), truth_clouds.to(device)
+    n_gids = split_x0.shape[0]
+    print(f"split {args.split}: {n_gids} conditions, {args.n_samples} samples each")
+
+    out = Path(args.out).resolve()
+    out.mkdir(parents=True, exist_ok=True)
+    split_offset = {"train": 1000, "val": 2000, "test": 3000}[args.split]
+    variants = method.variants_for_eval()
+    for v_idx, (label, sample_kwargs) in enumerate(variants.items()):
+        bundle_dir = out / label
+        bundle_dir.mkdir(parents=True, exist_ok=True)
+        torch.manual_seed(seed + split_offset + 100 * v_idx)
+        with torch.no_grad():
+            for i in range(n_gids):
+                x0_i, gid = split_x0[i], int(gids[i])
+                torch.manual_seed(int(args.seed_base) * 100003 + gid)
+                if args.chunk is not None and args.n_samples > int(args.chunk):
+                    parts, left = [], int(args.n_samples)
+                    while left > 0:
+                        m = min(int(args.chunk), left)
+                        parts.append(method.sample(model, x0_i, n_samples=m, **sample_kwargs).squeeze(0))
+                        left -= m
+                    gen_cloud = torch.cat(parts, dim=0)
+                else:
+                    gen_cloud = method.sample(model, x0_i, n_samples=int(args.n_samples), **sample_kwargs).squeeze(0)
+                np.savez(bundle_dir / f"group_{gid:04d}.npz", true_xT_norm=truth_clouds[i].cpu().numpy(),
+                         gen_xT_norm=gen_cloud.cpu().numpy(), x0_norm=x0_i.cpu().numpy(), gid=np.int64(gid))
+                print(f"  {label}: condition {gid} ({i + 1}/{n_gids})", flush=True)
+
+    record = {"config": str(Path(args.config).resolve()), "checkpoint": str(ckpt_path), "checkpoint_sha256": _sha256(ckpt_path),
+              "weights": weights, "split": args.split, "gids": list(gids), "n_samples": int(args.n_samples),
+              "sampling": {k: cfg["method"]["sampling"].get(k) for k in ("n_sde_steps", "n_sde_steps_per_seg", "n_steps",
+                                                                        "n_segments", "noise_scale") if k in cfg["method"]["sampling"]},
+              "seed_base": int(args.seed_base), "chunk": args.chunk, "variants": list(variants), "device": device,
+              "allow_tf32": allow_tf32, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
+    (out / "SAMPLES.json").write_text(json.dumps(record, indent=1) + "\n")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/bifurcation/train.py

@@ -0,0 +1,115 @@
+"""Train one model from a run configuration.
+
+    python -m bifurcation.train --config records/euler/runs/cascaded_si_S10/config.yaml --data $EULER_DATA \\
+        --out runs/euler_cascaded_si_S10 [--seed 0] [--device cuda] [--resume]
+
+The configuration names the dataset, the model, the method with its hyperparameters and the training schedule
+(records/<system>/runs/<model>/config.yaml are the configurations of the paper's models). Before training, the
+training pairs are written to OUT/manifest/: Cascaded SI with S segments reads the first 200 / S realizations of each
+of the 80 training conditions, cut into S segments (16,000 pairs for every S; S = 1 is Direct SI); EDM reads the
+endpoint pair (x_0, x_T) of all 200 realizations. Checkpoints are saved every ``training.val_every_updates`` updates
+under OUT/checkpoints/ (bifurcation.trainer); ``--resume`` continues an interrupted run exactly from its
+``restart_latest.pt``.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .data import manifests as MF
+from .data.loaders import get_dataloaders
+from .registry import build_method, build_model
+from .trainer import Trainer
+from .utils import load_yaml, save_resolved_config, seed_everything
+
+# a restart may only continue a run with the same data, method, model and seed
+RESUME_IDENTITY_KEYS = (("data", "name"), ("data", "packaged_dataset_path"), ("method", "name"), ("method", "data_mode"),
+                        ("method", "sigma_train"), ("method", "P_mean"), ("method", "lr"), ("training", "batch_size"),
+                        ("model", "name"))
+
+
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--config", required=True, help="run configuration (YAML)")
+    p.add_argument("--data", required=True, help="dataset package directory")
+    p.add_argument("--out", required=True, help="run directory (new, or the one to resume)")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--device", default=None, help="default: cuda if available, else cpu")
+    p.add_argument("--resume", action="store_true", help="continue from OUT/checkpoints/restart_latest.pt")
+    return p.parse_args(argv)
+
+
+def write_training_manifest(cfg: dict, out: Path) -> Path:
+    """The training pairs of the run (see the module docstring); returns the manifest file."""
+    packaged = cfg["data"]["packaged_dataset_path"]
+    n_ic, R = int(cfg["data"]["train_ics"]), int(cfg["data"]["train_realizations"])
+    g = MF.fine_grid(packaged)
+    if len(g["gids_train"]) < n_ic or g["vecs_train_shape"][1] < R:
+        raise ValueError(f"the package holds {len(g['gids_train'])} training conditions x {g['vecs_train_shape'][1]} "
+                         f"realizations, the configuration needs {n_ic} x {R}")
+    m = cfg["method"]
+    T = float(m["sampling"]["t_values"][-1])
+    if m["data_mode"] == "fine_segment_pairs":
+        S = int(m["sampling"].get("n_segments", 1))
+        if R % S:
+            raise ValueError(f"S = {S} must divide the {R} realizations per training condition")
+        geo = MF.segment_geometry(T, S, g["t_values"])
+        rows = MF.block_rows_segment(g["gids_train"], n_ic, R // S)
+        name = f"segments_S{S:02d}"
+        path = MF.write_segment_manifest(out / "manifest", rows, geo, name, packaged,
+                                         extra={"realizations_per_condition": R // S})
+    elif m["data_mode"] == "fine_endpoint_pairs":
+        geo = MF.segment_geometry(T, 1, g["t_values"])
+        rows = MF.block_rows_endpoint(g["gids_train"], n_ic, R)
+        path = MF.write_endpoint_manifest(out / "manifest", rows, geo["n_T"], g["t_values"], "endpoints", packaged,
+                                          extra={"realizations_per_condition": R})
+    else:
+        raise ValueError(f"unknown method.data_mode {m['data_mode']!r}")
+    if int(m["fine_horizon_idx"]) != int(geo["n_T"]):
+        raise ValueError(f"method.fine_horizon_idx {m['fine_horizon_idx']} != the index {geo['n_T']} of T = {T}")
+    return path
+
+
+def main(argv=None) -> int:
+    args = parse_args(argv)
+    cfg = load_yaml(args.config)
+    cfg["data"]["packaged_dataset_path"] = str(Path(args.data).resolve()) + "/"
+    cfg["seed"] = int(args.seed)
+    run_dir = Path(args.out).resolve()
+    run_dir.mkdir(parents=True, exist_ok=True)
+    cfg["method"]["fine_slice_manifest"] = str(write_training_manifest(cfg, run_dir))
+    cfg["training"]["device"] = str(args.device or os.environ.get("BIFURCATION_DEVICE")
+                                    or ("cuda" if torch.cuda.is_available() else "cpu"))
+
+    resume_payload = None
+    if args.resume:
+        restart_path = run_dir / "checkpoints" / "restart_latest.pt"
+        if not restart_path.is_file():
+            raise FileNotFoundError(f"--resume: no restart checkpoint at {restart_path}")
+        resume_payload = torch.load(restart_path, map_location="cpu", weights_only=False)
+        old = resume_payload.get("config", {})
+        for section, key in RESUME_IDENTITY_KEYS:
+            a, b = old.get(section, {}).get(key), cfg.get(section, {}).get(key)
+            if a is not None and a != b:
+                raise AssertionError(f"--resume: {section}.{key} was {a!r}, now {b!r}; refusing to continue a different run")
+        if int(old.get("seed", args.seed)) != int(args.seed):
+            raise AssertionError(f"--resume: the run has seed {old.get('seed')}, --seed {args.seed}")
+        print(f"resume: restart_latest.pt at step {resume_payload['step']}")
+    save_resolved_config(cfg, run_dir / "config_resolved.yaml")
+
+    seed_everything(int(args.seed))
+    train_loader, val_loader = get_dataloaders(cfg, seed=int(args.seed))
+    model = build_model(cfg["model"], cfg["method"], cfg["data"])
+    method = build_method(cfg["method"])
+    Trainer(model=model, method=method, train_loader=train_loader, val_loader=val_loader, cfg=cfg, run_dir=run_dir,
+            device=str(cfg["training"]["device"])).fit(resume=resume_payload)
+    print(f"\nrun:         {run_dir}\ncheckpoints: {run_dir / 'checkpoints'}\nlog:         {run_dir / 'train_log.csv'}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

code/bifurcation/trainer.py

@@ -0,0 +1,371 @@
+"""Training loop: exactly ``training.max_steps`` optimizer updates, with EMA weights, gradient clipping and
+validation events.
+
+A validation event fires after global update ``step`` when ``step % val_every_updates == 0``, at ``max_steps`` and
+at every step listed in ``val_extra_steps`` (0 = the untrained model). An event evaluates the training objective on
+the validation pairs (``val_loss``, and ``val_loss_ema`` for the EMA weights), writes one row of ``train_log.csv``
+(``event`` = 1), saves ``checkpoints/checkpoint_step<NNNNNNNN>.pt`` and the full-state ``restart_latest.pt``. With
+``log_loss_every_epochs: k`` a loss-only row (``event`` = 0, no checkpoint) is written at the end of every k-th epoch.
+``train_loss`` of a row is the mean over the updates since the previous row. The validation passes run under a
+forked RNG with a fixed seed and restore the loader's generator state, so they never perturb the training
+trajectory. ``checkpoints/checkpoint.pt`` holds the final weights; every checkpoint stores the raw and the EMA
+weights and the run configuration.
+
+EMA and LR ramp-up follow the convention of NVlabs/edm (``training/training_loop.py``): ``ema_halflife_kimg``,
+``ema_rampup_ratio`` and ``lr_rampup_kimg`` come from the training section (overrides) or the method section.
+"""
+from __future__ import annotations
+
+import copy
+import os
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import torch
+
+RESTART_SCHEMA = "restart_v1"
+RESTART_NAME = "restart_latest.pt"
+LOG_COLUMNS = "step,epoch,batch_index,event,train_loss,val_loss,val_loss_ema,train_eval_loss,train_eval_loss_ema,lr,elapsed_s"
+
+
+def _ema_beta(cur_nimg: int, batch_size: int, ema_halflife_kimg: float, ema_rampup_ratio: float | None) -> float:
+    ema_halflife_nimg = ema_halflife_kimg * 1000
+    if ema_rampup_ratio is not None:
+        ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * ema_rampup_ratio)
+    return 0.5 ** (batch_size / max(ema_halflife_nimg, 1e-8))
+
+
+def _set_rampup_lr(optimizer: torch.optim.Optimizer, base_lr: float, cur_nimg: int, lr_rampup_kimg: float) -> None:
+    factor = min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1.0)
+    for g in optimizer.param_groups:
+        g["lr"] = base_lr * factor
+
+
+def _eval_loss(model, loader, method, device, max_batches=None) -> float:
+    losses = []
+    with torch.no_grad():
+        for i, batch in enumerate(loader):
+            if max_batches is not None and i >= max_batches:
+                break
+            batch = tuple(t.to(device) for t in batch)
+            losses.append(method.loss(model, batch).item())
+    return float(np.mean(losses)) if losses else float("nan")
+
+
+def _save_checkpoint(path: Path, model, ema_model, cfg: Dict[str, Any], epoch: int, step: int) -> None:
+    payload = {"model_state_dict": model.state_dict(), "epoch": epoch, "config": cfg, "step": int(step)}
+    if ema_model is not None:
+        payload["ema_state_dict"] = ema_model.state_dict()
+    torch.save(payload, path)
+
+
+def _save_restart(path: Path, payload: Dict[str, Any]) -> None:
+    tmp = path.with_name(path.name + ".tmp")            # atomic: a job killed mid-save leaves no truncated file
+    torch.save(payload, tmp)
+    os.replace(tmp, path)
+
+
+def _truncate_log(log_path: Path, keep_step: int) -> int:
+    """Drop the rows after ``keep_step`` (and a partial final line)."""
+    lines = log_path.read_text().splitlines()
+    header, rows = lines[0], lines[1:]
+    kept = []
+    for line in rows:
+        try:
+            if int(line.split(",", 1)[0]) <= keep_step:
+                kept.append(line)
+        except ValueError:
+            continue
+    log_path.write_text("\n".join([header] + kept) + "\n")
+    return len(rows) - len(kept)
+
+
+class Trainer:
+    def __init__(self, model, method, train_loader, val_loader, cfg: Dict[str, Any], run_dir: Path, device: str):
+        self.model = model
+        self.method = method
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.cfg = cfg
+        self.run_dir = Path(run_dir)
+        self.device = device
+
+        tr, m = cfg["training"], cfg["method"]
+        allow_tf32 = tr.get("allow_tf32")
+        if allow_tf32 is not None:                    # one arithmetic class on every GPU (False = strict FP32)
+            torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
+            torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
+            print(f"precision: allow_tf32={bool(allow_tf32)} (matmul + cudnn)")
+        if bool(tr.get("deterministic", False)):      # a resumed run is bit-identical to an uninterrupted one
+            os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+            torch.use_deterministic_algorithms(True)
+            print("determinism: use_deterministic_algorithms(True), cudnn.deterministic, benchmark off, "
+                  "CUBLAS_WORKSPACE_CONFIG=" + os.environ["CUBLAS_WORKSPACE_CONFIG"])
+
+        self.batch_size = int(tr["batch_size"])
+        if "lr" in tr:
+            self.base_lr = float(tr["lr"])
+        elif "lr" in m:
+            self.base_lr = float(m["lr"])
+        else:
+            raise KeyError("'lr' must be set in the training or the method section")
+        self.val_seed = int(tr.get("val_seed", 777))
+        self.restart_checkpointing = bool(tr.get("restart_checkpointing", True))
+
+        if tr.get("max_steps") is None or tr.get("val_every_updates") is None:
+            raise KeyError("training.max_steps and training.val_every_updates are required")
+        self.max_steps = int(tr["max_steps"])
+        self.val_every_updates = int(tr["val_every_updates"])
+        if self.max_steps <= 0 or self.val_every_updates <= 0:
+            raise ValueError("max_steps and val_every_updates must be positive")
+        self.val_extra_steps = tr.get("val_extra_steps")
+        if self.val_extra_steps is not None:
+            if not isinstance(self.val_extra_steps, (list, tuple)) or \
+                    any(isinstance(s, bool) or int(s) != s for s in self.val_extra_steps):
+                raise ValueError("training.val_extra_steps must be a list of integers")
+            xs = sorted({int(s) for s in self.val_extra_steps})
+            bad = [s for s in xs if s < 0 or s >= self.max_steps]
+            if bad:
+                raise ValueError(f"training.val_extra_steps must lie in [0, max_steps): {bad}")
+            self.val_extra_steps = xs
+        self.log_loss_every_epochs = tr.get("log_loss_every_epochs")
+        if self.log_loss_every_epochs is not None:
+            k = self.log_loss_every_epochs
+            if isinstance(k, bool) or int(k) != k or int(k) <= 0:
+                raise ValueError("training.log_loss_every_epochs must be a positive integer")
+            self.log_loss_every_epochs = int(k)
+
+        def _opt(key, default=None):                  # training section overrides the method section
+            return tr[key] if key in tr else m.get(key, default)
+
+        self.ema_halflife_kimg = _opt("ema_halflife_kimg")           # None: no EMA
+        self.ema_rampup_ratio = _opt("ema_rampup_ratio", 0.05)
+        self.lr_rampup_kimg = _opt("lr_rampup_kimg")                 # None: no LR ramp
+        self.max_grad_norm = _opt("max_grad_norm")                   # None: no clipping
+        self.weight_decay = _opt("weight_decay")
+        self.use_ema = self.ema_halflife_kimg is not None and float(self.ema_halflife_kimg) > 0
+        self.use_lr_rampup = self.lr_rampup_kimg is not None and float(self.lr_rampup_kimg) > 0
+
+    def fit(self, resume: Dict[str, Any] | None = None) -> None:
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+        ckpt_dir = self.run_dir / "checkpoints"
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+        model = self.model.to(self.device)
+        if self.use_ema:
+            ema_model = copy.deepcopy(model).eval().to(self.device)
+            for p in ema_model.parameters():
+                p.requires_grad_(False)
+        else:
+            ema_model = None
+        optimizer = self.method.configure_optimizer(model, {"lr": self.base_lr, "weight_decay": self.weight_decay})
+
+        start_epoch = 1
+        if resume is not None:
+            if resume.get("schema") != RESTART_SCHEMA:
+                raise AssertionError(f"restart checkpoint schema {resume.get('schema')!r} != {RESTART_SCHEMA}")
+            model.load_state_dict(resume["model_state_dict"])
+            if ema_model is not None:
+                if resume.get("ema_state_dict") is None:
+                    raise AssertionError("EMA is on but the restart checkpoint has no EMA state")
+                ema_model.load_state_dict(resume["ema_state_dict"])
+            optimizer.load_state_dict(resume["optimizer_state_dict"])
+            start_epoch = int(resume["epoch"])        # the interrupted epoch is replayed
+
+        n_train_pairs = len(self.train_loader.dataset)
+        steps_per_epoch = len(self.train_loader)
+        n_epochs = max(1, -(-self.max_steps // steps_per_epoch))
+        print(f"schedule: max_steps = {self.max_steps}, validation + checkpoint every {self.val_every_updates} updates "
+              f"(+ at max_steps); {n_train_pairs} training pairs, batch {self.batch_size}, {steps_per_epoch} steps/epoch -> "
+              f"{n_epochs} epochs (last one partial: {self.max_steps - (n_epochs - 1) * steps_per_epoch} steps)")
+        if self.use_ema:
+            seen_kimg = (self.max_steps - 1) * self.batch_size / 1000.0
+            eff_kimg = min(float(self.ema_halflife_kimg), float(self.ema_rampup_ratio) * seen_kimg)
+            print(f"EMA: half-life {self.ema_halflife_kimg} kimg, ramp-up ratio {self.ema_rampup_ratio} "
+                  f"(effective half-life at the last update ~= {eff_kimg:.1f} kimg)")
+        else:
+            print("EMA: disabled")
+        print(f"LR ramp-up: {self.lr_rampup_kimg if self.use_lr_rampup else 'disabled'}; "
+              f"gradient clipping: {self.max_grad_norm if self.max_grad_norm is not None else 'disabled'}")
+
+        log_path = self.run_dir / "train_log.csv"
+        if resume is not None:
+            if int(resume["step"]) >= self.max_steps:
+                raise AssertionError(f"resume step {int(resume['step'])} >= max_steps {self.max_steps} (nothing to continue)")
+            if not log_path.exists():
+                raise FileNotFoundError(f"--resume requires the existing {log_path}")
+            dropped = _truncate_log(log_path, int(resume["step"]))
+            print(f"resume: continuing at step {int(resume['step'])} (epoch {start_epoch}, batch "
+                  f"{int(resume['batch_index'])}; {dropped} log row(s) dropped)")
+        else:
+            log_path.write_text(LOG_COLUMNS + "\n")
+
+        cur_nimg = 0
+        if resume is not None:                        # RNG streams restored last, after every setup step
+            cur_nimg = int(resume["cur_nimg"])
+            torch.set_rng_state(resume["torch_rng_state"])
+            if str(self.device).startswith("cuda"):
+                torch.cuda.set_rng_state_all(resume["cuda_rng_state_all"])
+            np.random.set_state(resume["numpy_rng_state"])
+            random.setstate(resume["python_rng_state"])
+            gen = self.train_loader.generator
+            if resume.get("loader_generator_state_epoch_start") is not None:
+                if gen is None:
+                    raise AssertionError("restart checkpoint has a loader generator state but the train loader has no generator")
+                gen.set_state(resume["loader_generator_state_epoch_start"])
+            print(f"resume: restored optimizer, EMA, counters (cur_nimg={cur_nimg}), RNG streams and loader generator", flush=True)
+        self._fit_updates(model=model, ema_model=ema_model, optimizer=optimizer, ckpt_dir=ckpt_dir, log_path=log_path,
+                          resume=resume, start_epoch=start_epoch, cur_nimg=cur_nimg, n_val_batches=len(self.val_loader))
+
+    def _fit_updates(self, *, model, ema_model, optimizer, ckpt_dir, log_path, resume, start_epoch, cur_nimg, n_val_batches):
+        max_steps, every = self.max_steps, self.val_every_updates
+        extra_steps = set(self.val_extra_steps or [])
+        loss_every = int(self.log_loss_every_epochs) if self.log_loss_every_epochs else 0
+        devices_all = [self.device] if str(self.device).startswith("cuda") else []
+        losses_since_epoch_start: list[float] = []
+        last_event_step = -1
+        t_row = time.time()
+        if resume is not None:
+            losses_since_epoch_start = [float(x) for x in resume.get("losses_since_epoch_start", [])]
+            last_event_step = int(resume["step"])
+        step, start_batch = 0, 0
+        if resume is not None:
+            step, start_batch = int(resume["step"]), int(resume["batch_index"])
+        epoch = start_epoch
+        losses_since_event: list[float] = []
+        t_event = time.time()
+        done = False
+        gen = self.train_loader.generator
+
+        def _event(epoch_i: int, batch_i: int, gen_state_epoch_start):
+            nonlocal losses_since_event, t_event, last_event_step, t_row
+            train_loss = float(np.mean(losses_since_event)) if losses_since_event else float("nan")
+            cur_lr = optimizer.param_groups[0]["lr"]
+            model.eval()
+            g_state = gen.get_state() if gen is not None else None
+            with torch.random.fork_rng(devices=devices_all):
+                torch.manual_seed(self.val_seed)
+                val_loss = _eval_loss(model, self.val_loader, self.method, self.device)
+                train_eval = _eval_loss(model, self.train_loader, self.method, self.device, max_batches=n_val_batches)
+                if self.use_ema:
+                    val_loss_ema = _eval_loss(ema_model, self.val_loader, self.method, self.device)
+                    train_eval_ema = _eval_loss(ema_model, self.train_loader, self.method, self.device, max_batches=n_val_batches)
+                else:
+                    val_loss_ema = train_eval_ema = float("nan")
+            if g_state is not None:
+                gen.set_state(g_state)
+            elapsed = time.time() - t_event
+            t_event = t_row = time.time()
+            last_event_step = step
+            print(f"step {step} (epoch {epoch_i}, batch {batch_i}): train_loss={train_loss:.6f}  val={val_loss:.6f} "
+                  f"val_ema={val_loss_ema:.6f}  lr={cur_lr:.2e}  elapsed={elapsed:.2f}s", flush=True)
+            with open(log_path, "a") as f:
+                f.write(f"{step},{epoch_i},{batch_i},1,{train_loss},{val_loss},{val_loss_ema},"
+                        f"{train_eval},{train_eval_ema},{cur_lr},{elapsed}\n")
+            losses_since_event = []
+            ckpt = ckpt_dir / f"checkpoint_step{step:08d}.pt"
+            _save_checkpoint(ckpt, model, ema_model, self.cfg, epoch_i, step=step)
+            print(f"  saved: {ckpt.name}", flush=True)
+            if self.restart_checkpointing:
+                _save_restart(ckpt_dir / RESTART_NAME, {
+                    "schema": RESTART_SCHEMA, "epoch": epoch_i, "step": step, "batch_index": batch_i, "cur_nimg": cur_nimg,
+                    "model_state_dict": model.state_dict(),
+                    "ema_state_dict": (ema_model.state_dict() if ema_model is not None else None),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "torch_rng_state": torch.get_rng_state(),
+                    "cuda_rng_state_all": (torch.cuda.get_rng_state_all() if str(self.device).startswith("cuda") else []),
+                    "numpy_rng_state": np.random.get_state(), "python_rng_state": random.getstate(),
+                    "loader_generator_state_epoch_start": gen_state_epoch_start,
+                    "losses_since_epoch_start": list(losses_since_epoch_start),
+                    "config": self.cfg})
+            model.train()
+
+        def _loss_row(epoch_i: int, batch_i: int):
+            nonlocal losses_since_epoch_start, t_row
+            if not losses_since_epoch_start or step == last_event_step:
+                losses_since_epoch_start = []
+                return                                # nothing to log, or an event just logged this step
+            train_loss = float(np.mean(losses_since_epoch_start))
+            cur_lr = optimizer.param_groups[0]["lr"]
+            model.eval()
+            g_state = gen.get_state() if gen is not None else None
+            with torch.random.fork_rng(devices=devices_all):
+                torch.manual_seed(self.val_seed)
+                val_loss = _eval_loss(model, self.val_loader, self.method, self.device)
+                val_loss_ema = (_eval_loss(ema_model, self.val_loader, self.method, self.device) if self.use_ema else float("nan"))
+            if g_state is not None:
+                gen.set_state(g_state)
+            model.train()
+            nan = float("nan")
+            elapsed = time.time() - t_row
+            t_row = time.time()
+            with open(log_path, "a") as f:
+                f.write(f"{step},{epoch_i},{batch_i},0,{train_loss},{val_loss},{val_loss_ema},{nan},{nan},{cur_lr},{elapsed}\n")
+            losses_since_epoch_start = []
+
+        if 0 in extra_steps and resume is None and step == 0:
+            # the untrained model: its train_loss is a forward-only loss on the first batch; RNG and loader untouched
+            g_state = gen.get_state() if gen is not None else None
+            model.eval()
+            with torch.random.fork_rng(devices=devices_all):
+                torch.manual_seed(self.val_seed)
+                with torch.no_grad():
+                    b0 = tuple(t.to(self.device) for t in next(iter(self.train_loader)))
+                    loss0 = float(self.method.loss(model, b0).item())
+            model.train()
+            if g_state is not None:
+                gen.set_state(g_state)
+            losses_since_event = [loss0]
+            _event(epoch, 0, gen.get_state() if gen is not None else None)
+
+        while not done:
+            gen_state_epoch_start = gen.get_state() if gen is not None else None
+            replay = resume is not None and epoch == start_epoch and start_batch > 0
+            if replay:                                # the interrupted epoch: same permutation, consumed batches skipped
+                gen_state_epoch_start = resume["loader_generator_state_epoch_start"]
+            model.train()
+            if not replay:
+                losses_since_epoch_start = []
+            it = iter(self.train_loader)
+            skip = start_batch if (resume is not None and epoch == start_epoch) else 0
+            for _ in range(skip):
+                next(it)
+            batch_i = skip
+            for batch in it:
+                batch = tuple(t.to(self.device) for t in batch)
+                if self.use_lr_rampup:
+                    _set_rampup_lr(optimizer, self.base_lr, cur_nimg, float(self.lr_rampup_kimg))
+                optimizer.zero_grad()
+                loss = self.method.loss(model, batch)
+                loss.backward()
+                if self.max_grad_norm is not None:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), float(self.max_grad_norm))
+                optimizer.step()
+                if self.use_ema:
+                    beta = _ema_beta(cur_nimg, self.batch_size, float(self.ema_halflife_kimg), self.ema_rampup_ratio)
+                    for p_ema, p_net in zip(ema_model.parameters(), model.parameters()):
+                        p_ema.copy_(p_net.detach().lerp(p_ema, beta))
+                    for b_ema, b_net in zip(ema_model.buffers(), model.buffers()):
+                        b_ema.copy_(b_net)
+                cur_nimg += batch[0].shape[0]
+                _li = loss.item()
+                losses_since_event.append(_li)
+                losses_since_epoch_start.append(_li)
+                step += 1
+                batch_i += 1
+                if step % every == 0 or step == max_steps or step in extra_steps:
+                    _event(epoch, batch_i, gen_state_epoch_start)
+                if step >= max_steps:
+                    done = True
+                    break
+            if not done and loss_every and epoch % loss_every == 0:
+                _loss_row(epoch, batch_i)
+            epoch += 1
+
+        print(f"trained exactly {max_steps} updates", flush=True)
+        _save_checkpoint(ckpt_dir / "checkpoint.pt", model, ema_model, self.cfg, epoch - 1, step=max_steps)

code/bifurcation/utils.py

@@ -0,0 +1,27 @@
+"""Seeding and YAML helpers."""
+from __future__ import annotations
+
+import os
+import random
+from typing import Any, Dict
+
+import numpy as np
+import torch
+import yaml
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def load_yaml(path: str | os.PathLike) -> Dict[str, Any]:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def save_resolved_config(cfg: Dict[str, Any], path: str | os.PathLike) -> None:
+    with open(path, "w") as f:
+        yaml.safe_dump(cfg, f, sort_keys=False)

code/configs/euler/dataset.yaml

@@ -0,0 +1,21 @@
+# 2-D Euler (incompressible, periodic): the package `euler_v2_uniform`. Set EULER_DATA (the package directory:
+# trajectories_fine.npz, split.json, normalization.npz) and EULER_TRUTH_BANK (group_<gid>.npz of the ten test
+# conditions, 10,000 realizations each).
+name: euler_v2_uniform
+system: euler
+data: ${EULER_DATA}
+truth_bank: ${EULER_TRUTH_BANK}
+expected_sha256: {split.json: 0040ecf525ba7c6b5cfe7d874812e04c69ccc4396ddbdc2db09e5489d8280b6e, normalization.npz: 06573da0f073b5bfabf884ac497f642c5d717824afd5644272239030dd7d4a7f,
+  trajectories_fine.npz: 664c3be66eee8b63dbad1fb379d72f6c7dd9a1094e3125c1e8c4223a4e2e7e66}
+state:
+  dim: 196
+  observed: coarse spectral vorticity, |k| <= 8, Hermitian half (98 complex modes), omega_hat / N_sim^2
+  hidden: the unresolved vorticity 8 < |k| <= 64 of the initial condition
+time_grid: {n_snapshots: 41, dt: 0.025, T: 1.0}
+package: {vecs_train: [80, 200, 41, 196], vecs_eval: [20, 100, 41, 196], vecs_new: [720, 20, 41, 196]}
+splits:
+  train: 80 conditions x 200 realizations (vecs_train); vecs_new holds 720 further conditions x 20 realizations, unused by the paper's models
+  val: 10 conditions x 100 realizations
+  test: 10 conditions x 100 realizations in the package; 10,000 realizations each in the truth bank
+normalization: per-coordinate mean and std of the 80 training conditions x realizations j < 20 x all 41 snapshots
+fid_adapter: {name: euler_coarse_vorticity, params: {n_disp: 128, k_c: 8}}

code/configs/evaluation.yaml

@@ -0,0 +1,30 @@
+# Evaluation settings shared by the five systems: csFID-10k (the paper's image metric) and the sampling of the test read.
+
+fid:
+  n_truth: 10000                   # truth realizations per test condition (the first 10,000 members of the truth bank)
+  n_generated: 10000               # generated samples per model and test condition
+  backend: clean-fid
+  clean_fid_version: "0.1.35"      # checked at run time
+  mode: clean                      # clean-fid "clean" mode: one PIL bicubic resize to 299 x 299
+  inception:
+    file: inception-2015-12-05.pt  # the TF Inception-v3 weights distributed with clean-fid (sha256 checked)
+    sha256: f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4
+    feature_layer: pool3           # 2048-D, used by cFID / cKID
+  resize: {size: 299, filter: bicubic, backend: PIL, antialias: true, once: true}
+  kid: {num_subsets: 100, max_subset_size: 1000}
+  clip_quantile: 0.9999            # image range: c = Q_0.9999(|x|) over the training fields (records/<system>/fid_adapter.json)
+  feature_batch: 200
+  primary_metric: csfid            # csFID: the Frechet distance of the spatial Inception feature below
+  secondary_metrics: [cfid, ckid]
+  spatial:
+    layer: mixed_6/conv            # the first branch of the mixed_6 block, before concatenation (192 x 17 x 17)
+    channels: 7                    # the first 7 channels
+    grid: 17
+    dim: 2023                      # 17 * 17 * 7, flattened NHWC, no spatial pooling
+    flatten: NHWC
+    pooling: none
+    reference: "Nash et al., ICML 2021; openai/guided-diffusion evaluations/evaluator.py FID_SPATIAL_NAME"
+
+sampling:
+  crn_base_seed: 777               # the noise of test condition gid is drawn from the stream seeded with 777 * 100003 + gid
+  sample_chunk: 2000               # the 10,000 samples of a condition are drawn in chunks of 2,000 from that stream

code/configs/nlkg/dataset.yaml

@@ -0,0 +1,23 @@
+# Nonlinear Klein-Gordon (2-D, cubic, defocusing; Dedalus): the package `nlkg_v3_candidate`, 392-D coarse state of
+# (u, v/4). Set NLKG_DATA and NLKG_TRUTH_BANK. The ten test conditions are only in the truth bank (their x_0 is read
+# from it).
+name: nlkg_v3_candidate
+system: nlkg
+data: ${NLKG_DATA}
+truth_bank: ${NLKG_TRUTH_BANK}
+expected_sha256: {trajectories_fine.npz: 258cdda0a2707a494d2f4bb34ae96b2074ebbe70c0960e7f586702d771ac8a22, split.json: 95c8d315d18dd94a8cc32c2d810e00d69e374a338a315ce2dab9858d77566de3,
+  normalization.npz: 94e78f3fc04a40b3f71c883b97cf492c3763853c8432736112ec0167e30543ca}
+state:
+  dim: 392
+  observed: the coarse operator (|k| <= 8, Hermitian half, [Re, Im]) applied to u and to v/4; the image metric reads u
+  hidden: equipartition modes 8 < |k| <= 24 of the initial displacement (rms 0.20)
+physics: {equation: 'u_t = v, v_t = c^2 lap(u) - m^2 u - lambda u^3', c: 1.0, m: 1.0, lambda: 64.0, L: 6.283185307179586,
+  n: 128, dt: 0.0025, hidden_rms: 0.2, v_scale: 4.0, T: 12.0}
+time_grid: {n_snapshots: 41, dt: 0.3, T: 12.0}
+package: {vecs_train: [80, 200, 41, 392], vecs_eval: [20, 200, 41, 392]}
+splits:
+  train: conditions 0-79 x 200 realizations
+  val: conditions 80-89 x 200 realizations
+  test: conditions 100-109, only in the truth bank (10,000 realizations each)
+normalization: per-coordinate mean and std of the 80 training conditions x realizations j < 20 x all 41 snapshots
+fid_adapter: {name: nlkg_coarse_u, params: {n_disp: 128, k_c: 8}}

code/configs/passive_scalar/dataset.yaml

@@ -0,0 +1,21 @@
+# Passive scalar (advection-diffusion in a fixed cellular flow): the package `passive_scalar_v1`. Set PASSIVE_SCALAR_DATA
+# and PASSIVE_SCALAR_TRUTH_BANK.
+name: passive_scalar_v1
+system: passive_scalar
+data: ${PASSIVE_SCALAR_DATA}
+truth_bank: ${PASSIVE_SCALAR_TRUTH_BANK}
+expected_sha256: {split.json: fb8ca7914199c2ae19ce61c0a6baed8736c98f2db77e19f5fa128003a2979081, normalization.npz: 7f79196a94bba0509e1a307df39a4f572a1f8100dbd5ba9aad82f53aae02095e,
+  trajectories_fine.npz: 7a810b232414307ade1cebc0ecff4961d381561e400e034fc356cbf12f764a52}
+state:
+  dim: 196
+  observed: coarse spectral scalar, |k| <= 8, Hermitian half (98 complex modes), theta_hat / N^2
+  hidden: the unresolved part of the initial scalar field, advected by the fixed cellular flow
+physics: {K: 2, V: 0.5, kappa: 6.18415427504503e-06, N: 512, m_substeps: 6}
+time_grid: {n_snapshots: 41, dt: 0.025, T: 1.0}
+package: {vecs_train: [80, 400, 41, 196], vecs_eval: [20, 4096, 41, 196], vecs_new: [1520, 20, 41, 196]}
+splits:
+  train: 80 conditions x 400 realizations (the paper's models read the first 200); vecs_new holds 1,520 further conditions x 20 realizations, unused
+  val: 10 conditions x 4,096 realizations
+  test: 10 conditions (gid 0 is one of them) x 4,096 realizations in the package; 10,000 each in the truth bank
+normalization: per-coordinate mean and std of the 80 training conditions x realizations j < 20 x all 41 snapshots
+fid_adapter: {name: passive_scalar_coarse, params: {n_disp: 128, k_c: 8}}

code/configs/sqg/dataset.yaml

@@ -0,0 +1,21 @@
+# Surface quasi-geostrophy (pyqg): the package `sqg_v3_candidate_v2`. Set SQG_DATA and SQG_TRUTH_BANK.
+#
+name: sqg_v3_candidate_v2
+system: sqg
+data: ${SQG_DATA}
+truth_bank: ${SQG_TRUTH_BANK}
+expected_sha256: {trajectories_fine.npz: 510b027125b30065489edbae1d0ac23c3cb15765630be4113c119af993c0dee4, split.json: fb482ffd0353ba47dfca67a34466f7026278962ca9c8fb44839aa039935d08e2,
+  normalization.npz: 46f94fc794069c9ad985d23ae4cc2b8ed866fe13d6a0a5463f1c8384352898f8}
+state:
+  dim: 196
+  observed: coarse spectral surface buoyancy b, |k| <= 8, Hermitian half (98 complex modes), b_hat / N^2
+  hidden: a Gaussian buoyancy field on 8 < |k| <= 32 with |b_k|^2 ~ k^-5/3 and rms 0.1, drawn per realization
+physics: {solver: pyqg.SQGModel, nx: 256, dt: 0.005, L: 2*pi, H: 1.0, Nb: 1.0, f_0: 1.0, beta: 0.0, rek: 0.0, T_phys: 20.0}
+time_grid: {n_snapshots: 41, dt: 0.025, T: 1.0, note: benchmark time 1.0 = 20 SQG time units}
+package: {vecs_train: [80, 200, 41, 196], vecs_eval: [20, 200, 41, 196]}
+splits:
+  train: 80 conditions x 200 realizations
+  val: 10 conditions x 200 realizations
+  test: conditions 100-109, 200 realizations in the package; 10,000 each in the truth bank
+normalization: per-coordinate mean and std of the 80 training conditions x realizations j < 20 x all 41 snapshots
+fid_adapter: {name: passive_scalar_coarse, params: {n_disp: 128, k_c: 8}}

code/configs/waves/dataset.yaml

@@ -0,0 +1,23 @@
+# Heterogeneous waves: the package `heterogeneous_waves_v3_candidate` (augmented state (u, v), 392-D). Set WAVES_DATA
+# and WAVES_TRUTH_BANK. The ten test conditions are only in the truth bank (their x_0 is read from it).
+#
+name: heterogeneous_waves_v3_candidate
+system: heterogeneous_waves
+data: ${WAVES_DATA}
+truth_bank: ${WAVES_TRUTH_BANK}
+expected_sha256: {trajectories_fine.npz: 4211aa62b2e6e61066519396e20bdf0a3b002b02a27641dd39031cbb680c9e53, split.json: ccb5fb3fba13a8f0e5528c3757fb54ba4044ca180328da99a53b58fbf6b8e88c,
+  normalization.npz: 23123f0b0547197bbd6275ac4a34d894a91ab98e4f1810b80416f6b891d67cb7, u_norm/normalization.npz: 55e9a00ddcf36eece80a88bf1f012026f4934989ae8a2969b1cdcb5c78228d2a,
+  v_norm/normalization.npz: 54441e3917067932261ac67dacb3b15dd76458147d9ff41775d1148b20b06ddf}
+state:
+  dim: 392
+  observed: coarse displacement u (196-D, |k| <= 8) and the finite-difference coarse velocity v (196-D); the image metric reads u
+  hidden: the velocity and the band 8 < |k| <= 16 of the initial condition
+physics: {equation: 'u_t = v, v_t = div(c^2 grad u) - gamma v', c2: '1 + 0.5 sin(2 pi x) sin(2 pi y)'}
+time_grid: {n_snapshots: 41, dt: 0.025, T: 1.0}
+package: {vecs_train: [80, 200, 41, 392], vecs_eval: [10, 200, 41, 392]}
+splits:
+  train: 80 conditions x 200 realizations
+  val: 10 conditions x 200 realizations
+  test: 10 conditions, only in the truth bank (10,000 realizations each)
+normalization: per-coordinate mean and std of the 80 training conditions x realizations j < 20 x all 41 snapshots, fitted separately for u (u_norm/) and v (v_norm/); normalization.npz is their concatenation
+fid_adapter: {name: heterogeneous_waves_augmented_coarse, params: {}}