| """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: $<SYSTEM>_DATA)") |
| ap.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $<SYSTEM>_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") |
| 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 |
|
|
|
|
| |
|
|
| 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) |
|
|
|
|
| |
|
|
| SFID_CHANNELS = 7 |
| SFID_SPATIAL = 17 |
| SFID_DIM = SFID_SPATIAL * SFID_SPATIAL * SFID_CHANNELS |
| 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") |
|
|
