File size: 8,189 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """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 # 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())
|