"""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())