| """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<R>].{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 |
| from bifurcation.fid import backend as BE |
| from bifurcation.fid.evaluate import bank_members |
| from bifurcation.models.spectral import K_C, mode_int_wavenumbers |
|
|
|
|
| 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} |
| |
| 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()) |
|
|
