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