| """Print Tables 2, 5 and 6 of the paper from records/<system>/units.json. No data, no GPU. |
| |
| python -m reproduce.tables [table2 | table5 | table6 | all] |
| |
| `python -m reproduce.table2` regenerates any entry of Table 2 or Table 6 from a checkpoint (or rescores the archived |
| samples) and compares it with these records. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import sys |
| from pathlib import Path |
|
|
| RECORDS = Path(__file__).resolve().parents[1] / "records" |
| ORDER = ("euler", "sqg", "waves", "nlkg", "passive_scalar") |
|
|
|
|
| def _rec(system): |
| return json.loads((RECORDS / system / "units.json").read_text()) |
|
|
|
|
| def _unit(rec, key): |
| return next(u for u in rec["units"] if u["key"] == key) |
|
|
|
|
| def _fmt(e): |
| return f"{e['mean']:.2f} +- {e['sem']:.2f}" |
|
|
|
|
| def table2(): |
| print("Table 2: test csFID-10k, mean +- SEM over the ten test conditions (+: a selected hyperparameter lies on the boundary of its searched range)") |
| print(f"{'system':16s} {'Direct SI':>18s} {'EDM':>18s} {'Cascaded SI':>18s} {'S*':>4s}") |
| for s in ORDER: |
| r = _rec(s) |
| cells = [] |
| for m in ("direct_si", "edm", "cascaded_si"): |
| u = _unit(r, r["table2"][m]) |
| flag = " +" if u.get("selection_flags") else " " |
| cells.append(_fmt(u["expected_csfid"]) + flag) |
| print(f"{s:16s} {cells[0]:>18s} {cells[1]:>18s} {cells[2]:>18s} {r['s_star']:4d}") |
|
|
|
|
| def table5(): |
| print("Table 5: selected configurations (S, learning rate, noise parameter, selected update, sampler steps per segment, NFE)") |
| for s in ORDER: |
| r = _rec(s) |
| for m in ("direct_si", "edm", "cascaded_si"): |
| if s == "passive_scalar" and m == "cascaded_si": |
| continue |
| u = _unit(r, r["table2"][m]); sel = u["selection"] |
| noise = f"P_mean={sel['P_mean']:g}" if "P_mean" in sel else f"sigma={sel['sigma']:g}" |
| print(f"{s:16s} {m:12s} S={u['S']:<3d} lr={sel['lr']:<8g} {noise:16s} update={sel['selected_update']:<6d} " |
| f"steps={u['sampler']['n_steps']:<4d} NFE={u['sampler']['nfe']}") |
|
|
|
|
| def table6(): |
| print("Table 6 / Fig. 3: test csFID-10k of the selected Cascaded SI model at each depth S (S = 1 is Direct SI), and EDM") |
| for s in ORDER: |
| r = _rec(s) |
| row = [f"S={u['S']}: {_fmt(u['expected_csfid'])}" for u in sorted((u for u in r["units"] if u["key"] != "edm"), key=lambda u: u["S"])] |
| print(f"{s:16s} " + " | ".join(row) + f" | EDM: {_fmt(_unit(r, 'edm')['expected_csfid'])}") |
|
|
|
|
| if __name__ == "__main__": |
| what = sys.argv[1] if len(sys.argv) > 1 else "all" |
| for name, fn in (("table2", table2), ("table5", table5), ("table6", table6)): |
| if what in (name, "all"): |
| fn(); print() |
|
|
