| """Reproduce one entry of Table 2, or one depth point of Table 6 / Fig. 3: csFID-10k of a model on the ten test conditions. |
| |
| python -m reproduce.table2 list |
| python -m reproduce.table2 check --system SYS --method M [--S S] [--checkpoint FILE] [--samples DIR] [--hash-large] |
| python -m reproduce.table2 run --system SYS --method M [--S S] --checkpoint FILE --out DIR [--device cuda] |
| python -m reproduce.table2 generate --system SYS --method M [--S S] --checkpoint FILE --out DIR [--gids g1,g2] [--n-samples N] |
| python -m reproduce.table2 score --system SYS --method M [--S S] --out DIR [--samples DIR] |
| python -m reproduce.table2 compare --system SYS --method M [--S S] --out DIR [--tolerance REL] |
| |
| SYS is euler | sqg | waves | nlkg | passive_scalar; M is direct_si | edm | cascaded_si (without --S the selected depth S*, |
| with --S the Cascaded SI model of that depth). The package and the truth bank are taken from <SYS>_DATA and |
| <SYS>_TRUTH_BANK or given with --data and --truth-bank. Everything else comes from records/<SYS>/: the run configuration |
| of the model, its sampler resolution n*, the image adapter and the recorded values. |
| |
| generate 10,000 samples per test condition (bifurcation.sample with the settings of configs/evaluation.yaml: |
| n* sampler steps, noise stream 777 * 100003 + gid, chunks of 2,000) -> OUT/samples/<variant>/group_<gid>.npz |
| score csFID per condition (bifurcation.fid.evaluate.evaluate_csfid_run) of OUT/samples, or of the archived samples |
| given with --samples -> OUT/csfid/{csfid_per_ic.csv, csfid_summary.json, MANIFEST_CSFID.json} |
| compare per-condition values, mean and SEM against the recorded values -> OUT/COMPARE.json |
| run generate + score + compare |
| check the package, truth bank, checkpoint and archived samples against the recorded sha256 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import hashlib |
| import json |
| import sys |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| from bifurcation.config import RECORDS, REPO, SYSTEMS, data_paths, dataset_card, evaluation_settings, records, run_config, unit |
|
|
|
|
| def sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 24), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def asset_hashes(system: str) -> dict: |
| """{section: {relative path: sha256}} from records/<system>/assets.sha256.""" |
| out, sec = {}, None |
| for line in (RECORDS / system / "assets.sha256").read_text().splitlines(): |
| if line.startswith("# "): |
| sec = line[2:].split(":")[0].strip() |
| continue |
| h, q = line.split(" ", 1) |
| out.setdefault(sec, {})[q] = h |
| return out |
|
|
|
|
| def cmd_list(a): |
| print(f"{'system':15s} {'model':16s} {'paper':22s} {'n*':>4s} {'NFE':>6s} csFID-10k") |
| for s in SYSTEMS: |
| rec = records(s) |
| for u in rec["units"]: |
| names = ", ".join({"direct_si": "Direct SI", "edm": "EDM", "cascaded_si": f"Cascaded SI (S*={rec['s_star']})"}[n] |
| for n in u["paper_names"]) or f"depth S={u['S']}" |
| e = u["expected_csfid"] |
| print(f"{s:15s} {u['key']:16s} {names:22s} {u['sampler']['n_steps']:4d} {u['sampler']['nfe']:6d} {e['mean']:.2f} +- {e['sem']:.2f}") |
|
|
|
|
| def cmd_check(a) -> int: |
| rec, u = unit(a.system, a.method, a.S) |
| data, bank = data_paths(a.system, a.data, a.truth_bank) |
| card = dataset_card(a.system) |
| problems, checked = [], 0 |
| big = 200 * 1024 * 1024 |
| for rel, want in card["expected_sha256"].items(): |
| p = data / rel |
| if not p.is_file(): |
| problems.append(f"--data: missing {rel}"); continue |
| if p.stat().st_size > big and not a.hash_large: |
| continue |
| checked += 1 |
| if sha256(p) != want: |
| problems.append(f"--data: {rel} sha256 differs from configs/{a.system}/dataset.yaml") |
| hashes = asset_hashes(a.system) |
| for gid in rec["test_gids"]: |
| rel = f"group_{gid:04d}.npz" |
| p = bank / rel |
| if not p.is_file(): |
| problems.append(f"--truth-bank: missing {rel}"); continue |
| if a.hash_large: |
| checked += 1 |
| if sha256(p) != hashes["truth_bank"][rel]: |
| problems.append(f"--truth-bank: {rel} sha256 differs") |
| if a.checkpoint: |
| got = sha256(Path(a.checkpoint)); checked += 1 |
| if u["checkpoint"] is None: |
| problems.append(f"--checkpoint: no checkpoint of {u['key']} is distributed") |
| elif got != u["checkpoint"]["sha256"]: |
| problems.append(f"--checkpoint: sha256 {got[:16]} is not the recorded {u['checkpoint']['file']}") |
| if a.samples: |
| v = u["sampler"]["variant"] |
| sec = f"samples/{u['paper_names'][0]}" if u["paper_names"] else None |
| for gid in rec["test_gids"]: |
| rel = f"{v}/group_{gid:04d}.npz" |
| p = Path(a.samples) / rel |
| if not p.is_file(): |
| problems.append(f"--samples: missing {rel}"); continue |
| want = (hashes.get(sec) or {}).get(rel) if sec else None |
| if want and a.hash_large: |
| checked += 1 |
| if sha256(p) != want: |
| problems.append(f"--samples: {rel} sha256 differs from the archived samples") |
| print(json.dumps({"system": a.system, "model": u["key"], "files_hashed": checked, "problems": problems}, indent=1)) |
| return 2 if problems else 0 |
|
|
|
|
| def _fresh(out: Path, sub: str) -> Path: |
| d = out / sub |
| if d.exists() and any(d.iterdir()): |
| raise SystemExit(f"refusing: {d} exists and is not empty (use a fresh --out)") |
| d.mkdir(parents=True, exist_ok=True) |
| return d |
|
|
|
|
| def cmd_generate(a): |
| from bifurcation import sample |
| import yaml |
| rec, u = unit(a.system, a.method, a.S) |
| data, bank = data_paths(a.system, a.data, a.truth_bank) |
| ev = evaluation_settings() |
| out = Path(a.out).resolve() |
| out.mkdir(parents=True, exist_ok=True) |
| cfg_path = out / "config.yaml" |
| cfg_path.write_text(yaml.safe_dump(run_config(a.system, u, data), sort_keys=False)) |
| n_samples = int(a.n_samples) if getattr(a, "n_samples", None) else int(ev["fid"]["n_generated"]) |
| argv = ["--config", str(cfg_path), "--data", str(data), "--checkpoint", str(Path(a.checkpoint).resolve()), |
| "--out", str(out / "samples"), "--split", "test", "--n-samples", str(n_samples), |
| "--n-steps", str(u["sampler"]["n_steps"]), "--seed-base", str(ev["sampling"]["crn_base_seed"]), |
| "--chunk", str(ev["sampling"]["sample_chunk"]), "--device", a.device] |
| if rec["x0_from_truth_bank"]: |
| argv += ["--truth-bank", str(bank)] |
| if a.gids: |
| argv += ["--gids", a.gids] |
| print("$ python -m bifurcation.sample " + " ".join(argv), flush=True) |
| sample.main(argv) |
| log = out / "GENERATED.json" |
| entries = json.loads(log.read_text()) if log.is_file() else [] |
| entries.append({"system": a.system, "model": u["key"], "checkpoint": str(Path(a.checkpoint).resolve()), "gids": a.gids or "all", |
| "argv": argv, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")}) |
| log.write_text(json.dumps(entries, indent=1) + "\n") |
|
|
|
|
| def cmd_score(a): |
| from bifurcation.fid.adapters import FIDAdapter |
| from bifurcation.fid.evaluate import evaluate_csfid_run, generated_states |
| rec, u = unit(a.system, a.method, a.S) |
| data, bank = data_paths(a.system, a.data, a.truth_bank) |
| card, fid_cfg = dataset_card(a.system), evaluation_settings()["fid"] |
| out = Path(a.out).resolve() |
| samples = Path(a.samples).resolve() if a.samples else out / "samples" |
| variant = u["sampler"]["variant"] |
| missing = [g for g in rec["test_gids"] if not (samples / variant / f"group_{g:04d}.npz").is_file()] |
| if missing: |
| raise SystemExit(f"{samples / variant}: no samples for test conditions {missing}") |
| n_gen = int(fid_cfg["n_generated"]) |
| gen_raw = generated_states(samples, variant, data, n_gen) |
| gen_raw = {g: x for g, x in gen_raw.items() if g in set(rec["test_gids"])} |
| for g, x in gen_raw.items(): |
| if len(x) != n_gen: |
| raise SystemExit(f"condition {g}: {len(x)} generated states, the paper's read uses {n_gen}") |
| adapter = FIDAdapter.load(RECORDS / a.system / "fid_adapter.json", **(card["fid_adapter"].get("params") or {})) |
| csdir = _fresh(out, "csfid") |
| summary = evaluate_csfid_run(gen_raw, bank, int(rec["fine_horizon_idx"]), int(fid_cfg["n_truth"]), adapter, csdir, |
| out / "fid_cache", fid_cfg, a.device, |
| provenance={"system": a.system, "model": u["key"], "samples": str(samples), "source": "reproduce.table2 score"}) |
| print(f"csFID-10k {a.system} {u['key']}: {summary['csFID']['mean']:.6f} +- {summary['csFID']['sem']:.6f} (n = {summary['csFID']['n']})") |
|
|
|
|
| def cmd_compare(a) -> int: |
| rec, u = unit(a.system, a.method, a.S) |
| out = Path(a.out).resolve() |
| got = {int(r["gid"]): float(r["sfid"]) for r in csv.DictReader(open(out / "csfid" / "csfid_per_ic.csv"))} |
| summ = json.loads((out / "csfid" / "csfid_summary.json").read_text())["csFID"] |
| exp = u["expected_csfid"] |
| rows, worst = [], 0.0 |
| for g in rec["test_gids"]: |
| e, v = exp["per_condition"][str(g)], got.get(g, float("nan")) |
| rel = abs(v - e) / abs(e) |
| worst = max(worst, rel) |
| rows.append({"gid": g, "recorded": e, "this_run": v, "rel_diff": rel}) |
| agg = {k: {"recorded": exp[k], "this_run": summ[k], "rel_diff": abs(summ[k] - exp[k]) / abs(exp[k])} for k in ("mean", "sem")} |
| for r in rows: |
| print(f" gid {r['gid']:4d} recorded {r['recorded']:.6f} this run {r['this_run']:.6f} rel diff {r['rel_diff']:.2e}") |
| print(f" mean {summ['mean']:.6f} (recorded {exp['mean']:.6f}, rel diff {agg['mean']['rel_diff']:.2e}); " |
| f"SEM {summ['sem']:.6f} (recorded {exp['sem']:.6f}, rel diff {agg['sem']['rel_diff']:.2e})") |
| verdict = {"system": a.system, "model": u["key"], "max_rel_diff_per_condition": worst, "aggregate": agg, "per_condition": rows, |
| "bitwise_equal": all(r["rel_diff"] == 0.0 for r in rows) and all(v["rel_diff"] == 0.0 for v in agg.values())} |
| if a.tolerance is not None: |
| verdict["tolerance"] = a.tolerance |
| verdict["within_tolerance"] = worst <= a.tolerance and all(v["rel_diff"] <= a.tolerance for v in agg.values()) |
| (out / "COMPARE.json").write_text(json.dumps(verdict, indent=1) + "\n") |
| print(f" max relative difference per condition: {worst:.3e}" + (f"; within {a.tolerance:g}: {verdict['within_tolerance']}" if a.tolerance is not None else "")) |
| return 1 if (a.tolerance is not None and not verdict["within_tolerance"]) else 0 |
|
|
|
|
| def cmd_run(a) -> int: |
| cmd_generate(a) |
| a.samples = None |
| cmd_score(a) |
| return cmd_compare(a) |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(prog="python -m reproduce.table2", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| sub.add_parser("list") |
|
|
| def unit_args(p, need_data=True): |
| p.add_argument("--system", required=True, choices=SYSTEMS) |
| p.add_argument("--method", required=True, choices=("direct_si", "edm", "cascaded_si")) |
| p.add_argument("--S", type=int, default=None, help="depth of a Cascaded SI model (Table 6); default: the selected S*") |
| if need_data: |
| p.add_argument("--data", default=None, help="dataset package directory (default: $<SYSTEM>_DATA)") |
| p.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $<SYSTEM>_TRUTH_BANK)") |
| c = sub.add_parser("check"); unit_args(c) |
| c.add_argument("--checkpoint"); c.add_argument("--samples") |
| c.add_argument("--hash-large", action="store_true", help="also hash the large files (trajectories, truth banks, samples)") |
| for name in ("generate", "run"): |
| g = sub.add_parser(name); unit_args(g) |
| g.add_argument("--checkpoint", required=True); g.add_argument("--out", required=True) |
| g.add_argument("--gids", default=None, help="comma-separated subset of the test conditions (split generation across jobs)") |
| g.add_argument("--device", default="cuda") |
| if name == "generate": |
| g.add_argument("--n-samples", type=int, default=None, help="fewer samples per condition (a smoke test; `score` refuses fewer than 10,000)") |
| if name == "run": |
| g.add_argument("--tolerance", type=float, default=None) |
| s = sub.add_parser("score"); unit_args(s) |
| s.add_argument("--out", required=True); s.add_argument("--samples", default=None, help="archived samples (default OUT/samples)") |
| s.add_argument("--device", default="cuda") |
| m = sub.add_parser("compare"); unit_args(m, need_data=False) |
| m.add_argument("--out", required=True); m.add_argument("--tolerance", type=float, default=None) |
| a = ap.parse_args(argv) |
| fn = {"list": cmd_list, "check": cmd_check, "generate": cmd_generate, "score": cmd_score, "compare": cmd_compare, "run": cmd_run}[a.cmd] |
| return int(fn(a) or 0) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
