diff --git a/code/.gitignore b/code/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..8140fd9db427f07628920a752d515d9676849c17 --- /dev/null +++ b/code/.gitignore @@ -0,0 +1,7 @@ +__pycache__/ +*.pyc +.pytest_cache/ +.venv/ +*.egg-info/ +/runs/ +/outputs/ diff --git a/code/LICENSE b/code/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..1843191013c30bf7e17d8f30fbd1111496b0d019 --- /dev/null +++ b/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. diff --git a/code/README.md b/code/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cd016def843bd0355dbc8dd7bf129e35bf5bc7c8 --- /dev/null +++ b/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). diff --git a/code/bifurcation/__init__.py b/code/bifurcation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4f53f7354604580fd6c8c34178c46e7116e7eb48 --- /dev/null +++ b/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" diff --git a/code/bifurcation/config.py b/code/bifurcation/config.py new file mode 100644 index 0000000000000000000000000000000000000000..d89235673594872e5276bed109c1c9158d253dad --- /dev/null +++ b/code/bifurcation/config.py @@ -0,0 +1,100 @@ +"""Configuration files of the repository: the dataset cards, the evaluation settings and the recorded runs. + +* ``configs//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//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 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 diff --git a/code/bifurcation/data/__init__.py b/code/bifurcation/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/bifurcation/data/fine.py b/code/bifurcation/data/fine.py new file mode 100644 index 0000000000000000000000000000000000000000..7903b7dee14f11a780fa989a3aea3399cfd088bc --- /dev/null +++ b/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_.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_.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 diff --git a/code/bifurcation/data/loaders.py b/code/bifurcation/data/loaders.py new file mode 100644 index 0000000000000000000000000000000000000000..830c26e09140a17d73e4b39da967315e9257fe4e --- /dev/null +++ b/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) diff --git a/code/bifurcation/data/manifests.py b/code/bifurcation/data/manifests.py new file mode 100644 index 0000000000000000000000000000000000000000..c09f9796f0fc70e47a70fa13597e27def02be934 --- /dev/null +++ b/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_.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 diff --git a/code/bifurcation/fid/__init__.py b/code/bifurcation/fid/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e6b0d9d5a859b21ebd1cd581dcb00e4f73eadb20 --- /dev/null +++ b/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//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) diff --git a/code/bifurcation/fid/adapters.py b/code/bifurcation/fid/adapters.py new file mode 100644 index 0000000000000000000000000000000000000000..5c0baefac47572dfb6406d0ac5743d95e8b0bd70 --- /dev/null +++ b/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//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//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} diff --git a/code/bifurcation/fid/backend.py b/code/bifurcation/fid/backend.py new file mode 100644 index 0000000000000000000000000000000000000000..4712f761d34979ed198fddc49ee347dd3a6ac204 --- /dev/null +++ b/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 diff --git a/code/bifurcation/fid/euler.py b/code/bifurcation/fid/euler.py new file mode 100644 index 0000000000000000000000000000000000000000..39ef9145748ab9a85ed39833dfee91c57d80ec15 --- /dev/null +++ b/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")} diff --git a/code/bifurcation/fid/evaluate.py b/code/bifurcation/fid/evaluate.py new file mode 100644 index 0000000000000000000000000000000000000000..a48a0d475d1cefb622b9a4953355e956477e92af --- /dev/null +++ b/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 ``/group_.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//group_.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.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"}) diff --git a/code/bifurcation/fid/heterogeneous_waves.py b/code/bifurcation/fid/heterogeneous_waves.py new file mode 100644 index 0000000000000000000000000000000000000000..65ff28050ca2bfc8c712d2e20ac9b76eceae9633 --- /dev/null +++ b/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")} diff --git a/code/bifurcation/fid/heterogeneous_waves_augmented.py b/code/bifurcation/fid/heterogeneous_waves_augmented.py new file mode 100644 index 0000000000000000000000000000000000000000..f1e9001fde29fc5cdaf916f3e77beee803957118 --- /dev/null +++ b/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 diff --git a/code/bifurcation/fid/nlkg.py b/code/bifurcation/fid/nlkg.py new file mode 100644 index 0000000000000000000000000000000000000000..dacbe34e3300be7633b9f99423c9781b6633b547 --- /dev/null +++ b/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 diff --git a/code/bifurcation/fid/passive_scalar.py b/code/bifurcation/fid/passive_scalar.py new file mode 100644 index 0000000000000000000000000000000000000000..60466653640251e780bc78971416e1b7e4b05401 --- /dev/null +++ b/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")} diff --git a/code/bifurcation/methods/__init__.py b/code/bifurcation/methods/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c64ed4ea77056095cde852d8aabe9f7b2b31b5e --- /dev/null +++ b/code/bifurcation/methods/__init__.py @@ -0,0 +1,3 @@ +from .base import Method + +__all__ = ["Method"] diff --git a/code/bifurcation/methods/base.py b/code/bifurcation/methods/base.py new file mode 100644 index 0000000000000000000000000000000000000000..d1e62af4df6ef6b942d72aa366f39d19d562c388 --- /dev/null +++ b/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}") diff --git a/code/bifurcation/methods/cascaded_si/__init__.py b/code/bifurcation/methods/cascaded_si/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6e821cca0a536d7c9918bb6f368427a9009380 --- /dev/null +++ b/code/bifurcation/methods/cascaded_si/__init__.py @@ -0,0 +1,3 @@ +from .method import CascadedSIMethod + +__all__ = ["CascadedSIMethod"] diff --git a/code/bifurcation/methods/cascaded_si/core.py b/code/bifurcation/methods/cascaded_si/core.py new file mode 100644 index 0000000000000000000000000000000000000000..283fd273fe0eba2a2e91c7ebf5abef9c407b8eb6 --- /dev/null +++ b/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 diff --git a/code/bifurcation/methods/cascaded_si/method.py b/code/bifurcation/methods/cascaded_si/method.py new file mode 100644 index 0000000000000000000000000000000000000000..d0846ab01645adcf1106ad29e79a8b286f8196fb --- /dev/null +++ b/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") diff --git a/code/bifurcation/methods/direct_si/__init__.py b/code/bifurcation/methods/direct_si/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0043ceea60c35b0807bd8e5cc7f8e8bdd126d736 --- /dev/null +++ b/code/bifurcation/methods/direct_si/__init__.py @@ -0,0 +1,3 @@ +from .method import DirectSIMethod + +__all__ = ["DirectSIMethod"] diff --git a/code/bifurcation/methods/direct_si/core.py b/code/bifurcation/methods/direct_si/core.py new file mode 100644 index 0000000000000000000000000000000000000000..741ec4890fb2d4595ad9cfe403b26ecf5e549070 --- /dev/null +++ b/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 diff --git a/code/bifurcation/methods/direct_si/method.py b/code/bifurcation/methods/direct_si/method.py new file mode 100644 index 0000000000000000000000000000000000000000..a9d8fa64cc1ddc64524be605d26a4e44650e8207 --- /dev/null +++ b/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") diff --git a/code/bifurcation/methods/edm/__init__.py b/code/bifurcation/methods/edm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..133845af935e9f9b667db8d1fe43e996caaea248 --- /dev/null +++ b/code/bifurcation/methods/edm/__init__.py @@ -0,0 +1,3 @@ +from .method import EDMMethod + +__all__ = ["EDMMethod"] diff --git a/code/bifurcation/methods/edm/core.py b/code/bifurcation/methods/edm/core.py new file mode 100644 index 0000000000000000000000000000000000000000..20ff724176d30f34a54280c70481707d551af773 --- /dev/null +++ b/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 diff --git a/code/bifurcation/methods/edm/method.py b/code/bifurcation/methods/edm/method.py new file mode 100644 index 0000000000000000000000000000000000000000..1ffe07b64687584221e216cdc401144dbc891d2a --- /dev/null +++ b/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) diff --git a/code/bifurcation/metrics.py b/code/bifurcation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..23d1720b4b63a7dd204fd3949f414334c9853654 --- /dev/null +++ b/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") diff --git a/code/bifurcation/models/__init__.py b/code/bifurcation/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/code/bifurcation/models/_common.py b/code/bifurcation/models/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..3d12ce91a222dad65d38989c065c4b52ab8c191f --- /dev/null +++ b/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) diff --git a/code/bifurcation/models/spectral.py b/code/bifurcation/models/spectral.py new file mode 100644 index 0000000000000000000000000000000000000000..53e33bdc31003b01387961e362cb102e173af1ad --- /dev/null +++ b/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)) diff --git a/code/bifurcation/models/unet2.py b/code/bifurcation/models/unet2.py new file mode 100644 index 0000000000000000000000000000000000000000..78cbd971f61f34c2b6d2199db5a2b7634aaf3fff --- /dev/null +++ b/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))) diff --git a/code/bifurcation/models/unet2_2comp.py b/code/bifurcation/models/unet2_2comp.py new file mode 100644 index 0000000000000000000000000000000000000000..9e8f62155b50819d27fe62587d205b49dcb95468 --- /dev/null +++ b/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") diff --git a/code/bifurcation/models/unet2_5ch.py b/code/bifurcation/models/unet2_5ch.py new file mode 100644 index 0000000000000000000000000000000000000000..7e504378e8c0a6a9ef761a88d0c22045f4ff60f2 --- /dev/null +++ b/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.") diff --git a/code/bifurcation/models/unet2_5ch_coord.py b/code/bifurcation/models/unet2_5ch_coord.py new file mode 100644 index 0000000000000000000000000000000000000000..da8f0f43f4a51ec8df89ac20477ddaf20454b14e --- /dev/null +++ b/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).") diff --git a/code/bifurcation/models/unet2_5ch_coord_aug.py b/code/bifurcation/models/unet2_5ch_coord_aug.py new file mode 100644 index 0000000000000000000000000000000000000000..c948a8694aa23165bce71dbfa9bd153ce8fd6322 --- /dev/null +++ b/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``, + + /u_norm/normalization.npz (196-D, displacement) + /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 /u_norm and /v_norm; " + "otherwise identical trunk and parameter count to unet2_5ch_coord.") diff --git a/code/bifurcation/models/unet2_coord.py b/code/bifurcation/models/unet2_coord.py new file mode 100644 index 0000000000000000000000000000000000000000..f228a45596d0e597de32e9611ffddf40a707d786 --- /dev/null +++ b/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)") diff --git a/code/bifurcation/registry.py b/code/bifurcation/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..50090db979a0482ea68ebfecc8f7fa34df008afd --- /dev/null +++ b/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//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 diff --git a/code/bifurcation/sample.py b/code/bifurcation/sample.py new file mode 100644 index 0000000000000000000000000000000000000000..92749079328c5cf462be2ef73fff0a77c453c5e5 --- /dev/null +++ b/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//runs//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//group_.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()) diff --git a/code/bifurcation/train.py b/code/bifurcation/train.py new file mode 100644 index 0000000000000000000000000000000000000000..3b04bd66bc5ab44104720a71a668a98602716410 --- /dev/null +++ b/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//runs//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()) diff --git a/code/bifurcation/trainer.py b/code/bifurcation/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..8513cf468febe150bf561416c9fa7e4b13de8299 --- /dev/null +++ b/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.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) diff --git a/code/bifurcation/utils.py b/code/bifurcation/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cca566f536d29a03c29d5f9a365e903fd45ff42a --- /dev/null +++ b/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) diff --git a/code/configs/euler/dataset.yaml b/code/configs/euler/dataset.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7764d2ee3608d932b896f44d26b29e15f51b0d04 --- /dev/null +++ b/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_.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}} diff --git a/code/configs/evaluation.yaml b/code/configs/evaluation.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5c200a85eaff4f68df948af50bf9049e3a89cbeb --- /dev/null +++ b/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//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 diff --git a/code/configs/nlkg/dataset.yaml b/code/configs/nlkg/dataset.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c8cedf0269f3b07d7a4718a944fcddc32751e0c8 --- /dev/null +++ b/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}} diff --git a/code/configs/passive_scalar/dataset.yaml b/code/configs/passive_scalar/dataset.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eadecedced60f22ea1e2de5e5051833418ae28be --- /dev/null +++ b/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}} diff --git a/code/configs/sqg/dataset.yaml b/code/configs/sqg/dataset.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c72315a19b7d47448abd62bf3e6f6ec8f289491a --- /dev/null +++ b/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}} diff --git a/code/configs/waves/dataset.yaml b/code/configs/waves/dataset.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2c56803d4d21ae73907cc65b545b491dfcf85b1b --- /dev/null +++ b/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: {}} diff --git a/code/pyproject.toml b/code/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..3120a31a43e7365d72bb6a7288e3c60bdbcfede2 --- /dev/null +++ b/code/pyproject.toml @@ -0,0 +1,42 @@ +[project] +name = "coarse-grained-forecasting" +version = "1.0.0" +description = "Direct SI, Cascaded SI and EDM for conditional forecasting from coarse-grained observations: the code of the paper's models, training, sampling and csFID evaluation." +readme = "README.md" +license = { file = "LICENSE" } +requires-python = ">=3.10,<3.11" +dependencies = [ + "torch==2.1.2", + "torchvision==0.16.2", + "numpy==1.26.2", + "scipy==1.15.3", + "pyyaml==6.0.1", + "pandas==2.2.1", + "matplotlib==3.8.2", + "scikit-learn==1.3.2", + "h5py==3.15.1", + "einops==0.8.1", + "Pillow==10.0.1", + "clean-fid==0.1.35", + "requests", +] + +[tool.uv] +# PyTorch 2.1.2 built for CUDA 11.8 (the environment of the recorded Euler, NLKG and passive-scalar values); +# on a machine without CUDA the same wheels run on the CPU. +[[tool.uv.index]] +name = "pytorch-cu118" +url = "https://download.pytorch.org/whl/cu118" +explicit = true + +[tool.uv.sources] +torch = { index = "pytorch-cu118" } +torchvision = { index = "pytorch-cu118" } + +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["bifurcation*", "reproduce*"] + diff --git a/code/records/depth_sweep/depth_scores.csv b/code/records/depth_sweep/depth_scores.csv new file mode 100644 index 0000000000000000000000000000000000000000..afabb1d8857a7fef9898993e76759370912b521e --- /dev/null +++ b/code/records/depth_sweep/depth_scores.csv @@ -0,0 +1,36 @@ +dataset,method,S,is_selected,mean,sem,sd,n,n_gen_nonfinite +euler,Direct SI,1,False,57.26483068072441,3.349242766217627,10.591235578090549,10,0 +euler,EDM,1,False,56.32467439472261,3.7438824883413795,11.839195955177548,10,0 +euler,Cascaded SI,10,True,16.32376118294813,1.140960328430044,3.608033357732705,10,0 +euler,Cascaded SI,2,False,53.59809032958364,3.2837081054094,10.38399678424998,10,0 +euler,Cascaded SI,40,False,23.290646768253133,1.4956564897631588,4.729681104863894,10,0 +euler,Cascaded SI,5,False,22.944786616995856,2.554606140179877,8.078373927619797,10,0 +euler,Cascaded SI,20,False,17.861238826203806,0.8661358704549127,2.7389621138100644,10,0 +passive_scalar,Direct SI,1,True,8.671174542198553,0.27900852753571304,0.8823024336226594,10,0 +passive_scalar,EDM,1,False,7.881187747848367,0.3104017524493821,0.9815765274477966,10,0 +passive_scalar,Cascaded SI,2,False,18.753759972929185,1.1744716754985969,3.714005541929739,10,0 +passive_scalar,Cascaded SI,40,False,75.20627377548983,4.738450881551495,14.984297366535456,10,0 +passive_scalar,Cascaded SI,5,False,32.973458058652966,1.6136256279837646,5.102732275248231,10,0 +passive_scalar,Cascaded SI,10,False,41.14088092045725,2.3611758048704354,7.46669349947187,10,0 +passive_scalar,Cascaded SI,20,False,56.945176663275745,2.791231902897987,8.82665029088358,10,0 +waves,Direct SI,1,False,115.86330141284854,6.763827859735928,21.389101738067428,10,0 +waves,Cascaded SI,2,False,71.23029867133164,3.911905547976706,12.370531523055481,10,0 +waves,Cascaded SI,5,True,36.66245678459255,1.7764049801888095,5.617485784262924,10,0 +waves,Cascaded SI,10,False,44.513957438707095,2.335970099052613,7.386986059055395,10,0 +waves,Cascaded SI,20,False,58.27589480277599,2.5799720893000146,8.158588101851374,10,0 +waves,Cascaded SI,40,False,62.81217572244084,3.248342568310332,10.272161136341742,10,0 +waves,EDM,1,False,123.04089328034769,7.067465996801545,22.349289835685173,10,0 +sqg,Direct SI,1,False,138.47532203507993,23.08644501940587,73.00574933757272,10,0 +sqg,EDM,1,False,118.62393662746766,21.672506523637665,68.53448321955285,10,0 +sqg,Cascaded SI,10,True,51.24497917459168,8.31809762870687,26.304134306359305,10,0 +sqg,Cascaded SI,2,False,83.2523176072445,16.314034007096794,51.58950528786942,10,0 +sqg,Cascaded SI,5,False,43.56877605319788,8.790125902695719,27.79681877216208,10,0 +sqg,Cascaded SI,20,False,89.08168688252472,7.557464374313741,23.898800758410744,10,1 +sqg,Cascaded SI,40,False,86.97838872337205,11.997151993142262,37.93832573355832,10,0 +nlkg,Direct SI,1,False,98.4767962962498,6.1787087828142875,19.53879275257978,10,0 +nlkg,EDM,1,False,92.36841011398715,5.428927787320806,17.167777060511934,10,0 +nlkg,Cascaded SI,20,True,16.24257143793273,1.5538946704054246,4.913846402477781,10,0 +nlkg,Cascaded SI,2,False,99.75472041704076,5.894287147084769,18.639372567843775,10,0 +nlkg,Cascaded SI,40,False,18.227291851587534,2.0481309411127504,6.47675872018049,10,0 +nlkg,Cascaded SI,5,False,53.09719674051487,7.301625987850265,23.08976874428377,10,0 +nlkg,Cascaded SI,10,False,22.504916177811378,5.395604960058663,17.062401028287212,10,0 diff --git a/code/records/euler/assets.sha256 b/code/records/euler/assets.sha256 new file mode 100644 index 0000000000000000000000000000000000000000..4a3bcfa053339556a6b1a65b41873f59d89b4ec5 --- /dev/null +++ b/code/records/euler/assets.sha256 @@ -0,0 +1,59 @@ +# sha256 of the inputs of the Table 2 entries; paths are relative to the directory named in the heading +# data: the dataset package (--data) +0040ecf525ba7c6b5cfe7d874812e04c69ccc4396ddbdc2db09e5489d8280b6e split.json +06573da0f073b5bfabf884ac497f642c5d717824afd5644272239030dd7d4a7f normalization.npz +664c3be66eee8b63dbad1fb379d72f6c7dd9a1094e3125c1e8c4223a4e2e7e66 trajectories_fine.npz +# truth_bank: the truth banks of the test conditions (--truth-bank) +edd0d44aa5af0385b87079ed73ea48c7644b87e89a90e4eba88d4a126680efc1 group_0005.npz +eaf5e92d04b1e4d26437679b25ece9a7312688d794ef299d4b9051a4478835cd group_0016.npz +431edb5d45e4e5060c199d192b4d0cb7814ee0c347b3ffed6b648ec0e1ca7e97 group_0020.npz +6e3970df538b8b859839276c7e3b2d5c865b61587c525ee6f1d8c84182318abe group_0036.npz +5554c43f9e77a3d99dcaf40c3934991a0f7d04222d8ae20753f694e938f61aaf group_0052.npz +86185b8f497482ebd47342cb6287923205bd8f3cb115fccb1437c3a96b19b51a group_0072.npz +d06812553b8823345dd903c2f57c5f619bbef4e526b2902230ba96486c8a1777 group_0082.npz +d6b8e8ddf67ddf319bf2b86728249e4708dae86d77c3e8416b37a71d05ea7ebe group_0090.npz +a336510e86fa648fcc2a933a05ca0a04a078e56fe83ab871737e41847ee30768 group_0093.npz +154b1c2433cfaa8471d9c1b1aa36b238ab0243cfd03a5114966da687e1a6aba9 group_0094.npz +# inception: the Inception-v3 weights in $CLEAN_FID_CACHE +f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4 inception-2015-12-05.pt +# samples/direct_si: the archived test samples (--samples) +50b2121c12712abbce4e362dd04876a3424f31e70b052879b7c8667c7caf906a noise_scale_1p0/group_0005.npz +744a563d4b992b6909ec6f998a99f4225998a90701fbece23d2e3d47cb370e47 noise_scale_1p0/group_0016.npz +4b90bee0507c762ed25542d9ec4931f7733db5fa0829d6e28522a715f7f96343 noise_scale_1p0/group_0020.npz +7022292c22adea25b57949daf001871f1ce623f004b9470912cf58a5ffc64ca7 noise_scale_1p0/group_0036.npz +dd556193e23362d646d53d1bbf818c485f3abe03a5cb617322fcb374b5b28ef2 noise_scale_1p0/group_0052.npz +7540d7ddc6590451ac201ee0d451c4a6a0f00afa3bf4e9d08e2f023d5c45c856 noise_scale_1p0/group_0072.npz +11aa9b80b21fe340d5fb2771894795e9e74ac42a61bcfd99bac9197c3d51ff03 noise_scale_1p0/group_0082.npz +e7e6754b5b7cbcdc658f2ac6997010fd9409c84769bbb9d241a70e5efa589da5 noise_scale_1p0/group_0090.npz +7df03b8776d8491569283dd938ce504d75f223152ad5887558f49e90483fc8ea noise_scale_1p0/group_0093.npz +df4dc3407b18a219b8b5e14af98ba7c3e00b3cfd2d55d0782d772c854bcbd8dd noise_scale_1p0/group_0094.npz +# samples/edm: the archived test samples (--samples) +2af7b7bb29c2fa0db253d384e751e01960e33947c8948854f3b3333717c37127 default/group_0005.npz +2825db6e19c7a68b265d76a9171c442188cd7620cb2047d43724f45653a2819c default/group_0016.npz +ce12e84d2ec6f314f3bad060e47fe821bec6c69539b158d862cb6896977a65b1 default/group_0020.npz +093adeb9bedff5fe6da840df215f381a7160c1a97a087f139cd0ecadfb91914d default/group_0036.npz +50e806f0cb281dfeb07eee5520d04f91cd284e1a0ee9a823a9df6a128061a80a default/group_0052.npz +edfcb254804cd3de9fe552fa5b1c1c58a791b9c684310b258c07844ec4f55c66 default/group_0072.npz +7b6851fcd70f7e5d3bf0d80b70585fa3dd9b85a49527e83c71896903f02752cb default/group_0082.npz +706e4201b78b7ef75e8833cc6f6f8c77edde489a23be629a995d5e79696de4ae default/group_0090.npz +cba651fa9cd75dacd13e831acfb9ad66b057de74d5098b56f898873d7a6ea622 default/group_0093.npz +6d1f3a79250d39e77b39ff4c03ed54723363246b1dc134ca6cc5632adf51d8ba default/group_0094.npz +# samples/cascaded_si: the archived test samples (--samples) +26af6f213e665d83f7568a09bc8db835da0030bd16c3fa55cf6e3bd78ce61bcb noise_scale_1p0/group_0005.npz +326799e4dbf4520d047138c4480ce3c837b90c1d7e13823c77ecb63a7796d166 noise_scale_1p0/group_0016.npz +ceb22a9ac9e9856398d95bf0deee54ae9da81b20a92c0e546fca43488bbb9350 noise_scale_1p0/group_0020.npz +8a8fad5be136e8867d7c58b25fdea46c52cd3b052c9d03862d96c3e125d83ad2 noise_scale_1p0/group_0036.npz +4183f5eb915354358926ff7972a2a5fc5e1fca620dcec5e35c6c6621e33da2d8 noise_scale_1p0/group_0052.npz +3e866e546acf19e67d6213a1dcb0669d48d2c3482c984cf895df3c5d83fe7eec noise_scale_1p0/group_0072.npz +d69f9a1b070a2a3f39cad1bc0aac7508a1b255a8615b75f907d3532eeea99de4 noise_scale_1p0/group_0082.npz +2d00be3f29b480124cd5f7eb01f30bf052d2e751b6aadfc1df81e15614d7d760 noise_scale_1p0/group_0090.npz +b499f492f29f41031aa3dbc57c29a689df067766bd4e312d1ec76656cf654587 noise_scale_1p0/group_0093.npz +512907c5d93bdfdc51366228b2da1956557b20fc3271ccb0b663b247de58e38b noise_scale_1p0/group_0094.npz +# checkpoints: the released checkpoints (--checkpoint), relative to the downloaded models repository +d670d439ea0e36979cef8bcee46bc78fd0b47b4fcc4cc517ffe0219b2d255c32 euler/direct-si/model.pt +3b6e0173fabe9528bf63a1bb13b05454630f046f05d17a94b9fbc3467272b4a0 euler/cascaded-si-S02/model.pt +9833a124ffaa155d569140630eb8bd91f7854c4965eb0042a437462a4f0c8407 euler/cascaded-si-S05/model.pt +d467e250984f326c17ab3daf4b6a956e47976ddd9948dc56ab87c00b96308dbb euler/cascaded-si-S10/model.pt +64edb1d25448691ecae8edb69a1090300bb30f56504ebe4ea8b60ee25103b466 euler/cascaded-si-S20/model.pt +3c7d254bda7e5f9d528666be727a1c2825c7bede04500212e12b560f75e842d6 euler/cascaded-si-S40/model.pt +3d1d855977e423cc7929d7c2c2be749b2651142bb7d0d8a1a479cee3d0aad715 euler/edm/model.pt diff --git a/code/records/euler/fid_adapter.json b/code/records/euler/fid_adapter.json new file mode 100644 index 0000000000000000000000000000000000000000..c5196164cde50e2e3d6a1c3d4a0bdcc92a4598b7 --- /dev/null +++ b/code/records/euler/fid_adapter.json @@ -0,0 +1,35 @@ +{ + "adapter": "euler_coarse_vorticity", + "field_kind": "signed", + "quantile": 0.9999, + "clip_c": 52.105898142242836, + "range": [ + -52.105898142242836, + 52.105898142242836 + ], + "n_train_states": 65600, + "field_shape": [ + 128, + 128 + ], + "pixel_stride": 4, + "n_pixel_values": 67174400, + "map": "uint8 = rint((clip(w,-c,c)+c)/(2c)*255)", + "channels": "3 identical (grayscale scalar field), no colormap", + "source": { + "packaged_dataset_path": "", + "block": "vecs_train", + "ics": "canonical rows 0..79", + "members": "j < 20", + "snapshots": "all", + "pixel_stride": 4, + "n_states": 65600, + "rule": "training data only; validation, test and generated data never enter" + }, + "state": "196-D coarse spectral vorticity (98 complex modes |k| <= 8, Hermitian half, omega_hat / N_sim^2)", + "reconstruction": "exact inverse FFT of the |k| <= 8 modes on a 128x128 grid (alias-free)", + "n_disp": 128, + "k_c": 8, + "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" +} diff --git a/code/records/euler/runs/cascaded_si_S01/config.yaml b/code/records/euler/runs/cascaded_si_S01/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7225a6cce639929f73c28d24f69c587a7530c752 --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S01/config.yaml @@ -0,0 +1,51 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 1.0 + n_segments: 1 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 8.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/cascaded_si_S02/config.yaml b/code/records/euler/runs/cascaded_si_S02/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8d2399a4469c1ca60894cbdc792339362a423b68 --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S02/config.yaml @@ -0,0 +1,52 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.5 + - 1.0 + n_segments: 2 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/cascaded_si_S05/config.yaml b/code/records/euler/runs/cascaded_si_S05/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dbd0189936f087f63aa14909ebb991aa9b0d601c --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S05/config.yaml @@ -0,0 +1,55 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.2 + - 0.4 + - 0.6 + - 0.8 + - 1.0 + n_segments: 5 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 2.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/cascaded_si_S10/config.yaml b/code/records/euler/runs/cascaded_si_S10/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..48e0905b36d287c472ca1de0f5661d12085de3a3 --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S10/config.yaml @@ -0,0 +1,60 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.1 + - 0.2 + - 0.3 + - 0.4 + - 0.5 + - 0.6 + - 0.7 + - 0.8 + - 0.9 + - 1.0 + n_segments: 10 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + sigma_train: 8.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/cascaded_si_S20/config.yaml b/code/records/euler/runs/cascaded_si_S20/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..13bb782fc2291581d68aaa13f6153536e73c3d14 --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S20/config.yaml @@ -0,0 +1,70 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.05 + - 0.1 + - 0.15 + - 0.2 + - 0.25 + - 0.3 + - 0.35 + - 0.4 + - 0.45 + - 0.5 + - 0.55 + - 0.6 + - 0.65 + - 0.7 + - 0.75 + - 0.8 + - 0.85 + - 0.9 + - 0.95 + - 1.0 + n_segments: 20 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 4.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/cascaded_si_S40/config.yaml b/code/records/euler/runs/cascaded_si_S40/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3d4b01c0986c313652ab04ae01842c7342b4f3e9 --- /dev/null +++ b/code/records/euler/runs/cascaded_si_S40/config.yaml @@ -0,0 +1,90 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.025 + - 0.05 + - 0.075 + - 0.1 + - 0.125 + - 0.15 + - 0.175 + - 0.2 + - 0.225 + - 0.25 + - 0.275 + - 0.3 + - 0.325 + - 0.35 + - 0.375 + - 0.4 + - 0.425 + - 0.45 + - 0.475 + - 0.5 + - 0.525 + - 0.55 + - 0.575 + - 0.6 + - 0.625 + - 0.65 + - 0.675 + - 0.7 + - 0.725 + - 0.75 + - 0.775 + - 0.8 + - 0.825 + - 0.85 + - 0.875 + - 0.9 + - 0.925 + - 0.95 + - 0.975 + - 1.0 + n_segments: 40 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 8.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/runs/edm/config.yaml b/code/records/euler/runs/edm/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f06b2fa81b595c6bf7b51be870ece4f9495588d --- /dev/null +++ b/code/records/euler/runs/edm/config.yaml @@ -0,0 +1,52 @@ +data: + name: euler_v2_uniform + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: edm + sigma_data: 1.0 + P_std: 1.2 + sampling: + sigma_min: 0.002 + sigma_max: 80.0 + rho: 7 + t_values: + - 0.0 + - 1.0 + n_steps: 20 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + lr_rampup_kimg: 0 + data_mode: fine_endpoint_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + P_mean: -1.6 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/euler/scores.json b/code/records/euler/scores.json new file mode 100644 index 0000000000000000000000000000000000000000..d5913651618966430d614ae6d7925beafc96ebe5 --- /dev/null +++ b/code/records/euler/scores.json @@ -0,0 +1,135 @@ +{ + "cascaded_si_S01": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,65.22805372184507,10000,10000,0\n16,74.21144398449542,10000,10000,0\n20,57.722502872449866,10000,10000,0\n36,67.79368464405024,10000,10000,0\n52,65.97300778843123,10000,10000,0\n72,55.176617517055035,10000,10000,0\n82,45.79409540054547,10000,10000,0\n90,47.49815183296431,10000,10000,0\n93,49.303691655995465,10000,10000,0\n94,43.94705738941195,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 57.2648306807244, + "sem": 3.349242766217627, + "sd": 10.591235578090549, + "n": 10 + }, + "csFID": { + "mean": 57.2648306807244, + "sem": 3.349242766217627, + "sd": 10.591235578090549, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S02": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,54.1538648581913,10000,10000,0\n16,67.98451096899879,10000,10000,0\n20,57.22661189268939,10000,10000,0\n36,59.757197222057584,10000,10000,0\n52,68.41507622194683,10000,10000,0\n72,55.252674268824535,10000,10000,0\n82,46.365731642060894,10000,10000,0\n90,39.6683529152707,10000,10000,0\n93,48.196651131611475,10000,10000,0\n94,38.96023217418491,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 53.59809032958365, + "sem": 3.2837081054094, + "sd": 10.38399678424998, + "n": 10 + }, + "csFID": { + "mean": 53.59809032958365, + "sem": 3.2837081054094, + "sd": 10.38399678424998, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S05": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,30.410875968219045,10000,10000,0\n16,19.064416322877435,10000,10000,0\n20,24.489591370006337,10000,10000,0\n36,42.52903215094415,10000,10000,0\n52,22.474309040249864,10000,10000,0\n72,20.04828380768828,10000,10000,0\n82,18.07861395489266,10000,10000,0\n90,17.54196528074999,10000,10000,0\n93,15.435461749067315,10000,10000,0\n94,19.375316525263486,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 22.944786616995856, + "sem": 2.554606140179877, + "sd": 8.078373927619797, + "n": 10 + }, + "csFID": { + "mean": 22.944786616995856, + "sem": 2.554606140179877, + "sd": 8.078373927619797, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S10": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,17.125925224700666,10000,10000,0\n16,15.044156088762293,10000,10000,0\n20,17.154272915168804,10000,10000,0\n36,25.737405958101135,10000,10000,0\n52,14.770237920268642,10000,10000,0\n72,16.922398737731328,10000,10000,0\n82,13.987127562070441,10000,10000,0\n90,15.370199565153257,10000,10000,0\n93,12.792013760495365,10000,10000,0\n94,14.333874097029366,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 16.32376118294813, + "sem": 1.1409603284300438, + "sd": 3.608033357732705, + "n": 10 + }, + "csFID": { + "mean": 16.32376118294813, + "sem": 1.1409603284300438, + "sd": 3.608033357732705, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S20": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,19.55871022475492,10000,10000,0\n16,15.659390557453492,10000,10000,0\n20,22.415182482281978,10000,10000,0\n36,20.763638233740835,10000,10000,0\n52,18.640790936192616,10000,10000,0\n72,16.307478807910684,10000,10000,0\n82,19.055494277409764,10000,10000,0\n90,15.106865511296462,10000,10000,0\n93,13.561268269729396,10000,10000,0\n94,17.54356896126791,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 17.861238826203806, + "sem": 0.8661358704549128, + "sd": 2.7389621138100644, + "n": 10 + }, + "csFID": { + "mean": 17.861238826203806, + "sem": 0.8661358704549128, + "sd": 2.7389621138100644, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S40": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,28.60062549215712,10000,10000,0\n16,20.295543199685767,10000,10000,0\n20,19.887490471116678,10000,10000,0\n36,33.02334450956084,10000,10000,0\n52,24.391798822395174,10000,10000,0\n72,25.587415111124756,10000,10000,0\n82,23.19201478612888,10000,10000,0\n90,19.838081079198332,10000,10000,0\n93,17.743499150036826,10000,10000,0\n94,20.346655061126967,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 23.290646768253133, + "sem": 1.4956564897631588, + "sd": 4.729681104863894, + "n": 10 + }, + "csFID": { + "mean": 23.290646768253133, + "sem": 1.4956564897631588, + "sd": 4.729681104863894, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "edm": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n5,62.9826373415093,10000,10000,0\n16,78.67235255538606,10000,10000,0\n20,53.11195109011504,10000,10000,0\n36,69.99068046524508,10000,10000,0\n52,61.51092491805218,10000,10000,0\n72,55.1068677983626,10000,10000,0\n82,46.46768540666159,10000,10000,0\n90,46.34913439335213,10000,10000,0\n93,47.46498014473616,10000,10000,0\n94,41.589529833805955,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 56.324674394722614, + "sem": 3.7438824883413795, + "sd": 11.839195955177548, + "n": 10 + }, + "csFID": { + "mean": 56.324674394722614, + "sem": 3.7438824883413795, + "sd": 11.839195955177548, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + } +} diff --git a/code/records/euler/units.json b/code/records/euler/units.json new file mode 100644 index 0000000000000000000000000000000000000000..0f09062bbd0d526d44d14c1766895470425679b8 --- /dev/null +++ b/code/records/euler/units.json @@ -0,0 +1,347 @@ +{ + "system": "euler", + "paper_name": "2-D Euler", + "dataset": "euler_v2_uniform", + "test_gids": [ + 5, + 16, + 20, + 36, + 52, + 72, + 82, + 90, + 93, + 94 + ], + "val_gids": [ + 8, + 9, + 10, + 11, + 13, + 27, + 37, + 75, + 81, + 83 + ], + "fine_horizon_idx": 40, + "x0_from_truth_bank": false, + "s_star": 10, + "table2": { + "direct_si": "cascaded_si_S01", + "edm": "edm", + "cascaded_si": "cascaded_si_S10" + }, + "validation_energy_distance_by_S": { + "1": 1.2814549667065398, + "2": 1.2434741392067776, + "5": 0.4341282997258393, + "10": 0.19754292785567965, + "20": 0.2174059261514088, + "40": 0.34261693919405695 + }, + "units": [ + { + "key": "cascaded_si_S01", + "S": 1, + "method": "cascaded_si", + "paper_names": [ + "direct_si" + ], + "config": "records/euler/runs/cascaded_si_S01/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 8.0, + "selected_update": 5000, + "validation_energy_distance": 1.2814549667065398 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 40 + }, + "checkpoint": { + "file": "euler/direct-si/model.pt", + "sha256": "d670d439ea0e36979cef8bcee46bc78fd0b47b4fcc4cc517ffe0219b2d255c32", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "5": 65.22805372184507, + "16": 74.21144398449542, + "20": 57.722502872449866, + "36": 67.79368464405024, + "52": 65.97300778843123, + "72": 55.176617517055035, + "82": 45.79409540054547, + "90": 47.49815183296431, + "93": 49.303691655995465, + "94": 43.94705738941195 + }, + "n_gen_nonfinite": {}, + "mean": 57.2648306807244, + "sem": 3.349242766217627, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S02", + "S": 2, + "method": "cascaded_si", + "paper_names": [], + "config": "records/euler/runs/cascaded_si_S02/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 7000, + "validation_energy_distance": 1.2434741392067776 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 160 + }, + "checkpoint": { + "file": "euler/cascaded-si-S02/model.pt", + "sha256": "3b6e0173fabe9528bf63a1bb13b05454630f046f05d17a94b9fbc3467272b4a0", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "5": 54.1538648581913, + "16": 67.98451096899879, + "20": 57.22661189268939, + "36": 59.757197222057584, + "52": 68.41507622194683, + "72": 55.252674268824535, + "82": 46.365731642060894, + "90": 39.6683529152707, + "93": 48.196651131611475, + "94": 38.96023217418491 + }, + "n_gen_nonfinite": {}, + "mean": 53.59809032958365, + "sem": 3.2837081054094, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S05", + "S": 5, + "method": "cascaded_si", + "paper_names": [], + "config": "records/euler/runs/cascaded_si_S05/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 2.0, + "selected_update": 25000, + "validation_energy_distance": 0.4341282997258393 + }, + "sampler": { + "n_steps": 160, + "variant": "noise_scale_1p0", + "nfe": 800 + }, + "checkpoint": { + "file": "euler/cascaded-si-S05/model.pt", + "sha256": "9833a124ffaa155d569140630eb8bd91f7854c4965eb0042a437462a4f0c8407", + "bytes": 18578218 + }, + "expected_csfid": { + "per_condition": { + "5": 30.410875968219045, + "16": 19.064416322877435, + "20": 24.489591370006337, + "36": 42.52903215094415, + "52": 22.474309040249864, + "72": 20.04828380768828, + "82": 18.07861395489266, + "90": 17.54196528074999, + "93": 15.435461749067315, + "94": 19.375316525263486 + }, + "n_gen_nonfinite": {}, + "mean": 22.944786616995856, + "sem": 2.554606140179877, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S10", + "S": 10, + "method": "cascaded_si", + "paper_names": [ + "cascaded_si" + ], + "config": "records/euler/runs/cascaded_si_S10/config.yaml", + "selection": { + "lr": 0.0003, + "sigma": 8.0, + "selected_update": 17000, + "validation_energy_distance": 0.19754292785567965 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 800 + }, + "checkpoint": { + "file": "euler/cascaded-si-S10/model.pt", + "sha256": "d467e250984f326c17ab3daf4b6a956e47976ddd9948dc56ab87c00b96308dbb", + "bytes": 18578282 + }, + "expected_csfid": { + "per_condition": { + "5": 17.125925224700666, + "16": 15.044156088762293, + "20": 17.154272915168804, + "36": 25.737405958101135, + "52": 14.770237920268642, + "72": 16.922398737731328, + "82": 13.987127562070441, + "90": 15.370199565153257, + "93": 12.792013760495365, + "94": 14.333874097029366 + }, + "n_gen_nonfinite": {}, + "mean": 16.32376118294813, + "sem": 1.1409603284300438, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S20", + "S": 20, + "method": "cascaded_si", + "paper_names": [], + "config": "records/euler/runs/cascaded_si_S20/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 4.0, + "selected_update": 27000, + "validation_energy_distance": 0.2174059261514088 + }, + "sampler": { + "n_steps": 160, + "variant": "noise_scale_1p0", + "nfe": 3200 + }, + "checkpoint": { + "file": "euler/cascaded-si-S20/model.pt", + "sha256": "64edb1d25448691ecae8edb69a1090300bb30f56504ebe4ea8b60ee25103b466", + "bytes": 18578346 + }, + "expected_csfid": { + "per_condition": { + "5": 19.55871022475492, + "16": 15.659390557453492, + "20": 22.415182482281978, + "36": 20.763638233740835, + "52": 18.640790936192616, + "72": 16.307478807910684, + "82": 19.055494277409764, + "90": 15.106865511296462, + "93": 13.561268269729396, + "94": 17.54356896126791 + }, + "n_gen_nonfinite": {}, + "mean": 17.861238826203806, + "sem": 0.8661358704549128, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S40", + "S": 40, + "method": "cascaded_si", + "paper_names": [], + "config": "records/euler/runs/cascaded_si_S40/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 8.0, + "selected_update": 12000, + "validation_energy_distance": 0.34261693919405695 + }, + "sampler": { + "n_steps": 160, + "variant": "noise_scale_1p0", + "nfe": 6400 + }, + "checkpoint": { + "file": "euler/cascaded-si-S40/model.pt", + "sha256": "3c7d254bda7e5f9d528666be727a1c2825c7bede04500212e12b560f75e842d6", + "bytes": 18578538 + }, + "expected_csfid": { + "per_condition": { + "5": 28.60062549215712, + "16": 20.295543199685767, + "20": 19.887490471116678, + "36": 33.02334450956084, + "52": 24.391798822395174, + "72": 25.587415111124756, + "82": 23.19201478612888, + "90": 19.838081079198332, + "93": 17.743499150036826, + "94": 20.346655061126967 + }, + "n_gen_nonfinite": {}, + "mean": 23.290646768253133, + "sem": 1.4956564897631588, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "edm", + "S": 1, + "method": "edm", + "paper_names": [ + "edm" + ], + "config": "records/euler/runs/edm/config.yaml", + "selection": { + "lr": 0.001, + "P_mean": -1.6, + "selected_update": 300, + "validation_energy_distance": 1.2337307787158724 + }, + "sampler": { + "n_steps": 40, + "variant": "default", + "nfe": 79 + }, + "checkpoint": { + "file": "euler/edm/model.pt", + "sha256": "3d1d855977e423cc7929d7c2c2be749b2651142bb7d0d8a1a479cee3d0aad715", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "5": 62.9826373415093, + "16": 78.67235255538606, + "20": 53.11195109011504, + "36": 69.99068046524508, + "52": 61.51092491805218, + "72": 55.1068677983626, + "82": 46.46768540666159, + "90": 46.34913439335213, + "93": 47.46498014473616, + "94": 41.589529833805955 + }, + "n_gen_nonfinite": {}, + "mean": 56.324674394722614, + "sem": 3.7438824883413795, + "n": 10 + }, + "selection_flags": [] + } + ] +} diff --git a/code/records/nlkg/assets.sha256 b/code/records/nlkg/assets.sha256 new file mode 100644 index 0000000000000000000000000000000000000000..a5e5aac6dfe246627c1bd87ebcae3a7fcb08959d --- /dev/null +++ b/code/records/nlkg/assets.sha256 @@ -0,0 +1,59 @@ +# sha256 of the inputs of the Table 2 entries; paths are relative to the directory named in the heading +# data: the dataset package (--data) +258cdda0a2707a494d2f4bb34ae96b2074ebbe70c0960e7f586702d771ac8a22 trajectories_fine.npz +95c8d315d18dd94a8cc32c2d810e00d69e374a338a315ce2dab9858d77566de3 split.json +94e78f3fc04a40b3f71c883b97cf492c3763853c8432736112ec0167e30543ca normalization.npz +# truth_bank: the truth banks of the test conditions (--truth-bank) +9a6757c2736ae9ad755964d1c341735f56dd1eddbd89b971ae731c8e588c717a group_0100.npz +98e703513e8d873b43e8988eead9509c4eed66ab55cd91aa4da56f107d792dbc group_0101.npz +f4b821f6910f5b892508769ac69fbff97f45ea7e62aa5c051e147246a664863c group_0102.npz +ec6736f898793004c2064443f8cbdb830eba022e7860a280c0891e3dfc23306c group_0103.npz +9a8d7933620e902cc3b1486650f79a86b1a614d38009412f26b2f64a25053115 group_0104.npz +a53d95ba7f7ef5891a513efad434fb8b81c0d4e8cdd692bd67ab0892226268f9 group_0105.npz +19ba8d2bd73ff40b1d74720cea5a12bfb17c8c7510bf7614d1363f67356d30c3 group_0106.npz +ae11d493e3da9161c625d557c889a62fdb58bbd60dc2692a6e0845629dfb4b12 group_0107.npz +b353eb2c85a6cead3fa7fe67c5a0a64ee1cf894f50ce0f669a39a7b583dceaf6 group_0108.npz +d3241ab7af85bd49e0f21f83edf4da5a50756df8015ef63d951ce2825b763daf group_0109.npz +# inception: the Inception-v3 weights in $CLEAN_FID_CACHE +f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4 inception-2015-12-05.pt +# samples/direct_si: the archived test samples (--samples) +26274c091c3145f93d6cefe67e42c05bba3c394dca28010dc0e08980155598ba noise_scale_1p0/group_0100.npz +8bd498c3dfde60aa76cd11981f95f4dbe8d99a4ad6221692aaed2bbe0b5645fd noise_scale_1p0/group_0101.npz +4ddfbfaa7aef088286347b327b515b1f40ee9f6d053403882f73aba84d4280cc noise_scale_1p0/group_0102.npz +82f720f72dcd3b6ec360df49c2113da1dd3a66d1cc151a3a0db818cdf2d62608 noise_scale_1p0/group_0103.npz +c29b609d71bd19b41adc9155b8bfe68c63b15e137230ab060805d276254e36ef noise_scale_1p0/group_0104.npz +c5694a44b1d320c35e241da1cd8cfbae94ddc08ba68da6112a28b8ab8b231e6d noise_scale_1p0/group_0105.npz +4a6f1c0cc2adfb785fe54d6e3a3d1f0229742d839f2d44e76251d484d208aefa noise_scale_1p0/group_0106.npz +cab72be6fd2604b2071375191bc973172a9e8bcb410cf8b1072a6a5401618902 noise_scale_1p0/group_0107.npz +da178e706e7ebcbde3ac1158fca6df15cbb2bdb8ede0e50ff60f40db51363e1a noise_scale_1p0/group_0108.npz +c189b96bc233fb4ac4baa2e096dc246bc3bc03cdd313adeb9361ff4bd7d4f4d1 noise_scale_1p0/group_0109.npz +# samples/edm: the archived test samples (--samples) +9b2756e21d62dcc83722977edd003bcc748d2512d545f8892a15907005decdee default/group_0100.npz +22646ed678556eb606d8d02a7fcd4d3d5c5cc0b2d3b0505e8c7cb4b91e778e53 default/group_0101.npz +7241b031324525982a88fae233538db1d8c70d7d0f6003c41bc621c09b27d92b default/group_0102.npz +f163b63fd5ce37167bb6a54ffdce1f970441625cb8b17939c376940d3f99fc70 default/group_0103.npz +a324ae1e92d193dfdfb2f4ce0bd5d7a20913f4e3e8d55ac5e969c0dd8b9f130e default/group_0104.npz +325ebd77f518e95512772cd4c4948e5995bdedc610d8b6f5bab5d416cf7a4a42 default/group_0105.npz +a19d96fc5f5237ed7ec93de24692af0986fafde1a67a48c37c53b41e3182ca45 default/group_0106.npz +0727c56939ee69f89de99345e539c376b3d5b5055d42286486892640f6de72f2 default/group_0107.npz +a1dc0fdabab2a6772b9947bb66dfa574ffdad4a3adec9b0140dabf4af23104f2 default/group_0108.npz +58f7e6b8c6af539b26208a356b53d92aa1af1377fc418d078db794843303ffb1 default/group_0109.npz +# samples/cascaded_si: the archived test samples (--samples) +2a8310f1662828f673d27d9399f844362d78600203675fab53ebfa5f08d3302c noise_scale_1p0/group_0100.npz +921ff32e7122b836f142ee8b7b1a056fdc3918c0b5edbf43385fab452e2a1b60 noise_scale_1p0/group_0101.npz +a8a7497ad27ca4de5000d0e651c70b2040e6431782461ea55e125ec2a94d9d43 noise_scale_1p0/group_0102.npz +3b1b01360180d69feebb835bea7b87486d370004dcc33b8e4d90708d70c87711 noise_scale_1p0/group_0103.npz +15fc57d6e693a3e381273ba1c5e7fe1971c8ac948d2a02f746f89a449f05f404 noise_scale_1p0/group_0104.npz +c3f9d99c6aa5202734d2baa8918e2fcda98843a1ce7a161008a3735f968c7f4e noise_scale_1p0/group_0105.npz +08a3b172eaa6575f77cdf24ee84d0bb6b5582454bad2b17140e67728964d1629 noise_scale_1p0/group_0106.npz +17fd3e6d75bd11bb65c0ae1201ee7f99242a8a585729415fa3b0f71135af150e noise_scale_1p0/group_0107.npz +8b624ed8bcf53e6954aeb3fa26a0f450a40549a680e84161465164ab53e8d5b7 noise_scale_1p0/group_0108.npz +931bbf6e638b9964570ba365dc786268134499ae26a7e2c2da2e05ceb9ec36ec noise_scale_1p0/group_0109.npz +# checkpoints: the released checkpoints (--checkpoint), relative to the downloaded models repository +1c4ee6e6261d46f193cf328a3e5198df3e7e1bcce29ab31a91ca567737dd9646 nonlinear-kg/direct-si/model.pt +30c313ffaf84bec673566d35268f152e88dccd32686890f70bee271b79727b10 nonlinear-kg/cascaded-si-S02/model.pt +b01eeee29892c90a77ccb156738e1d4024167f2cbd30b2f624124eb0d19342c8 nonlinear-kg/cascaded-si-S05/model.pt +4c21f902a8a64fa6e7e03bf8aa2a5d0d1d647fd02bf0cf0d9520651bede97185 nonlinear-kg/cascaded-si-S10/model.pt +546ee87205aaf7ac0b7f688abf0de10309b212781d5f4c6c5f945323fb243c4d nonlinear-kg/cascaded-si-S20/model.pt +356ffe38e95681e7c84dec6927e5fb6da2032816ec7ab77118e4223640439bd9 nonlinear-kg/cascaded-si-S40/model.pt +633bbcce5f6d8f39ccb29e6c674bd497cdc4337e05d51de9d2a6db90a64e5843 nonlinear-kg/edm/model.pt diff --git a/code/records/nlkg/fid_adapter.json b/code/records/nlkg/fid_adapter.json new file mode 100644 index 0000000000000000000000000000000000000000..266f3ed3a2cd75ff9a9c44c66d804157e9b5f8a2 --- /dev/null +++ b/code/records/nlkg/fid_adapter.json @@ -0,0 +1,37 @@ +{ + "adapter": "nlkg_coarse_u", + "field_kind": "signed", + "quantile": 0.9999, + "clip_c": 0.5955045945108044, + "range": [ + -0.5955045945108044, + 0.5955045945108044 + ], + "n_train_states": 1600, + "field_shape": [ + 128, + 128 + ], + "pixel_stride": 1, + "n_pixel_values": 26214400, + "map": "uint8 = rint((clip(w,-c,c)+c)/(2c)*255)", + "channels": "3 identical (grayscale scalar field), no colormap", + "source": { + "packaged_dataset_path": "", + "block": "vecs_train", + "ics": "canonical rows 0..79", + "members": "j < 20", + "snapshots": [ + 40 + ], + "pixel_stride": 1, + "n_states": 1600, + "rule": "training data only; validation, test and generated data never enter" + }, + "state": "392-D coarse state (196-D u + 196-D v/4, house coarse operator |k| <= 8 on each); the image reads ONLY the first component u -- identical reconstruction to euler_coarse_vorticity", + "reconstruction": "exact inverse FFT of the |k| <= 8 modes on a 128x128 grid (alias-free)", + "n_disp": 128, + "k_c": 8, + "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)" +} diff --git a/code/records/nlkg/runs/cascaded_si_S01/config.yaml b/code/records/nlkg/runs/cascaded_si_S01/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a5a2c06265c4888b3089742761bd3e0f0fee9340 --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S01/config.yaml @@ -0,0 +1,51 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 12.0 + n_segments: 1 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 128.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/cascaded_si_S02/config.yaml b/code/records/nlkg/runs/cascaded_si_S02/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..26312198558226cf8c892c87b3295d1c61a66399 --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S02/config.yaml @@ -0,0 +1,52 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 6.0 + - 12.0 + n_segments: 2 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 3.0e-05 + sigma_train: 8.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/cascaded_si_S05/config.yaml b/code/records/nlkg/runs/cascaded_si_S05/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c3ebb6f10e5a9a5339b9afbc368db5ee1dd2e3a --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S05/config.yaml @@ -0,0 +1,55 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 2.4000000000000004 + - 4.800000000000001 + - 7.2 + - 9.600000000000001 + - 12.0 + n_segments: 5 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 0.25 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/cascaded_si_S10/config.yaml b/code/records/nlkg/runs/cascaded_si_S10/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ae1f8afdeab918229498670c8df3954b628ead70 --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S10/config.yaml @@ -0,0 +1,60 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 1.2000000000000002 + - 2.4000000000000004 + - 3.6 + - 4.800000000000001 + - 6.0 + - 7.2 + - 8.4 + - 9.600000000000001 + - 10.8 + - 12.0 + n_segments: 10 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.003 + sigma_train: 0.5 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/cascaded_si_S20/config.yaml b/code/records/nlkg/runs/cascaded_si_S20/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..67b6d7d56f3ab1c8da4dfc72133cdf76d125edd2 --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S20/config.yaml @@ -0,0 +1,70 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.6000000000000001 + - 1.2000000000000002 + - 1.8 + - 2.4000000000000004 + - 3.0 + - 3.6 + - 4.2 + - 4.800000000000001 + - 5.4 + - 6.0 + - 6.6000000000000005 + - 7.2 + - 7.800000000000001 + - 8.4 + - 9.0 + - 9.600000000000001 + - 10.200000000000001 + - 10.8 + - 11.4 + - 12.0 + n_segments: 20 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 1.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/cascaded_si_S40/config.yaml b/code/records/nlkg/runs/cascaded_si_S40/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b73ec0719c49304324ec5e5ae118935b5b2acfaa --- /dev/null +++ b/code/records/nlkg/runs/cascaded_si_S40/config.yaml @@ -0,0 +1,90 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.30000000000000004 + - 0.6000000000000001 + - 0.9 + - 1.2000000000000002 + - 1.5 + - 1.8 + - 2.1 + - 2.4000000000000004 + - 2.7 + - 3.0 + - 3.3000000000000003 + - 3.6 + - 3.9000000000000004 + - 4.2 + - 4.5 + - 4.800000000000001 + - 5.1000000000000005 + - 5.4 + - 5.7 + - 6.0 + - 6.300000000000001 + - 6.6000000000000005 + - 6.9 + - 7.2 + - 7.5 + - 7.800000000000001 + - 8.1 + - 8.4 + - 8.700000000000001 + - 9.0 + - 9.3 + - 9.600000000000001 + - 9.9 + - 10.200000000000001 + - 10.5 + - 10.8 + - 11.100000000000001 + - 11.4 + - 11.700000000000001 + - 12.0 + n_segments: 40 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + sigma_train: 2.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/runs/edm/config.yaml b/code/records/nlkg/runs/edm/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..deb2bb7ac52e0566238106789008f42a11d4ad7e --- /dev/null +++ b/code/records/nlkg/runs/edm/config.yaml @@ -0,0 +1,52 @@ +data: + name: nlkg_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_2comp + D: 392 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 12.0 +method: + name: edm + sigma_data: 1.0 + P_std: 1.2 + sampling: + sigma_min: 0.002 + sigma_max: 80.0 + rho: 7 + t_values: + - 0.0 + - 12.0 + n_steps: 20 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + lr_rampup_kimg: 0 + data_mode: fine_endpoint_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + P_mean: -4.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/nlkg/scores.json b/code/records/nlkg/scores.json new file mode 100644 index 0000000000000000000000000000000000000000..ab3ec13843e9b9f59ef66c48352424c73c02e9cf --- /dev/null +++ b/code/records/nlkg/scores.json @@ -0,0 +1,135 @@ +{ + "cascaded_si_S01": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,86.50339206862986,10000,10000,0\n101,82.17386645891077,10000,10000,0\n102,85.08862443439534,10000,10000,0\n103,105.96661352064586,10000,10000,0\n104,88.66118022487461,10000,10000,0\n105,107.9384631741924,10000,10000,0\n106,146.80570228809182,10000,10000,0\n107,104.12904226827146,10000,10000,0\n108,83.81643790189827,10000,10000,0\n109,93.68464062258761,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 98.4767962962498, + "sem": 6.1787087828142875, + "sd": 19.53879275257978, + "n": 10 + }, + "csFID": { + "mean": 98.4767962962498, + "sem": 6.1787087828142875, + "sd": 19.53879275257978, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S02": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,87.30338201802135,10000,10000,0\n101,83.25596893728516,10000,10000,0\n102,84.38834295153731,10000,10000,0\n103,104.20679677950466,10000,10000,0\n104,95.16958079739976,10000,10000,0\n105,104.74304727808482,10000,10000,0\n106,146.66769326658493,10000,10000,0\n107,107.77579599865578,10000,10000,0\n108,89.48169875816228,10000,10000,0\n109,94.5548973851715,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 99.75472041704074, + "sem": 5.894287147084769, + "sd": 18.639372567843775, + "n": 10 + }, + "csFID": { + "mean": 99.75472041704074, + "sem": 5.894287147084769, + "sd": 18.639372567843775, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S05": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,92.63917294930161,10000,10000,0\n101,49.86823139250902,10000,10000,0\n102,47.46047500644772,10000,10000,0\n103,91.45318172827865,10000,10000,0\n104,56.347194609326436,10000,10000,0\n105,23.913809957644332,10000,10000,0\n106,56.97006944240752,10000,10000,0\n107,44.54986523952408,10000,10000,0\n108,37.93354061755201,10000,10000,0\n109,29.836426462157306,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 53.09719674051487, + "sem": 7.301625987850265, + "sd": 23.08976874428377, + "n": 10 + }, + "csFID": { + "mean": 53.09719674051487, + "sem": 7.301625987850265, + "sd": 23.08976874428377, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S10": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,31.385955252590236,10000,10000,0\n101,12.348480970992114,10000,10000,0\n102,14.496443143531451,10000,10000,0\n103,66.2757090789733,10000,10000,0\n104,26.915870774500263,10000,10000,0\n105,9.220095330626975,10000,10000,0\n106,24.368288260847805,10000,10000,0\n107,15.194038156306192,10000,10000,0\n108,13.39604081556297,10000,10000,0\n109,11.448239994182472,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 22.504916177811378, + "sem": 5.395604960058663, + "sd": 17.062401028287212, + "n": 10 + }, + "csFID": { + "mean": 22.504916177811378, + "sem": 5.395604960058663, + "sd": 17.062401028287212, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S20": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,19.46689280522986,10000,10000,0\n101,12.170450558502239,10000,10000,0\n102,14.166808232319966,10000,10000,0\n103,22.766638013605927,10000,10000,0\n104,15.439152288157743,10000,10000,0\n105,11.638441653905943,10000,10000,0\n106,25.92547770681159,10000,10000,0\n107,15.615507640775547,10000,10000,0\n108,11.845380452458926,10000,10000,0\n109,13.39096502755956,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 16.24257143793273, + "sem": 1.5538946704054242, + "sd": 4.91384640247778, + "n": 10 + }, + "csFID": { + "mean": 16.24257143793273, + "sem": 1.5538946704054242, + "sd": 4.91384640247778, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S40": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,18.152889464190935,10000,10000,0\n101,11.520560759428633,10000,10000,0\n102,14.366384927800482,10000,10000,0\n103,22.63483205135941,10000,10000,0\n104,18.517078785282024,10000,10000,0\n105,12.129927401498435,10000,10000,0\n106,31.301773541882937,10000,10000,0\n107,25.390870435077545,10000,10000,0\n108,12.531017617367638,10000,10000,0\n109,15.727583531987307,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 18.227291851587534, + "sem": 2.0481309411127504, + "sd": 6.47675872018049, + "n": 10 + }, + "csFID": { + "mean": 18.227291851587534, + "sem": 2.0481309411127504, + "sd": 6.47675872018049, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "edm": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,76.58221938982206,10000,10000,0\n101,75.14550751056231,10000,10000,0\n102,79.71499245058573,10000,10000,0\n103,96.73988122315717,10000,10000,0\n104,89.98229585007118,10000,10000,0\n105,109.45974001585324,10000,10000,0\n106,128.81857368748194,10000,10000,0\n107,101.06772607934283,10000,10000,0\n108,77.93207366082618,10000,10000,0\n109,88.24109127216889,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 92.36841011398715, + "sem": 5.428927787320806, + "sd": 17.167777060511934, + "n": 10 + }, + "csFID": { + "mean": 92.36841011398715, + "sem": 5.428927787320806, + "sd": 17.167777060511934, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + } +} diff --git a/code/records/nlkg/units.json b/code/records/nlkg/units.json new file mode 100644 index 0000000000000000000000000000000000000000..e7011ee3538e506a1001dd60d7d0e99ccf192c6c --- /dev/null +++ b/code/records/nlkg/units.json @@ -0,0 +1,351 @@ +{ + "system": "nlkg", + "paper_name": "Nonlinear Klein-Gordon (NLKG)", + "dataset": "nlkg_v3_candidate", + "test_gids": [ + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109 + ], + "val_gids": [ + 80, + 81, + 82, + 83, + 84, + 85, + 86, + 87, + 88, + 89 + ], + "fine_horizon_idx": 40, + "x0_from_truth_bank": true, + "s_star": 20, + "table2": { + "direct_si": "cascaded_si_S01", + "edm": "edm", + "cascaded_si": "cascaded_si_S20" + }, + "validation_energy_distance_by_S": { + "1": 10.650600428383587, + "2": 10.686098889729823, + "5": 4.469934852532934, + "10": 1.5280741957806083, + "20": 1.3706359090695461, + "40": 1.4106361928761457 + }, + "units": [ + { + "key": "cascaded_si_S01", + "S": 1, + "method": "cascaded_si", + "paper_names": [ + "direct_si" + ], + "config": "records/nlkg/runs/cascaded_si_S01/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 128.0, + "selected_update": 15000, + "validation_energy_distance": 10.650600428383587 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 40 + }, + "checkpoint": { + "file": "nonlinear-kg/direct-si/model.pt", + "sha256": "1c4ee6e6261d46f193cf328a3e5198df3e7e1bcce29ab31a91ca567737dd9646", + "bytes": 25032338 + }, + "expected_csfid": { + "per_condition": { + "100": 86.50339206862986, + "101": 82.17386645891077, + "102": 85.08862443439534, + "103": 105.96661352064586, + "104": 88.66118022487461, + "105": 107.9384631741924, + "106": 146.80570228809182, + "107": 104.12904226827146, + "108": 83.81643790189827, + "109": 93.68464062258761 + }, + "n_gen_nonfinite": {}, + "mean": 98.4767962962498, + "sem": 6.1787087828142875, + "n": 10 + }, + "selection_flags": [ + "selected_hyperparameter_on_search_boundary" + ] + }, + { + "key": "cascaded_si_S02", + "S": 2, + "method": "cascaded_si", + "paper_names": [], + "config": "records/nlkg/runs/cascaded_si_S02/config.yaml", + "selection": { + "lr": 3e-05, + "sigma": 8.0, + "selected_update": 750, + "validation_energy_distance": 10.686098889729823 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 80 + }, + "checkpoint": { + "file": "nonlinear-kg/cascaded-si-S02/model.pt", + "sha256": "30c313ffaf84bec673566d35268f152e88dccd32686890f70bee271b79727b10", + "bytes": 25032338 + }, + "expected_csfid": { + "per_condition": { + "100": 87.30338201802135, + "101": 83.25596893728516, + "102": 84.38834295153731, + "103": 104.20679677950466, + "104": 95.16958079739976, + "105": 104.74304727808482, + "106": 146.66769326658493, + "107": 107.77579599865578, + "108": 89.48169875816228, + "109": 94.5548973851715 + }, + "n_gen_nonfinite": {}, + "mean": 99.75472041704074, + "sem": 5.894287147084769, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S05", + "S": 5, + "method": "cascaded_si", + "paper_names": [], + "config": "records/nlkg/runs/cascaded_si_S05/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 0.25, + "selected_update": 26000, + "validation_energy_distance": 4.469934852532934 + }, + "sampler": { + "n_steps": 320, + "variant": "noise_scale_1p0", + "nfe": 1600 + }, + "checkpoint": { + "file": "nonlinear-kg/cascaded-si-S05/model.pt", + "sha256": "b01eeee29892c90a77ccb156738e1d4024167f2cbd30b2f624124eb0d19342c8", + "bytes": 25032338 + }, + "expected_csfid": { + "per_condition": { + "100": 92.63917294930161, + "101": 49.86823139250902, + "102": 47.46047500644772, + "103": 91.45318172827865, + "104": 56.347194609326436, + "105": 23.913809957644332, + "106": 56.97006944240752, + "107": 44.54986523952408, + "108": 37.93354061755201, + "109": 29.836426462157306 + }, + "n_gen_nonfinite": {}, + "mean": 53.09719674051487, + "sem": 7.301625987850265, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S10", + "S": 10, + "method": "cascaded_si", + "paper_names": [], + "config": "records/nlkg/runs/cascaded_si_S10/config.yaml", + "selection": { + "lr": 0.003, + "sigma": 0.5, + "selected_update": 17000, + "validation_energy_distance": 1.5280741957806083 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 400 + }, + "checkpoint": { + "file": "nonlinear-kg/cascaded-si-S10/model.pt", + "sha256": "4c21f902a8a64fa6e7e03bf8aa2a5d0d1d647fd02bf0cf0d9520651bede97185", + "bytes": 25032402 + }, + "expected_csfid": { + "per_condition": { + "100": 31.385955252590236, + "101": 12.348480970992114, + "102": 14.496443143531451, + "103": 66.2757090789733, + "104": 26.915870774500263, + "105": 9.220095330626975, + "106": 24.368288260847805, + "107": 15.194038156306192, + "108": 13.39604081556297, + "109": 11.448239994182472 + }, + "n_gen_nonfinite": {}, + "mean": 22.504916177811378, + "sem": 5.395604960058663, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S20", + "S": 20, + "method": "cascaded_si", + "paper_names": [ + "cascaded_si" + ], + "config": "records/nlkg/runs/cascaded_si_S20/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 1.0, + "selected_update": 30000, + "validation_energy_distance": 1.3706359090695461 + }, + "sampler": { + "n_steps": 320, + "variant": "noise_scale_1p0", + "nfe": 6400 + }, + "checkpoint": { + "file": "nonlinear-kg/cascaded-si-S20/model.pt", + "sha256": "546ee87205aaf7ac0b7f688abf0de10309b212781d5f4c6c5f945323fb243c4d", + "bytes": 25032466 + }, + "expected_csfid": { + "per_condition": { + "100": 19.46689280522986, + "101": 12.170450558502239, + "102": 14.166808232319966, + "103": 22.766638013605927, + "104": 15.439152288157743, + "105": 11.638441653905943, + "106": 25.92547770681159, + "107": 15.615507640775547, + "108": 11.845380452458926, + "109": 13.39096502755956 + }, + "n_gen_nonfinite": {}, + "mean": 16.24257143793273, + "sem": 1.5538946704054242, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S40", + "S": 40, + "method": "cascaded_si", + "paper_names": [], + "config": "records/nlkg/runs/cascaded_si_S40/config.yaml", + "selection": { + "lr": 0.0003, + "sigma": 2.0, + "selected_update": 22000, + "validation_energy_distance": 1.4106361928761457 + }, + "sampler": { + "n_steps": 320, + "variant": "noise_scale_1p0", + "nfe": 12800 + }, + "checkpoint": { + "file": "nonlinear-kg/cascaded-si-S40/model.pt", + "sha256": "356ffe38e95681e7c84dec6927e5fb6da2032816ec7ab77118e4223640439bd9", + "bytes": 25032658 + }, + "expected_csfid": { + "per_condition": { + "100": 18.152889464190935, + "101": 11.520560759428633, + "102": 14.366384927800482, + "103": 22.63483205135941, + "104": 18.517078785282024, + "105": 12.129927401498435, + "106": 31.301773541882937, + "107": 25.390870435077545, + "108": 12.531017617367638, + "109": 15.727583531987307 + }, + "n_gen_nonfinite": {}, + "mean": 18.227291851587534, + "sem": 2.0481309411127504, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "edm", + "S": 1, + "method": "edm", + "paper_names": [ + "edm" + ], + "config": "records/nlkg/runs/edm/config.yaml", + "selection": { + "lr": 0.001, + "P_mean": -4.0, + "selected_update": 17000, + "validation_energy_distance": 10.465239160940088 + }, + "sampler": { + "n_steps": 40, + "variant": "default", + "nfe": 79 + }, + "checkpoint": { + "file": "nonlinear-kg/edm/model.pt", + "sha256": "633bbcce5f6d8f39ccb29e6c674bd497cdc4337e05d51de9d2a6db90a64e5843", + "bytes": 25032274 + }, + "expected_csfid": { + "per_condition": { + "100": 76.58221938982206, + "101": 75.14550751056231, + "102": 79.71499245058573, + "103": 96.73988122315717, + "104": 89.98229585007118, + "105": 109.45974001585324, + "106": 128.81857368748194, + "107": 101.06772607934283, + "108": 77.93207366082618, + "109": 88.24109127216889 + }, + "n_gen_nonfinite": {}, + "mean": 92.36841011398715, + "sem": 5.428927787320806, + "n": 10 + }, + "selection_flags": [ + "selected_hyperparameter_on_search_boundary" + ] + } + ] +} diff --git a/code/records/passive_scalar/assets.sha256 b/code/records/passive_scalar/assets.sha256 new file mode 100644 index 0000000000000000000000000000000000000000..8a0ae0e7add59638d94a5bccc0c9f9bf4727e8c1 --- /dev/null +++ b/code/records/passive_scalar/assets.sha256 @@ -0,0 +1,49 @@ +# sha256 of the inputs of the Table 2 entries; paths are relative to the directory named in the heading +# data: the dataset package (--data) +fb8ca7914199c2ae19ce61c0a6baed8736c98f2db77e19f5fa128003a2979081 split.json +7f79196a94bba0509e1a307df39a4f572a1f8100dbd5ba9aad82f53aae02095e normalization.npz +7a810b232414307ade1cebc0ecff4961d381561e400e034fc356cbf12f764a52 trajectories_fine.npz +# truth_bank: the truth banks of the test conditions (--truth-bank) +331e4840052421a1aa1dbf89917229b3603319983103f413ee67384f44529363 group_0000.npz +0d337d5ae77421d95f8c9a444f47cee5c302a6db4f89f1f10bf119996348b8e5 group_0007.npz +822b862678411231d36ba5a8bca17c52ee2d1f3a3be8ce1ec357dd49830016af group_0017.npz +978d41b073d0e13a89f358e1a6ebd6cf7e6528e21ae291905692d0c5a1023c60 group_0019.npz +45f10476eb9db10c3a55c580bc6a726f4409261ecaba504b5fe1680ea502810e group_0045.npz +64b6143d49437e427e453edcee94bbd783af8b9ce7d509418ab768cc3304ff14 group_0054.npz +d64473363f82717e67e9d1fb12fa36198b1f181e93e6493516299311b4ab97b7 group_0071.npz +06cf8258d5f5f9df14c186088b5cceca76c7e3a705740d2dfb11854baed0b4ec group_0078.npz +c0c1023204edd56ee74b7c8306510ffea4ff613a0a72e7f063383993a8c3ea4a group_0085.npz +999d10d4cc6d3e3dc4165d3d1474aa8cf65aeaf5958bf06ead66e4a8cb4b321e group_0094.npz +# inception: the Inception-v3 weights in $CLEAN_FID_CACHE +f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4 inception-2015-12-05.pt +# samples/direct_si: the archived test samples (--samples) +3daacb35f62eb6163f84b9b8a136b70e2b259a7458f7eb44eb661b4f41df5c41 noise_scale_1p0/group_0000.npz +f872a66742aba024d31af68f0a4f8a00e08bb27f4d5b12ceb07542223b7b9115 noise_scale_1p0/group_0007.npz +7356a371fa75ccdde134f7f487a8e646e2ca7a11ee4fdd4167e31ae0ed86e762 noise_scale_1p0/group_0017.npz +1433caef41297e1a4f28391cc302b5b5c8caad1033648fb50840887f6f8f45c6 noise_scale_1p0/group_0019.npz +bc2199bc7c43b9749eff1f753ba7f270d85c3cce7da463d015e7261074de4845 noise_scale_1p0/group_0045.npz +5c005a217d1c9406045315ca72fca0725ad6d0f9369e6040d2c74f75b2ac3213 noise_scale_1p0/group_0054.npz +b123a34c0076bf3fd7ec29026354c6dfede0f795e2317c8442751bb96b5516ae noise_scale_1p0/group_0071.npz +1eeadbaa9c4d3f7531d3ab92e4c717dc9b417029653612cc4fdda8eb031147a6 noise_scale_1p0/group_0078.npz +b8cb3b70a92ae0d9fbc5b8352f20af799db8653d8bd8a14d84bc2e7729a4dfd1 noise_scale_1p0/group_0085.npz +1714586f9df7b5343dfb0bf90d61e24ace6ab4f8b6cb23d481e278b5b66da363 noise_scale_1p0/group_0094.npz +# samples/edm: the archived test samples (--samples) +32c9b1067800fc29cb16e3d4690cdb9509cb904cd6e2a1e26dbfbdb001ef8902 default/group_0000.npz +ab1fec53800f89620473f27e70d1a40d5bfc124caafe47eb9ffb37b0df353a18 default/group_0007.npz +35bc946934825ccbb65f8f344140a0e091a03c9be76046da6f310c2a35f3fac2 default/group_0017.npz +497bc06058fdb730f8f05d31ef426dd13da8abf7550a32d83f4940469dae9cf5 default/group_0019.npz +f6008be676e440b253e1acc10bc94b2eb93217cb0f33e5d657f05cef7de938a8 default/group_0045.npz +10b392c68e7347b5fe887dcdf208b4330d8807ee39584801fc8622574359d0e2 default/group_0054.npz +31136aa90e92edeabe598ed946385e5158cf84bab87248b5ef59a2414733cc14 default/group_0071.npz +b86182932bb591900a518ee5b2bbe6bcd39b147c14d41ca7327da867542a9aa4 default/group_0078.npz +72873e563da1d3b053d99d04c6eb9dc9b8022535f41bf0146d6d746d227573fd default/group_0085.npz +7f4f4449608c92b9e051d26cc65684a3cccf27e82a87198d2f9a86cb811de8b0 default/group_0094.npz +# samples/cascaded_si: S* = 1, the same model and samples as direct_si +# checkpoints: the released checkpoints (--checkpoint), relative to the downloaded models repository +7ac8fbd0edb3658952fcafe1c8e10feec069bae40df67bdb73ec3b2c5fd27f3f passive-scalar/direct-si/model.pt +7faa6a6466ddcae430d77e501e0f76f3b92836baa415dfd606425ed863493b9f passive-scalar/cascaded-si-S02/model.pt +41be255ea7bb0beaef16caa37d2bc408c6c5b525a6c26958916f5e27618e801b passive-scalar/cascaded-si-S05/model.pt +f730fbbc1bfd042b771cee1500b863c191bfb3e41722a19859987683fb12180b passive-scalar/cascaded-si-S10/model.pt +1cbcf0ed47cc2423f20d97480c8ecdba416bab1f155c3d0bdf4a9ca2f34cdbb9 passive-scalar/cascaded-si-S20/model.pt +5d3c6ff0b138298dc2c0af3c116316909d6aa10d15ae9ee85c71906aa6d5e21c passive-scalar/cascaded-si-S40/model.pt +c74afc28febc20af1221ed8cc580b8214e6372a9d9f67d909b94cb8365a5c7bb passive-scalar/edm/model.pt diff --git a/code/records/passive_scalar/fid_adapter.json b/code/records/passive_scalar/fid_adapter.json new file mode 100644 index 0000000000000000000000000000000000000000..579f0cecdf2af37295c98527f4b266fda2b50024 --- /dev/null +++ b/code/records/passive_scalar/fid_adapter.json @@ -0,0 +1,35 @@ +{ + "adapter": "passive_scalar_coarse", + "field_kind": "signed", + "quantile": 0.9999, + "clip_c": 3.8628467146160794, + "range": [ + -3.8628467146160794, + 3.8628467146160794 + ], + "n_train_states": 65600, + "field_shape": [ + 128, + 128 + ], + "pixel_stride": 4, + "n_pixel_values": 67174400, + "map": "uint8 = rint((clip(w,-c,c)+c)/(2c)*255)", + "channels": "3 identical (grayscale scalar field), no colormap", + "source": { + "packaged_dataset_path": "", + "block": "vecs_train", + "ics": "canonical rows 0..79", + "members": "j < 20", + "snapshots": "all", + "pixel_stride": 4, + "n_states": 65600, + "rule": "training data only; validation, test and generated data never enter" + }, + "state": "196-D coarse spectral passive scalar (98 complex modes |k| <= 8, Hermitian half, theta_hat / N^2)", + "reconstruction": "exact inverse FFT of the |k| <= 8 modes on a 128x128 grid (alias-free)", + "n_disp": 128, + "k_c": 8, + "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" +} diff --git a/code/records/passive_scalar/runs/cascaded_si_S01/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S01/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73a8735484c8bdb8766cf1362be02596e5dafa3e --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S01/config.yaml @@ -0,0 +1,51 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 1.0 + n_segments: 1 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.003 + sigma_train: 1.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/cascaded_si_S02/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S02/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ebd40366b50276b1693e85b70422bc0212df934 --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S02/config.yaml @@ -0,0 +1,52 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.5 + - 1.0 + n_segments: 2 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + sigma_train: 4.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/cascaded_si_S05/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S05/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..88e0f466ff3eec5c374af336a79cd83b94791532 --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S05/config.yaml @@ -0,0 +1,55 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.2 + - 0.4 + - 0.6 + - 0.8 + - 1.0 + n_segments: 5 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/cascaded_si_S10/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S10/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3762bd955230a3650c825ae4293ac4d0c6eb1d94 --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S10/config.yaml @@ -0,0 +1,60 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.1 + - 0.2 + - 0.3 + - 0.4 + - 0.5 + - 0.6 + - 0.7000000000000001 + - 0.8 + - 0.9 + - 1.0 + n_segments: 10 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/cascaded_si_S20/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S20/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5b1a79a6c484293fc48b66d446f6c2f9e77ac4dd --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S20/config.yaml @@ -0,0 +1,70 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.05 + - 0.1 + - 0.15 + - 0.2 + - 0.25 + - 0.3 + - 0.35000000000000003 + - 0.4 + - 0.45 + - 0.5 + - 0.55 + - 0.6 + - 0.65 + - 0.7000000000000001 + - 0.75 + - 0.8 + - 0.85 + - 0.9 + - 0.9500000000000001 + - 1.0 + n_segments: 20 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/cascaded_si_S40/config.yaml b/code/records/passive_scalar/runs/cascaded_si_S40/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f6a88cb5e32583cd68b6adc02710e8fe460e2120 --- /dev/null +++ b/code/records/passive_scalar/runs/cascaded_si_S40/config.yaml @@ -0,0 +1,90 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.025 + - 0.05 + - 0.075 + - 0.1 + - 0.125 + - 0.15 + - 0.17500000000000002 + - 0.2 + - 0.225 + - 0.25 + - 0.275 + - 0.3 + - 0.325 + - 0.35000000000000003 + - 0.375 + - 0.4 + - 0.425 + - 0.45 + - 0.47500000000000003 + - 0.5 + - 0.525 + - 0.55 + - 0.5750000000000001 + - 0.6 + - 0.625 + - 0.65 + - 0.675 + - 0.7000000000000001 + - 0.725 + - 0.75 + - 0.775 + - 0.8 + - 0.8250000000000001 + - 0.85 + - 0.875 + - 0.9 + - 0.925 + - 0.9500000000000001 + - 0.975 + - 1.0 + n_segments: 40 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/runs/edm/config.yaml b/code/records/passive_scalar/runs/edm/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..64d7fa45384a48b75331c65b2d3ba22ab343fc59 --- /dev/null +++ b/code/records/passive_scalar/runs/edm/config.yaml @@ -0,0 +1,52 @@ +data: + name: passive_scalar_v1 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: edm + sigma_data: 1.0 + P_std: 1.2 + sampling: + sigma_min: 0.002 + sigma_max: 80.0 + rho: 7 + t_values: + - 0.0 + - 1.0 + n_steps: 20 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + lr_rampup_kimg: 0 + data_mode: fine_endpoint_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + P_mean: -1.2 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/passive_scalar/scores.json b/code/records/passive_scalar/scores.json new file mode 100644 index 0000000000000000000000000000000000000000..7c82426a50937e6fde2a036274a4881c01e3a880 --- /dev/null +++ b/code/records/passive_scalar/scores.json @@ -0,0 +1,135 @@ +{ + "cascaded_si_S01": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,9.301505129800887,10000,10000,0\n7,6.999326082120092,10000,10000,0\n17,7.988786201364974,10000,10000,0\n19,8.764889866887358,10000,10000,0\n45,9.098692218227313,10000,10000,0\n54,8.834131472642184,10000,10000,0\n71,9.85266447519433,10000,10000,0\n78,9.691036184475479,10000,10000,0\n85,8.188405591947514,10000,10000,0\n94,7.992308199325407,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 8.671174542198553, + "sem": 0.279008527535713, + "sd": 0.8823024336226594, + "n": 10 + }, + "csFID": { + "mean": 8.671174542198553, + "sem": 0.279008527535713, + "sd": 0.8823024336226594, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S02": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,18.644554989597793,10000,10000,0\n7,15.1181824349909,10000,10000,0\n17,15.397152953385728,10000,10000,0\n19,16.79350199533451,10000,10000,0\n45,26.08087482672596,10000,10000,0\n54,19.623429040131953,10000,10000,0\n71,24.385423667695306,10000,10000,0\n78,18.070791837513184,10000,10000,0\n85,17.3606518007266,10000,10000,0\n94,16.063036183189922,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 18.753759972929185, + "sem": 1.1744716754985969, + "sd": 3.714005541929739, + "n": 10 + }, + "csFID": { + "mean": 18.753759972929185, + "sem": 1.1744716754985969, + "sd": 3.714005541929739, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S05": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,33.008599859598746,10000,10000,0\n7,25.33035308765824,10000,10000,0\n17,27.029751435494006,10000,10000,0\n19,36.74900160599293,10000,10000,0\n45,40.36190972639173,10000,10000,0\n54,34.068754837466884,10000,10000,0\n71,40.39227733968147,10000,10000,0\n78,32.0267967976792,10000,10000,0\n85,29.60857688494275,10000,10000,0\n94,31.15855901162371,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 32.973458058652966, + "sem": 1.6136256279837646, + "sd": 5.102732275248231, + "n": 10 + }, + "csFID": { + "mean": 32.973458058652966, + "sem": 1.6136256279837646, + "sd": 5.102732275248231, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S10": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,39.40223997774541,10000,10000,0\n7,28.391132538688794,10000,10000,0\n17,33.66047867345088,10000,10000,0\n19,50.70222471216181,10000,10000,0\n45,46.69638187161013,10000,10000,0\n54,43.83844979928182,10000,10000,0\n71,51.80822874686169,10000,10000,0\n78,41.560247566900614,10000,10000,0\n85,40.191818714206875,10000,10000,0\n94,35.157606603664476,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 41.14088092045725, + "sem": 2.3611758048704354, + "sd": 7.46669349947187, + "n": 10 + }, + "csFID": { + "mean": 41.14088092045725, + "sem": 2.3611758048704354, + "sd": 7.46669349947187, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S20": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,55.523564819455146,10000,10000,0\n7,41.609638471631285,10000,10000,0\n17,52.22519174041531,10000,10000,0\n19,69.32357767903466,10000,10000,0\n45,67.22082543260842,10000,10000,0\n54,61.02834986212878,10000,10000,0\n71,65.35893570806462,10000,10000,0\n78,55.95454194548762,10000,10000,0\n85,47.95601314330446,10000,10000,0\n94,53.25112783062713,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 56.945176663275745, + "sem": 2.791231902897987, + "sd": 8.82665029088358, + "n": 10 + }, + "csFID": { + "mean": 56.945176663275745, + "sem": 2.791231902897987, + "sd": 8.82665029088358, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S40": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,75.8534706135913,10000,10000,0\n7,54.137818501023844,10000,10000,0\n17,57.83665352795538,10000,10000,0\n19,93.75504399670626,10000,10000,0\n45,96.6819009445029,10000,10000,0\n54,75.58561441286702,10000,10000,0\n71,93.10959723101439,10000,10000,0\n78,70.75124986539777,10000,10000,0\n85,67.16704531254987,10000,10000,0\n94,67.18434334928963,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 75.20627377548985, + "sem": 4.738450881551494, + "sd": 14.984297366535456, + "n": 10 + }, + "csFID": { + "mean": 75.20627377548985, + "sem": 4.738450881551494, + "sd": 14.984297366535456, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "edm": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n0,8.214064883145397,10000,10000,0\n7,6.340900551288996,10000,10000,0\n17,6.946395490153236,10000,10000,0\n19,8.17481947293578,10000,10000,0\n45,8.22899421983783,10000,10000,0\n54,7.849720606557412,10000,10000,0\n71,9.988332341720536,10000,10000,0\n78,8.283878193030546,10000,10000,0\n85,7.2782537114106844,10000,10000,0\n94,7.506518008403248,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 7.881187747848367, + "sem": 0.31040175244938206, + "sd": 0.9815765274477966, + "n": 10 + }, + "csFID": { + "mean": 7.881187747848367, + "sem": 0.31040175244938206, + "sd": 0.9815765274477966, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + } +} diff --git a/code/records/passive_scalar/units.json b/code/records/passive_scalar/units.json new file mode 100644 index 0000000000000000000000000000000000000000..7084f86b7f9b1837c56b3417d3afcb3e21a2a05b --- /dev/null +++ b/code/records/passive_scalar/units.json @@ -0,0 +1,346 @@ +{ + "system": "passive_scalar", + "paper_name": "Passive scalar", + "dataset": "passive_scalar_v1", + "test_gids": [ + 0, + 7, + 17, + 19, + 45, + 54, + 71, + 78, + 85, + 94 + ], + "val_gids": [ + 25, + 32, + 43, + 44, + 52, + 55, + 58, + 69, + 89, + 93 + ], + "fine_horizon_idx": 40, + "x0_from_truth_bank": false, + "s_star": 1, + "table2": { + "direct_si": "cascaded_si_S01", + "edm": "edm", + "cascaded_si": "cascaded_si_S01" + }, + "validation_energy_distance_by_S": { + "1": 0.17306941650503, + "2": 0.5197569839991051, + "5": 0.8252763928969088, + "10": 1.3195428451076956, + "20": 2.128299478057891, + "40": 3.269881230277261 + }, + "units": [ + { + "key": "cascaded_si_S01", + "S": 1, + "method": "cascaded_si", + "paper_names": [ + "direct_si", + "cascaded_si" + ], + "config": "records/passive_scalar/runs/cascaded_si_S01/config.yaml", + "selection": { + "lr": 0.003, + "sigma": 1.0, + "selected_update": 6000, + "validation_energy_distance": 0.17306941650503 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 40 + }, + "checkpoint": { + "file": "passive-scalar/direct-si/model.pt", + "sha256": "7ac8fbd0edb3658952fcafe1c8e10feec069bae40df67bdb73ec3b2c5fd27f3f", + "bytes": 18587370 + }, + "expected_csfid": { + "per_condition": { + "0": 9.301505129800887, + "7": 6.999326082120092, + "17": 7.988786201364974, + "19": 8.764889866887358, + "45": 9.098692218227313, + "54": 8.834131472642184, + "71": 9.85266447519433, + "78": 9.691036184475479, + "85": 8.188405591947514, + "94": 7.992308199325407 + }, + "n_gen_nonfinite": {}, + "mean": 8.671174542198553, + "sem": 0.279008527535713, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S02", + "S": 2, + "method": "cascaded_si", + "paper_names": [], + "config": "records/passive_scalar/runs/cascaded_si_S02/config.yaml", + "selection": { + "lr": 0.0003, + "sigma": 4.0, + "selected_update": 3000, + "validation_energy_distance": 0.5197569839991051 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 80 + }, + "checkpoint": { + "file": "passive-scalar/cascaded-si-S02/model.pt", + "sha256": "7faa6a6466ddcae430d77e501e0f76f3b92836baa415dfd606425ed863493b9f", + "bytes": 18587370 + }, + "expected_csfid": { + "per_condition": { + "0": 18.644554989597793, + "7": 15.1181824349909, + "17": 15.397152953385728, + "19": 16.79350199533451, + "45": 26.08087482672596, + "54": 19.623429040131953, + "71": 24.385423667695306, + "78": 18.070791837513184, + "85": 17.3606518007266, + "94": 16.063036183189922 + }, + "n_gen_nonfinite": {}, + "mean": 18.753759972929185, + "sem": 1.1744716754985969, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S05", + "S": 5, + "method": "cascaded_si", + "paper_names": [], + "config": "records/passive_scalar/runs/cascaded_si_S05/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 8000, + "validation_energy_distance": 0.8252763928969088 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 200 + }, + "checkpoint": { + "file": "passive-scalar/cascaded-si-S05/model.pt", + "sha256": "41be255ea7bb0beaef16caa37d2bc408c6c5b525a6c26958916f5e27618e801b", + "bytes": 18587434 + }, + "expected_csfid": { + "per_condition": { + "0": 33.008599859598746, + "7": 25.33035308765824, + "17": 27.029751435494006, + "19": 36.74900160599293, + "45": 40.36190972639173, + "54": 34.068754837466884, + "71": 40.39227733968147, + "78": 32.0267967976792, + "85": 29.60857688494275, + "94": 31.15855901162371 + }, + "n_gen_nonfinite": {}, + "mean": 32.973458058652966, + "sem": 1.6136256279837646, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S10", + "S": 10, + "method": "cascaded_si", + "paper_names": [], + "config": "records/passive_scalar/runs/cascaded_si_S10/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 11000, + "validation_energy_distance": 1.3195428451076956 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 800 + }, + "checkpoint": { + "file": "passive-scalar/cascaded-si-S10/model.pt", + "sha256": "f730fbbc1bfd042b771cee1500b863c191bfb3e41722a19859987683fb12180b", + "bytes": 18587434 + }, + "expected_csfid": { + "per_condition": { + "0": 39.40223997774541, + "7": 28.391132538688794, + "17": 33.66047867345088, + "19": 50.70222471216181, + "45": 46.69638187161013, + "54": 43.83844979928182, + "71": 51.80822874686169, + "78": 41.560247566900614, + "85": 40.191818714206875, + "94": 35.157606603664476 + }, + "n_gen_nonfinite": {}, + "mean": 41.14088092045725, + "sem": 2.3611758048704354, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S20", + "S": 20, + "method": "cascaded_si", + "paper_names": [], + "config": "records/passive_scalar/runs/cascaded_si_S20/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 17000, + "validation_energy_distance": 2.128299478057891 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 1600 + }, + "checkpoint": { + "file": "passive-scalar/cascaded-si-S20/model.pt", + "sha256": "1cbcf0ed47cc2423f20d97480c8ecdba416bab1f155c3d0bdf4a9ca2f34cdbb9", + "bytes": 18587562 + }, + "expected_csfid": { + "per_condition": { + "0": 55.523564819455146, + "7": 41.609638471631285, + "17": 52.22519174041531, + "19": 69.32357767903466, + "45": 67.22082543260842, + "54": 61.02834986212878, + "71": 65.35893570806462, + "78": 55.95454194548762, + "85": 47.95601314330446, + "94": 53.25112783062713 + }, + "n_gen_nonfinite": {}, + "mean": 56.945176663275745, + "sem": 2.791231902897987, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S40", + "S": 40, + "method": "cascaded_si", + "paper_names": [], + "config": "records/passive_scalar/runs/cascaded_si_S40/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 16.0, + "selected_update": 13000, + "validation_energy_distance": 3.269881230277261 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 3200 + }, + "checkpoint": { + "file": "passive-scalar/cascaded-si-S40/model.pt", + "sha256": "5d3c6ff0b138298dc2c0af3c116316909d6aa10d15ae9ee85c71906aa6d5e21c", + "bytes": 18587754 + }, + "expected_csfid": { + "per_condition": { + "0": 75.8534706135913, + "7": 54.137818501023844, + "17": 57.83665352795538, + "19": 93.75504399670626, + "45": 96.6819009445029, + "54": 75.58561441286702, + "71": 93.10959723101439, + "78": 70.75124986539777, + "85": 67.16704531254987, + "94": 67.18434334928963 + }, + "n_gen_nonfinite": {}, + "mean": 75.20627377548985, + "sem": 4.738450881551494, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "edm", + "S": 1, + "method": "edm", + "paper_names": [ + "edm" + ], + "config": "records/passive_scalar/runs/edm/config.yaml", + "selection": { + "lr": 0.001, + "P_mean": -1.2, + "selected_update": 16000, + "validation_energy_distance": 0.10606232925071438 + }, + "sampler": { + "n_steps": 40, + "variant": "default", + "nfe": 79 + }, + "checkpoint": { + "file": "passive-scalar/edm/model.pt", + "sha256": "c74afc28febc20af1221ed8cc580b8214e6372a9d9f67d909b94cb8365a5c7bb", + "bytes": 18587370 + }, + "expected_csfid": { + "per_condition": { + "0": 8.214064883145397, + "7": 6.340900551288996, + "17": 6.946395490153236, + "19": 8.17481947293578, + "45": 8.22899421983783, + "54": 7.849720606557412, + "71": 9.988332341720536, + "78": 8.283878193030546, + "85": 7.2782537114106844, + "94": 7.506518008403248 + }, + "n_gen_nonfinite": {}, + "mean": 7.881187747848367, + "sem": 0.31040175244938206, + "n": 10 + }, + "selection_flags": [] + } + ] +} diff --git a/code/records/sqg/assets.sha256 b/code/records/sqg/assets.sha256 new file mode 100644 index 0000000000000000000000000000000000000000..cae84e6db132ca198c9df043f99775c2e959d3e7 --- /dev/null +++ b/code/records/sqg/assets.sha256 @@ -0,0 +1,59 @@ +# sha256 of the inputs of the Table 2 entries; paths are relative to the directory named in the heading +# data: the dataset package (--data) +510b027125b30065489edbae1d0ac23c3cb15765630be4113c119af993c0dee4 trajectories_fine.npz +fb482ffd0353ba47dfca67a34466f7026278962ca9c8fb44839aa039935d08e2 split.json +46f94fc794069c9ad985d23ae4cc2b8ed866fe13d6a0a5463f1c8384352898f8 normalization.npz +# truth_bank: the truth banks of the test conditions (--truth-bank) +72af84372082450a9d2c228879aec31879378ae308e912bc7a4f4e5a05847086 group_0100.npz +c5c37a9973f23b88a61c1589da467020e4b59037d7e67945d80cadc3152c52f0 group_0101.npz +4962c55bc4196a779441762f74517c886d75901056a4c6944b13c4fac32987ec group_0102.npz +2191a3cc2c01d66182bfcfe39e682bf321c00634f3f8f0f6f4f5c28727696247 group_0103.npz +49b85a98af0f3946a7d9dc46d22f9410a2682742ead691ac6430419ceb2575f2 group_0104.npz +b451244846039c6db72f1e635ce64ef6e0d876b89a045d796a203ae8a188c991 group_0105.npz +44c8f010d6f155f2df554055862b2647118cc4b382e61ab5ab9bc0e43e26c9df group_0106.npz +ab4918c1de1d5de08f4dd70bdcd5e7714bc838d54affc2f6874f2ae2405597b3 group_0107.npz +5e07235b934cee9f968073d679259db3ac65d7f5c0ffb21b6c1e328b4d7f64f6 group_0108.npz +dc9144acbe3965e19257fcafeb4cb86b7d20b1c8ac4ddde9402352c178fe739a group_0109.npz +# inception: the Inception-v3 weights in $CLEAN_FID_CACHE +f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4 inception-2015-12-05.pt +# samples/direct_si: the archived test samples (--samples) +f94895658eaf3020e5f982fee696e743f8a3ad195376261958500c910f88af59 noise_scale_1p0/group_0100.npz +192fa6429a2bbf28663c3bcdece645e2524ee6b760e0fb61625ca0fcb3c9ed23 noise_scale_1p0/group_0101.npz +660c2daa79b45bffc2116eb50702d8322f461f7a8264d2bc4081ea7fc8c9970d noise_scale_1p0/group_0102.npz +58348ae771f43810e79d840965a92927a0af627ed2c4036c8288b68419f06193 noise_scale_1p0/group_0103.npz +47ce8ab6495bf0348d2bd31887c828c823c94e80520d0ae204933ace802c8c2c noise_scale_1p0/group_0104.npz +83a7b7af9b74b1c651cf3dfd5cee26ac381ac66b1a6223948391840bf62255be noise_scale_1p0/group_0105.npz +792d5221c31f6634cc398a4dbcc454876a27d4d4687a3a15e2dffbe87f4880d6 noise_scale_1p0/group_0106.npz +92671e6e63235942ecea84d51820c1235e192ca04ce93c0bb717f726bfcdc1a5 noise_scale_1p0/group_0107.npz +a7f0781ff67bbe247b7598866914e39387be5f67eac2eab7a73237715f15ab31 noise_scale_1p0/group_0108.npz +a26263a3c16f7d52a123fae422326db3797607103111001606d35244aa130606 noise_scale_1p0/group_0109.npz +# samples/edm: the archived test samples (--samples) +79c9161aacdf0091ecddbfda48b83c436a5ee677092ea8fc647546fe6f127c5d default/group_0100.npz +e292619a79ef7901ffd669e8ae0fd2a1833fa7ccea688978d8e10cfb53088bea default/group_0101.npz +90eb1472eafd3f70a354fcf8733acbbb3529f4c3137f5971e5c1d4bd5b9c23a6 default/group_0102.npz +cb2b0c8fb992466a4eca71f9acac5a2aae62cea19e6b5e0c9b42dce17c278416 default/group_0103.npz +78c24d081e91ae39692e83548533090fa092b152ad0105f2cab1cf86e1866f20 default/group_0104.npz +0bdbf136760efdbc6f62f2c82676d1eb8b8b4b46ff55df0377cf12ff6684c3fa default/group_0105.npz +bc36aaaf7d516d2a2a48da98281d2bbf266629c48b7050347ab3ff95856e4e3d default/group_0106.npz +a08412a19c19dabe9b5303594541f37b7d0fc9e3d234c19d2fc40404b150c2dc default/group_0107.npz +e12c8eb8a2b763bad217f7174d1ba96f5e19dae65e447f008a365a03322e7284 default/group_0108.npz +e68913c05dfd8abf9480a8f2df0af2c12c440bc8198392c9e49b56218887d0ab default/group_0109.npz +# samples/cascaded_si: the archived test samples (--samples) +450176f1bf2e25743b5f17e0514fc58164eb6b0ce84fcc3641908185b3bfe890 noise_scale_1p0/group_0100.npz +d55900a5382059ed02a4a0c91fb5ca9cd7bb1a7839f54a23f12cf9b61ea8a53a noise_scale_1p0/group_0101.npz +2b1ba2ba490f65fe3e7fb7832ac8842a356a1ffe72a75e928db43872c48f7b1b noise_scale_1p0/group_0102.npz +ff653d19d67c9b9c70e0a9e7e6839ad81d9ee66b21813d4070889a0e42d39e05 noise_scale_1p0/group_0103.npz +c9d5d8d717b277af12bdf0c194af91198114787c48bec38803bc0d8d510af984 noise_scale_1p0/group_0104.npz +f93d4e260db29103fb0d0e5e7af486534de5eb65831c0c7d62eb5121ccd43115 noise_scale_1p0/group_0105.npz +e5ec180a96303ec62a8a543662acdc0609bc6d5916cd9527dab65ea2bdd09817 noise_scale_1p0/group_0106.npz +b3f2eb65b12d07d168893ae38379b0128fec963ed42c084f6c0fe6eddd768a7b noise_scale_1p0/group_0107.npz +79ac519c839c7a435f8739b13137da714f3a38838d91cb91d2158b4a98d5ccf4 noise_scale_1p0/group_0108.npz +a62df35501a20ad7d6027480bd44b67cacd2a9324b1ed1fb5bd995073fdb29ca noise_scale_1p0/group_0109.npz +# checkpoints: the released checkpoints (--checkpoint), relative to the downloaded models repository +043edcf9b92836ab185e2400f38af9da51217ee9314513546c560139c833dd4b sqg/direct-si/model.pt +7897199322040cd8e416c47d52f9c627a897ab3f61315c159c3fc21524dfbf95 sqg/cascaded-si-S02/model.pt +a37023ac4e537ad368207189137a505301f460b1e745288752686cd48837be3b sqg/cascaded-si-S05/model.pt +aa94d2608a7aa21ec81cb421692ce4d8cd456d9925f428692003238f55d4cb8d sqg/cascaded-si-S10/model.pt +52cb13be41c51528ca28ba9b2e1ea00c5f8fab9e86c68812d94432a009df4883 sqg/cascaded-si-S20/model.pt +9add06517f1ff2a6f0f0a22bc52553323926f7102524ae13c5feac13a7a264ba sqg/cascaded-si-S40/model.pt +7387034ad6413ff6073a115602533645c6f31b8005ee1a1a89ea9579e6297a99 sqg/edm/model.pt diff --git a/code/records/sqg/fid_adapter.json b/code/records/sqg/fid_adapter.json new file mode 100644 index 0000000000000000000000000000000000000000..5e14f8f593bc793667b53b1efd657ea1ce39745a --- /dev/null +++ b/code/records/sqg/fid_adapter.json @@ -0,0 +1,35 @@ +{ + "adapter": "passive_scalar_coarse", + "field_kind": "signed", + "quantile": 0.9999, + "clip_c": 1.1251138597965493, + "range": [ + -1.1251138597965493, + 1.1251138597965493 + ], + "n_train_states": 65600, + "field_shape": [ + 128, + 128 + ], + "pixel_stride": 4, + "n_pixel_values": 67174400, + "map": "uint8 = rint((clip(w,-c,c)+c)/(2c)*255)", + "channels": "3 identical (grayscale scalar field), no colormap", + "source": { + "packaged_dataset_path": "", + "block": "vecs_train", + "ics": "canonical rows 0..79", + "members": "j < 20", + "snapshots": "all", + "pixel_stride": 4, + "n_states": 65600, + "rule": "training data only; validation, test and generated data never enter" + }, + "state": "196-D coarse spectral passive scalar (98 complex modes |k| <= 8, Hermitian half, theta_hat / N^2)", + "reconstruction": "exact inverse FFT of the |k| <= 8 modes on a 128x128 grid (alias-free)", + "n_disp": 128, + "k_c": 8, + "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" +} diff --git a/code/records/sqg/runs/cascaded_si_S01/config.yaml b/code/records/sqg/runs/cascaded_si_S01/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c9210b287f9edac185af042098783e02b179d982 --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S01/config.yaml @@ -0,0 +1,51 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 1.0 + n_segments: 1 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/cascaded_si_S02/config.yaml b/code/records/sqg/runs/cascaded_si_S02/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ce6a050c230d8f328ae87499fbab0b958c7e25b4 --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S02/config.yaml @@ -0,0 +1,52 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.5 + - 1.0 + n_segments: 2 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 2.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/cascaded_si_S05/config.yaml b/code/records/sqg/runs/cascaded_si_S05/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8095f2246726acf42bf4664e664feca29ee7b484 --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S05/config.yaml @@ -0,0 +1,55 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.2 + - 0.4 + - 0.6000000000000001 + - 0.8 + - 1.0 + n_segments: 5 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + sigma_train: 1.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/cascaded_si_S10/config.yaml b/code/records/sqg/runs/cascaded_si_S10/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..41064e204d068c6ef30d518a3ff3b6879108d34e --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S10/config.yaml @@ -0,0 +1,60 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.1 + - 0.2 + - 0.30000000000000004 + - 0.4 + - 0.5 + - 0.6000000000000001 + - 0.7000000000000001 + - 0.8 + - 0.9 + - 1.0 + n_segments: 10 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/cascaded_si_S20/config.yaml b/code/records/sqg/runs/cascaded_si_S20/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b3524c3b4494b03d4b76e82209a072c20d81c11e --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S20/config.yaml @@ -0,0 +1,70 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.05 + - 0.1 + - 0.15000000000000002 + - 0.2 + - 0.25 + - 0.30000000000000004 + - 0.35000000000000003 + - 0.4 + - 0.45 + - 0.5 + - 0.55 + - 0.6000000000000001 + - 0.65 + - 0.7000000000000001 + - 0.75 + - 0.8 + - 0.8500000000000001 + - 0.9 + - 0.9500000000000001 + - 1.0 + n_segments: 20 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/cascaded_si_S40/config.yaml b/code/records/sqg/runs/cascaded_si_S40/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c039cc2941cfe8daec979145c7b94adba940a54d --- /dev/null +++ b/code/records/sqg/runs/cascaded_si_S40/config.yaml @@ -0,0 +1,90 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.025 + - 0.05 + - 0.07500000000000001 + - 0.1 + - 0.125 + - 0.15000000000000002 + - 0.17500000000000002 + - 0.2 + - 0.225 + - 0.25 + - 0.275 + - 0.30000000000000004 + - 0.325 + - 0.35000000000000003 + - 0.375 + - 0.4 + - 0.42500000000000004 + - 0.45 + - 0.47500000000000003 + - 0.5 + - 0.525 + - 0.55 + - 0.5750000000000001 + - 0.6000000000000001 + - 0.625 + - 0.65 + - 0.675 + - 0.7000000000000001 + - 0.7250000000000001 + - 0.75 + - 0.775 + - 0.8 + - 0.8250000000000001 + - 0.8500000000000001 + - 0.875 + - 0.9 + - 0.925 + - 0.9500000000000001 + - 0.9750000000000001 + - 1.0 + n_segments: 40 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/runs/edm/config.yaml b/code/records/sqg/runs/edm/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cd3b55623a3d5a78480d3994b098eaa472c88b2f --- /dev/null +++ b/code/records/sqg/runs/edm/config.yaml @@ -0,0 +1,52 @@ +data: + name: sqg_v3_candidate_v2 + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: edm + sigma_data: 1.0 + P_std: 1.2 + sampling: + sigma_min: 0.002 + sigma_max: 80.0 + rho: 7 + t_values: + - 0.0 + - 1.0 + n_steps: 20 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + lr_rampup_kimg: 0 + data_mode: fine_endpoint_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + P_mean: -3.2 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/sqg/scores.json b/code/records/sqg/scores.json new file mode 100644 index 0000000000000000000000000000000000000000..acc9ff233ec0e16c26f75d1a6593196ff1c67412 --- /dev/null +++ b/code/records/sqg/scores.json @@ -0,0 +1,135 @@ +{ + "cascaded_si_S01": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,174.77764334659662,10000,10000,0\n101,190.22545507837197,10000,10000,0\n102,32.19350595285715,10000,10000,0\n103,259.54823436220556,10000,10000,0\n104,187.9102583014046,10000,10000,0\n105,144.12298637015994,10000,10000,0\n106,69.83315328274227,10000,10000,0\n107,180.1658685174503,10000,10000,0\n108,89.04146761014198,10000,10000,0\n109,56.93464752886899,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 138.47532203507996, + "sem": 23.08644501940587, + "sd": 73.00574933757272, + "n": 10 + }, + "csFID": { + "mean": 138.47532203507996, + "sem": 23.08644501940587, + "sd": 73.00574933757272, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S02": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,108.64610599112586,10000,10000,0\n101,148.213528710336,10000,10000,0\n102,12.249737372758716,10000,10000,0\n103,166.43311996183718,10000,10000,0\n104,101.97250494419791,10000,10000,0\n105,81.0061190095802,10000,10000,0\n106,55.60906851984424,10000,10000,0\n107,97.34060878782299,10000,10000,0\n108,41.70360215641068,10000,10000,0\n109,19.348780618531293,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 83.2523176072445, + "sem": 16.314034007096794, + "sd": 51.58950528786942, + "n": 10 + }, + "csFID": { + "mean": 83.2523176072445, + "sem": 16.314034007096794, + "sd": 51.58950528786942, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S05": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,47.44390232255884,10000,10000,0\n101,97.64966657928815,10000,10000,0\n102,13.897894548199758,10000,10000,0\n103,78.32269942558258,10000,10000,0\n104,53.961217997417634,10000,10000,0\n105,33.35838313262059,10000,10000,0\n106,23.07217369951556,10000,10000,0\n107,50.69801284998158,10000,10000,0\n108,18.393931957047016,10000,10000,0\n109,18.889878019767053,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 43.56877605319787, + "sem": 8.790125902695719, + "sd": 27.796818772162084, + "n": 10 + }, + "csFID": { + "mean": 43.56877605319787, + "sem": 8.790125902695719, + "sd": 27.796818772162084, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S10": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,64.94213015904387,10000,10000,0\n101,106.07373925577525,10000,10000,0\n102,28.498497011053587,10000,10000,0\n103,78.46324992969,10000,10000,0\n104,53.92776605817534,10000,10000,0\n105,32.34964086870488,10000,10000,0\n106,34.21413790129114,10000,10000,0\n107,57.82609257156645,10000,10000,0\n108,26.65452682986927,10000,10000,0\n109,29.50001116074702,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 51.244979174591684, + "sem": 8.318097628706868, + "sd": 26.304134306359302, + "n": 10 + }, + "csFID": { + "mean": 51.244979174591684, + "sem": 8.318097628706868, + "sd": 26.304134306359302, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S20": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,106.47051789719507,10000,10000,0\n101,98.1515112149221,10000,9999,1\n102,67.8273652273723,10000,10000,0\n103,142.1986562667293,10000,10000,0\n104,94.2926867903455,10000,10000,0\n105,88.04353243050758,10000,10000,0\n106,60.465059315902636,10000,10000,0\n107,90.76624161381369,10000,10000,0\n108,70.9330081827298,10000,10000,0\n109,71.66828988572928,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 89.08168688252472, + "sem": 7.557464374313741, + "sd": 23.898800758410744, + "n": 10 + }, + "csFID": { + "mean": 89.08168688252472, + "sem": 7.557464374313741, + "sd": 23.898800758410744, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S40": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,103.311354260536,10000,10000,0\n101,138.43431083940442,10000,10000,0\n102,49.069336701801376,10000,10000,0\n103,151.12233778182332,10000,10000,0\n104,82.95322414622876,10000,10000,0\n105,61.61246415076869,10000,10000,0\n106,55.004062919050455,10000,10000,0\n107,116.30086955370848,10000,10000,0\n108,52.43701041972753,10000,10000,0\n109,59.53891646067143,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 86.97838872337205, + "sem": 11.997151993142262, + "sd": 37.93832573355832, + "n": 10 + }, + "csFID": { + "mean": 86.97838872337205, + "sem": 11.997151993142262, + "sd": 37.93832573355832, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "edm": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n100,138.3704066650215,10000,10000,0\n101,177.4114321188299,10000,10000,0\n102,26.592048990966646,10000,10000,0\n103,231.37837575590652,10000,10000,0\n104,173.26949861737825,10000,10000,0\n105,126.41908743961136,10000,10000,0\n106,57.35276085330884,10000,10000,0\n107,151.24564285169487,10000,10000,0\n108,68.33837760817943,10000,10000,0\n109,35.8617353737792,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 118.62393662746763, + "sem": 21.67250652363767, + "sd": 68.53448321955287, + "n": 10 + }, + "csFID": { + "mean": 118.62393662746763, + "sem": 21.67250652363767, + "sd": 68.53448321955287, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + } +} diff --git a/code/records/sqg/units.json b/code/records/sqg/units.json new file mode 100644 index 0000000000000000000000000000000000000000..106a57005fc315c983204927616efb04cfb578d1 --- /dev/null +++ b/code/records/sqg/units.json @@ -0,0 +1,349 @@ +{ + "system": "sqg", + "paper_name": "Surface quasi-geostrophy (SQG)", + "dataset": "sqg_v3_candidate_v2", + "test_gids": [ + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109 + ], + "val_gids": [ + 3, + 6, + 9, + 10, + 11, + 17, + 28, + 50, + 70, + 73 + ], + "fine_horizon_idx": 40, + "x0_from_truth_bank": false, + "s_star": 10, + "table2": { + "direct_si": "cascaded_si_S01", + "edm": "edm", + "cascaded_si": "cascaded_si_S10" + }, + "validation_energy_distance_by_S": { + "1": 8.574122041243665, + "2": 6.64245033209117, + "5": 4.160166093654704, + "10": 3.2081736575575603, + "20": 3.9763564860504417, + "40": 4.650876671501445 + }, + "units": [ + { + "key": "cascaded_si_S01", + "S": 1, + "method": "cascaded_si", + "paper_names": [ + "direct_si" + ], + "config": "records/sqg/runs/cascaded_si_S01/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 16.0, + "selected_update": 4000, + "validation_energy_distance": 8.574122041243665 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 40 + }, + "checkpoint": { + "file": "sqg/direct-si/model.pt", + "sha256": "043edcf9b92836ab185e2400f38af9da51217ee9314513546c560139c833dd4b", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "100": 174.77764334659662, + "101": 190.22545507837197, + "102": 32.19350595285715, + "103": 259.54823436220556, + "104": 187.9102583014046, + "105": 144.12298637015994, + "106": 69.83315328274227, + "107": 180.1658685174503, + "108": 89.04146761014198, + "109": 56.93464752886899 + }, + "n_gen_nonfinite": {}, + "mean": 138.47532203507996, + "sem": 23.08644501940587, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S02", + "S": 2, + "method": "cascaded_si", + "paper_names": [], + "config": "records/sqg/runs/cascaded_si_S02/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 2.0, + "selected_update": 500, + "validation_energy_distance": 6.64245033209117 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 80 + }, + "checkpoint": { + "file": "sqg/cascaded-si-S02/model.pt", + "sha256": "7897199322040cd8e416c47d52f9c627a897ab3f61315c159c3fc21524dfbf95", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "100": 108.64610599112586, + "101": 148.213528710336, + "102": 12.249737372758716, + "103": 166.43311996183718, + "104": 101.97250494419791, + "105": 81.0061190095802, + "106": 55.60906851984424, + "107": 97.34060878782299, + "108": 41.70360215641068, + "109": 19.348780618531293 + }, + "n_gen_nonfinite": {}, + "mean": 83.2523176072445, + "sem": 16.314034007096794, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S05", + "S": 5, + "method": "cascaded_si", + "paper_names": [], + "config": "records/sqg/runs/cascaded_si_S05/config.yaml", + "selection": { + "lr": 0.0003, + "sigma": 1.0, + "selected_update": 8000, + "validation_energy_distance": 4.160166093654704 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 400 + }, + "checkpoint": { + "file": "sqg/cascaded-si-S05/model.pt", + "sha256": "a37023ac4e537ad368207189137a505301f460b1e745288752686cd48837be3b", + "bytes": 18578218 + }, + "expected_csfid": { + "per_condition": { + "100": 47.44390232255884, + "101": 97.64966657928815, + "102": 13.897894548199758, + "103": 78.32269942558258, + "104": 53.961217997417634, + "105": 33.35838313262059, + "106": 23.07217369951556, + "107": 50.69801284998158, + "108": 18.393931957047016, + "109": 18.889878019767053 + }, + "n_gen_nonfinite": {}, + "mean": 43.56877605319787, + "sem": 8.790125902695719, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S10", + "S": 10, + "method": "cascaded_si", + "paper_names": [ + "cascaded_si" + ], + "config": "records/sqg/runs/cascaded_si_S10/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 16.0, + "selected_update": 11000, + "validation_energy_distance": 3.2081736575575603 + }, + "sampler": { + "n_steps": 160, + "variant": "noise_scale_1p0", + "nfe": 1600 + }, + "checkpoint": { + "file": "sqg/cascaded-si-S10/model.pt", + "sha256": "aa94d2608a7aa21ec81cb421692ce4d8cd456d9925f428692003238f55d4cb8d", + "bytes": 18578282 + }, + "expected_csfid": { + "per_condition": { + "100": 64.94213015904387, + "101": 106.07373925577525, + "102": 28.498497011053587, + "103": 78.46324992969, + "104": 53.92776605817534, + "105": 32.34964086870488, + "106": 34.21413790129114, + "107": 57.82609257156645, + "108": 26.65452682986927, + "109": 29.50001116074702 + }, + "n_gen_nonfinite": {}, + "mean": 51.244979174591684, + "sem": 8.318097628706868, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S20", + "S": 20, + "method": "cascaded_si", + "paper_names": [], + "config": "records/sqg/runs/cascaded_si_S20/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 16.0, + "selected_update": 17000, + "validation_energy_distance": 3.9763564860504417 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 800 + }, + "checkpoint": { + "file": "sqg/cascaded-si-S20/model.pt", + "sha256": "52cb13be41c51528ca28ba9b2e1ea00c5f8fab9e86c68812d94432a009df4883", + "bytes": 18578346 + }, + "expected_csfid": { + "per_condition": { + "100": 106.47051789719507, + "101": 98.1515112149221, + "102": 67.8273652273723, + "103": 142.1986562667293, + "104": 94.2926867903455, + "105": 88.04353243050758, + "106": 60.465059315902636, + "107": 90.76624161381369, + "108": 70.9330081827298, + "109": 71.66828988572928 + }, + "n_gen_nonfinite": { + "101": 1 + }, + "mean": 89.08168688252472, + "sem": 7.557464374313741, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S40", + "S": 40, + "method": "cascaded_si", + "paper_names": [], + "config": "records/sqg/runs/cascaded_si_S40/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 16.0, + "selected_update": 11000, + "validation_energy_distance": 4.650876671501445 + }, + "sampler": { + "n_steps": 640, + "variant": "noise_scale_1p0", + "nfe": 25600 + }, + "checkpoint": { + "file": "sqg/cascaded-si-S40/model.pt", + "sha256": "9add06517f1ff2a6f0f0a22bc52553323926f7102524ae13c5feac13a7a264ba", + "bytes": 18578538 + }, + "expected_csfid": { + "per_condition": { + "100": 103.311354260536, + "101": 138.43431083940442, + "102": 49.069336701801376, + "103": 151.12233778182332, + "104": 82.95322414622876, + "105": 61.61246415076869, + "106": 55.004062919050455, + "107": 116.30086955370848, + "108": 52.43701041972753, + "109": 59.53891646067143 + }, + "n_gen_nonfinite": {}, + "mean": 86.97838872337205, + "sem": 11.997151993142262, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "edm", + "S": 1, + "method": "edm", + "paper_names": [ + "edm" + ], + "config": "records/sqg/runs/edm/config.yaml", + "selection": { + "lr": 0.0003, + "P_mean": -3.2, + "selected_update": 4000, + "validation_energy_distance": 8.014866943636475 + }, + "sampler": { + "n_steps": 40, + "variant": "default", + "nfe": 79 + }, + "checkpoint": { + "file": "sqg/edm/model.pt", + "sha256": "7387034ad6413ff6073a115602533645c6f31b8005ee1a1a89ea9579e6297a99", + "bytes": 18578154 + }, + "expected_csfid": { + "per_condition": { + "100": 138.3704066650215, + "101": 177.4114321188299, + "102": 26.592048990966646, + "103": 231.37837575590652, + "104": 173.26949861737825, + "105": 126.41908743961136, + "106": 57.35276085330884, + "107": 151.24564285169487, + "108": 68.33837760817943, + "109": 35.8617353737792 + }, + "n_gen_nonfinite": {}, + "mean": 118.62393662746763, + "sem": 21.67250652363767, + "n": 10 + }, + "selection_flags": [] + } + ] +} diff --git a/code/records/waves/assets.sha256 b/code/records/waves/assets.sha256 new file mode 100644 index 0000000000000000000000000000000000000000..2b348f94fb5768191159e00a7dfb96fee90f8745 --- /dev/null +++ b/code/records/waves/assets.sha256 @@ -0,0 +1,61 @@ +# sha256 of the inputs of the Table 2 entries; paths are relative to the directory named in the heading +# data: the dataset package (--data) +4211aa62b2e6e61066519396e20bdf0a3b002b02a27641dd39031cbb680c9e53 trajectories_fine.npz +ccb5fb3fba13a8f0e5528c3757fb54ba4044ca180328da99a53b58fbf6b8e88c split.json +23123f0b0547197bbd6275ac4a34d894a91ab98e4f1810b80416f6b891d67cb7 normalization.npz +55e9a00ddcf36eece80a88bf1f012026f4934989ae8a2969b1cdcb5c78228d2a u_norm/normalization.npz +54441e3917067932261ac67dacb3b15dd76458147d9ff41775d1148b20b06ddf v_norm/normalization.npz +# truth_bank: the truth banks of the test conditions (--truth-bank) +d977cadfd45ec116f1d8c8f0e26d2e8943f00c62acf2a5ad9ca16a45b9d26e2c group_0016.npz +e19de5ab4a05cb5cc774c0a97e052e9cb01085715a035bb8c5801f5659c5249a group_0031.npz +f546597cec69e0b63e71865889e7d84d167979952f15ae9fa77372dfd04a518f group_0032.npz +c59b0648942adabe26476342aabc7f65bdfdf8b519f12edae9f657860c87d34c group_0043.npz +b73a3da843b2d7a9c51e5a241d77608440e028b5f8e3362c1678be8abd01d7c5 group_0050.npz +2954d92b8cc05f6ddf04852ecb80d53a31b825234f15010e136928125892dbbb group_0056.npz +beedd6caf6eafd7b2df793d2ad6248c1cd43b3c1b1b3ce80816b0694be356b5a group_0058.npz +e3d95e740542311e9d682048a3de1de6bb084dd7a1edf9cb2c1f0ab92e5c396f group_0060.npz +af9b221c4035f3e128e662bed2efa1c383b984be6fe4db69301580143eeaf0b1 group_0083.npz +a37b73f7d741093021f494987e7a2152b11f4f1eea1a4dbf3a2af544e97421d8 group_0094.npz +# inception: the Inception-v3 weights in $CLEAN_FID_CACHE +f58cb9b6ec323ed63459aa4fb441fe750cfe39fafad6da5cb504a16f19e958f4 inception-2015-12-05.pt +# samples/direct_si: the archived test samples (--samples) +074759fd1554871f4ad31b9635dacbd0c9423b97380ffa19e03a2bfb8161dfff noise_scale_1p0/group_0016.npz +c7bb5769360ee79a106e160fd0e9eac488ae91229ee9fe95a563781782484dba noise_scale_1p0/group_0031.npz +f06e5d235bb43cf556527edf1d60b881f8d82edc5f0c5d9ee7a5c360f9aa0f7a noise_scale_1p0/group_0032.npz +14906cb3df1fdab94efef83d7b2b8a4686fe1ee7ff6b1991e20df1935f48fb3d noise_scale_1p0/group_0043.npz +5e537397b6cdb057cef62478849167b487cdd600d90d33e66e4161013ab4ac8d noise_scale_1p0/group_0050.npz +50739bb903a8c1aea7bb52f93042e85ab5be8f2ebc5a40638340d42401e0f490 noise_scale_1p0/group_0056.npz +b0c757effb47885445aea8ed8e5ab3c7bb2e6ce41f1214e0bad6b5ebfce12593 noise_scale_1p0/group_0058.npz +43c62c4530fb8669131d352f5fd1760c407b6a105717de2423815021b8362d1b noise_scale_1p0/group_0060.npz +0f95f25d29c223f45cc13a7898ee1f24e8fce57f4c5898bb65ceba16f15115c0 noise_scale_1p0/group_0083.npz +b5b61b909f5974c96adede397654a8261db8f776d935d223a7e4ae60e621b866 noise_scale_1p0/group_0094.npz +# samples/edm: the archived test samples (--samples) +15e059ec0d797d8e6348f5d34be5f80aeb9d1467d955c9aeeed97d7d35fd248d default/group_0016.npz +465c560b5d609558c7941759a88178e6733d7a84521ef8a0eab2dc36a04dd24b default/group_0031.npz +ce568e331161290b6530c2a03c32b4318e43bc9af374ef048dfe700285082a5c default/group_0032.npz +087d0c7970b9132fe3cfb3b98f14cfa2f62ae76b8d9ba4dc9a577bc83da3ce6a default/group_0043.npz +707888f07deaf5b99ee6e81a417ae8269a7feb671e4f3fc01e00dbf3da31659f default/group_0050.npz +9ba345c595a601a58cf26c66aff4515c54a0d24a4fa616197c65cdefc2221a32 default/group_0056.npz +2bf37e56ee5c16f879ff3c3262f2c7eacc900dcec0b47fba58838f8a54147582 default/group_0058.npz +0f8d90544dcbb73e18df6dafa3476bf07c772e9cd934f5b8ce29c85040d0980e default/group_0060.npz +6ffb1c4c40b791651aa06d143d133eb3647668cd7406d113555b7453eb453f9b default/group_0083.npz +5c46aa0e1e581c53f3bbd1bc0d15625b3da727e2b0eb886205a9c3c1ad7b9264 default/group_0094.npz +# samples/cascaded_si: the archived test samples (--samples) +118767e4519a3fb6fe7e4c8492d8ada3a1dcb90be788c22e731b97b7671262ec noise_scale_1p0/group_0016.npz +396685b2a1f047513d726ceb794ea892f5507038e5a1304bed59899e8344b529 noise_scale_1p0/group_0031.npz +c16a3b3376080ecfa7a9640cd66f94f2b9811a529bd5f13d218ccb1e0c4af4fc noise_scale_1p0/group_0032.npz +c31d97dbde26db29141601a520ff15bc1a5f89ca1efc35b7a78205f4cdb7fc5d noise_scale_1p0/group_0043.npz +ecd8655f82819a35b1311a7f25205a99a1dc5803d5f399df1d433b26d4753132 noise_scale_1p0/group_0050.npz +2da6b346c735ca40383f19bbd2cb55fa82ce2939551ac1b0af4124b343dd2ae4 noise_scale_1p0/group_0056.npz +d14e7cef851e3debbda81a8d4d468913c6ced3f244d0cd3ac6e3af2c14f96633 noise_scale_1p0/group_0058.npz +c1b20fd344203b2afcaf6eb3816ee5eddba0510b44d3c48f84de00e83581c8e3 noise_scale_1p0/group_0060.npz +69e55b8194f7df16f41b938bbc6be47ca4c1c56901ca48b4c110516f0df40ac1 noise_scale_1p0/group_0083.npz +62165d8b3427ad4942e66c09e4b9b3e5f4a3d4a6124d1c79c32dd383f04bdfb0 noise_scale_1p0/group_0094.npz +# checkpoints: the released checkpoints (--checkpoint), relative to the downloaded models repository +f13c4ecaef486f3cc746dc4d1ac8c17af474e88863de88b99f2adbca34c0f449 het-waves/direct-si/model.pt +81ad0a0ffd6813fdf0d29049755205130592d292168db6eb405815b6752d9843 het-waves/cascaded-si-S02/model.pt +55d7268475ea02239b848aea1f0e848303bed5fce1b942db4d39d2acaafdd1a3 het-waves/cascaded-si-S05/model.pt +19ae726f34a4633f2b0cff8c7e18bf77838798647ab4100d786f9c13f9cc62f0 het-waves/cascaded-si-S10/model.pt +8af387e8de03c0137dc29786fad6b7838e3a870fa31d3051725211c56981204f het-waves/cascaded-si-S20/model.pt +6a01579cc242c5c0af7bff9f23263e47177e0db5a9b6dd211c99d0230899faf7 het-waves/cascaded-si-S40/model.pt +239c3b5a212dc128d49607c7d3dd701116b5419ba8c0e58f085dace1b6b24327 het-waves/edm/model.pt diff --git a/code/records/waves/fid_adapter.json b/code/records/waves/fid_adapter.json new file mode 100644 index 0000000000000000000000000000000000000000..9a1c05e09f8f87895d835e43b97d0dca32c74cb1 --- /dev/null +++ b/code/records/waves/fid_adapter.json @@ -0,0 +1,37 @@ +{ + "adapter": "heterogeneous_waves_augmented_coarse", + "field_kind": "signed", + "quantile": 0.9999, + "clip_c": 2.3817124402523167, + "range": [ + -2.3817124402523167, + 2.3817124402523167 + ], + "n_train_states": 1600, + "field_shape": [ + 128, + 128 + ], + "pixel_stride": 1, + "n_pixel_values": 26214400, + "map": "uint8 = rint((clip(w,-c,c)+c)/(2c)*255)", + "channels": "3 identical (grayscale scalar field), no colormap", + "source": { + "packaged_dataset_path": "", + "block": "vecs_train", + "ics": "canonical rows 0..79", + "members": "j < 20", + "snapshots": [ + 40 + ], + "pixel_stride": 1, + "n_states": 1600, + "rule": "training data only; validation, test and generated data never enter" + }, + "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", + "reconstruction": "exact inverse FFT of the |k| <= 8 modes on a 128x128 grid (alias-free)", + "n_disp": 128, + "k_c": 8, + "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" +} diff --git a/code/records/waves/runs/cascaded_si_S01/config.yaml b/code/records/waves/runs/cascaded_si_S01/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ea5cb772fbb6e84c0deb59630d2fce57485bb2e1 --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S01/config.yaml @@ -0,0 +1,51 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 1.0 + n_segments: 1 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/cascaded_si_S02/config.yaml b/code/records/waves/runs/cascaded_si_S02/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..279c2afa48745dd84a1fa8bb1b2f33cdcac005d1 --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S02/config.yaml @@ -0,0 +1,52 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.5 + - 1.0 + n_segments: 2 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 32.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/cascaded_si_S05/config.yaml b/code/records/waves/runs/cascaded_si_S05/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5c9286199903eb52f5a12068dfca63e092ee91e8 --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S05/config.yaml @@ -0,0 +1,55 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.2 + - 0.4 + - 0.6 + - 0.8 + - 1.0 + n_segments: 5 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.003 + sigma_train: 2.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/cascaded_si_S10/config.yaml b/code/records/waves/runs/cascaded_si_S10/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..89882245dbfe5c9343178dc32a9af92fbab2afeb --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S10/config.yaml @@ -0,0 +1,60 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.1 + - 0.2 + - 0.3 + - 0.4 + - 0.5 + - 0.6 + - 0.7 + - 0.8 + - 0.9 + - 1.0 + n_segments: 10 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0001 + sigma_train: 2.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/cascaded_si_S20/config.yaml b/code/records/waves/runs/cascaded_si_S20/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a933861a2e9b41a29f8a045cf63e55b18ea6c3b8 --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S20/config.yaml @@ -0,0 +1,70 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.05 + - 0.1 + - 0.15 + - 0.2 + - 0.25 + - 0.3 + - 0.35 + - 0.4 + - 0.45 + - 0.5 + - 0.55 + - 0.6 + - 0.65 + - 0.7 + - 0.75 + - 0.8 + - 0.85 + - 0.9 + - 0.95 + - 1.0 + n_segments: 20 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.0003 + sigma_train: 16.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/cascaded_si_S40/config.yaml b/code/records/waves/runs/cascaded_si_S40/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2063e0253a9a0b489fda00e3f19fc5c6c8e260cd --- /dev/null +++ b/code/records/waves/runs/cascaded_si_S40/config.yaml @@ -0,0 +1,90 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: cascaded_si + t_eps: 0.0 + beta_power: 1 + max_grad_norm: 1.0 + sampling: + noise_scale: 1.0 + t_values: + - 0.0 + - 0.025 + - 0.05 + - 0.075 + - 0.1 + - 0.125 + - 0.15 + - 0.175 + - 0.2 + - 0.225 + - 0.25 + - 0.275 + - 0.3 + - 0.325 + - 0.35 + - 0.375 + - 0.4 + - 0.425 + - 0.45 + - 0.475 + - 0.5 + - 0.525 + - 0.55 + - 0.575 + - 0.6 + - 0.625 + - 0.65 + - 0.675 + - 0.7 + - 0.725 + - 0.75 + - 0.775 + - 0.8 + - 0.825 + - 0.85 + - 0.875 + - 0.9 + - 0.925 + - 0.95 + - 0.975 + - 1.0 + n_segments: 40 + n_sde_steps_per_seg: 20 + eval_noise_scales: + - 1.0 + data_mode: fine_segment_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 0.001 + sigma_train: 1.0 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/runs/edm/config.yaml b/code/records/waves/runs/edm/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9cb81f319b8588d015d4fc45dff348cd84cc7980 --- /dev/null +++ b/code/records/waves/runs/edm/config.yaml @@ -0,0 +1,52 @@ +data: + name: heterogeneous_waves_v3_candidate + packaged_dataset_path: / + train_ics: 80 + train_realizations: 200 +model: + name: unet2_5ch_coord_aug + D: 196 + base_channels: 64 + n_grid: 32 + k_c: 8 + t_start: 0.0 + t_end: 1.0 +method: + name: edm + sigma_data: 1.0 + P_std: 1.2 + sampling: + sigma_min: 0.002 + sigma_max: 80.0 + rho: 7 + t_values: + - 0.0 + - 1.0 + n_steps: 20 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + lr_rampup_kimg: 0 + data_mode: fine_endpoint_pairs + fine_val_members: 20 + fine_horizon_idx: 40 + lr: 3.0e-05 + P_mean: -0.4 +training: + batch_size: 256 + max_steps: 30000 + val_every_updates: 1000 + val_extra_steps: + - 0 + - 100 + - 200 + - 300 + - 500 + - 750 + log_loss_every_epochs: 1 + ema_halflife_kimg: 500 + ema_rampup_ratio: 0.05 + allow_tf32: false + deterministic: true + num_workers: 0 + val_seed: 777 +seed: 0 diff --git a/code/records/waves/scores.json b/code/records/waves/scores.json new file mode 100644 index 0000000000000000000000000000000000000000..e5e841f5ac1287889c5fb05378fa7491e2adc8b5 --- /dev/null +++ b/code/records/waves/scores.json @@ -0,0 +1,135 @@ +{ + "cascaded_si_S01": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,111.57677687508175,10000,10000,0\n31,130.47201849091908,10000,10000,0\n32,113.68342065291057,10000,10000,0\n43,154.5689081717628,10000,10000,0\n50,126.61445237825666,10000,10000,0\n56,103.40463274616053,10000,10000,0\n58,117.3358626128699,10000,10000,0\n60,91.44539732427745,10000,10000,0\n83,129.9377790048021,10000,10000,0\n94,79.59376587144453,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 115.86330141284854, + "sem": 6.763827859735928, + "sd": 21.389101738067428, + "n": 10 + }, + "csFID": { + "mean": 115.86330141284854, + "sem": 6.763827859735928, + "sd": 21.389101738067428, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S02": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,71.59043316880343,10000,10000,0\n31,76.96027415153361,10000,10000,0\n32,62.30672053737942,10000,10000,0\n43,75.17741286596083,10000,10000,0\n50,87.62135749251325,10000,10000,0\n56,61.317991409406545,10000,10000,0\n58,67.51171426289727,10000,10000,0\n60,64.57135125750239,10000,10000,0\n83,92.87292253550976,10000,10000,0\n94,52.37280903180988,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 71.23029867133165, + "sem": 3.911905547976706, + "sd": 12.370531523055481, + "n": 10 + }, + "csFID": { + "mean": 71.23029867133165, + "sem": 3.911905547976706, + "sd": 12.370531523055481, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S05": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,43.277408000709556,10000,10000,0\n31,31.478095865981516,10000,10000,0\n32,32.7243260764908,10000,10000,0\n43,37.05278122099753,10000,10000,0\n50,42.025427704230765,10000,10000,0\n56,36.5351355911925,10000,10000,0\n58,40.55751135611888,10000,10000,0\n60,38.089085921104015,10000,10000,0\n83,40.06167408683427,10000,10000,0\n94,24.823122022265693,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 36.66245678459255, + "sem": 1.7764049801888093, + "sd": 5.617485784262924, + "n": 10 + }, + "csFID": { + "mean": 36.66245678459255, + "sem": 1.7764049801888093, + "sd": 5.617485784262924, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S10": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,51.42688856934524,10000,10000,0\n31,36.93803459490931,10000,10000,0\n32,38.37838964203388,10000,10000,0\n43,50.53022064717459,10000,10000,0\n50,54.05021536909919,10000,10000,0\n56,44.34002012135983,10000,10000,0\n58,47.00815975206427,10000,10000,0\n60,43.66229700421155,10000,10000,0\n83,48.40036733819585,10000,10000,0\n94,30.40498134867721,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 44.513957438707095, + "sem": 2.335970099052613, + "sd": 7.386986059055395, + "n": 10 + }, + "csFID": { + "mean": 44.513957438707095, + "sem": 2.335970099052613, + "sd": 7.386986059055395, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S20": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,66.53229366559452,10000,10000,0\n31,45.68399569936173,10000,10000,0\n32,53.02356625090238,10000,10000,0\n43,65.26812035429066,10000,10000,0\n50,70.10738302005865,10000,10000,0\n56,57.76848661963652,10000,10000,0\n58,60.14032272987788,10000,10000,0\n60,54.80270287224158,10000,10000,0\n83,62.290169858514275,10000,10000,0\n94,47.14190695728166,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 58.27589480277599, + "sem": 2.5799720893000138, + "sd": 8.158588101851374, + "n": 10 + }, + "csFID": { + "mean": 58.27589480277599, + "sem": 2.5799720893000138, + "sd": 8.158588101851374, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "cascaded_si_S40": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,74.40973311154386,10000,10000,0\n31,54.97513717069876,10000,10000,0\n32,52.62212580463097,10000,10000,0\n43,69.59406161768857,10000,10000,0\n50,79.10202339196337,10000,10000,0\n56,61.4363669411012,10000,10000,0\n58,60.99766409970283,10000,10000,0\n60,58.90514788190768,10000,10000,0\n83,69.71323604885907,10000,10000,0\n94,46.366261156312135,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 62.81217572244085, + "sem": 3.248342568310332, + "sd": 10.272161136341742, + "n": 10 + }, + "csFID": { + "mean": 62.81217572244085, + "sem": 3.248342568310332, + "sd": 10.272161136341742, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + }, + "edm": { + "per_condition_csv": "gid,sfid,n_truth,n_gen,n_gen_nonfinite\n16,114.77058356503778,10000,10000,0\n31,139.43023051445556,10000,10000,0\n32,116.59819284001048,10000,10000,0\n43,158.88143885433328,10000,10000,0\n50,137.1861054276901,10000,10000,0\n56,108.73299349922962,10000,10000,0\n58,127.65699132725206,10000,10000,0\n60,100.70698060403907,10000,10000,0\n83,142.50279305131977,10000,10000,0\n94,83.94262312010909,10000,10000,0\n", + "summary": { + "n_ic": 10, + "sfid": { + "mean": 123.0408932803477, + "sem": 7.067465996801545, + "sd": 22.349289835685173, + "n": 10 + }, + "csFID": { + "mean": 123.0408932803477, + "sem": 7.067465996801545, + "sd": 22.349289835685173, + "n": 10 + }, + "n_truth_per_ic": 10000 + } + } +} diff --git a/code/records/waves/units.json b/code/records/waves/units.json new file mode 100644 index 0000000000000000000000000000000000000000..0f23729a47352578e54999d2a5b0d7570a97eb8f --- /dev/null +++ b/code/records/waves/units.json @@ -0,0 +1,349 @@ +{ + "system": "waves", + "paper_name": "Heterogeneous waves", + "dataset": "heterogeneous_waves_v3_candidate", + "test_gids": [ + 16, + 31, + 32, + 43, + 50, + 56, + 58, + 60, + 83, + 94 + ], + "val_gids": [ + 3, + 29, + 37, + 72, + 76, + 79, + 81, + 82, + 91, + 92 + ], + "fine_horizon_idx": 40, + "x0_from_truth_bank": true, + "s_star": 5, + "table2": { + "direct_si": "cascaded_si_S01", + "edm": "edm", + "cascaded_si": "cascaded_si_S05" + }, + "validation_energy_distance_by_S": { + "1": 10.124500017288877, + "2": 6.614841657404936, + "5": 2.804017431478681, + "10": 3.804343481341147, + "20": 4.814091804528606, + "40": 5.055250247057583 + }, + "units": [ + { + "key": "cascaded_si_S01", + "S": 1, + "method": "cascaded_si", + "paper_names": [ + "direct_si" + ], + "config": "records/waves/runs/cascaded_si_S01/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 6000, + "validation_energy_distance": 10.124500017288877 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 40 + }, + "checkpoint": { + "file": "het-waves/direct-si/model.pt", + "sha256": "f13c4ecaef486f3cc746dc4d1ac8c17af474e88863de88b99f2adbca34c0f449", + "bytes": 25042002 + }, + "expected_csfid": { + "per_condition": { + "16": 111.57677687508175, + "31": 130.47201849091908, + "32": 113.68342065291057, + "43": 154.5689081717628, + "50": 126.61445237825666, + "56": 103.40463274616053, + "58": 117.3358626128699, + "60": 91.44539732427745, + "83": 129.9377790048021, + "94": 79.59376587144453 + }, + "n_gen_nonfinite": {}, + "mean": 115.86330141284854, + "sem": 6.763827859735928, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S02", + "S": 2, + "method": "cascaded_si", + "paper_names": [], + "config": "records/waves/runs/cascaded_si_S02/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 32.0, + "selected_update": 30000, + "validation_energy_distance": 6.614841657404936 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 80 + }, + "checkpoint": { + "file": "het-waves/cascaded-si-S02/model.pt", + "sha256": "81ad0a0ffd6813fdf0d29049755205130592d292168db6eb405815b6752d9843", + "bytes": 25042002 + }, + "expected_csfid": { + "per_condition": { + "16": 71.59043316880343, + "31": 76.96027415153361, + "32": 62.30672053737942, + "43": 75.17741286596083, + "50": 87.62135749251325, + "56": 61.317991409406545, + "58": 67.51171426289727, + "60": 64.57135125750239, + "83": 92.87292253550976, + "94": 52.37280903180988 + }, + "n_gen_nonfinite": {}, + "mean": 71.23029867133165, + "sem": 3.911905547976706, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S05", + "S": 5, + "method": "cascaded_si", + "paper_names": [ + "cascaded_si" + ], + "config": "records/waves/runs/cascaded_si_S05/config.yaml", + "selection": { + "lr": 0.003, + "sigma": 2.0, + "selected_update": 30000, + "validation_energy_distance": 2.804017431478681 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 200 + }, + "checkpoint": { + "file": "het-waves/cascaded-si-S05/model.pt", + "sha256": "55d7268475ea02239b848aea1f0e848303bed5fce1b942db4d39d2acaafdd1a3", + "bytes": 25042002 + }, + "expected_csfid": { + "per_condition": { + "16": 43.277408000709556, + "31": 31.478095865981516, + "32": 32.7243260764908, + "43": 37.05278122099753, + "50": 42.025427704230765, + "56": 36.5351355911925, + "58": 40.55751135611888, + "60": 38.089085921104015, + "83": 40.06167408683427, + "94": 24.823122022265693 + }, + "n_gen_nonfinite": {}, + "mean": 36.66245678459255, + "sem": 1.7764049801888093, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S10", + "S": 10, + "method": "cascaded_si", + "paper_names": [], + "config": "records/waves/runs/cascaded_si_S10/config.yaml", + "selection": { + "lr": 0.0001, + "sigma": 2.0, + "selected_update": 23000, + "validation_energy_distance": 3.804343481341147 + }, + "sampler": { + "n_steps": 40, + "variant": "noise_scale_1p0", + "nfe": 400 + }, + "checkpoint": { + "file": "het-waves/cascaded-si-S10/model.pt", + "sha256": "19ae726f34a4633f2b0cff8c7e18bf77838798647ab4100d786f9c13f9cc62f0", + "bytes": 25042066 + }, + "expected_csfid": { + "per_condition": { + "16": 51.42688856934524, + "31": 36.93803459490931, + "32": 38.37838964203388, + "43": 50.53022064717459, + "50": 54.05021536909919, + "56": 44.34002012135983, + "58": 47.00815975206427, + "60": 43.66229700421155, + "83": 48.40036733819585, + "94": 30.40498134867721 + }, + "n_gen_nonfinite": {}, + "mean": 44.513957438707095, + "sem": 2.335970099052613, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S20", + "S": 20, + "method": "cascaded_si", + "paper_names": [], + "config": "records/waves/runs/cascaded_si_S20/config.yaml", + "selection": { + "lr": 0.0003, + "sigma": 16.0, + "selected_update": 27000, + "validation_energy_distance": 4.814091804528606 + }, + "sampler": { + "n_steps": 80, + "variant": "noise_scale_1p0", + "nfe": 1600 + }, + "checkpoint": { + "file": "het-waves/cascaded-si-S20/model.pt", + "sha256": "8af387e8de03c0137dc29786fad6b7838e3a870fa31d3051725211c56981204f", + "bytes": 25042130 + }, + "expected_csfid": { + "per_condition": { + "16": 66.53229366559452, + "31": 45.68399569936173, + "32": 53.02356625090238, + "43": 65.26812035429066, + "50": 70.10738302005865, + "56": 57.76848661963652, + "58": 60.14032272987788, + "60": 54.80270287224158, + "83": 62.290169858514275, + "94": 47.14190695728166 + }, + "n_gen_nonfinite": {}, + "mean": 58.27589480277599, + "sem": 2.5799720893000138, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "cascaded_si_S40", + "S": 40, + "method": "cascaded_si", + "paper_names": [], + "config": "records/waves/runs/cascaded_si_S40/config.yaml", + "selection": { + "lr": 0.001, + "sigma": 1.0, + "selected_update": 500, + "validation_energy_distance": 5.055250247057583 + }, + "sampler": { + "n_steps": 160, + "variant": "noise_scale_1p0", + "nfe": 6400 + }, + "checkpoint": { + "file": "het-waves/cascaded-si-S40/model.pt", + "sha256": "6a01579cc242c5c0af7bff9f23263e47177e0db5a9b6dd211c99d0230899faf7", + "bytes": 25042322 + }, + "expected_csfid": { + "per_condition": { + "16": 74.40973311154386, + "31": 54.97513717069876, + "32": 52.62212580463097, + "43": 69.59406161768857, + "50": 79.10202339196337, + "56": 61.4363669411012, + "58": 60.99766409970283, + "60": 58.90514788190768, + "83": 69.71323604885907, + "94": 46.366261156312135 + }, + "n_gen_nonfinite": {}, + "mean": 62.81217572244085, + "sem": 3.248342568310332, + "n": 10 + }, + "selection_flags": [] + }, + { + "key": "edm", + "S": 1, + "method": "edm", + "paper_names": [ + "edm" + ], + "config": "records/waves/runs/edm/config.yaml", + "selection": { + "lr": 3e-05, + "P_mean": -0.4, + "selected_update": 3000, + "validation_energy_distance": 10.684587013096882 + }, + "sampler": { + "n_steps": 40, + "variant": "default", + "nfe": 79 + }, + "checkpoint": { + "file": "het-waves/edm/model.pt", + "sha256": "239c3b5a212dc128d49607c7d3dd701116b5419ba8c0e58f085dace1b6b24327", + "bytes": 25042002 + }, + "expected_csfid": { + "per_condition": { + "16": 114.77058356503778, + "31": 139.43023051445556, + "32": 116.59819284001048, + "43": 158.88143885433328, + "50": 137.1861054276901, + "56": 108.73299349922962, + "58": 127.65699132725206, + "60": 100.70698060403907, + "83": 142.50279305131977, + "94": 83.94262312010909 + }, + "n_gen_nonfinite": {}, + "mean": 123.0408932803477, + "sem": 7.067465996801545, + "n": 10 + }, + "selection_flags": [ + "selected_hyperparameter_on_search_boundary" + ] + } + ] +} diff --git a/code/reproduce/__init__.py b/code/reproduce/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94558162eaeb48ed9fe02407ef58374d57bf9df8 --- /dev/null +++ b/code/reproduce/__init__.py @@ -0,0 +1 @@ +"""Entry points that reproduce the paper's tables and figures (README, Reproduction).""" diff --git a/code/reproduce/appendix/common.py b/code/reproduce/appendix/common.py new file mode 100644 index 0000000000000000000000000000000000000000..10fe2aab3ce71ccd6bd6814f7cd65faab2071e9f --- /dev/null +++ b/code/reproduce/appendix/common.py @@ -0,0 +1,240 @@ +"""Shared pieces of the appendix diagnostics: the inputs of one system, exact spectral transforms of coarse states, +and the joint extraction of the pool3 (cFID) and spatial (csFID) Inception features.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +from bifurcation.config import RECORDS, SYSTEMS, data_paths, dataset_card, evaluation_settings, records + +METHODS = ("direct_si", "edm", "cascaded_si") +LABELS = {"direct_si": "Direct SI", "edm": "EDM", "cascaded_si": "Cascaded SI"} + + +def add_system_args(ap: argparse.ArgumentParser, samples: bool = True) -> None: + ap.add_argument("--system", required=True, choices=SYSTEMS) + ap.add_argument("--data", default=None, help="dataset package directory (default: $_DATA)") + ap.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $_TRUTH_BANK)") + if samples: + ap.add_argument("--samples", action="append", default=[], metavar="METHOD=DIR", + help="test samples of a Table 2 model (direct_si, edm, cascaded_si); repeatable") + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--device", default="cuda") + + +class SystemInputs: + """The package, truth bank, image adapter, normalization and recorded models of one system.""" + + def __init__(self, a): + from bifurcation.fid.adapters import FIDAdapter + self.system = a.system + self.data, self.bank = data_paths(a.system, a.data, a.truth_bank) + self.card = dataset_card(a.system) + self.rec = records(a.system) + self.fid_cfg = evaluation_settings()["fid"] + self.adapter = FIDAdapter.load(RECORDS / a.system / "fid_adapter.json", **(self.card["fid_adapter"].get("params") or {})) + z = np.load(self.data / "normalization.npz") + self.nmean, self.nstd = np.asarray(z["mean"], np.float64), np.asarray(z["std"], np.float64) + self.t_idx = int(self.rec["fine_horizon_idx"]) + self.gids = sorted(int(g) for g in self.rec["test_gids"]) + self.samples = {} + for s in getattr(a, "samples", []) or []: + m, _, d = s.partition("=") + if m not in METHODS or not d: + raise SystemExit(f"--samples expects METHOD=DIR with METHOD in {METHODS}") + self.samples[m] = Path(d).resolve() + + def unit(self, method: str) -> dict: + key = self.rec["table2"][method] + return next(u for u in self.rec["units"] if u["key"] == key) + + def generated(self, method: str, gid: int, n: int) -> np.ndarray: + """RAW generated states (n, D) of test condition ``gid`` from the samples of ``method``.""" + u = self.unit(method) + d = self.samples.get(method) + if d is None and method == "cascaded_si" and self.rec["table2"]["cascaded_si"] == self.rec["table2"]["direct_si"]: + d = self.samples.get("direct_si") # S* = 1: the same model and samples as Direct SI + if d is None: + raise SystemExit(f"--samples {method}=DIR is needed") + p = d / u["sampler"]["variant"] / f"group_{gid:04d}.npz" + if not p.is_file(): + raise FileNotFoundError(f"{method}: no samples of condition {gid} at {p}") + with np.load(p) as z: + return np.asarray(z["gen_xT_norm"], np.float64)[:n] * self.nstd + self.nmean + + def models(self) -> list[dict]: + """The distinct Table 2 models with samples, in the order direct_si, edm, cascaded_si (a shared model once).""" + out, seen = [], {} + for m in METHODS: + key = self.rec["table2"][m] + if key in seen: + seen[key]["methods"].append(m); continue + if m not in self.samples and not (m == "cascaded_si" and key == self.rec["table2"]["direct_si"] and "direct_si" in self.samples): + continue + seen[key] = {"key": key, "methods": [m], "n_steps": int(self.unit(m)["sampler"]["n_steps"])} + out.append(seen[key]) + for m in out: + m["label"] = "/".join(LABELS[x] for x in m["methods"]) + return out + + +# ── exact transforms of coarse spectral states (no interpolation) ────────────────────────────────────────────── + +def shifted(coeff_raw, kx, ky, dx, dy, n_disp): + """Periodic translation by (dx, dy) pixels of the n_disp grid: a phase factor on every mode.""" + M = len(kx) + c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:] + ph = np.exp(-2j * np.pi * (kx * dx + ky * dy) / n_disp) + c = c * ph + return np.concatenate([c.real, c.imag], axis=1) + + +def rot90(coeff_raw, kx, ky, ix): + """Rotation by ix * 90 degrees: a permutation of the Hermitian-half modes (conjugated where R k leaves the half).""" + M = len(kx) + c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:] + key = {(int(a), int(b)): i for i, (a, b) in enumerate(zip(kx, ky))} + out = np.zeros_like(c) + for i, (a, b) in enumerate(zip(kx, ky)): + for _ in range(ix): + a, b = -b, a + if (int(a), int(b)) in key: + out[:, key[(int(a), int(b))]] = c[:, i] + elif (-int(a), -int(b)) in key: + out[:, key[(-int(a), -int(b))]] = np.conj(c[:, i]) + else: + raise KeyError((a, b)) + return np.concatenate([out.real, out.imag], axis=1) + + +def flip(coeff_raw, kx, ky, axis): + """Reflection k_x -> -k_x (axis 0) or k_y -> -k_y (axis 1), with the Hermitian partner conjugated.""" + M = len(kx) + c = coeff_raw[:, :M] + 1j * coeff_raw[:, M:] + key = {(int(a), int(b)): i for i, (a, b) in enumerate(zip(kx, ky))} + out = np.zeros_like(c) + for i, (a, b) in enumerate(zip(kx, ky)): + a2, b2 = (-int(a), int(b)) if axis == 0 else (int(a), -int(b)) + if (a2, b2) in key: + out[:, key[(a2, b2)]] = c[:, i] + else: + out[:, key[(-a2, -b2)]] = np.conj(c[:, i]) + return np.concatenate([out.real, out.imag], axis=1) + + +# ── cFID and csFID features from one forward pass ────────────────────────────────────────────────────────────── + +SFID_CHANNELS = 7 +SFID_SPATIAL = 17 +SFID_DIM = SFID_SPATIAL * SFID_SPATIAL * SFID_CHANNELS # 2023 +CFID_DIM = 2048 +CANON_MIXED5_SHAPE = (768, SFID_SPATIAL, SFID_SPATIAL) + + +def _stages(L): + pre = [L.conv, L.conv_1, L.conv_2, L.pool0, L.conv_3, L.conv_4, L.pool1, + L.mixed, L.mixed_1, L.mixed_2, L.mixed_3, L.mixed_4, L.mixed_5] + post = [L.mixed_6, L.mixed_7, L.mixed_8, L.mixed_9, L.mixed_10, L.pool2] + return pre, post + + +def verify_layer_identity(extractor, device, cpu_extractor=None) -> dict: + """Check that the stepwise traversal reproduces the pool3 feature (bit-exactly on CPU, so the spatial tap sits on + the cFID path) and that the tapped tensor is the 17 x 17 mixed_6/conv activation. Returns the record.""" + import torch + L = extractor.layers + if not hasattr(L, "mixed_5") or not hasattr(L, "mixed_6") or not hasattr(L.mixed_6, "conv"): + raise RuntimeError("this Inception build does not expose layers.mixed_6.conv") + pre, post = _stages(L) + x = (torch.rand(3, 3, 299, 299, device=device) * 255.0 - 128) / 128 + with torch.no_grad(): + h = x + for st in pre: + h = st.forward(h) + mixed5 = h + spatial = L.mixed_6.conv.forward(mixed5) + for st in post: + h = st.forward(h) + gpu_diff = float((h.view(3, CFID_DIM) - L.forward(x).view(3, CFID_DIM)).abs().max()) + cpu_diff = None + if cpu_extractor is not None: + Lc = cpu_extractor.layers + pre_c, post_c = _stages(Lc) + xc = (torch.rand(3, 3, 299, 299) * 255.0 - 128) / 128 + with torch.no_grad(): + hc = xc + for st in pre_c: + hc = st.forward(hc) + for st in post_c: + hc = st.forward(hc) + cpu_diff = float((hc.view(3, CFID_DIM) - Lc.forward(xc).view(3, CFID_DIM)).abs().max()) + flat = spatial[:, :SFID_CHANNELS].permute(0, 2, 3, 1).reshape(3, -1) + rec = {"submodule": "layers.mixed_6.conv", "mixed_5_output_shape_CHW": list(mixed5.shape[1:]), + "mixed_6_conv_output_shape_CHW": list(spatial.shape[1:]), "channels_total": int(spatial.shape[1]), + "channels_kept": SFID_CHANNELS, "spatial_kept": [int(spatial.shape[2]), int(spatial.shape[3])], + "feature_dim": int(flat.shape[1]), "flatten_order": "NHWC", + "split_traversal_matches_layers_forward_bitexact_cpu": bool(cpu_diff == 0.0) if cpu_diff is not None else None, + "split_traversal_max_abs_diff_cpu": cpu_diff, "split_traversal_max_abs_diff_cuda": gpu_diff} + problems = [] + if tuple(mixed5.shape[1:]) != CANON_MIXED5_SHAPE: + problems.append(f"mixed_5 output {tuple(mixed5.shape[1:])} != {CANON_MIXED5_SHAPE}") + if tuple(spatial.shape[2:]) != (SFID_SPATIAL, SFID_SPATIAL): + problems.append(f"mixed_6/conv spatial {tuple(spatial.shape[2:])} != {SFID_SPATIAL}x{SFID_SPATIAL}") + if int(flat.shape[1]) != SFID_DIM: + problems.append(f"feature dim {int(flat.shape[1])} != {SFID_DIM}") + if cpu_diff is not None and cpu_diff != 0.0: + problems.append(f"CPU split traversal differs from layers.forward by {cpu_diff}") + if gpu_diff > 1e-2: + problems.append(f"device split traversal differs from layers.forward by {gpu_diff}") + rec["problems"] = problems + if problems: + raise RuntimeError("spatial feature verification failed: " + "; ".join(problems)) + print(f"Inception features verified: mixed_6/conv {list(spatial.shape[1:])}, {SFID_CHANNELS} channels kept -> " + f"{int(flat.shape[1])}-D; stepwise traversal vs layers.forward: cpu {cpu_diff}, device {gpu_diff}\n") + return rec + + +def both_features(imgs_u8, extractor, device, batch: int = 200): + """uint8 (N, H, W, 3) -> (pool3 features (N, 2048), spatial features (N, 2023)) from one forward pass; the same + resize and normalization as bifurcation.fid.backend.""" + import torch + from cleanfid.resize import build_resizer + imgs_u8 = np.asarray(imgs_u8) + 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}") + L = extractor.layers + pre, _ = _stages(L) + resizer = build_resizer("clean") + fc, fs = [], [] + 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:]) != (299, 299): + raise AssertionError("resizer did not produce 299x299") + fc.append(extractor(x).cpu().numpy()) + h = (x - 128) / 128 + for st in pre: + h = st.forward(h) + sp = L.mixed_6.conv.forward(h)[:, :SFID_CHANNELS] + fs.append(sp.permute(0, 2, 3, 1).reshape(sp.shape[0], -1).cpu().numpy()) + return np.concatenate(fc).astype(np.float32), np.concatenate(fs).astype(np.float32) + + +def extractors(fid_cfg: dict, device: str): + """(extractor on the device, the layer-identity record); the pins of the fid settings are checked first.""" + import torch + from bifurcation.fid import backend as BE + BE.check_pins(fid_cfg) + dev = torch.device(device if torch.cuda.is_available() else "cpu") + ext = BE.build_extractor(dev) + cpu_ext = BE.build_extractor(torch.device("cpu")) if dev.type != "cpu" else ext + rec = verify_layer_identity(ext, dev, cpu_extractor=cpu_ext) + return ext, dev, rec + + +def write_json(path: Path, obj) -> None: + path.write_text(json.dumps(obj, indent=1, default=float) + "\n") diff --git a/code/reproduce/appendix/fid_sample_stability.py b/code/reproduce/appendix/fid_sample_stability.py new file mode 100644 index 0000000000000000000000000000000000000000..21f4f694aef944acecdbfd39ebf542b477edc004 --- /dev/null +++ b/code/reproduce/appendix/fid_sample_stability.py @@ -0,0 +1,107 @@ +"""Appendix B.3: finite-sample stability of cFID and csFID (the ordering of the models and their relative gaps at N = 10,000). + + python reproduce/appendix/fid_sample_stability.py --system euler --samples direct_si=DIR --samples edm=DIR \\ + --samples cascaded_si=DIR --out DIR [--floor-ns 1000,2500,5000] [--model-ns 2000,5000,10000] [--device cuda] + + floor rows FID(A_N, B_N), A and B disjoint halves of the truth ensemble (N <= 5,000: the bank holds 10,000 members) + model rows FID(gen_N, truth_N) for each Table 2 model +Members are the first N in the stored order, so every N is a nested subset; features are extracted once at the largest +N. Writes OUT/sample_stability.{csv,json}. +""" +from __future__ import annotations + +import argparse +import csv +import math +import sys +from collections import defaultdict +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import SystemInputs, add_system_args, both_features, extractors, write_json # noqa: E402 +from bifurcation.fid import backend as BE # noqa: E402 +from bifurcation.fid.evaluate import bank_members # noqa: E402 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap) + ap.add_argument("--floor-ns", default="1000,2500,5000") + ap.add_argument("--model-ns", default="2000,5000,10000") + a = ap.parse_args(argv) + + inp = SystemInputs(a) + outdir = a.out.resolve(); outdir.mkdir(parents=True, exist_ok=True) + ext, dev, layer_rec = extractors(inp.fid_cfg, a.device) + batch = int(inp.fid_cfg.get("feature_batch", 200)) + adapter, gids, t_idx = inp.adapter, inp.gids, inp.t_idx + models = inp.models() + floor_ns = [int(x) for x in a.floor_ns.split(",")] + model_ns = [int(x) for x in a.model_ns.split(",")] + nmax = max(model_ns + [2 * max(floor_ns)]) + print(f"{a.system}: {len(gids)} test conditions, features at N = {nmax}") + print("models: " + "; ".join(f"{m['label']} ({m['key']}, n* = {m['n_steps']})" for m in models)) + print(f"floor N {floor_ns} (at most 5000: an independent floor at N needs 2N of the 10,000 members); model N {model_ns}\n") + + acc = defaultdict(lambda: defaultdict(list)) # acc[(row, N)][metric] -> per-condition values + for g in gids: + tv = bank_members(inp.bank, g, t_idx, nmax) + imgs, _ = adapter.to_images(tv) + tc, ts = both_features(imgs, ext, dev, batch) + for N in floor_ns: + acc[("floor (truth vs truth)", N)]["cfid"].append(BE.fid_from_features(tc[:N], tc[N:2 * N])) + acc[("floor (truth vs truth)", N)]["sfid"].append(BE.fid_from_features(ts[:N], ts[N:2 * N])) + line = (f" gid {g:4d} floor({floor_ns[-1]}) c{acc[('floor (truth vs truth)', floor_ns[-1])]['cfid'][-1]:6.2f}" + f"/s{acc[('floor (truth vs truth)', floor_ns[-1])]['sfid'][-1]:7.2f}") + for m in models: + gi, _ = adapter.to_images(inp.generated(m["methods"][0], g, nmax)) + gc, gs = both_features(gi, ext, dev, batch) + for N in model_ns: + acc[(m["label"], N)]["cfid"].append(BE.fid_from_features(gc[:N], tc[:N])) + acc[(m["label"], N)]["sfid"].append(BE.fid_from_features(gs[:N], ts[:N])) + line += f" | {m['label'][:20]} c{acc[(m['label'], model_ns[-1])]['cfid'][-1]:6.2f}/s{acc[(m['label'], model_ns[-1])]['sfid'][-1]:7.2f}" + print(line, flush=True) + + def ms(v): + x = np.asarray(v, float); return x.mean(), x.std(ddof=1) / math.sqrt(len(x)) + + rows = [] + print(f"\n{'row':28s} {'N':>6s} {'cFID':>17s} {'csFID':>18s}") + for (label, N), d in sorted(acc.items(), key=lambda kv: (kv[0][0], kv[0][1])): + mc, sc = ms(d["cfid"]); msf, ssf = ms(d["sfid"]) + rows.append({"row": label, "N": N, "n_ic": len(d["cfid"]), "cfid": mc, "cfid_sem": sc, "sfid": msf, "sfid_sem": ssf, + "per_gid_cfid": {int(g): float(v) for g, v in zip(gids, d["cfid"])}, + "per_gid_sfid": {int(g): float(v) for g, v in zip(gids, d["sfid"])}}) + print(f"{label:28s} {N:6d} {mc:8.3f} +- {sc:6.3f} {msf:9.3f} +- {ssf:6.3f}") + + print("\nordering and relative gaps across N (model rows only)") + summary = {} + for metric, key in (("cFID", "cfid"), ("csFID", "sfid")): + summary[metric] = {} + for N in model_ns: + vals = {r["row"]: r[key] for r in rows if r["N"] == N and not r["row"].startswith("floor")} + order = sorted(vals, key=vals.get) + best = vals[order[0]] + summary[metric][N] = {"order": order, "values": vals, "ratio_to_best": {k: vals[k] / best for k in order}} + print(f" {metric} N={N:6d} order: {' < '.join(order)} ratios to best: " + + ", ".join(f"{k}={v:.2f}" for k, v in summary[metric][N]["ratio_to_best"].items())) + + write_json(outdir / "sample_stability.json", {"at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "system": a.system, + "gids": gids, "t_idx": t_idx, "floor_ns": floor_ns, "model_ns": model_ns, + "models": models, "adapter": adapter.params, "layer": layer_rec, "pins": BE.pins(), + "rows": rows, "stability": summary}) + with open(outdir / "sample_stability.csv", "w", newline="") as fh: + w = csv.writer(fh); w.writerow(["row", "N", "n_ic", "cfid", "cfid_sem", "sfid", "sfid_sem"]) + for r in rows: + w.writerow([r["row"], r["N"], r["n_ic"], r["cfid"], r["cfid_sem"], r["sfid"], r["sfid_sem"]]) + print(f"\nwrote {outdir / 'sample_stability.csv'}\n {outdir / 'sample_stability.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/appendix/fid_spatial_sensitivity.py b/code/reproduce/appendix/fid_spatial_sensitivity.py new file mode 100644 index 0000000000000000000000000000000000000000..c3ffd0897aba3219821230cc048a7c2558c26c65 --- /dev/null +++ b/code/reproduce/appendix/fid_spatial_sensitivity.py @@ -0,0 +1,138 @@ +"""Table 8: how much cFID and csFID respond to an exact spatial rearrangement of the truth ensemble (no model output). + + python reproduce/appendix/fid_spatial_sensitivity.py --system passive_scalar --out DIR [--half 5000] [--shifts 2,8,16,32] + [--replicates 1] [--device cuda] + +Each test condition's 2 * half truth members are split into two disjoint halves A and B: + floor row FID(A, B) + perturbation row FID(T(A), B) T = a periodic shift by (d, d) pixels, a rotation or a reflection, exact in the + coarse spectral representation +With --replicates R > 1 every row is re-estimated on R random disjoint splits of the same members (replicate 0 is the +first/second-half split), and the per-condition mean +- SEM over the splits is written beside the across-condition +values. Writes OUT/spatial_sensitivity[_rep].{csv,json} (and _per_ic.csv for R > 1). +""" +from __future__ import annotations + +import argparse +import csv +import math +import sys +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import SystemInputs, add_system_args, both_features, extractors, flip, rot90, shifted, write_json # noqa: E402 +from bifurcation.fid import backend as BE # noqa: E402 +from bifurcation.fid.evaluate import bank_members # noqa: E402 +from bifurcation.models.spectral import K_C, mode_int_wavenumbers # noqa: E402 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap, samples=False) + ap.add_argument("--half", type=int, default=5000, help="members per side; 2 * half must not exceed the bank") + ap.add_argument("--shifts", default="2,8,16,32") + ap.add_argument("--replicates", type=int, default=1) + ap.add_argument("--split-seed", type=int, default=0) + a = ap.parse_args(argv) + + inp = SystemInputs(a) + outdir = a.out.resolve(); outdir.mkdir(parents=True, exist_ok=True) + ext, dev, layer_rec = extractors(inp.fid_cfg, a.device) + batch = int(inp.fid_cfg.get("feature_batch", 200)) + adapter, gids, t_idx = inp.adapter, inp.gids, inp.t_idx + kx, ky = mode_int_wavenumbers(K_C) + n_disp = int(adapter.params.get("n_disp", 128)) + h = int(a.half) + print(f"{a.system}: {len(gids)} test conditions, snapshot {t_idx}, {h} vs {h} independent truth members per row\n") + + transforms = [(f"shift {d} px", (lambda v, d=d: shifted(v, kx, ky, d, d, n_disp))) + for d in [int(x) for x in a.shifts.split(",")]] + transforms += [("rot 90", lambda v: rot90(v, kx, ky, 1)), ("rot 180", lambda v: rot90(v, kx, ky, 2)), + ("flip x", lambda v: flip(v, kx, ky, 0))] + + def feats(states_raw): + imgs, _ = adapter.to_images(np.asarray(states_raw, np.float64)) + return both_features(imgs, ext, dev, batch) + + R = max(1, int(a.replicates)) + rng_split = np.random.default_rng(int(a.split_seed)) + row_names = ["truth vs truth (floor)"] + [n for n, _ in transforms] + per_gid = {name: {k: {} for k in ("cfid", "sfid", "cfid_sem", "sfid_sem", "cfid_reps", "sfid_reps")} for name in row_names} + # features of the whole 2 * half pool are extracted once per row; every replicate is a row selection + for g in gids: + v = bank_members(inp.bank, g, t_idx, 2 * h) + fc0, fs0 = feats(v) + fcT = {}; fsT = {} + for name, fn in transforms: + fcT[name], fsT[name] = feats(fn(v)) + reps = {name: {"cfid": [], "sfid": []} for name in row_names} + for r in range(R): + idx = np.arange(2 * h) if r == 0 else rng_split.permutation(2 * h) + A, B = idx[:h], idx[h:] + reps["truth vs truth (floor)"]["cfid"].append(BE.fid_from_features(fc0[A], fc0[B])) + reps["truth vs truth (floor)"]["sfid"].append(BE.fid_from_features(fs0[A], fs0[B])) + for name, _ in transforms: + reps[name]["cfid"].append(BE.fid_from_features(fcT[name][A], fc0[B])) + reps[name]["sfid"].append(BE.fid_from_features(fsT[name][A], fs0[B])) + del fc0, fs0, fcT, fsT + for name in row_names: + for m in ("cfid", "sfid"): + x = np.asarray(reps[name][m], float) + per_gid[name][m][g] = float(x.mean()) + per_gid[name][f"{m}_sem"][g] = float(x.std(ddof=1) / math.sqrt(len(x))) if len(x) > 1 else float("nan") + per_gid[name][f"{m}_reps"][g] = [float(t) for t in x] + fl = per_gid["truth vs truth (floor)"] + print(f" gid {g:4d} floor cFID {fl['cfid'][g]:7.3f} csFID {fl['sfid'][g]:8.3f} | " + " ".join( + f"{n}: {per_gid[n]['cfid'][g]:6.2f}/{per_gid[n]['sfid'][g]:7.2f}" for n, _ in transforms), flush=True) + + def ms(d): + x = np.array([d[g] for g in gids], float) + return x.mean(), x.std(ddof=1) / math.sqrt(len(x)) + + rows = [] + fl_c, _ = ms(per_gid["truth vs truth (floor)"]["cfid"]); fl_s, _ = ms(per_gid["truth vs truth (floor)"]["sfid"]) + print(f"\n{'row':24s} {'cFID':>18s} {'csFID':>19s} {'cFID/floor':>11s} {'csFID/floor':>12s}") + for name in row_names: + mc, sc = ms(per_gid[name]["cfid"]); msf, ssf = ms(per_gid[name]["sfid"]) + row = {"row": name, "n_ic": len(gids), "cfid": mc, "cfid_sem": sc, "sfid": msf, "sfid_sem": ssf, + "cfid_over_floor": mc / fl_c, "sfid_over_floor": msf / fl_s, + "per_gid_cfid": {int(g): float(per_gid[name]["cfid"][g]) for g in gids}, + "per_gid_sfid": {int(g): float(per_gid[name]["sfid"][g]) for g in gids}} + if R > 1: + csem = np.array([per_gid[name]["cfid_sem"][g] for g in gids], float) + ssem = np.array([per_gid[name]["sfid_sem"][g] for g in gids], float) + row.update({"per_gid_cfid_sem": {int(g): float(per_gid[name]["cfid_sem"][g]) for g in gids}, + "per_gid_sfid_sem": {int(g): float(per_gid[name]["sfid_sem"][g]) for g in gids}, + "per_gid_cfid_replicates": {int(g): per_gid[name]["cfid_reps"][g] for g in gids}, + "per_gid_sfid_replicates": {int(g): per_gid[name]["sfid_reps"][g] for g in gids}, + "median_per_ic_cfid_sem": float(np.median(csem)), "median_per_ic_sfid_sem": float(np.median(ssem)), + "max_per_ic_cfid_sem": float(csem.max()), "max_per_ic_sfid_sem": float(ssem.max())}) + rows.append(row) + print(f"{name:24s} {mc:9.3f} +- {sc:6.3f} {msf:10.3f} +- {ssf:6.3f} {mc/fl_c:11.2f} {msf/fl_s:12.2f}") + + stem = "spatial_sensitivity" if R == 1 else f"spatial_sensitivity_rep{R}" + write_json(outdir / f"{stem}.json", {"at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "system": a.system, + "half": h, "t_idx": t_idx, "gids": gids, "replicates": R, "split_seed": int(a.split_seed), + "design": "floor = FID(A,B); perturbation = FID(T(A),B); A,B disjoint halves of equal size", + "adapter": adapter.params, "layer": layer_rec, "pins": BE.pins(), "rows": rows}) + if R > 1: + with open(outdir / f"{stem}_per_ic.csv", "w", newline="") as fh: + w = csv.writer(fh); w.writerow(["row", "gid", "cfid", "cfid_sem", "sfid", "sfid_sem", "replicates"]) + for r in rows: + for g in gids: + w.writerow([r["row"], g, r["per_gid_cfid"][g], r["per_gid_cfid_sem"][g], r["per_gid_sfid"][g], r["per_gid_sfid_sem"][g], R]) + with open(outdir / f"{stem}.csv", "w", newline="") as fh: + w = csv.writer(fh); w.writerow(["row", "n_ic", "cfid", "cfid_sem_across_ic", "sfid", "sfid_sem_across_ic", "cfid_over_floor", "sfid_over_floor"]) + for r in rows: + w.writerow([r["row"], r["n_ic"], r["cfid"], r["cfid_sem"], r["sfid"], r["sfid_sem"], r["cfid_over_floor"], r["sfid_over_floor"]]) + print(f"\nwrote {outdir / f'{stem}.csv'}\n {outdir / f'{stem}.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/appendix/sfid_probe.py b/code/reproduce/appendix/sfid_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..286a7ccf17e373828c6107c854bddab13c7ae7d5 --- /dev/null +++ b/code/reproduce/appendix/sfid_probe.py @@ -0,0 +1,129 @@ +"""Table 7: cFID and csFID of the Euler models beside their reference points. + + python reproduce/appendix/sfid_probe.py --system euler --samples direct_si=DIR --samples edm=DIR --samples cascaded_si=DIR \\ + --out DIR [--n 10000] [--shifts 1,2,4,8,16,32,64] [--sections floor,models,climatology,symmetry] [--device cuda] + +Rows (each a mean +- SEM over the ten test conditions, both metrics from one Inception forward pass): + floor FID between two disjoint halves of the truth ensemble of a condition + models the Table 2 models (their archived or regenerated test samples) against the truth ensemble + climatology a no-skill reference: truth members pooled from the other nine conditions + symmetry the truth ensemble against itself after an exact periodic shift, rotation or reflection +Writes OUT/sfid_probe.csv and OUT/sfid_probe.json; the truth features are cached under OUT/cache/. +""" +from __future__ import annotations + +import argparse +import sys +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import (SystemInputs, add_system_args, both_features, extractors, flip, rot90, # noqa: E402 + shifted, write_json) +from bifurcation.fid import backend as BE # noqa: E402 +from bifurcation.fid.evaluate import bank_members # noqa: E402 +from bifurcation.models.spectral import K_C, mode_int_wavenumbers # noqa: E402 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap) + ap.add_argument("--n", type=int, default=10000, help="members per ensemble") + ap.add_argument("--shifts", default="1,2,4,8,16,32,64") + ap.add_argument("--sections", default="floor,models,climatology,symmetry") + ap.add_argument("--no-cache", action="store_true", help="recompute the truth features") + a = ap.parse_args(argv) + import pandas as pd + + inp = SystemInputs(a) + outdir = a.out.resolve(); (outdir / "cache").mkdir(parents=True, exist_ok=True) + ext, dev, layer_rec = extractors(inp.fid_cfg, a.device) + batch = int(inp.fid_cfg.get("feature_batch", 200)) + adapter, gids, t_idx = inp.adapter, inp.gids, inp.t_idx + kx, ky = mode_int_wavenumbers(K_C) + n_disp = int(adapter.params.get("n_disp", 128)) + + def feats(states_raw): + imgs, _ = adapter.to_images(np.asarray(states_raw, np.float64)) + return both_features(imgs, ext, dev, batch) + + rows = [] + + def record(probe, detail, per_gid_c, per_gid_s): + c, s = np.array(per_gid_c, float), np.array(per_gid_s, float) + rows.append({"probe": probe, "detail": detail, "n_ic": len(c), + "cfid": c.mean(), "cfid_sem": c.std(ddof=1) / np.sqrt(len(c)), + "sfid": s.mean(), "sfid_sem": s.std(ddof=1) / np.sqrt(len(s)), + "sfid_over_cfid": s.mean() / c.mean() if c.mean() else float("nan"), + "per_gid_cfid": {int(g): float(v) for g, v in zip(gids, c)}, + "per_gid_sfid": {int(g): float(v) for g, v in zip(gids, s)}}) + print(f" {probe:24s} {detail:14s} cFID {c.mean():8.3f} +- {c.std(ddof=1)/np.sqrt(len(c)):5.3f}" + f" csFID {s.mean():9.3f} +- {s.std(ddof=1)/np.sqrt(len(s)):6.3f}", flush=True) + + sections = {x.strip() for x in a.sections.split(",") if x.strip()} + print(f"truth features ({len(gids)} conditions x {a.n} members)") + FC, FS = {}, {} + for g in gids: + p = outdir / "cache" / f"truth_gid{g:04d}_t{t_idx}_n{a.n}.npz" + if not a.no_cache and p.is_file(): + with np.load(p) as z: + FC[g], FS[g] = np.asarray(z["cfid"], np.float32), np.asarray(z["sfid"], np.float32) + else: + FC[g], FS[g] = feats(bank_members(inp.bank, g, t_idx, a.n)) + np.savez_compressed(p, cfid=FC[g], sfid=FS[g]) + print(f" pool3 feature {FC[gids[0]].shape}, spatial feature {FS[gids[0]].shape}\n") + + if "floor" in sections: + print("FLOOR") + h = a.n // 2 + record("floor (truth vs truth)", f"{h}+{h}", + [BE.fid_from_features(FC[g][:h], FC[g][h:]) for g in gids], + [BE.fid_from_features(FS[g][:h], FS[g][h:]) for g in gids]) + if "models" in sections: + print("\nMODELS") + for title, method in (("Cascaded SI", "cascaded_si"), ("Direct SI", "direct_si"), ("EDM", "edm")): + cc, ss = [], [] + for g in gids: + fc, fs = feats(inp.generated(method, g, a.n)) + cc.append(BE.fid_from_features(fc, FC[g])); ss.append(BE.fid_from_features(fs, FS[g])) + record("model", title, cc, ss) + if "climatology" in sections: + print("\nNO-SKILL REFERENCE") + rng = np.random.default_rng(0) + cc, ss = [], [] + for g in gids: + others = [x for x in gids if x != g] + per = a.n // len(others) + idx = [rng.permutation(len(FC[o]))[: per + (1 if i < a.n - per * len(others) else 0)] + for i, o in enumerate(others)] + pc = np.concatenate([FC[o][ix] for o, ix in zip(others, idx)]) + ps = np.concatenate([FS[o][ix] for o, ix in zip(others, idx)]) + cc.append(BE.fid_from_features(pc, FC[g])); ss.append(BE.fid_from_features(ps, FS[g])) + record("climatology (LOO)", f"{a.n} members", cc, ss) + if "symmetry" in sections: + print("\nSYMMETRY / LAYOUT FLOOR (truth vs exactly transformed truth)") + transforms = [(f"shift {d} px", (lambda v, d=d: shifted(v, kx, ky, d, d, n_disp))) + for d in [int(x) for x in a.shifts.split(",")]] + transforms += [("rot 90", lambda v: rot90(v, kx, ky, 1)), ("rot 180", lambda v: rot90(v, kx, ky, 2)), + ("flip x", lambda v: flip(v, kx, ky, 0))] + for name, fn in transforms: + cc, ss = [], [] + for g in gids: + fc, fs = feats(fn(bank_members(inp.bank, g, t_idx, a.n))) + cc.append(BE.fid_from_features(fc, FC[g])); ss.append(BE.fid_from_features(fs, FS[g])) + record("symmetry floor", name, cc, ss) + + pd.DataFrame([{k: v for k, v in r.items() if not k.startswith("per_gid")} for r in rows]).to_csv(outdir / "sfid_probe.csv", index=False) + write_json(outdir / "sfid_probe.json", {"at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "system": a.system, + "n": a.n, "gids": gids, "t_idx": t_idx, "layer": layer_rec, "adapter": adapter.params, + "pins": BE.pins(), "rows": rows}) + print(f"\nwrote {outdir / 'sfid_probe.csv'}\n {outdir / 'sfid_probe.json'}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/figures/dataset_overview_panels.py b/code/reproduce/figures/dataset_overview_panels.py new file mode 100644 index 0000000000000000000000000000000000000000..a78fcb72f79994733dcbfc0c4e89260a75996479 --- /dev/null +++ b/code/reproduce/figures/dataset_overview_panels.py @@ -0,0 +1,157 @@ +"""Figure 4, waves panel: five test conditions x (x_0 | five true end states at the forecast horizon), from the truth +bank, as one grid and as single tiles. No axes, text or colorbar inside the figure. + + python reproduce/figures/dataset_overview_panels.py --system waves --out DIR [--gids g1 g2 g3 g4 g5] [--vmax-from endstates|all] + +Conditions: the test conditions ranked by the spread of their truth ensemble at the horizon (mean pointwise variance +over the first 200 members), at the ranks round(q (n - 1)) for q = 0, 1/4, 1/2, 3/4, 1, narrowest to widest. End +states: members 0..4. One symmetric colour scale, Q_0.999 |field| over the displayed end states. Writes +OUT/_ic_endstates.{png,pdf,svg}, OUT/tiles/, OUT/_colorbar.*, OUT/_fields.npz, metadata.json. +""" +from __future__ import annotations + +import argparse +import json +import math +import sys +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import SystemInputs, add_system_args # noqa: E402 +from reproduce.figures.qualitative_panels import fmt_tick, nice_ticks # noqa: E402 + +N_SPREAD = 200 +MEMBERS = (0, 1, 2, 3, 4) +QUANTILES = (0.0, 0.25, 0.5, 0.75, 1.0) + + +def _mpl(): + import matplotlib + matplotlib.use("Agg") + matplotlib.rcParams["pdf.fonttype"] = 42 + matplotlib.rcParams["ps.fonttype"] = 42 + matplotlib.rcParams["svg.fonttype"] = "path" + import matplotlib.pyplot as plt + return matplotlib, plt + + +def render_grid(x0s, ends, vmax, stem: Path, P=1.0, gap=0.04, x0_gap=0.16, cmap="RdBu_r", dpi=300): + _, plt = _mpl() + nr, nc = len(x0s), 1 + ends.shape[1] + W = P * nc + gap * (nc - 2) + x0_gap + H = P * nr + gap * (nr - 1) + fig = plt.figure(figsize=(W, H)) + for i in range(nr): + y = H - (i + 1) * P - i * gap + for j in range(nc): + x = 0.0 if j == 0 else P + x0_gap + (j - 1) * (P + gap) + ax = fig.add_axes([x / W, y / H, P / W, P / H]) + f = x0s[i] if j == 0 else ends[i, j - 1] + ax.imshow(f.T, cmap=cmap, vmin=-vmax, vmax=vmax, origin="lower", extent=(0, 1, 0, 1), + interpolation="bilinear", aspect="equal") + ax.set_axis_off() + for ext in ("png", "pdf", "svg"): + fig.savefig(stem.with_suffix(f".{ext}"), dpi=dpi, pad_inches=0, facecolor="white") + plt.close(fig) + + +def save_tile(field, stem: Path, vmax, cmap="RdBu_r", px=800): + _, plt = _mpl() + fig = plt.figure(figsize=(px / 300, px / 300), dpi=300) + ax = fig.add_axes([0, 0, 1, 1]) + ax.imshow(field.T, cmap=cmap, vmin=-vmax, vmax=vmax, origin="lower", extent=(0, 1, 0, 1), + interpolation="bilinear", aspect="equal") + ax.set_axis_off() + for ext in ("png", "pdf", "svg"): + fig.savefig(stem.with_suffix(f".{ext}"), dpi=300, pad_inches=0) + plt.close(fig) + + +def save_colorbar(stem: Path, vmax, cmap="RdBu_r"): + matplotlib, plt = _mpl() + ticks = nice_ticks(-vmax, vmax) + fig = plt.figure(figsize=(1.1, 4.0)) + ax = fig.add_axes([0.15, 0.04, 0.22, 0.92]) + cb = matplotlib.colorbar.ColorbarBase(ax, cmap=plt.get_cmap(cmap), norm=matplotlib.colors.Normalize(vmin=-vmax, vmax=vmax)) + cb.set_ticks(ticks) + cb.set_ticklabels([fmt_tick(t) for t in ticks]) + cb.ax.tick_params(labelsize=8, length=2) + cb.outline.set_linewidth(0.6) + for ext in ("png", "pdf", "svg"): + fig.savefig(stem.with_suffix(f".{ext}"), dpi=300, bbox_inches="tight", pad_inches=0.02) + plt.close(fig) + return ticks + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap, samples=False) + ap.add_argument("--gids", type=int, nargs="*", default=None, help="override the five conditions") + ap.add_argument("--vmax-from", choices=["endstates", "all"], default="endstates", + help="colour scale from the end states (default; x0 shares it) or from x0 and end states together") + a = ap.parse_args(argv) + from bifurcation.fid.evaluate import bank_members + inp = SystemInputs(a) + adapter, cand = inp.adapter, inp.gids + out = a.out.resolve() + (out / "tiles").mkdir(parents=True, exist_ok=True) + + def get(gid, members, t_idx): + return bank_members(inp.bank, gid, t_idx, max(members) + 1)[list(members)] + + spreads = {} + for g in cand: + f = np.asarray(adapter.to_fields(get(g, range(N_SPREAD), inp.t_idx)), np.float64) + spreads[g] = float(np.mean(f.var(axis=0))) + order = sorted(cand, key=lambda g: spreads[g]) + if a.gids: + bad = [g for g in a.gids if g not in cand] + if bad: + raise SystemExit(f"{bad} are not test conditions of {a.system}: {cand}") + chosen, rule = list(a.gids), "user-specified (--gids)" + else: + idx = [int(math.floor(q * (len(order) - 1) + 0.5)) for q in QUANTILES] + chosen = [order[k] for k in idx] + rule = (f"test conditions ranked by truth spread at the horizon (mean pointwise variance over the first {N_SPREAD} " + f"members); ranks round(q*(n-1)), q = {list(QUANTILES)}; rows narrowest to widest") + assert len(set(chosen)) == len(chosen) + + x0s, ends, x0_dev = [], [], {} + for g in chosen: + s0 = get(g, range(N_SPREAD), 0) + x0_dev[g] = float(np.max(np.abs(s0 - s0[0])) / max(np.max(np.abs(s0[0])), 1e-300)) + x0s.append(np.asarray(adapter.to_fields(s0[:1]), np.float64)[0]) + ends.append(np.asarray(adapter.to_fields(get(g, MEMBERS, inp.t_idx)), np.float64)) + x0s, ends = np.stack(x0s), np.stack(ends) + pool = np.concatenate([x0s.ravel(), ends.ravel()]) if a.vmax_from == "all" else ends.ravel() + vmax = float(np.quantile(np.abs(pool), 0.999)) + clip = {"x0_frac_clipped": float(np.mean(np.abs(x0s) > vmax)), "end_frac_clipped": float(np.mean(np.abs(ends) > vmax)), + "x0_absmax": float(np.abs(x0s).max()), "end_absmax": float(np.abs(ends).max())} + + stem = out / f"{a.system}_ic_endstates" + render_grid(x0s, ends, vmax, stem) + for i, g in enumerate(chosen): + save_tile(x0s[i], out / "tiles" / f"row{i}_gid{g:04d}_x0", vmax) + for j, m in enumerate(MEMBERS): + save_tile(ends[i, j], out / "tiles" / f"row{i}_gid{g:04d}_end{m}", vmax) + ticks = save_colorbar(out / f"{a.system}_colorbar", vmax) + np.savez_compressed(out / f"{a.system}_fields.npz", x0=x0s.astype(np.float32), ends=ends.astype(np.float32), + gids=np.asarray(chosen), members=np.asarray(MEMBERS), vmax=np.float32(vmax)) + meta = {"system": a.system, "t_idx": inp.t_idx, "candidate_gids": cand, "truth_spread_by_gid_sorted": {str(g): spreads[g] for g in order}, + "selection_rule": rule, "rows_gids": chosen, "columns": ["x0"] + [f"truth member {m} at the horizon" for m in MEMBERS], + "x0_identity_max_rel_dev": {str(g): d for g, d in x0_dev.items()}, + "colour": {"cmap": "RdBu_r", "vmin": -vmax, "vmax": vmax, "colorbar_ticks": ticks, **clip}, + "written": datetime.now(timezone.utc).isoformat(timespec="seconds")} + (out / "metadata.json").write_text(json.dumps(meta, indent=1) + "\n") + print(f"[{a.system}] rows (gids): {chosen} vmax +-{vmax:.4g} x0 clipped {clip['x0_frac_clipped']:.2%}") + print(f"[{a.system}] wrote {stem}.{{png,pdf,svg}}, {len(chosen) * 6} tiles, colorbar, metadata -> {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/figures/dataset_overview_tiles.py b/code/reproduce/figures/dataset_overview_tiles.py new file mode 100644 index 0000000000000000000000000000000000000000..b2a5d3022ab14ff81c6891fcd5eaa293b0542f70 --- /dev/null +++ b/code/reproduce/figures/dataset_overview_tiles.py @@ -0,0 +1,82 @@ +"""Figure 4, Euler and passive-scalar panels: for five test conditions, the initial condition x_0 and five true end +states at the forecast horizon, each as its own square image (no axes, border, colorbar or text). + + python reproduce/figures/dataset_overview_tiles.py --system euler --out DIR [--gids 5 16 ...] + +Default conditions: the first five test conditions in numeric order (the paper shows Euler 5 and passive scalar 0). +End states: truth-bank members 0..4. One symmetric colour scale per system, +-min(adapter clip c, Q_0.999 |field| +over the 30 images). Writes OUT//ic_x0.{png,pdf}, ic_end.{png,pdf} and metadata.json. +""" +from __future__ import annotations + +import argparse +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +import matplotlib +import numpy as np + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import SystemInputs, add_system_args # noqa: E402 +from bifurcation.fid.evaluate import bank_members # noqa: E402 + +STYLE = dict(tile_px=800, cmap="RdBu_r", interpolation="bilinear", n_endstates=5) + + +def save_tile(field, stem, vmax, style=STYLE): + dpi = 100 + fig = plt.figure(figsize=(style["tile_px"] / dpi, style["tile_px"] / dpi), dpi=dpi) + ax = fig.add_axes([0, 0, 1, 1]) + ax.imshow(field.T, cmap=style["cmap"], vmin=-vmax, vmax=vmax, origin="lower", extent=(0, 1, 0, 1), + interpolation=style["interpolation"], aspect="equal") + ax.set_axis_off() + fig.savefig(stem.with_suffix(".png"), dpi=dpi, pad_inches=0) + fig.savefig(stem.with_suffix(".pdf"), pad_inches=0) + plt.close(fig) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap, samples=False) + ap.add_argument("--gids", type=int, nargs="*", default=None, help="test conditions to show (default: the first five)") + a = ap.parse_args(argv) + inp = SystemInputs(a) + adapter, clip_c = inp.adapter, float(inp.adapter.params["clip_c"]) + gids = inp.gids[:5] if not a.gids else [int(g) for g in a.gids] + bad = [g for g in gids if g not in inp.gids] + if bad: + raise SystemExit(f"not test conditions of {a.system}: {bad} (test conditions: {inp.gids})") + n_end = STYLE["n_endstates"] + tiles = {} + for g in gids: + tiles[g] = {"x0": adapter.to_fields(bank_members(inp.bank, g, 0, 1))[0], + "end": adapter.to_fields(bank_members(inp.bank, g, inp.t_idx, n_end))} + shown = np.concatenate([np.abs(np.concatenate([t["x0"][None], t["end"]])).ravel() for t in tiles.values()]) + vmax = float(min(clip_c, np.quantile(shown, 0.999))) + out = a.out.resolve() / a.system + out.mkdir(parents=True, exist_ok=True) + files = [] + for g in gids: + save_tile(tiles[g]["x0"], out / f"ic{g}_x0", vmax); files.append(f"ic{g}_x0") + for k in range(n_end): + save_tile(tiles[g]["end"][k], out / f"ic{g}_end{k}", vmax); files.append(f"ic{g}_end{k}") + meta = {"system": a.system, "written": datetime.now(timezone.utc).isoformat(timespec="seconds"), "gids": gids, + "endstates": {"source": "truth bank members 0..4", "snapshot_index": inp.t_idx}, + "x0": "truth bank member 0 at snapshot 0 (every member shares the initial condition)", + "colour": {"cmap": STYLE["cmap"], "vmin": -vmax, "vmax": vmax, "adapter_clip_c": clip_c}, + "image": {"tile_px": STYLE["tile_px"], "orientation": "imshow(field.T, origin='lower')", "interpolation": STYLE["interpolation"]}, + "files": [f + ext for f in files for ext in (".png", ".pdf")]} + (out / "metadata.json").write_text(json.dumps(meta, indent=1) + "\n") + print(f"{a.system}: conditions {gids}, snapshot {inp.t_idx}, colour +-{vmax:.4g} -> {len(files)} images x (png, pdf) in {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/figures/fig3_depth_sweeps.py b/code/reproduce/figures/fig3_depth_sweeps.py new file mode 100644 index 0000000000000000000000000000000000000000..349f248ae433d9ea75d057a86f20b7f42ea26f4d --- /dev/null +++ b/code/reproduce/figures/fig3_depth_sweeps.py @@ -0,0 +1,84 @@ +"""Figure 3: test csFID-10k against the depth S of Cascaded SI, with EDM as a band, from records/depth_sweep/depth_scores.csv. + + python reproduce/figures/fig3_depth_sweeps.py --out DIR +""" +from pathlib import Path +import argparse +import csv + +import numpy as np +from plot_style import plt, DIRECT, CASCADED, EDM, INK +from matplotlib.lines import Line2D + +ROOT = Path(__file__).resolve().parents[2] + + +def build(rows, datasets, output): + five = len(datasets) == 5 + fig, axes = plt.subplots(1, len(datasets), figsize=(5.5, 2.05)) + fig.subplots_adjust(left=.084, right=.970 if five else .985, bottom=.24, + top=.78, + wspace=.35 if five else .33) + titles = {'euler': 'euler', 'sqg': 'sqg', 'waves': 'het-waves', + 'nlkg': 'nonlinear-kg', 'passive_scalar': 'passive-scalar'} + label_size = plt.rcParams['axes.labelsize'] + for i, (ax, ds) in enumerate(zip(axes, datasets)): + rr = [r for r in rows if r['dataset'] == ds] + si = sorted([r for r in rr if r['method'] != 'EDM'], key=lambda r: int(r['S'])) + assert [int(r['S']) for r in si] == [1, 2, 5, 10, 20, 40], ds + selected = [r for r in si if r['is_selected'].lower() == 'true'] + assert len(selected) == 1, (ds, selected) + x = np.array([int(r['S']) for r in si]) + y = np.array([float(r['mean']) for r in si]) + e = np.array([float(r['sem']) for r in si]) + assert np.isfinite(y).all() and np.isfinite(e).all() + assert (e >= 0).all() + ax.errorbar(x, y, yerr=e, fmt='o-', color=CASCADED, + markersize=3, capsize=2, lw=1, elinewidth=.7, zorder=3) + ax.plot(x[0], y[0], 'o', ms=4, color=DIRECT, zorder=4) + r = selected[0] + ax.plot(int(r['S']), float(r['mean']), 'o', ms=7, + mfc='none', mec=INK, mew=.8, zorder=5) + edm_rows = [r for r in rr if r['method'] == 'EDM'] + assert len(edm_rows) == 1, ds + m, se = float(edm_rows[0]['mean']), float(edm_rows[0]['sem']) + ax.axhspan(m-se, m+se, color=EDM, alpha=.25, lw=0) + ax.axhline(m, color=EDM, ls='--', lw=.75) + ax.set_xscale('log') + ax.set_xlim(.78, 51) + ax.set_ylim(bottom=0) + ax.set_xticks(x, [str(t) for t in x]) + ax.minorticks_off() + ax.set_xlabel(r'Segments $S$', labelpad=2) + # Center unbroken identifiers instead of adding panel letters or wrapping + # names. Courier matches the manuscript's times-package \texttt family. + ax.set_title(titles[ds], loc='center', fontsize=8, pad=5, + fontfamily='Courier') + ax.spines[['top', 'right']].set_visible(False) + ax.tick_params(length=2, pad=2) + ax.grid(axis='y', color='.9', lw=.4) + ax.set_axisbelow(True) + axes[0].set_ylabel(r'csFID-10k $\downarrow$', labelpad=2) + handles = [ + Line2D([], [], marker='o', color=DIRECT, lw=0, ms=3, label=r'Direct SI ($S=1$)'), + Line2D([], [], marker='o', color=CASCADED, ms=3, label='Cascaded SI'), + Line2D([], [], color=EDM, ls='--', label='EDM'), + Line2D([], [], marker='o', color=INK, mfc='none', ms=6, lw=0, label=r'Selected $S^\star$')] + legend = fig.legend(handles=handles, ncol=4, loc='upper center', bbox_to_anchor=(.52, 1), + frameon=False, fontsize=label_size, columnspacing=1.2, handlelength=1.3) + output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output, dpi=400) + fig.savefig(output.with_suffix('.png'), dpi=240) + plt.close(fig) + + +if __name__ == '__main__': + p = argparse.ArgumentParser() + p.add_argument('--data', type=Path, default=ROOT/'records'/'depth_sweep'/'depth_scores.csv') + p.add_argument('--out', type=Path, required=True) + args = p.parse_args() + with args.data.open() as f: + rows = list(csv.DictReader(f)) + build(rows, ['euler', 'sqg', 'waves', 'nlkg', 'passive_scalar'], + args.out/'depth_sweeps_all_datasets.pdf') + build(rows, ['sqg', 'nlkg'], args.out/'depth_sweeps_sqg_nlkg.pdf') diff --git a/code/reproduce/figures/plot_style.py b/code/reproduce/figures/plot_style.py new file mode 100644 index 0000000000000000000000000000000000000000..b5d3288e041a49711ec58a286124f59cf5b27459 --- /dev/null +++ b/code/reproduce/figures/plot_style.py @@ -0,0 +1,27 @@ +"""Shared publication style. Figure dimensions are final printed inches.""" +import os +from pathlib import Path +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +DIRECT = '#356a98' +CASCADED = '#b74445' +EDM = '#777777' +TRUTH = '#222222' +INK = '#222222' +COLORS = {'Direct SI': DIRECT, 'Cascaded SI': CASCADED, 'EDM': EDM, 'Truth': TRUTH} +plt.rcParams.update({ + 'font.family': 'serif', 'font.serif': ['STIXGeneral'], + 'mathtext.fontset': 'stix', 'font.size': 8, + 'axes.labelsize': 8, 'axes.titlesize': 8.5, + 'xtick.labelsize': 7, 'ytick.labelsize': 7, + 'axes.linewidth': .6, 'lines.linewidth': 1.2, + 'pdf.fonttype': 42, 'ps.fonttype': 42, + 'savefig.facecolor': 'white', +}) + +def save(fig, path): + fig.savefig(path.with_suffix('.pdf'), dpi=400) + fig.savefig(path.with_suffix('.png'), dpi=240) + plt.close(fig) diff --git a/code/reproduce/figures/qualitative_panels.py b/code/reproduce/figures/qualitative_panels.py new file mode 100644 index 0000000000000000000000000000000000000000..7b8b1339db3b7ea15ba8c8fd0ccdda2ee8015bf4 --- /dev/null +++ b/code/reproduce/figures/qualitative_panels.py @@ -0,0 +1,186 @@ +"""Figures 2 and 5: one row of panels per system, x_0 | Truth | Direct SI | EDM | Cascaded SI (three samples each, with +the ensemble mean and standard deviation), panels and colorbar ticks only. + + python reproduce/figures/qualitative_panels.py --system euler --gid 36 --samples direct_si=DIR --samples edm=DIR \\ + --samples cascaded_si=DIR --out DIR + +The paper shows Euler 36, waves 16 (Fig. 2), SQG 103, NLKG 107 and passive scalar 78 (Fig. 5). Samples are the +members 0, 1, 2 of the truth bank and of each model's test samples; mean and standard deviation are over the full +10,000-member ensembles. Colour limits: the field panels share +-min(adapter clip c, Q_0.999 |truth|), the standard +deviation panels [0, Q_0.999 over all methods]. Writes OUT/.{png,pdf,svg}. +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import matplotlib +import numpy as np + +matplotlib.use("Agg") +matplotlib.rcParams["pdf.fonttype"] = 42 # TrueType outlines in the PDF; SVG text as paths +matplotlib.rcParams["ps.fonttype"] = 42 +matplotlib.rcParams["svg.fonttype"] = "path" +import matplotlib.pyplot as plt # noqa: E402 + +REPO = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO)) + +from reproduce.appendix.common import SystemInputs, add_system_args # noqa: E402 +from bifurcation.fid.evaluate import bank_members # noqa: E402 + +SAMPLE_INDICES = (0, 1, 2) +METHOD_ORDER = [("Truth", None), ("Direct SI", "direct_si"), ("EDM", "edm"), ("Cascaded SI", "cascaded_si")] +STYLE = dict(panel_in=0.95, gap_in=0.045, block_gap_in=0.30, x0_gap_in=0.30, row_gap_in=0.06, top_margin_in=0.05, + bottom_margin_in=0.05, left_margin_in=0.08, cbar_margin_in=0.38, cbar_width_in=0.09, cbar_gap_in=0.14, + cbar_tick_fontsize=7.5, cbar_tick_len=2, cbar_tick_width=0.4, cbar_n_ticks=(4, 5), axis_linewidth=0.4, + axis_edgecolor="0.85", cmap_field="RdBu_r", cmap_std="magma", dpi_png=300) + + +def _streaming_moments(raw, adapter, chunk=500): + n = len(raw) + s = s2 = None + for i in range(0, n, chunk): + f = adapter.to_fields(raw[i:i + chunk]) + s = f.sum(0) if s is None else s + f.sum(0) + s2 = (f ** 2).sum(0) if s2 is None else s2 + (f ** 2).sum(0) + mean = s / n + var = np.maximum(s2 / n - mean ** 2, 0.0) + return mean, np.sqrt(var) + + +def nice_ticks(vmin, vmax, n_ok=(4, 5)): + """Round-number ticks inside [vmin, vmax] (step m * 10^e, m in 1, 2, 5 preferred), covering most of the bar.""" + span = float(vmax - vmin) + e0 = int(np.floor(np.log10(span))) + rank = {1: 0, 2: 0, 5: 0, 2.5: 1, 4: 1, 1.5: 2, 3: 2} + best, best_score = None, None + for e in range(e0 - 2, e0 + 2): + for m, r in rank.items(): + step = m * 10.0 ** e + lo = np.ceil(vmin / step - 1e-9) * step + hi = np.floor(vmax / step + 1e-9) * step + if hi < lo: + continue + ticks = np.arange(lo, hi + step / 2, step) + n = len(ticks) + coverage = (ticks[-1] - ticks[0]) / span if n > 1 else 0.0 + ok = n in n_ok + score = (ok, coverage - 0.08 * r - (0.0 if ok else 0.1 * abs(n - np.mean(n_ok)))) + if best_score is None or score > best_score: + best, best_score = ticks, score + return [float(t) for t in best] + + +def fmt_tick(t): + t = round(float(t), 10) + return "0" if abs(t) < 1e-12 else f"{t:g}" + + +def load_panels(inp: SystemInputs, gid: int, n_ensemble: int = 10000) -> dict: + adapter = inp.adapter + clip_c = float(adapter.params["clip_c"]) + if gid not in inp.gids: + raise ValueError(f"gid {gid} is not a test condition of {inp.system} ({inp.gids})") + panels = {} + truth_raw = bank_members(inp.bank, gid, inp.t_idx, n_ensemble) + panels["Truth"] = {"samples": adapter.to_fields(truth_raw[list(SAMPLE_INDICES)])} + panels["Truth"]["mean"], panels["Truth"]["std"] = _streaming_moments(truth_raw, adapter) + x0_field = None + for key, method in METHOD_ORDER[1:]: + gen = inp.generated(method, gid, n_ensemble) + if x0_field is None: + u = inp.unit(method) + d = inp.samples.get(method) or inp.samples.get("direct_si") + with np.load(d / u["sampler"]["variant"] / f"group_{gid:04d}.npz") as z: + x0_field = adapter.to_fields((np.asarray(z["x0_norm"], np.float64) * inp.nstd + inp.nmean)[None])[0] + panels[key] = {"samples": adapter.to_fields(gen[list(SAMPLE_INDICES)])} + panels[key]["mean"], panels[key]["std"] = _streaming_moments(gen, adapter) + truth_disp = np.concatenate([panels["Truth"]["samples"].ravel(), panels["Truth"]["mean"].ravel()]) + vmax = float(min(clip_c, np.quantile(np.abs(truth_disp), 0.999))) + smax = float(max(np.quantile(p["std"], 0.999) for p in panels.values())) + return dict(x0=x0_field, panels=panels, vmax=vmax, smax=smax, gid=gid) + + +def draw_row(fig, fig_w, y_bottom, row_h, data, style): + S = style + P, GIN, BG = S["panel_in"], S["gap_in"], S["block_gap_in"] + RG = S["row_gap_in"] + GW = 3 * P + 2 * GIN + pair_xoff = (GW - (2 * P + GIN)) / 2 + fig_h = fig.get_figheight() + + def ax_at(x_in, y_in, w_in=P, h_in=P): + return fig.add_axes([x_in / fig_w, y_in / fig_h, w_in / fig_w, h_in / fig_h]) + + def show(ax, f, vmin, vmax, cmap): + im = ax.imshow(f.T, cmap=cmap, vmin=vmin, vmax=vmax, origin="lower", extent=(0, 1, 0, 1), + interpolation="bilinear", aspect="equal") + ax.set_xticks([]); ax.set_yticks([]) + for sp in ax.spines.values(): + sp.set_color(S["axis_edgecolor"]); sp.set_linewidth(S["axis_linewidth"]) + return im + + vmax, smax = data["vmax"], data["smax"] + top = y_bottom + row_h + y_samples = top - S["top_margin_in"] - P + y_meanstd = y_samples - RG - P + x = S["left_margin_in"] + show(ax_at(x, (y_samples + y_meanstd + P) / 2 - P / 2), data["x0"], -vmax, vmax, S["cmap_field"]) + x += P + S["x0_gap_in"] + im_field = im_std = None + for title, _ in METHOD_ORDER: + pnl = data["panels"][title] + for k in range(3): + im_field = show(ax_at(x + k * (P + GIN), y_samples), pnl["samples"][k], -vmax, vmax, S["cmap_field"]) + show(ax_at(x + pair_xoff, y_meanstd), pnl["mean"], -vmax, vmax, S["cmap_field"]) + im_std = show(ax_at(x + pair_xoff + P + GIN, y_meanstd), pnl["std"], 0.0, smax, S["cmap_std"]) + x += GW + BG + content_top, content_bot = y_samples + P, y_meanstd + cbh = (content_top - content_bot - S["cbar_gap_in"]) / 2 + cbx = fig_w - S["cbar_margin_in"] + 0.10 + cbw = S["cbar_width_in"] + cb1_ax = fig.add_axes([cbx / fig_w, (content_bot + cbh + S["cbar_gap_in"]) / fig_h, cbw / fig_w, cbh / fig_h]) + cb2_ax = fig.add_axes([cbx / fig_w, content_bot / fig_h, cbw / fig_w, cbh / fig_h]) + for ax, im, vmin_, vmax_ in ((cb1_ax, im_field, -vmax, vmax), (cb2_ax, im_std, 0.0, smax)): + cc = fig.colorbar(im, cax=ax) + cc.outline.set_linewidth(S["axis_linewidth"]); cc.outline.set_edgecolor("0.7") + ax.tick_params(labelsize=S["cbar_tick_fontsize"], length=S["cbar_tick_len"], width=S["cbar_tick_width"], pad=2, colors="0.35") + ticks = nice_ticks(vmin_, vmax_, S["cbar_n_ticks"]) + ax.set_yticks(ticks) + ax.set_yticklabels([fmt_tick(t) for t in ticks]) + + +def render(data, style=STYLE): + row_h = style["top_margin_in"] + 2 * style["panel_in"] + style["row_gap_in"] + style["bottom_margin_in"] + fig_w = (style["left_margin_in"] + style["panel_in"] + style["x0_gap_in"] + + 4 * (3 * style["panel_in"] + 2 * style["gap_in"]) + 3 * style["block_gap_in"] + style["cbar_margin_in"]) + fig = plt.figure(figsize=(fig_w, row_h + style["bottom_margin_in"])) + draw_row(fig, fig_w, style["bottom_margin_in"], row_h, data, style) + return fig + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + add_system_args(ap) + ap.add_argument("--gid", type=int, required=True, help="the test condition shown") + a = ap.parse_args(argv) + inp = SystemInputs(a) + data = load_panels(inp, a.gid) + print(f"{a.system}: gid {a.gid} field +-{data['vmax']:.4g} std [0, {data['smax']:.4g}] colorbar ticks " + f"{[fmt_tick(t) for t in nice_ticks(-data['vmax'], data['vmax'], STYLE['cbar_n_ticks'])]} / " + f"{[fmt_tick(t) for t in nice_ticks(0.0, data['smax'], STYLE['cbar_n_ticks'])]}") + out = a.out.resolve(); out.mkdir(parents=True, exist_ok=True) + fig = render(data) + stem = out / a.system + fig.savefig(stem.with_suffix(".png"), dpi=STYLE["dpi_png"], bbox_inches="tight", pad_inches=0.02) + fig.savefig(stem.with_suffix(".pdf"), bbox_inches="tight", pad_inches=0.02) + fig.savefig(stem.with_suffix(".svg"), bbox_inches="tight", pad_inches=0.02) + plt.close(fig) + print(f"wrote {stem}.png / .pdf / .svg") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/reproduce/table2.py b/code/reproduce/table2.py new file mode 100644 index 0000000000000000000000000000000000000000..16d6042ab7007dc923c2753f94489ec7743380f0 --- /dev/null +++ b/code/reproduce/table2.py @@ -0,0 +1,246 @@ +"""Reproduce one entry of Table 2, or one depth point of Table 6 / Fig. 3: csFID-10k of a model on the ten test conditions. + + python -m reproduce.table2 list + python -m reproduce.table2 check --system SYS --method M [--S S] [--checkpoint FILE] [--samples DIR] [--hash-large] + python -m reproduce.table2 run --system SYS --method M [--S S] --checkpoint FILE --out DIR [--device cuda] + python -m reproduce.table2 generate --system SYS --method M [--S S] --checkpoint FILE --out DIR [--gids g1,g2] [--n-samples N] + python -m reproduce.table2 score --system SYS --method M [--S S] --out DIR [--samples DIR] + python -m reproduce.table2 compare --system SYS --method M [--S S] --out DIR [--tolerance REL] + +SYS is euler | sqg | waves | nlkg | passive_scalar; M is direct_si | edm | cascaded_si (without --S the selected depth S*, +with --S the Cascaded SI model of that depth). The package and the truth bank are taken from _DATA and +_TRUTH_BANK or given with --data and --truth-bank. Everything else comes from records//: the run configuration +of the model, its sampler resolution n*, the image adapter and the recorded values. + + generate 10,000 samples per test condition (bifurcation.sample with the settings of configs/evaluation.yaml: + n* sampler steps, noise stream 777 * 100003 + gid, chunks of 2,000) -> OUT/samples//group_.npz + score csFID per condition (bifurcation.fid.evaluate.evaluate_csfid_run) of OUT/samples, or of the archived samples + given with --samples -> OUT/csfid/{csfid_per_ic.csv, csfid_summary.json, MANIFEST_CSFID.json} + compare per-condition values, mean and SEM against the recorded values -> OUT/COMPARE.json + run generate + score + compare + check the package, truth bank, checkpoint and archived samples against the recorded sha256 +""" +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + +from bifurcation.config import RECORDS, REPO, SYSTEMS, data_paths, dataset_card, evaluation_settings, records, run_config, unit + + +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 asset_hashes(system: str) -> dict: + """{section: {relative path: sha256}} from records//assets.sha256.""" + out, sec = {}, None + for line in (RECORDS / system / "assets.sha256").read_text().splitlines(): + if line.startswith("# "): + sec = line[2:].split(":")[0].strip() + continue + h, q = line.split(" ", 1) + out.setdefault(sec, {})[q] = h + return out + + +def cmd_list(a): + print(f"{'system':15s} {'model':16s} {'paper':22s} {'n*':>4s} {'NFE':>6s} csFID-10k") + for s in SYSTEMS: + rec = records(s) + for u in rec["units"]: + names = ", ".join({"direct_si": "Direct SI", "edm": "EDM", "cascaded_si": f"Cascaded SI (S*={rec['s_star']})"}[n] + for n in u["paper_names"]) or f"depth S={u['S']}" + e = u["expected_csfid"] + print(f"{s:15s} {u['key']:16s} {names:22s} {u['sampler']['n_steps']:4d} {u['sampler']['nfe']:6d} {e['mean']:.2f} +- {e['sem']:.2f}") + + +def cmd_check(a) -> int: + rec, u = unit(a.system, a.method, a.S) + data, bank = data_paths(a.system, a.data, a.truth_bank) + card = dataset_card(a.system) + problems, checked = [], 0 + big = 200 * 1024 * 1024 + for rel, want in card["expected_sha256"].items(): + p = data / rel + if not p.is_file(): + problems.append(f"--data: missing {rel}"); continue + if p.stat().st_size > big and not a.hash_large: + continue + checked += 1 + if sha256(p) != want: + problems.append(f"--data: {rel} sha256 differs from configs/{a.system}/dataset.yaml") + hashes = asset_hashes(a.system) + for gid in rec["test_gids"]: + rel = f"group_{gid:04d}.npz" + p = bank / rel + if not p.is_file(): + problems.append(f"--truth-bank: missing {rel}"); continue + if a.hash_large: + checked += 1 + if sha256(p) != hashes["truth_bank"][rel]: + problems.append(f"--truth-bank: {rel} sha256 differs") + if a.checkpoint: + got = sha256(Path(a.checkpoint)); checked += 1 + if u["checkpoint"] is None: + problems.append(f"--checkpoint: no checkpoint of {u['key']} is distributed") + elif got != u["checkpoint"]["sha256"]: + problems.append(f"--checkpoint: sha256 {got[:16]} is not the recorded {u['checkpoint']['file']}") + if a.samples: + v = u["sampler"]["variant"] + sec = f"samples/{u['paper_names'][0]}" if u["paper_names"] else None # archived samples exist for the Table 2 models + for gid in rec["test_gids"]: + rel = f"{v}/group_{gid:04d}.npz" + p = Path(a.samples) / rel + if not p.is_file(): + problems.append(f"--samples: missing {rel}"); continue + want = (hashes.get(sec) or {}).get(rel) if sec else None + if want and a.hash_large: + checked += 1 + if sha256(p) != want: + problems.append(f"--samples: {rel} sha256 differs from the archived samples") + print(json.dumps({"system": a.system, "model": u["key"], "files_hashed": checked, "problems": problems}, indent=1)) + return 2 if problems else 0 + + +def _fresh(out: Path, sub: str) -> Path: + d = out / sub + if d.exists() and any(d.iterdir()): + raise SystemExit(f"refusing: {d} exists and is not empty (use a fresh --out)") + d.mkdir(parents=True, exist_ok=True) + return d + + +def cmd_generate(a): + from bifurcation import sample + import yaml + rec, u = unit(a.system, a.method, a.S) + data, bank = data_paths(a.system, a.data, a.truth_bank) + ev = evaluation_settings() + out = Path(a.out).resolve() + out.mkdir(parents=True, exist_ok=True) + cfg_path = out / "config.yaml" + cfg_path.write_text(yaml.safe_dump(run_config(a.system, u, data), sort_keys=False)) + n_samples = int(a.n_samples) if getattr(a, "n_samples", None) else int(ev["fid"]["n_generated"]) + argv = ["--config", str(cfg_path), "--data", str(data), "--checkpoint", str(Path(a.checkpoint).resolve()), + "--out", str(out / "samples"), "--split", "test", "--n-samples", str(n_samples), + "--n-steps", str(u["sampler"]["n_steps"]), "--seed-base", str(ev["sampling"]["crn_base_seed"]), + "--chunk", str(ev["sampling"]["sample_chunk"]), "--device", a.device] + if rec["x0_from_truth_bank"]: + argv += ["--truth-bank", str(bank)] + if a.gids: + argv += ["--gids", a.gids] + print("$ python -m bifurcation.sample " + " ".join(argv), flush=True) + sample.main(argv) + log = out / "GENERATED.json" + entries = json.loads(log.read_text()) if log.is_file() else [] + entries.append({"system": a.system, "model": u["key"], "checkpoint": str(Path(a.checkpoint).resolve()), "gids": a.gids or "all", + "argv": argv, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}) + log.write_text(json.dumps(entries, indent=1) + "\n") + + +def cmd_score(a): + from bifurcation.fid.adapters import FIDAdapter + from bifurcation.fid.evaluate import evaluate_csfid_run, generated_states + rec, u = unit(a.system, a.method, a.S) + data, bank = data_paths(a.system, a.data, a.truth_bank) + card, fid_cfg = dataset_card(a.system), evaluation_settings()["fid"] + out = Path(a.out).resolve() + samples = Path(a.samples).resolve() if a.samples else out / "samples" + variant = u["sampler"]["variant"] + missing = [g for g in rec["test_gids"] if not (samples / variant / f"group_{g:04d}.npz").is_file()] + if missing: + raise SystemExit(f"{samples / variant}: no samples for test conditions {missing}") + n_gen = int(fid_cfg["n_generated"]) + gen_raw = generated_states(samples, variant, data, n_gen) + gen_raw = {g: x for g, x in gen_raw.items() if g in set(rec["test_gids"])} + for g, x in gen_raw.items(): + if len(x) != n_gen: + raise SystemExit(f"condition {g}: {len(x)} generated states, the paper's read uses {n_gen}") + adapter = FIDAdapter.load(RECORDS / a.system / "fid_adapter.json", **(card["fid_adapter"].get("params") or {})) + csdir = _fresh(out, "csfid") + summary = evaluate_csfid_run(gen_raw, bank, int(rec["fine_horizon_idx"]), int(fid_cfg["n_truth"]), adapter, csdir, + out / "fid_cache", fid_cfg, a.device, + provenance={"system": a.system, "model": u["key"], "samples": str(samples), "source": "reproduce.table2 score"}) + print(f"csFID-10k {a.system} {u['key']}: {summary['csFID']['mean']:.6f} +- {summary['csFID']['sem']:.6f} (n = {summary['csFID']['n']})") + + +def cmd_compare(a) -> int: + rec, u = unit(a.system, a.method, a.S) + out = Path(a.out).resolve() + got = {int(r["gid"]): float(r["sfid"]) for r in csv.DictReader(open(out / "csfid" / "csfid_per_ic.csv"))} + summ = json.loads((out / "csfid" / "csfid_summary.json").read_text())["csFID"] + exp = u["expected_csfid"] + rows, worst = [], 0.0 + for g in rec["test_gids"]: + e, v = exp["per_condition"][str(g)], got.get(g, float("nan")) + rel = abs(v - e) / abs(e) + worst = max(worst, rel) + rows.append({"gid": g, "recorded": e, "this_run": v, "rel_diff": rel}) + agg = {k: {"recorded": exp[k], "this_run": summ[k], "rel_diff": abs(summ[k] - exp[k]) / abs(exp[k])} for k in ("mean", "sem")} + for r in rows: + print(f" gid {r['gid']:4d} recorded {r['recorded']:.6f} this run {r['this_run']:.6f} rel diff {r['rel_diff']:.2e}") + print(f" mean {summ['mean']:.6f} (recorded {exp['mean']:.6f}, rel diff {agg['mean']['rel_diff']:.2e}); " + f"SEM {summ['sem']:.6f} (recorded {exp['sem']:.6f}, rel diff {agg['sem']['rel_diff']:.2e})") + verdict = {"system": a.system, "model": u["key"], "max_rel_diff_per_condition": worst, "aggregate": agg, "per_condition": rows, + "bitwise_equal": all(r["rel_diff"] == 0.0 for r in rows) and all(v["rel_diff"] == 0.0 for v in agg.values())} + if a.tolerance is not None: + verdict["tolerance"] = a.tolerance + verdict["within_tolerance"] = worst <= a.tolerance and all(v["rel_diff"] <= a.tolerance for v in agg.values()) + (out / "COMPARE.json").write_text(json.dumps(verdict, indent=1) + "\n") + print(f" max relative difference per condition: {worst:.3e}" + (f"; within {a.tolerance:g}: {verdict['within_tolerance']}" if a.tolerance is not None else "")) + return 1 if (a.tolerance is not None and not verdict["within_tolerance"]) else 0 + + +def cmd_run(a) -> int: + cmd_generate(a) + a.samples = None + cmd_score(a) + return cmd_compare(a) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(prog="python -m reproduce.table2", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd", required=True) + sub.add_parser("list") + + def unit_args(p, need_data=True): + p.add_argument("--system", required=True, choices=SYSTEMS) + p.add_argument("--method", required=True, choices=("direct_si", "edm", "cascaded_si")) + p.add_argument("--S", type=int, default=None, help="depth of a Cascaded SI model (Table 6); default: the selected S*") + if need_data: + p.add_argument("--data", default=None, help="dataset package directory (default: $_DATA)") + p.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $_TRUTH_BANK)") + c = sub.add_parser("check"); unit_args(c) + c.add_argument("--checkpoint"); c.add_argument("--samples") + c.add_argument("--hash-large", action="store_true", help="also hash the large files (trajectories, truth banks, samples)") + for name in ("generate", "run"): + g = sub.add_parser(name); unit_args(g) + g.add_argument("--checkpoint", required=True); g.add_argument("--out", required=True) + g.add_argument("--gids", default=None, help="comma-separated subset of the test conditions (split generation across jobs)") + g.add_argument("--device", default="cuda") + if name == "generate": + g.add_argument("--n-samples", type=int, default=None, help="fewer samples per condition (a smoke test; `score` refuses fewer than 10,000)") + if name == "run": + g.add_argument("--tolerance", type=float, default=None) + s = sub.add_parser("score"); unit_args(s) + s.add_argument("--out", required=True); s.add_argument("--samples", default=None, help="archived samples (default OUT/samples)") + s.add_argument("--device", default="cuda") + m = sub.add_parser("compare"); unit_args(m, need_data=False) + m.add_argument("--out", required=True); m.add_argument("--tolerance", type=float, default=None) + a = ap.parse_args(argv) + fn = {"list": cmd_list, "check": cmd_check, "generate": cmd_generate, "score": cmd_score, "compare": cmd_compare, "run": cmd_run}[a.cmd] + return int(fn(a) or 0) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/code/reproduce/tables.py b/code/reproduce/tables.py new file mode 100644 index 0000000000000000000000000000000000000000..b69e517c94f3220dcde3e130f089ff262a901981 --- /dev/null +++ b/code/reproduce/tables.py @@ -0,0 +1,68 @@ +"""Print Tables 2, 5 and 6 of the paper from records//units.json. No data, no GPU. + + python -m reproduce.tables [table2 | table5 | table6 | all] + +`python -m reproduce.table2` regenerates any entry of Table 2 or Table 6 from a checkpoint (or rescores the archived +samples) and compares it with these records. +""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +RECORDS = Path(__file__).resolve().parents[1] / "records" +ORDER = ("euler", "sqg", "waves", "nlkg", "passive_scalar") + + +def _rec(system): + return json.loads((RECORDS / system / "units.json").read_text()) + + +def _unit(rec, key): + return next(u for u in rec["units"] if u["key"] == key) + + +def _fmt(e): + return f"{e['mean']:.2f} +- {e['sem']:.2f}" + + +def table2(): + print("Table 2: test csFID-10k, mean +- SEM over the ten test conditions (+: a selected hyperparameter lies on the boundary of its searched range)") + print(f"{'system':16s} {'Direct SI':>18s} {'EDM':>18s} {'Cascaded SI':>18s} {'S*':>4s}") + for s in ORDER: + r = _rec(s) + cells = [] + for m in ("direct_si", "edm", "cascaded_si"): + u = _unit(r, r["table2"][m]) + flag = " +" if u.get("selection_flags") else " " + cells.append(_fmt(u["expected_csfid"]) + flag) + print(f"{s:16s} {cells[0]:>18s} {cells[1]:>18s} {cells[2]:>18s} {r['s_star']:4d}") + + +def table5(): + print("Table 5: selected configurations (S, learning rate, noise parameter, selected update, sampler steps per segment, NFE)") + for s in ORDER: + r = _rec(s) + for m in ("direct_si", "edm", "cascaded_si"): + if s == "passive_scalar" and m == "cascaded_si": + continue + u = _unit(r, r["table2"][m]); sel = u["selection"] + noise = f"P_mean={sel['P_mean']:g}" if "P_mean" in sel else f"sigma={sel['sigma']:g}" + print(f"{s:16s} {m:12s} S={u['S']:<3d} lr={sel['lr']:<8g} {noise:16s} update={sel['selected_update']:<6d} " + f"steps={u['sampler']['n_steps']:<4d} NFE={u['sampler']['nfe']}") + + +def table6(): + print("Table 6 / Fig. 3: test csFID-10k of the selected Cascaded SI model at each depth S (S = 1 is Direct SI), and EDM") + for s in ORDER: + r = _rec(s) + row = [f"S={u['S']}: {_fmt(u['expected_csfid'])}" for u in sorted((u for u in r["units"] if u["key"] != "edm"), key=lambda u: u["S"])] + print(f"{s:16s} " + " | ".join(row) + f" | EDM: {_fmt(_unit(r, 'edm')['expected_csfid'])}") + + +if __name__ == "__main__": + what = sys.argv[1] if len(sys.argv) > 1 else "all" + for name, fn in (("table2", table2), ("table5", table5), ("table6", table6)): + if what in (name, "all"): + fn(); print() diff --git a/code/uv.lock b/code/uv.lock new file mode 100644 index 0000000000000000000000000000000000000000..45d15883a8c34b61505ac0f4cf237243fb7fc8c7 --- /dev/null +++ b/code/uv.lock @@ -0,0 +1,652 @@ +version = 1 +revision = 3 +requires-python = "==3.10.*" +resolution-markers = [ + "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'", + "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'", +] + +[[package]] +name = "certifi" +version = "2026.10.23" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/11/b524cbd9bbdaca58608e4cf4053fd463417549e1959fa8d3aa18319f9555/certifi-2026.10.23.tar.gz", hash = "sha256:a9895e6e2758b371e0545e321a3a8fd1fcfd9829ada8824a936dfb16aaaf36cd", size = 136225, upload-time = "2026-10-23T11:36:16.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/d1/a3652fe2c19134d0d15922177821d863cbd7e842d53a21bbe7e8e3a1b5ac/certifi-2026.10.23-py3-none-any.whl", hash = "sha256:e9b54de76e530bb6cf0ccd036787103acb58fd42ba930ea8bf948ddbaa2f61d1", size = 135082, upload-time = "2026-10-23T11:36:14.763Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/33/1c/f41d4e74c28ab327ff3acd36053f7ea506c55872d7a90b0fa71aa3ab0c89/charset_normalizer-3.5.2.tar.gz", hash = "sha256:39de2a259fc954455c57274dc94c79d5842774e1247a016aff30bc0efed0f4ef", size = 172659, upload-time = "2026-09-30T04:39:23.398Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/77/46e87bcfc45d25ab4db7cfc9bb544bfa3ffd302289ed31ae93f5433eb899/charset_normalizer-3.5.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:195c26fb65950f8fce54e26349852b7bdd7c5f120aeefbcc440b8a20faaed4a3", size = 373348, upload-time = "2026-09-30T04:34:42.843Z" }, + { url = "https://files.pythonhosted.org/packages/52/fc/e518013affcc43c9f919c3ba41bffe9b4ee4aceb0a6462a6243efcca5f2d/charset_normalizer-3.5.2-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9373ad13ef0d2c0fb761e04e55bfdee5a08b52cef2c882c8fbe9935b1517152e", size = 260860, upload-time = "2026-09-30T04:34:44.817Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ad/2a895c945ee61988dfd9ccee64f0b78dc09f29c9b34d1dd545246d78e0ad/charset_normalizer-3.5.2-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:ddf19c062bea7a0cc80f519243d2c01dd091be0cf952a0750d4ad576709559f5", size = 246491, upload-time = "2026-09-30T04:34:46.227Z" }, + { url = "https://files.pythonhosted.org/packages/c9/05/5d958bc8ea503e26be25ada5430fd409cfb45dc22ff33f9f96e649613c99/charset_normalizer-3.5.2-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3d14b50de6bf4d0edf857a9386836846f982b8f524e188e2e68b96d702bcf4aa", size = 290875, upload-time = "2026-09-30T04:34:47.691Z" }, + { url = "https://files.pythonhosted.org/packages/55/e2/06bad57dfdb49cad92c0ba85b6b4fa5827a67df37897287cfef0553843c2/charset_normalizer-3.5.2-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:28a15fdad492a99b6eccfaaed66ef3f74050680545ea61ec8b2f4c538f1f1320", size = 286987, upload-time = "2026-09-30T04:34:49.356Z" }, + { url = "https://files.pythonhosted.org/packages/c4/9f/9f52d2886d52645987d603425482c63c5045a3005db1354a7097e5ed1ae9/charset_normalizer-3.5.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8a893cc101149f80a653f82062ebc95b34525a2614382e1da5458fe7c6997249", size = 271105, upload-time = "2026-09-30T04:34:50.957Z" }, + { url = "https://files.pythonhosted.org/packages/4a/41/e05e19578b7b87e7db2c3ecb884bd09d065274ab41e1f535542e8bae9b06/charset_normalizer-3.5.2-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:619799369eeef6366ed3e8755a5670f4f2f0fb6b30a0fd7264dc0fdc2357058e", size = 269833, upload-time = "2026-09-30T04:34:52.623Z" }, + { url = "https://files.pythonhosted.org/packages/c2/cd/fc7152414561ff78f9c61a4a627025f65fd244abaefc99ac4e75c6169b33/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:447441e76ec720b15e64418d32e092297340387053047c7c694f579efb0ee1d9", size = 263048, upload-time = "2026-09-30T04:34:54.013Z" }, + { url = "https://files.pythonhosted.org/packages/71/71/fb379e399b1013962716a059d551e03bd80b3e02f05d81246886af6c0958/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:62588a277bfb59def052abd940703fa35107152bf479781a878617d60faf8fb5", size = 249237, upload-time = "2026-09-30T04:34:55.373Z" }, + { url = "https://files.pythonhosted.org/packages/33/4f/aeadccd6d20882909eb2597ec40eccc05d505e2262f14dc76d7620700657/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:44bd4fbb29dfbeba60e7d2bd000c59e4b21ddb3cc53912b14048d37092706d7c", size = 290756, upload-time = "2026-09-30T04:34:56.867Z" }, + { url = "https://files.pythonhosted.org/packages/b9/7b/e8a92613236b257d3b75d532496eb21f68ad8231d7d56df94497e460112e/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:30fcd120b732aa79317f08dee04d7de0847822e4cf7ee0e9f445bb958832252c", size = 269026, upload-time = "2026-09-30T04:34:58.278Z" }, + { url = "https://files.pythonhosted.org/packages/ff/ac/21d5c6b972285c5f095ff78afdc99f3539e3985dfcf7ff1cfbe9e772f529/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:50e3adfb96fc189eb27b1cf62d3b598b89b4bb0420d93a3d3e42e137409011be", size = 287601, upload-time = "2026-09-30T04:34:59.703Z" }, + { url = "https://files.pythonhosted.org/packages/59/72/263491ec1494a194b16fcbff88a0220f2af633738c79e92b4d7189226322/charset_normalizer-3.5.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b736353c0a625bbd5fcec108576e2385db3496f4f771f785ff32e108d3c3bc45", size = 272660, upload-time = "2026-09-30T04:35:01.26Z" }, + { url = "https://files.pythonhosted.org/packages/38/7a/81f1186c57e406233969ae96b5fd0c907254ab8b6e47da8ae24cd20fb9d7/charset_normalizer-3.5.2-cp310-cp310-win32.whl", hash = "sha256:f5833ad231be5eb6553de524a70f48d71b2c8563101750531e0b80184e175cd4", size = 187301, upload-time = "2026-09-30T04:35:02.659Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ab/fb8b5178fdf7bf68efa61cbe892286bc9f73846f26147d4725fb89388475/charset_normalizer-3.5.2-cp310-cp310-win_amd64.whl", hash = "sha256:1461ac396c4fdb983a675f20aa555624f0ee18ac83d832b9244ffff3d8055275", size = 215446, upload-time = "2026-09-30T04:35:04.011Z" }, + { url = "https://files.pythonhosted.org/packages/36/40/e270b74e0d3f583d78d902960845a149e80eb7c737ee5613ccc6efc7264b/charset_normalizer-3.5.2-cp310-cp310-win_arm64.whl", hash = "sha256:c6708715abcf3c73b99508253e961a9967f02fe536532834149574eda6de0d1c", size = 203389, upload-time = "2026-09-30T04:35:05.577Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ab/176fbfd5b64939c55d652366aa5b9ef1d767af207a3aa6ebeb0d226c484d/charset_normalizer-3.5.2-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:4275811936e2f06feff5e598fb42a1b7ae852da8e39605211892b56b81a34efd", size = 331815, upload-time = "2026-09-30T04:38:26.216Z" }, + { url = "https://files.pythonhosted.org/packages/7e/84/371eac6b30bdbcbf2d632a1a01809103459216fcaae61b8b8d922c1bfb8a/charset_normalizer-3.5.2-cp37-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:1c50fe28bbc2ced33386f298650d91218076c05420e6cbd790b913adc41659e7", size = 253276, upload-time = "2026-09-30T04:38:28.032Z" }, + { url = "https://files.pythonhosted.org/packages/43/6f/c4fbae58febff71709c51bc7e18fdfa55341dc382704740f9f0cbf03817b/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d19fbd981a488e22cd04883659ca6b08f50b5974f9fd7c95655ef6a043e5893f", size = 241239, upload-time = "2026-09-30T04:38:29.732Z" }, + { url = "https://files.pythonhosted.org/packages/61/71/458c3f42164a07d0c5210798e9e704b39e540a6793b05aba67f3a35243a9/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0fed1d06615f022ee3b13caf5e8b180cfea32bb2c5aded8a9d44277afc040f93", size = 231121, upload-time = "2026-09-30T04:38:31.462Z" }, + { url = "https://files.pythonhosted.org/packages/09/54/ab9e89367076f6331bb6c65c4bf14a5361fa5191cb6561bf534f18504e1b/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:838dcc90063569a0448120554591a1d6c4a4ffe11babf048908793154ab86ade", size = 260350, upload-time = "2026-09-30T04:38:33.239Z" }, + { url = "https://files.pythonhosted.org/packages/7c/c1/061431ecc688d9d76602502cb57cc01e691e682c18f1beb45f9673b5bbd2/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2ce45c6627b22c47e390bc91a41c3d13032192e699fa0bea96e9671b373d69b0", size = 255430, upload-time = "2026-09-30T04:38:34.865Z" }, + { url = "https://files.pythonhosted.org/packages/8d/1f/20c8949f0676f7ab811abdeb7f4d7f1cbc6e61ff20bef08b44edeb092bc8/charset_normalizer-3.5.2-cp37-abi3-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0774bf9bf620249fee3e0b8b9fd3065de213be30f3aa94ce2494b3b638949e26", size = 250612, upload-time = "2026-09-30T04:38:36.649Z" }, + { url = "https://files.pythonhosted.org/packages/2b/9e/46f2fa4c431fc98c4ae76a8cb5bdca54e0341e3cfc3fcfd8e82740250818/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1db38f4c5496827c1a501846d64d14c3b80c7e6714e406cd7dc36a9899fa1011", size = 242083, upload-time = "2026-09-30T04:38:38.26Z" }, + { url = "https://files.pythonhosted.org/packages/bd/39/559be29a0c0f086e0bba6922babd38916cc5e0b58ced4de13ee01ea05508/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:304d8e4d493af723536393eee0c689eb7813f4a474c8b479dee63f1fdd98f621", size = 232738, upload-time = "2026-09-30T04:38:39.81Z" }, + { url = "https://files.pythonhosted.org/packages/ff/6c/387b0e4f756a282831c1d9fc6aeb6c51ca4507ca202767c8de15ce9b12e2/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:9b7f416ff0978e2f2249330527f0ad6fa02f4932e6199692d3b52da2048c19e4", size = 260703, upload-time = "2026-09-30T04:38:41.346Z" }, + { url = "https://files.pythonhosted.org/packages/96/92/1fdf015f09ef449f50d3ac4b67c90887c9c318b727daa95cc4f866e6521d/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:01077390b03f7988f11d700a2194e69b119741a86b1a638b1db88891e3eced8e", size = 247622, upload-time = "2026-09-30T04:38:42.937Z" }, + { url = "https://files.pythonhosted.org/packages/dc/3c/8e7b8a5671ad5d433669fb2a76f1a0164df2d9b1718b0206bc2a16d840cc/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_s390x.whl", hash = "sha256:7e841fb9010836c992c9f12fcbd43a831de93a5f726fc1ccd8ca1d0268c5014c", size = 257500, upload-time = "2026-09-30T04:38:44.604Z" }, + { url = "https://files.pythonhosted.org/packages/b4/f0/45b579df5cabc1d5d53ea1cc35e8437d3ca768c0acccc7041517cb6fbb32/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9cae88599c7219005d879f98e5ed53341e9a122af585e1091200358a3003d2a0", size = 255100, upload-time = "2026-09-30T04:38:46.289Z" }, + { url = "https://files.pythonhosted.org/packages/31/68/fdec18a343f5fb3f310588dd478b09ac4799e0b187dbade3a8cd776f03ef/charset_normalizer-3.5.2-cp37-abi3-win32.whl", hash = "sha256:01b0c0d2262a9e28e8484a278c7e1b5d650e3ac8cf2683d2967e25899f208bdf", size = 174499, upload-time = "2026-09-30T04:38:47.999Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8a/b618149cc5207943a0242068d7a27897f56a62947b5a039085f2a22029f8/charset_normalizer-3.5.2-cp37-abi3-win_amd64.whl", hash = "sha256:9f56f72050826f63dcee7a7f55b0a77168cb3bfc553fd405e7f8f9ece75a4036", size = 200092, upload-time = "2026-09-30T04:38:49.707Z" }, + { url = "https://files.pythonhosted.org/packages/03/cf/4c66866fa9e2b1c78e3c911516d1de497a677b7ac60f1eceda74ce777ca3/charset_normalizer-3.5.2-cp37-abi3-win_arm64.whl", hash = "sha256:40ab6bffa02ae10a0581e6c198be7d2d8ca5c2a0c64e4ed3465d766df457573e", size = 294363, upload-time = "2026-09-30T04:38:51.312Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ad/d07d7862a62ffa6d79d68074d14823243dd235a77c45262acbf6adeb28bf/charset_normalizer-3.5.2-py3-none-any.whl", hash = "sha256:b6b751274acb69d77b3323d6b7dbaa3c7fdfc1eb829b7eb61d262f32e1af9685", size = 68872, upload-time = "2026-09-30T04:39:21.828Z" }, +] + +[[package]] +name = "clean-fid" +version = "0.1.35" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "requests" }, + { name = "scipy" }, + { name = "torch" }, + { name = "torchvision", version = "0.16.2", source = { registry = "https://download.pytorch.org/whl/cu118" }, marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "torchvision", version = "0.16.2+cu118", source = { registry = "https://download.pytorch.org/whl/cu118" }, marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, + { name = "tqdm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/56/0dfc83e0fe455cfe6272b23a65039b4101c63a4e7446801e26178b675fbf/clean_fid-0.1.35-py3-none-any.whl", hash = "sha256:757cf49d75debe9b07ab14955fe59c845a296deaf0616153b40c5e75b3cf87fb", size = 26008, upload-time = "2022-12-18T20:19:35.819Z" }, +] + +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + +[[package]] +name = "coarse-grained-forecasting" +version = "1.0.0" +source = { editable = "." } +dependencies = [ + { name = "clean-fid" }, + { name = "einops" }, + { name = "h5py" }, + { name = "matplotlib" }, + { name = "numpy" }, + { name = "pandas" }, + { name = "pillow" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "scikit-learn" }, + { name = "scipy" }, + { name = "torch" }, + { name = "torchvision", version = "0.16.2", source = { registry = "https://download.pytorch.org/whl/cu118" }, marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "torchvision", version = "0.16.2+cu118", source = { registry = "https://download.pytorch.org/whl/cu118" }, marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, +] + +[package.metadata] +requires-dist = [ + { name = "clean-fid", specifier = "==0.1.35" }, + { name = "einops", specifier = "==0.8.1" }, + { name = "h5py", specifier = "==3.15.1" }, + { name = "matplotlib", specifier = "==3.8.2" }, + { name = "numpy", specifier = "==1.26.2" }, + { name = "pandas", specifier = "==2.2.1" }, + { name = "pillow", specifier = "==10.0.1" }, + { name = "pyyaml", specifier = "==6.0.1" }, + { name = "requests" }, + { name = "scikit-learn", specifier = "==1.3.2" }, + { name = "scipy", specifier = "==1.15.3" }, + { name = "torch", specifier = "==2.1.2", index = "https://download.pytorch.org/whl/cu118" }, + { name = "torchvision", specifier = "==0.16.2", index = "https://download.pytorch.org/whl/cu118" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "contourpy" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/a3/da4153ec8fe25d263aa48c1a4cbde7f49b59af86f0b6f7862788c60da737/contourpy-1.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ba38e3f9f330af820c4b27ceb4b9c7feee5fe0493ea53a8720f4792667465934", size = 268551, upload-time = "2025-04-15T17:34:46.581Z" }, + { url = "https://files.pythonhosted.org/packages/2f/6c/330de89ae1087eb622bfca0177d32a7ece50c3ef07b28002de4757d9d875/contourpy-1.3.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dc41ba0714aa2968d1f8674ec97504a8f7e334f48eeacebcaa6256213acb0989", size = 253399, upload-time = "2025-04-15T17:34:51.427Z" }, + { url = "https://files.pythonhosted.org/packages/c1/bd/20c6726b1b7f81a8bee5271bed5c165f0a8e1f572578a9d27e2ccb763cb2/contourpy-1.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9be002b31c558d1ddf1b9b415b162c603405414bacd6932d031c5b5a8b757f0d", size = 312061, upload-time = "2025-04-15T17:34:55.961Z" }, + { url = "https://files.pythonhosted.org/packages/22/fc/a9665c88f8a2473f823cf1ec601de9e5375050f1958cbb356cdf06ef1ab6/contourpy-1.3.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8d2e74acbcba3bfdb6d9d8384cdc4f9260cae86ed9beee8bd5f54fee49a430b9", size = 351956, upload-time = "2025-04-15T17:35:00.992Z" }, + { url = "https://files.pythonhosted.org/packages/25/eb/9f0a0238f305ad8fb7ef42481020d6e20cf15e46be99a1fcf939546a177e/contourpy-1.3.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e259bced5549ac64410162adc973c5e2fb77f04df4a439d00b478e57a0e65512", size = 320872, upload-time = "2025-04-15T17:35:06.177Z" }, + { url = "https://files.pythonhosted.org/packages/32/5c/1ee32d1c7956923202f00cf8d2a14a62ed7517bdc0ee1e55301227fc273c/contourpy-1.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad687a04bc802cbe8b9c399c07162a3c35e227e2daccf1668eb1f278cb698631", size = 325027, upload-time = "2025-04-15T17:35:11.244Z" }, + { url = "https://files.pythonhosted.org/packages/83/bf/9baed89785ba743ef329c2b07fd0611d12bfecbedbdd3eeecf929d8d3b52/contourpy-1.3.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cdd22595308f53ef2f891040ab2b93d79192513ffccbd7fe19be7aa773a5e09f", size = 1306641, upload-time = "2025-04-15T17:35:26.701Z" }, + { url = "https://files.pythonhosted.org/packages/d4/cc/74e5e83d1e35de2d28bd97033426b450bc4fd96e092a1f7a63dc7369b55d/contourpy-1.3.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b4f54d6a2defe9f257327b0f243612dd051cc43825587520b1bf74a31e2f6ef2", size = 1374075, upload-time = "2025-04-15T17:35:43.204Z" }, + { url = "https://files.pythonhosted.org/packages/0c/42/17f3b798fd5e033b46a16f8d9fcb39f1aba051307f5ebf441bad1ecf78f8/contourpy-1.3.2-cp310-cp310-win32.whl", hash = "sha256:f939a054192ddc596e031e50bb13b657ce318cf13d264f095ce9db7dc6ae81c0", size = 177534, upload-time = "2025-04-15T17:35:46.554Z" }, + { url = "https://files.pythonhosted.org/packages/54/ec/5162b8582f2c994721018d0c9ece9dc6ff769d298a8ac6b6a652c307e7df/contourpy-1.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:c440093bbc8fc21c637c03bafcbef95ccd963bc6e0514ad887932c18ca2a759a", size = 221188, upload-time = "2025-04-15T17:35:50.064Z" }, + { url = "https://files.pythonhosted.org/packages/33/05/b26e3c6ecc05f349ee0013f0bb850a761016d89cec528a98193a48c34033/contourpy-1.3.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:fd93cc7f3139b6dd7aab2f26a90dde0aa9fc264dbf70f6740d498a70b860b82c", size = 265681, upload-time = "2025-04-15T17:44:59.314Z" }, + { url = "https://files.pythonhosted.org/packages/2b/25/ac07d6ad12affa7d1ffed11b77417d0a6308170f44ff20fa1d5aa6333f03/contourpy-1.3.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:107ba8a6a7eec58bb475329e6d3b95deba9440667c4d62b9b6063942b61d7f16", size = 315101, upload-time = "2025-04-15T17:45:04.165Z" }, + { url = "https://files.pythonhosted.org/packages/8f/4d/5bb3192bbe9d3f27e3061a6a8e7733c9120e203cb8515767d30973f71030/contourpy-1.3.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:ded1706ed0c1049224531b81128efbd5084598f18d8a2d9efae833edbd2b40ad", size = 220599, upload-time = "2025-04-15T17:45:08.456Z" }, +] + +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + +[[package]] +name = "einops" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/81/df4fbe24dff8ba3934af99044188e20a98ed441ad17a274539b74e82e126/einops-0.8.1.tar.gz", hash = "sha256:de5d960a7a761225532e0f1959e5315ebeafc0cd43394732f103ca44b9837e84", size = 54805, upload-time = "2025-02-09T03:17:00.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/62/9773de14fe6c45c23649e98b83231fffd7b9892b6cf863251dc2afa73643/einops-0.8.1-py3-none-any.whl", hash = "sha256:919387eb55330f5757c6bea9165c5ff5cfe63a642682ea788a6d472576d81737", size = 64359, upload-time = "2025-02-09T03:17:01.998Z" }, +] + +[[package]] +name = "filelock" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/58/6fd434bec86eff7c38a3168454cb132b762b2bea9b3ac094101a2f7bc32a/filelock-4.1.0.tar.gz", hash = "sha256:ad7f724afef953e731b1cc39bcd3a09166d72ed7fcdf29e6e88b1c3235c6715d", size = 561277, upload-time = "2026-10-09T19:57:20.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/86/032133892a5de43b5a98200b01aadcad68cc255e274a762f08b8a76d2912/filelock-4.1.0-py3-none-any.whl", hash = "sha256:2ce9818e3e2d8f284c1a964414447ef148d42a5fd5e2a477a7118e574b293ec1", size = 133003, upload-time = "2026-10-09T19:57:18.716Z" }, +] + +[[package]] +name = "fonttools" +version = "4.65.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/77/51/d63c7e52163ac14393a35bd14bd7c0da95f8f74be5d7cc988092f9965129/fonttools-4.65.0.tar.gz", hash = "sha256:762ba5431358d0dbd4a01982484a1d494fb267e91f974cdcf20b80eab8560f6f", size = 3674467, upload-time = "2026-09-10T15:35:54.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/d9/1caaa015dd207da7ccd3feba87289997bd88334ea3e74a367cdc7b0e50a3/fonttools-4.65.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:93a73af2075036d36d7fbf856779c56a1b3b86ffcdae6abede7596604c42c156", size = 3095449, upload-time = "2026-09-10T15:33:04.814Z" }, + { url = "https://files.pythonhosted.org/packages/20/d6/988cd9b33ae2d92b15a51d73eb77c991f7a0fe90a7bc74526e9669247789/fonttools-4.65.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c130be2232e3caf8d2b476854ea78421ec1642917ff5ab695284bac31bbb072b", size = 2587697, upload-time = "2026-09-10T15:33:07.653Z" }, + { url = "https://files.pythonhosted.org/packages/f8/6a/36f465a1c277131f9569f6a56fe99cf135391b4860b9c4f4b23e5b1cd5df/fonttools-4.65.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e3944e0bdba42effb71959e43d91b599326b02b59c78310d5675e8a75525e7d8", size = 5313308, upload-time = "2026-09-10T15:33:09.975Z" }, + { url = "https://files.pythonhosted.org/packages/ee/56/151b5e81d20c63834f48ad37a0cbbbe2f9b248e38f8d10387f0cf5219244/fonttools-4.65.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fb53892b570f7f1f0055e75fc4de32673e32f749c4c8a606b63d5c436650e634", size = 5253649, upload-time = "2026-09-10T15:33:12.225Z" }, + { url = "https://files.pythonhosted.org/packages/c8/22/6389215da9d4f98623aacdca9479c1030bbd516cb658e893bc54b61098ce/fonttools-4.65.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a6c8d184e523580a7c55d21cde37176a3c91cb539cf06c2aa36ffc634fd75296", size = 5277745, upload-time = "2026-09-10T15:33:14.783Z" }, + { url = "https://files.pythonhosted.org/packages/89/e3/c1037a1dfb7c8efe6f2a7d1951ebdde40cbbf82e9c5d796fcb03077e4790/fonttools-4.65.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:e5ceccaf2e57d83b753a2b5db5d94aa0a8071886d4afebd2d520c9683e6bef0e", size = 5413259, upload-time = "2026-09-10T15:33:17.619Z" }, + { url = "https://files.pythonhosted.org/packages/44/b9/7dd72330168d39635c329f23a98279393b1e825e42a6db362e09897aad7b/fonttools-4.65.0-cp310-cp310-win32.whl", hash = "sha256:aff640a4fcb021fa83f9879d5bfa115b6931522dae991a24faa75888bd6aeff6", size = 1577124, upload-time = "2026-09-10T15:33:19.864Z" }, + { url = "https://files.pythonhosted.org/packages/aa/c2/e959385b4626989b25b82b9f4f99e7c3ac68377d6846b376239b6f126966/fonttools-4.65.0-cp310-cp310-win_amd64.whl", hash = "sha256:5c1700a60e4ff23a0425d5a64abf43d092e6b55071354825781faf255904dcb4", size = 1633354, upload-time = "2026-09-10T15:33:21.963Z" }, + { url = "https://files.pythonhosted.org/packages/e6/35/f894ceb867118c0261d0f69a9bd516b045a3754238f76c88a49513ac7a83/fonttools-4.65.0-py3-none-any.whl", hash = "sha256:3060b8c1fc2329fa20265b7c138614143ea7c1624e26c5c180c76aeb74deae6f", size = 1196441, upload-time = "2026-09-10T15:35:52.347Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/77/cd/9be253869fc42e764de7f3dedd6969af7d44ff9c3375214a3442a6f3fc08/fsspec-2026.9.0.tar.gz", hash = "sha256:0f08147951c8cb31d844c3547d631053b127863b60be04cf06e121333ee0e2fe", size = 333545, upload-time = "2026-09-18T17:50:42.825Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/c0/a98505f18594f1bce828bb159cec0fcf9860562f1a2c85913409fc8f3d9e/fsspec-2026.9.0-py3-none-any.whl", hash = "sha256:8dd6e646e99ea382bd85f97a45e6b526a442d79423a7dc673f1e2756d05fcb5f", size = 221738, upload-time = "2026-09-18T17:50:41.341Z" }, +] + +[[package]] +name = "h5py" +version = "3.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/6a/0d79de0b025aa85dc8864de8e97659c94cf3d23148394a954dc5ca52f8c8/h5py-3.15.1.tar.gz", hash = "sha256:c86e3ed45c4473564de55aa83b6fc9e5ead86578773dfbd93047380042e26b69", size = 426236, upload-time = "2025-10-16T10:35:27.404Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/30/8fa61698b438dd751fa46a359792e801191dadab560d0a5f1c709443ef8e/h5py-3.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:67e59f6c2f19a32973a40f43d9a088ae324fe228c8366e25ebc57ceebf093a6b", size = 3414477, upload-time = "2025-10-16T10:33:24.201Z" }, + { url = "https://files.pythonhosted.org/packages/16/16/db2f63302937337c4e9e51d97a5984b769bdb7488e3d37632a6ac297f8ef/h5py-3.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0e2f471688402c3404fa4e13466e373e622fd4b74b47b56cfdff7cc688209422", size = 2850298, upload-time = "2025-10-16T10:33:27.747Z" }, + { url = "https://files.pythonhosted.org/packages/fc/2e/f1bb7de9b05112bfd14d5206090f0f92f1e75bbb412fbec5d4653c3d44dd/h5py-3.15.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4c45802bcb711e128a6839cb6c01e9ac648dc55df045c9542a675c771f15c8d5", size = 4523605, upload-time = "2025-10-16T10:33:31.168Z" }, + { url = "https://files.pythonhosted.org/packages/05/8a/63f4b08f3628171ce8da1a04681a65ee7ac338fde3cb3e9e3c9f7818e4da/h5py-3.15.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:64ce3f6470adb87c06e3a8dd1b90e973699f1759ad79bfa70c230939bff356c9", size = 4735346, upload-time = "2025-10-16T10:33:34.759Z" }, + { url = "https://files.pythonhosted.org/packages/74/48/f16d12d9de22277605bcc11c0dcab5e35f06a54be4798faa2636b5d44b3c/h5py-3.15.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:4411c1867b9899a25e983fff56d820a66f52ac326bbe10c7cdf7d832c9dcd883", size = 4175305, upload-time = "2025-10-16T10:33:38.83Z" }, + { url = "https://files.pythonhosted.org/packages/d6/2f/47cdbff65b2ce53c27458c6df63a232d7bb1644b97df37b2342442342c84/h5py-3.15.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:2cbc4104d3d4aca9d6db8c0c694555e255805bfeacf9eb1349bda871e26cacbe", size = 4653602, upload-time = "2025-10-16T10:33:42.188Z" }, + { url = "https://files.pythonhosted.org/packages/c3/28/dc08de359c2f43a67baa529cb70d7f9599848750031975eed92d6ae78e1d/h5py-3.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:01f55111ca516f5568ae7a7fc8247dfce607de331b4467ee8a9a6ed14e5422c7", size = 2873601, upload-time = "2025-10-16T10:33:45.323Z" }, +] + +[[package]] +name = "idna" +version = "3.20" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/08/8eea9d4b8302028f3abb2c0813953f7aec26d33b7a8960ed760e65ff29fa/idna-3.20.tar.gz", hash = "sha256:a7db850025b95ded1eae8a46181a1a6c56c92c96f0e2b005d9ff8dc0210cab44", size = 216463, upload-time = "2026-09-17T14:11:04.752Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/a2/bb081bab032533a855d44de1d56f8e8426114ff1ba5d1f07a438a0a654f8/idna-3.20-py3-none-any.whl", hash = "sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c", size = 69583, upload-time = "2026-09-17T14:11:03.168Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "joblib" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d5/1d/537ab090f302b838943a1b56497dd53059b9a9b46a074936470173a2e207/joblib-1.6.0.tar.gz", hash = "sha256:2ccc96785b12046c08fd6d55839c12857831b54a3c1673ffadd2f04bfc4eda03", size = 327903, upload-time = "2026-08-31T09:39:04.122Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/53/84099323c2ec4be98d935f63c033ac4151ee83836ca1050ede3b3aadf155/joblib-1.6.0-py3-none-any.whl", hash = "sha256:3dbbf9f6e4b592a2357b854608e980fe6390d131d7a82f011a377ef2ebef7aba", size = 306115, upload-time = "2026-08-31T09:39:02.298Z" }, +] + +[[package]] +name = "kiwisolver" +version = "1.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/07/bd78e6a8fae171ea041ef5bba3ed21a003522fa088834b069b1909981f30/kiwisolver-1.5.1.tar.gz", hash = "sha256:f1303ef2eec81262a4b708c3e858afe58d7c75ad91c1c05266eda7673369859a", size = 104395, upload-time = "2026-08-28T10:28:27.153Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/54/1b8d2bc580414cc75b3a5b3d195981a6facacd8a8d986e589d0a3b51a709/kiwisolver-1.5.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:1acc7e5b7ef05e9da8bb70cd6c7c4513090213d2e1ad9720f599f0bf6c52aec5", size = 123362, upload-time = "2026-08-28T10:24:43.402Z" }, + { url = "https://files.pythonhosted.org/packages/57/19/92c30f540dcfff86ff625389427c39b53a9aaea16420ddcc09b2ab8d1073/kiwisolver-1.5.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:bad20d4c69c851c982a1e3606f4c293edfd5a87885786c50082412240c4b1ffd", size = 66551, upload-time = "2026-08-28T10:24:44.768Z" }, + { url = "https://files.pythonhosted.org/packages/30/ae/30b94088722e7bce8b821cd5ee935a87f80023e7243125f1accf98e39bf7/kiwisolver-1.5.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0a4faea5c6db201c6a21391d2ac926ea97acf7dacdbc3c417189e1adb1a00837", size = 64076, upload-time = "2026-08-28T10:24:46.184Z" }, + { url = "https://files.pythonhosted.org/packages/f1/0e/9f394d693be549fa3fab62498c2778294595991047adc3dcf10aa99b91c0/kiwisolver-1.5.1-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:e05c2f7925f1d88778e53cb44f14e0223204a3bdd09a41664750363acfb1f2ef", size = 1628782, upload-time = "2026-08-28T10:24:47.91Z" }, + { url = "https://files.pythonhosted.org/packages/64/96/26efc04348c0f332b6e6c471dff132bb8b18b2df494a1b65826601674e62/kiwisolver-1.5.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3a4e41e3096bf1f0f1b76e2ffd6d828d6547f574f702d59bdbef7acfa59db9c", size = 1228112, upload-time = "2026-08-28T10:24:49.503Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ac/df9ddc19ec972cd97b60b262dd2c4be28c6192b5a779cad429c6657ee680/kiwisolver-1.5.1-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f1d56ec54d257d05e0b50f5780d967540cd07beeaf9e5f645b26d50cce79f4d8", size = 1246756, upload-time = "2026-08-28T10:24:51.225Z" }, + { url = "https://files.pythonhosted.org/packages/cf/de/c38a246b8a4f4b293c7291aa650917b05c5549b0d532749bc523296e90c4/kiwisolver-1.5.1-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8de6f2a4ce7e7bd27d23dd94abf0ccafe0e0e5cc9c764b0577191f2c25f08f26", size = 1295555, upload-time = "2026-08-28T10:24:52.953Z" }, + { url = "https://files.pythonhosted.org/packages/b6/48/3c95a5232d8b5783050d8c8089b58bf1f57018f9135232d72c5aa6b0c01b/kiwisolver-1.5.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:16895f553ee6620a827d2da56b871f835fb70b9216cca5d188e885caf6e3bd23", size = 2179476, upload-time = "2026-08-28T10:24:54.568Z" }, + { url = "https://files.pythonhosted.org/packages/e6/e9/b9e999336afc561b154e93682f597832bdc7446f2ee18b467c5b0e924867/kiwisolver-1.5.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:6b92f60017dda7d877fdc546438b5e28f31c523264f49cf5a48c1d0ce1a0dfbc", size = 2275258, upload-time = "2026-08-28T10:24:56.186Z" }, + { url = "https://files.pythonhosted.org/packages/76/cd/c1c550796ada4a59644b8264b425fe332686d2a086c4729fcc640ececdfe/kiwisolver-1.5.1-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:7824b5e8bdbf0bccb4ccd37bbb115849a1dc45437fb4de8351385ed07c437ee0", size = 2443424, upload-time = "2026-08-28T10:24:58.093Z" }, + { url = "https://files.pythonhosted.org/packages/b2/d4/30f641eb8e10f6ad5bb9af8a15ec8eea386162a6ef491e68995b8680ef4e/kiwisolver-1.5.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:aa7d00b1700966d2917e54d278aba86897890ca9276dd8b76cf6446b6c181b92", size = 2249457, upload-time = "2026-08-28T10:25:00.223Z" }, + { url = "https://files.pythonhosted.org/packages/1c/d9/211e016f8aa5b5bf496dd3bd7ebab4b73b22d91f025f25e20156a44a77fc/kiwisolver-1.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:96c30002424670b5e1e46495c2b8cbffef39cf77c1d79e76462029d50339785b", size = 70643, upload-time = "2026-08-28T10:25:01.888Z" }, + { url = "https://files.pythonhosted.org/packages/9e/8e/ae7007266dea1fd93fcc48541042c1f6583958055de3b424a59ec16def25/kiwisolver-1.5.1-cp310-cp310-win_arm64.whl", hash = "sha256:f0f4a42db92d6ec7677ab9d12830a2a8ec145a9c6d15db2b593466bc875c78d7", size = 68215, upload-time = "2026-08-28T10:25:03.136Z" }, + { url = "https://files.pythonhosted.org/packages/ba/f1/8bae9fac0f1837679ed59a5dd7e97e7eb943b738defa7cc0117ea0d107dc/kiwisolver-1.5.1-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:a5716a33bfabb2c6ce27b6cf03253467b3804f83e215f4d202685cf93c6c9874", size = 59543, upload-time = "2026-08-28T10:28:13.794Z" }, + { url = "https://files.pythonhosted.org/packages/c3/3e/42e3639d32ef3ac6ede80c39dc1d1df6224927409960d8f762b6cb504efe/kiwisolver-1.5.1-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:685929988b208a911f1285e2f8ed54210b0d681a3dc0f03e00d599d291986e7e", size = 57504, upload-time = "2026-08-28T10:28:15.074Z" }, + { url = "https://files.pythonhosted.org/packages/e1/12/f4beaaceb740b96c363a8fa6f34dbdf37a58d9e5f15416427d7bf89c552a/kiwisolver-1.5.1-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e4e49f7e1a4e7191bdf9dc67a974db714501b1fc52c24324103d06a86abd5c08", size = 79885, upload-time = "2026-08-28T10:28:16.345Z" }, + { url = "https://files.pythonhosted.org/packages/02/24/18d8a755acdae79c37ca2f1a925795aab67b9456c30c8569c377ea3abc77/kiwisolver-1.5.1-pp310-pypy310_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f0a887b6565bbfe80efde2b7f6e8890d7d9bbdb11bdb17028a3690c32fe0621f", size = 77582, upload-time = "2026-08-28T10:28:17.719Z" }, + { url = "https://files.pythonhosted.org/packages/df/a9/b86dde4f553ef74db0f9c32398614e24b2e8ec1d29a6ca04ae080d7dd29b/kiwisolver-1.5.1-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:1983f0974a750a6f6556f368ba11105d1d8369c735b944747c9f12ae5aea7aae", size = 98321, upload-time = "2026-08-28T10:28:19.087Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/38/9b/e422a865e1d5d57d0e509b4e0bf1c1a70a7f6382c29a5aa428df994c8bc8/markupsafe-3.0.4.tar.gz", hash = "sha256:2e9ad7dd851bf45fab9f75cbff4cb493fee9979e8d8c7c9c3ee119022518edd6", size = 153777, upload-time = "2026-10-02T23:07:22.29Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/aa/9a65962e364bf19745f6bec7bde398fb1f5ca53ad2e734258e6198cc32ba/markupsafe-3.0.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:dd8ea6ebee7aedbf7c749fa80521d9ccf1ba473e0d1e14805caafbaad281c889", size = 11521, upload-time = "2026-10-02T23:04:20.141Z" }, + { url = "https://files.pythonhosted.org/packages/14/36/927999a34b7d1def6957de89153d327fedc4030061940b5a468584e6c5a8/markupsafe-3.0.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dff05cb7016dff1e9fd68f4122c127b65dfc59de5306cfb7ad92f956f230bee2", size = 12066, upload-time = "2026-10-02T23:04:21.275Z" }, + { url = "https://files.pythonhosted.org/packages/17/54/69e7b0db9bd687bfd83451eed5666384d67cfed3be062bfc59a06a15474e/markupsafe-3.0.4-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cf63c214fe879a65e69a386f915e36104fc84254ab141240f8854602d8e0be2a", size = 22081, upload-time = "2026-10-02T23:04:22.221Z" }, + { url = "https://files.pythonhosted.org/packages/bd/14/f6f5c97903f7d2db76bbfaced31509a47d360a103b3aaf4849f6536591ec/markupsafe-3.0.4-cp310-cp310-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2a6ef68ae94aed8721934072b27a3b654ea2100b97e4ab864cf1489c90926fbc", size = 26079, upload-time = "2026-10-02T23:04:23.346Z" }, + { url = "https://files.pythonhosted.org/packages/d0/04/c3cc9b75f94f8b54d7e503c44cebd4b4a115ec1d6f1b996b999807daba99/markupsafe-3.0.4-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:fd9f8797427910198f95bced71ddfed61130d7e349213bfb8466c9c99e2c46a8", size = 22611, upload-time = "2026-10-02T23:04:24.401Z" }, + { url = "https://files.pythonhosted.org/packages/00/26/c4708ed3b0f08e8e6d7cbce3314ac130cb5352d1f62951b6fef7878f2b1e/markupsafe-3.0.4-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d1aca03ede943eb80ab3d63bb082c84b7aab85ea83bd0fd0c200260945fb49d9", size = 20712, upload-time = "2026-10-02T23:04:25.386Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b9/f3894d6aae3d7a52c9363317f4baf2fcadc052163d6dbc871266e32639ed/markupsafe-3.0.4-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0764a13d34cae40db7bbf3a09b7e9b491bf4603e20b263a7a9d6b8e324975d0a", size = 20742, upload-time = "2026-10-02T23:04:26.403Z" }, + { url = "https://files.pythonhosted.org/packages/59/7c/8e248ddbfe286ab6bddb462bdac0851cbb1510d2cbbb815a415fcb0511af/markupsafe-3.0.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:9388003072b95f2f1e3fd908604194d653ba21330d811961a78b7da1a77e9e36", size = 21560, upload-time = "2026-10-02T23:04:27.522Z" }, + { url = "https://files.pythonhosted.org/packages/df/26/2353fef7d4bcf2b18e16bad81979fcecff2915156ca6882447917205b8e2/markupsafe-3.0.4-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:8698d70a8081ee8c090dbb394768b5789a1da8b131b5499f89d071dd3cfaf6be", size = 20973, upload-time = "2026-10-02T23:04:28.573Z" }, + { url = "https://files.pythonhosted.org/packages/67/6f/a9561d98d9a6ee3494b0b970a1c766e58bab128cc84841d56ec009456dbd/markupsafe-3.0.4-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:bf053da3c97a4bc5ecfbb218cdd2983febd91c617be8367d139882aa11e490aa", size = 22036, upload-time = "2026-10-02T23:04:29.568Z" }, + { url = "https://files.pythonhosted.org/packages/55/83/6217df9192eca95af3ff0cad955c9854a93fa43b288cf7411ae24713eb52/markupsafe-3.0.4-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:9438a2648b2195980cb2dd8e53ed7b8df91319e2d0b70ae61a9e1d1bc8d3bec9", size = 20314, upload-time = "2026-10-02T23:04:30.881Z" }, + { url = "https://files.pythonhosted.org/packages/35/7c/9cd8081dae70e17fdab3558121fa618d459c7950c8abd8fc10fc024486d7/markupsafe-3.0.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:88d59b473bfb03259722600839af9bbd7fa13a2eb514beefeedb95997882f69a", size = 20780, upload-time = "2026-10-02T23:04:31.975Z" }, + { url = "https://files.pythonhosted.org/packages/44/37/c5f2f45f4d8f0c24e4af7b9dc711640c888044cbf5f7e2ba2c0bb74cbb55/markupsafe-3.0.4-cp310-cp310-win32.whl", hash = "sha256:4a540e2d3192792fc84eced57bef37851ccb2b41f73291bb17408eea77bcd278", size = 14096, upload-time = "2026-10-02T23:04:33.129Z" }, + { url = "https://files.pythonhosted.org/packages/50/50/394a1c61c9b9972cb4f76875c332af7109bd39bdad09d3a0ec5327a339dc/markupsafe-3.0.4-cp310-cp310-win_amd64.whl", hash = "sha256:5c22873ad1f0532ba40fa1727f3c0fc1bbbaab6d373d4cbe3f0dc74b2e2521c7", size = 14329, upload-time = "2026-10-02T23:04:34.168Z" }, + { url = "https://files.pythonhosted.org/packages/d1/9a/86d03d2f32ba41a57baa124a7bcb76d0a6d39b1622e9bd4d550074e8b61a/markupsafe-3.0.4-cp310-cp310-win_arm64.whl", hash = "sha256:3d23795802fc8bd72534836d64489bbf0f67c088959091bdb22e10735a5107bf", size = 14196, upload-time = "2026-10-02T23:04:35.201Z" }, +] + +[[package]] +name = "matplotlib" +version = "3.8.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy" }, + { name = "cycler" }, + { name = "fonttools" }, + { name = "kiwisolver" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "pyparsing" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fb/ab/38a0e94cb01dacb50f06957c2bed1c83b8f9dac6618988a37b2487862944/matplotlib-3.8.2.tar.gz", hash = "sha256:01a978b871b881ee76017152f1f1a0cbf6bd5f7b8ff8c96df0df1bd57d8755a1", size = 35866957, upload-time = "2023-11-17T21:16:40.15Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/d0/fc5f6796a1956f5b9a33555611d01a3cec038f000c3d70ecb051b1631ac4/matplotlib-3.8.2-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:09796f89fb71a0c0e1e2f4bdaf63fb2cefc84446bb963ecdeb40dfee7dfa98c7", size = 7590640, upload-time = "2023-11-17T21:17:02.834Z" }, + { url = "https://files.pythonhosted.org/packages/57/44/007b592809f50883c910db9ec4b81b16dfa0136407250fb581824daabf03/matplotlib-3.8.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6f9c6976748a25e8b9be51ea028df49b8e561eed7809146da7a47dbecebab367", size = 7484350, upload-time = "2023-11-17T21:17:12.281Z" }, + { url = "https://files.pythonhosted.org/packages/01/87/c7b24f3048234fe10184560263be2173311376dc3d1fa329de7f012d6ce5/matplotlib-3.8.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b78e4f2cedf303869b782071b55fdde5987fda3038e9d09e58c91cc261b5ad18", size = 11382388, upload-time = "2023-11-17T21:17:26.461Z" }, + { url = "https://files.pythonhosted.org/packages/19/e5/a4ea514515f270224435c69359abb7a3d152ed31b9ee3ba5e63017461945/matplotlib-3.8.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e208f46cf6576a7624195aa047cb344a7f802e113bb1a06cfd4bee431de5e31", size = 11611959, upload-time = "2023-11-17T21:17:40.541Z" }, + { url = "https://files.pythonhosted.org/packages/09/23/ab5a562c9acb81e351b084bea39f65b153918417fb434619cf5a19f44a55/matplotlib-3.8.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:46a569130ff53798ea5f50afce7406e91fdc471ca1e0e26ba976a8c734c9427a", size = 9536938, upload-time = "2023-11-17T21:17:49.925Z" }, + { url = "https://files.pythonhosted.org/packages/46/37/b5e27ab30ecc0a3694c8a78287b5ef35dad0c3095c144fcc43081170bfd6/matplotlib-3.8.2-cp310-cp310-win_amd64.whl", hash = "sha256:830f00640c965c5b7f6bc32f0d4ce0c36dfe0379f7dd65b07a00c801713ec40a", size = 7643836, upload-time = "2023-11-17T21:17:58.379Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "networkx" +version = "3.4.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" }, +] + +[[package]] +name = "numpy" +version = "1.26.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/dd/2b/205ddff2314d4eea852e31d53b8e55eb3f32b292efc3dd86bd827ab9019d/numpy-1.26.2.tar.gz", hash = "sha256:f65738447676ab5777f11e6bbbdb8ce11b785e105f690bc45966574816b6d3ea", size = 15664248, upload-time = "2023-11-12T23:17:31.386Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/ac/dea2939dfc3c591a2494121669455fd7d049248ef284c9542904ddbe05d5/numpy-1.26.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:3703fc9258a4a122d17043e57b35e5ef1c5a5837c3db8be396c82e04c1cf9b0f", size = 20618422, upload-time = "2023-11-12T22:51:49.818Z" }, + { url = "https://files.pythonhosted.org/packages/2f/ac/be1f2767b7222347d2fefc18d8d58e9febfd9919190cc6fbd8a4d22d6eab/numpy-1.26.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:cc392fdcbd21d4be6ae1bb4475a03ce3b025cd49a9be5345d76d7585aea69440", size = 13956966, upload-time = "2023-11-12T22:52:18.231Z" }, + { url = "https://files.pythonhosted.org/packages/38/39/f726e49ca91cbc336ff297d458dd20b4db2a4204198b075b7f7cc3d3c0ba/numpy-1.26.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:36340109af8da8805d8851ef1d74761b3b88e81a9bd80b290bbfed61bd2b4f75", size = 14211889, upload-time = "2023-11-12T22:52:50.334Z" }, + { url = "https://files.pythonhosted.org/packages/64/41/284783f1014685201e447ea976e85fed0e351f5debbaf3ee6d7645521f1d/numpy-1.26.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bcc008217145b3d77abd3e4d5ef586e3bdfba8fe17940769f8aa09b99e856c00", size = 18228265, upload-time = "2023-11-12T22:53:30.754Z" }, + { url = "https://files.pythonhosted.org/packages/b9/f0/01a7dade8233e6f3c380e2b271aedc98dd1902363661adfb8ab4364c5629/numpy-1.26.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:3ced40d4e9e18242f70dd02d739e44698df3dcb010d31f495ff00a31ef6014fe", size = 13867221, upload-time = "2023-11-12T22:54:06.325Z" }, + { url = "https://files.pythonhosted.org/packages/6c/89/0ef844673002e08444881bff4b6d2a940fbce934d8cd431aa76c8e46e42a/numpy-1.26.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:b272d4cecc32c9e19911891446b72e986157e6a1809b7b56518b4f3755267523", size = 18065825, upload-time = "2023-11-12T22:54:45.565Z" }, + { url = "https://files.pythonhosted.org/packages/15/b7/a7acec96d8c58bf40a24c9fccce988a819840692762e16bf1fc256e1c26a/numpy-1.26.2-cp310-cp310-win32.whl", hash = "sha256:22f8fc02fdbc829e7a8c578dd8d2e15a9074b630d4da29cda483337e300e3ee9", size = 20752812, upload-time = "2023-11-12T22:55:26.963Z" }, + { url = "https://files.pythonhosted.org/packages/24/b5/fed6f7e582937eb947369dccf6c94602598a25f23e482d1b1f2299159328/numpy-1.26.2-cp310-cp310-win_amd64.whl", hash = "sha256:26c9d33f8e8b846d5a65dd068c14e04018d05533b348d9eaeef6c1bd787f9919", size = 15795633, upload-time = "2023-11-12T22:56:49.532Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pandas" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "pytz" }, + { name = "tzdata" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3d/59/2afa81b9fb300c90531803c0fd43ff4548074fa3e8d0f747ef63b3b5e77a/pandas-2.2.1.tar.gz", hash = "sha256:0ab90f87093c13f3e8fa45b48ba9f39181046e8f3317d3aadb2fffbb1b978572", size = 4395256, upload-time = "2024-02-23T15:32:10.41Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/f4495f8ab5a58b1eeee06b5abd811e0a93f7b75acdc89380797f99bdf91a/pandas-2.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8df8612be9cd1c7797c93e1c5df861b2ddda0b48b08f2c3eaa0702cf88fb5f88", size = 12543124, upload-time = "2024-02-23T15:30:19.248Z" }, + { url = "https://files.pythonhosted.org/packages/4f/19/0ae5f1557badfcae1052c1397041a2c5441e9f31e1c7b0cce7f8bc585f4e/pandas-2.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0f573ab277252ed9aaf38240f3b54cfc90fff8e5cab70411ee1d03f5d51f3944", size = 11285572, upload-time = "2024-02-23T15:30:24.562Z" }, + { url = "https://files.pythonhosted.org/packages/5d/d2/df8047f8c3648eb6b3ee86ef7ee811ad01e55b47a14ea02fe36d601e12cd/pandas-2.2.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f02a3a6c83df4026e55b63c1f06476c9aa3ed6af3d89b4f04ea656ccdaaaa359", size = 15629656, upload-time = "2024-02-23T15:30:29.578Z" }, + { url = "https://files.pythonhosted.org/packages/19/df/8d789d96a9e338cf28cb7978fa93ef5da53137624b7ef032f30748421c2b/pandas-2.2.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c38ce92cb22a4bea4e3929429aa1067a454dcc9c335799af93ba9be21b6beb51", size = 13024911, upload-time = "2024-02-23T15:30:33.593Z" }, + { url = "https://files.pythonhosted.org/packages/11/a1/9d5505c6c56740f7ed8bd78c8756fb76aeff1c706b30e6930ddf90693aee/pandas-2.2.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:c2ce852e1cf2509a69e98358e8458775f89599566ac3775e70419b98615f4b06", size = 16275635, upload-time = "2024-02-23T15:30:38.629Z" }, + { url = "https://files.pythonhosted.org/packages/d6/99/378e9108cf3562c7c6294249f1bfd3be08325af5e96af435fb221dd1c320/pandas-2.2.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:53680dc9b2519cbf609c62db3ed7c0b499077c7fefda564e330286e619ff0dd9", size = 13880082, upload-time = "2024-02-23T15:30:42.111Z" }, + { url = "https://files.pythonhosted.org/packages/93/26/2a695303a4a3194014dca7cb5d5ce08f0d2c6baa344fb5f562c642e77b2b/pandas-2.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:94e714a1cca63e4f5939cdce5f29ba8d415d85166be3441165edd427dc9f6bc0", size = 11592785, upload-time = "2024-02-23T15:30:46.129Z" }, +] + +[[package]] +name = "pillow" +version = "10.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/64/9e/7e638579cce7dc346632f020914141a164a872be813481f058883ee8d421/Pillow-10.0.1.tar.gz", hash = "sha256:d72967b06be9300fed5cfbc8b5bafceec48bf7cdc7dab66b1d2549035287191d", size = 50528551, upload-time = "2023-09-15T14:00:45.533Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/34/4e02804420b1cd9371ce5f7c0da7024a8450e6b092c336d3233d6c6448d9/Pillow-10.0.1-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:8f06be50669087250f319b706decf69ca71fdecd829091a37cc89398ca4dc17a", size = 3730082, upload-time = "2023-09-15T13:56:57.712Z" }, + { url = "https://files.pythonhosted.org/packages/55/93/9a8b06ca5917f246d1cd31351d5ec38651f8f5f9619c8db5751583f49ba3/Pillow-10.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:50bd5f1ebafe9362ad622072a1d2f5850ecfa44303531ff14353a4059113b12d", size = 3281492, upload-time = "2023-09-15T13:57:03.209Z" }, + { url = "https://files.pythonhosted.org/packages/53/ac/3f9450af6583199b1c408f3a9a0f65e6dedc25aed15f691a74bcfa2d1077/Pillow-10.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e6a90167bcca1216606223a05e2cf991bb25b14695c518bc65639463d7db722d", size = 3411397, upload-time = "2023-09-15T13:57:07.418Z" }, + { url = "https://files.pythonhosted.org/packages/52/97/9fb1b66fbf98893722247fdf6ac39467a7b177483b546a94dbcff66653f0/Pillow-10.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f11c9102c56ffb9ca87134bd025a43d2aba3f1155f508eff88f694b33a9c6d19", size = 3486505, upload-time = "2023-09-15T13:57:10.787Z" }, + { url = "https://files.pythonhosted.org/packages/a1/a7/35db8e4d15e8f139f9eda3378260a91b2ef7342342769dbeaff8dbc29c26/Pillow-10.0.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:186f7e04248103482ea6354af6d5bcedb62941ee08f7f788a1c7707bc720c66f", size = 3462862, upload-time = "2023-09-15T13:57:14.17Z" }, + { url = "https://files.pythonhosted.org/packages/7a/07/e896b096a77375e78e02ce222ae4fd6014928cd76c691d312060a1645dfa/Pillow-10.0.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:0462b1496505a3462d0f35dc1c4d7b54069747d65d00ef48e736acda2c8cbdff", size = 3607904, upload-time = "2023-09-15T13:57:18.109Z" }, + { url = "https://files.pythonhosted.org/packages/d2/4f/9bb94120b219e050588684cc5f4d1d48d97e0e4a1ec6c1aa68369132d2c7/Pillow-10.0.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d889b53ae2f030f756e61a7bff13684dcd77e9af8b10c6048fb2c559d6ed6eaf", size = 3548900, upload-time = "2023-09-15T13:57:21.438Z" }, + { url = "https://files.pythonhosted.org/packages/9e/33/0140c7dbe509d4e765558c976bea108fd2f63d527a9e16860392ccbcae72/Pillow-10.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:552912dbca585b74d75279a7570dd29fa43b6d93594abb494ebb31ac19ace6bd", size = 3587084, upload-time = "2023-09-15T13:57:25.339Z" }, + { url = "https://files.pythonhosted.org/packages/23/ca/7296d769f62266c0f94bf76496bc77114e7a96d2de3d7bcba91d0ba2856f/Pillow-10.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:787bb0169d2385a798888e1122c980c6eff26bf941a8ea79747d35d8f9210ca0", size = 2513767, upload-time = "2023-09-15T13:57:28.215Z" }, + { url = "https://files.pythonhosted.org/packages/2c/30/458e634572b833d8208ee83ca4e43bb6c21ccdd1b6149f4baf072f8807e8/Pillow-10.0.1-pp310-pypy310_pp73-macosx_10_10_x86_64.whl", hash = "sha256:32bec7423cdf25c9038fef614a853c9d25c07590e1a870ed471f47fb80b244db", size = 3684471, upload-time = "2023-09-15T13:59:52.5Z" }, + { url = "https://files.pythonhosted.org/packages/af/66/d2aaf99819c7aef39f65b43c67ef968128aff33d44c230ec675b40a7ffb4/Pillow-10.0.1-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b7cf63d2c6928b51d35dfdbda6f2c1fddbe51a6bc4a9d4ee6ea0e11670dd981e", size = 3400133, upload-time = "2023-09-15T13:59:56.646Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ba/391fd1ff93e158cb34de3d2e003c070ee85b8840a4fe99a900463d75efad/Pillow-10.0.1-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:f6d3d4c905e26354e8f9d82548475c46d8e0889538cb0657aa9c6f0872a37aa4", size = 3489956, upload-time = "2023-09-15T14:00:00.343Z" }, + { url = "https://files.pythonhosted.org/packages/c8/d0/069cb663f6b809668e4ed6620dc56d518f5bd30134458f783b5553d2ba67/Pillow-10.0.1-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:847e8d1017c741c735d3cd1883fa7b03ded4f825a6e5fcb9378fd813edee995f", size = 2513935, upload-time = "2023-09-15T14:00:04.705Z" }, +] + +[[package]] +name = "pyparsing" +version = "3.3.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e4/11/b213bebff182584360cb8d17c72c1677fec5c5c228de439e63bcf8ab1c8f/pyparsing-3.3.3.tar.gz", hash = "sha256:928ae7e20211f3b6f3915a72f06a0cfd29ab9d24279dd6346b6b1a7146397d36", size = 1050487, upload-time = "2026-09-20T20:59:05.609Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/bb/d215ee7c73b61497b28a5503f9f53523f294fcc936762b7caf90e0c1c2b5/pyparsing-3.3.3-py3-none-any.whl", hash = "sha256:ece8c00a69cf01b45d0b1dedabb469c90d8caf996d4fda40f147627a122849a4", size = 126420, upload-time = "2026-09-20T20:59:04.025Z" }, +] + +[[package]] +name = "python-dateutil" +version = "3.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, + { name = "tzdata" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/df/f15f80383f152452d8985b36d95cb8751f901a3e0ff4cd27f7e118f168e2/python_dateutil-3.0.0.tar.gz", hash = "sha256:7ebcb3c29a92f3ceafe79a7a48a5c45677e9e004bbeec5742ac0bbb00560b7a2", size = 271066, upload-time = "2026-10-21T23:41:45.075Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/17/cebee01d1c0c1441174cf0456ddc91b280d79acb65340c05e41d77866ee8/python_dateutil-3.0.0-py2.py3-none-any.whl", hash = "sha256:3c79630202e4152c6add678446984236efe4f57e84d8e892aab0637bab82e9fd", size = 94659, upload-time = "2026-10-21T23:41:43.431Z" }, +] + +[[package]] +name = "pytz" +version = "2026.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/b0/168030b0bfcc334cf09fa3313afdb888d892c8a0027324683d1de0ec277e/pytz-2026.6.tar.gz", hash = "sha256:3b0fc2727707663d6e517936f444fd8a8b240f1fc32f4bf7ee03274245945284", size = 318520, upload-time = "2026-10-29T09:05:47.011Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b6/25/564f3fc5b70c1398d706c51f2b1470d0ea0b1b2f96fee67877b6f5bdf3b6/pytz-2026.6-py2.py3-none-any.whl", hash = "sha256:7fed3fee8fd6bd4a2b7b75a8d56021d0624127bf69bd5691f115db929ac570b3", size = 506321, upload-time = "2026-10-29T09:05:45.14Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/af35f7ea75cf72f2cd079c95ee16797de7cd71f29ea7c68ae5ce7be1eda0/PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43", size = 125201, upload-time = "2023-07-18T00:00:23.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/06/4beb652c0fe16834032e54f0956443d4cc797fe645527acee59e7deaa0a2/PyYAML-6.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a", size = 189447, upload-time = "2023-07-17T23:57:04.325Z" }, + { url = "https://files.pythonhosted.org/packages/5b/07/10033a403b23405a8fc48975444463d3d10a5c2736b7eb2550b07b367429/PyYAML-6.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f", size = 169264, upload-time = "2023-07-17T23:57:07.787Z" }, + { url = "https://files.pythonhosted.org/packages/f1/26/55e4f21db1f72eaef092015d9017c11510e7e6301c62a6cfee91295d13c6/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938", size = 677003, upload-time = "2023-07-17T23:57:13.144Z" }, + { url = "https://files.pythonhosted.org/packages/ba/91/090818dfa62e85181f3ae23dd1e8b7ea7f09684864a900cab72d29c57346/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d", size = 699070, upload-time = "2023-07-17T23:57:19.402Z" }, + { url = "https://files.pythonhosted.org/packages/29/61/bf33c6c85c55bc45a29eee3195848ff2d518d84735eb0e2d8cb42e0d285e/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515", size = 705525, upload-time = "2023-07-17T23:57:25.272Z" }, + { url = "https://files.pythonhosted.org/packages/07/91/45dfd0ef821a7f41d9d0136ea3608bb5b1653e42fd56a7970532cb5c003f/PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290", size = 707514, upload-time = "2023-08-28T18:43:20.945Z" }, + { url = "https://files.pythonhosted.org/packages/b6/a0/b6700da5d49e9fed49dc3243d3771b598dad07abb37cc32e524607f96adc/PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924", size = 130488, upload-time = "2023-07-17T23:57:28.144Z" }, + { url = "https://files.pythonhosted.org/packages/24/97/9b59b43431f98d01806b288532da38099cc6f2fea0f3d712e21e269c0279/PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d", size = 145338, upload-time = "2023-07-17T23:57:31.118Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "joblib" }, + { name = "numpy" }, + { name = "scipy" }, + { name = "threadpoolctl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/00/835e3d280fdd7784e76bdef91dd9487582d7951a7254f59fc8004fc8b213/scikit-learn-1.3.2.tar.gz", hash = "sha256:a2f54c76accc15a34bfb9066e6c7a56c1e7235dda5762b990792330b52ccfb05", size = 7510251, upload-time = "2023-10-23T13:47:55.287Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/53/570b55a6e10b8694ac1e3024d2df5cd443f1b4ff6d28430845da8b9019b3/scikit_learn-1.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e326c0eb5cf4d6ba40f93776a20e9a7a69524c4db0757e7ce24ba222471ee8a1", size = 10209999, upload-time = "2023-10-23T13:46:30.373Z" }, + { url = "https://files.pythonhosted.org/packages/70/d0/50ace22129f79830e3cf682d0a2bd4843ef91573299d43112d52790163a8/scikit_learn-1.3.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:535805c2a01ccb40ca4ab7d081d771aea67e535153e35a1fd99418fcedd1648a", size = 9479353, upload-time = "2023-10-23T13:46:34.368Z" }, + { url = "https://files.pythonhosted.org/packages/8f/46/fcc35ed7606c50d3072eae5a107a45cfa5b7f5fa8cc48610edd8cc8e8550/scikit_learn-1.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1215e5e58e9880b554b01187b8c9390bf4dc4692eedeaf542d3273f4785e342c", size = 10304705, upload-time = "2023-10-23T13:46:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/d0/0b/26ad95cf0b747be967b15fb71a06f5ac67aba0fd2f9cd174de6edefc4674/scikit_learn-1.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0ee107923a623b9f517754ea2f69ea3b62fc898a3641766cb7deb2f2ce450161", size = 10827807, upload-time = "2023-10-23T13:46:41.59Z" }, + { url = "https://files.pythonhosted.org/packages/69/8a/cf17d6443f5f537e099be81535a56ab68a473f9393fbffda38cd19899fc8/scikit_learn-1.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:35a22e8015048c628ad099da9df5ab3004cdbf81edc75b396fd0cff8699ac58c", size = 9255427, upload-time = "2023-10-23T13:46:44.826Z" }, +] + +[[package]] +name = "scipy" +version = "1.15.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/2f/4966032c5f8cc7e6a60f1b2e0ad686293b9474b65246b0c642e3ef3badd0/scipy-1.15.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:a345928c86d535060c9c2b25e71e87c39ab2f22fc96e9636bd74d1dbf9de448c", size = 38702770, upload-time = "2025-05-08T16:04:20.849Z" }, + { url = "https://files.pythonhosted.org/packages/a0/6e/0c3bf90fae0e910c274db43304ebe25a6b391327f3f10b5dcc638c090795/scipy-1.15.3-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ad3432cb0f9ed87477a8d97f03b763fd1d57709f1bbde3c9369b1dff5503b253", size = 30094511, upload-time = "2025-05-08T16:04:27.103Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b1/4deb37252311c1acff7f101f6453f0440794f51b6eacb1aad4459a134081/scipy-1.15.3-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:aef683a9ae6eb00728a542b796f52a5477b78252edede72b8327a886ab63293f", size = 22368151, upload-time = "2025-05-08T16:04:31.731Z" }, + { url = "https://files.pythonhosted.org/packages/38/7d/f457626e3cd3c29b3a49ca115a304cebb8cc6f31b04678f03b216899d3c6/scipy-1.15.3-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:1c832e1bd78dea67d5c16f786681b28dd695a8cb1fb90af2e27580d3d0967e92", size = 25121732, upload-time = "2025-05-08T16:04:36.596Z" }, + { url = "https://files.pythonhosted.org/packages/db/0a/92b1de4a7adc7a15dcf5bddc6e191f6f29ee663b30511ce20467ef9b82e4/scipy-1.15.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:263961f658ce2165bbd7b99fa5135195c3a12d9bef045345016b8b50c315cb82", size = 35547617, upload-time = "2025-05-08T16:04:43.546Z" }, + { url = "https://files.pythonhosted.org/packages/8e/6d/41991e503e51fc1134502694c5fa7a1671501a17ffa12716a4a9151af3df/scipy-1.15.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e2abc762b0811e09a0d3258abee2d98e0c703eee49464ce0069590846f31d40", size = 37662964, upload-time = "2025-05-08T16:04:49.431Z" }, + { url = "https://files.pythonhosted.org/packages/25/e1/3df8f83cb15f3500478c889be8fb18700813b95e9e087328230b98d547ff/scipy-1.15.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:ed7284b21a7a0c8f1b6e5977ac05396c0d008b89e05498c8b7e8f4a1423bba0e", size = 37238749, upload-time = "2025-05-08T16:04:55.215Z" }, + { url = "https://files.pythonhosted.org/packages/93/3e/b3257cf446f2a3533ed7809757039016b74cd6f38271de91682aa844cfc5/scipy-1.15.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5380741e53df2c566f4d234b100a484b420af85deb39ea35a1cc1be84ff53a5c", size = 40022383, upload-time = "2025-05-08T16:05:01.914Z" }, + { url = "https://files.pythonhosted.org/packages/d1/84/55bc4881973d3f79b479a5a2e2df61c8c9a04fcb986a213ac9c02cfb659b/scipy-1.15.3-cp310-cp310-win_amd64.whl", hash = "sha256:9d61e97b186a57350f6d6fd72640f9e99d5a4a2b8fbf4b9ee9a841eab327dc13", size = 41259201, upload-time = "2025-05-08T16:05:08.166Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "threadpoolctl" +version = "3.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/dc/6c58154c1c65f758ea979e7139cb76993a9cfc662d14e9be3c4a667cfb77/threadpoolctl-3.7.0.tar.gz", hash = "sha256:61348cfb77d53b9242e0017029244b559b810c142ced65b4e21eeca1843959a7", size = 31961, upload-time = "2026-09-15T15:46:20.263Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/3f/f88a53f60a472b46f4023f56d204dd7de33d34c5d2acbfa0d70a674e639e/threadpoolctl-3.7.0-py3-none-any.whl", hash = "sha256:cd8b60b5641b45c67bbf73c64c843235fc2d8a480c87389f52f5dbee893b86be", size = 26362, upload-time = "2026-09-15T15:46:19.168Z" }, +] + +[[package]] +name = "torch" +version = "2.1.2+cu118" +source = { registry = "https://download.pytorch.org/whl/cu118" } +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu118/torch-2.1.2%2Bcu118-cp310-cp310-linux_x86_64.whl", hash = "sha256:60396358193f238888540f4a38d78485f161e28ec17fa445f0373b5350ef21f0", upload-time = "2024-01-29T18:38:38Z" }, + { url = "https://download-r2.pytorch.org/whl/cu118/torch-2.1.2%2Bcu118-cp310-cp310-win_amd64.whl", hash = "sha256:0ddfa0336d678316ff4c35172d85cddab5aa5ded4f781158e725096926491db9", upload-time = "2024-01-29T18:40:05Z" }, +] + +[[package]] +name = "torchvision" +version = "0.16.2" +source = { registry = "https://download.pytorch.org/whl/cu118" } +resolution-markers = [ + "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'", +] +dependencies = [ + { name = "numpy", marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "pillow", marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/torchvision-0.16.2-cp310-cp310-linux_aarch64.whl", hash = "sha256:96c7583700112a410bdc4e1e4f118c429dab49c29c9a31a2cc3579bc9b08b19d", upload-time = "2024-01-30T03:20:51Z" }, +] + +[[package]] +name = "torchvision" +version = "0.16.2+cu118" +source = { registry = "https://download.pytorch.org/whl/cu118" } +resolution-markers = [ + "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'", +] +dependencies = [ + { name = "numpy", marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, + { name = "pillow", marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, + { name = "requests", marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, + { name = "torch", marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu118/torchvision-0.16.2%2Bcu118-cp310-cp310-linux_x86_64.whl", hash = "sha256:18470aef0bbde73f5a6a96135cd457f4d8be31f60be7ceae4ef5174f02f73add", upload-time = "2024-01-29T19:09:38Z" }, + { url = "https://download-r2.pytorch.org/whl/cu118/torchvision-0.16.2%2Bcu118-cp310-cp310-win_amd64.whl", hash = "sha256:689f2458e8924c47b7ba9f50dca353423b75214184b905d540f69d9b962b2fdf", upload-time = "2024-01-29T19:09:39Z" }, +] + +[[package]] +name = "tqdm" +version = "4.70.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0d/ea/b2a5bd54b28a324dae8211928b2d730b6547500342c7e6c6dea08bd0a485/tqdm-4.70.1.tar.gz", hash = "sha256:cefd0eca11b2a37a3aee776544d4f4ae913f02688135b5556b8788dfa474afc4", size = 171846, upload-time = "2026-09-11T07:25:16.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/03/921a3d3c75785aca9ebfbfcabfbc3a1be12e2ab5265deb026d55a5a3f83e/tqdm-4.70.1-py3-none-any.whl", hash = "sha256:c293e525e6fef9c20e8728fd4612df02a0aa31bb5fe91ecd93e123b1b7bffa73", size = 80199, upload-time = "2026-09-11T07:25:14.599Z" }, +] + +[[package]] +name = "triton" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock", marker = "platform_machine != 'aarch64' or platform_python_implementation != 'CPython' or sys_platform != 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4d/22/91a8af421c8a8902dde76e6ef3db01b258af16c53d81e8c0d0dc13900a9e/triton-2.1.0-0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:66439923a30d5d48399b08a9eae10370f6c261a5ec864a64983bae63152d39d7", size = 89215103, upload-time = "2023-09-01T07:26:13.625Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "tzdata" +version = "2026.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7a/29/6a1391d31cd00ff9e927215a224a880e1e84e826b64ba80a499e51c074d8/tzdata-2026.6.tar.gz", hash = "sha256:85d2cd9c09f81f0bb29db27759fc83af4f1e620efb814f8c5bf8b4cc4feece6d", size = 201017, upload-time = "2026-10-28T09:57:09.137Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0a/66/9a06c2227fedfd4919f0eaaa8b94377d9730882a7e4f95823c9a331ab5e1/tzdata-2026.6-py2.py3-none-any.whl", hash = "sha256:0a10eec6e13ef16b627fabed9a5f2ba61a62b1c67f327305f4e7e412a70b4aec", size = 348210, upload-time = "2026-10-28T09:57:07.276Z" }, +] + +[[package]] +name = "urllib3" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, +]