1import json, math, time, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import sys
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import make_model, train_model, evaluate, sweep_baseline, make_report
  8from bench.custom_tracks.poisson_jfb_short_trace import get_dataset as poisson_data
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12EPOCHS = 12
 13BATCH = 128
 14# The idea's a priori schedules at m=30, C_N=1.
 15KAPPAS = (0.5, 1.0, 1.5)
 16M_FINE = 30
 17IDEA_N = {k: int(math.ceil(M_FINE ** k)) for k in KAPPAS}
 18# Include every idea lr and every idea data budget in the baseline grid.
 19LRS = (0.0015, 0.003, 0.006)
 20BASE_GRID = [{"lr": lr, "n_train": n} for lr in LRS for n in sorted(set(IDEA_N.values()))]
 21
 22
 23def seed_all(seed):
 24    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 25    if torch.cuda.is_available():
 26        try: torch.cuda.manual_seed_all(seed)
 27        except Exception: pass
 28
 29
 30def dataset(seed, n_train, m=M_FINE, n_test=400):
 31    # Track has exact 30-point PDE observations. Coarser resolutions are
 32    # represented by nested spatial subsampling, preserving the same operator.
 33    d = poisson_data(seed, n_train, n_test)
 34    if m == 30:
 35        return d
 36    idx = np.linspace(0, d["out_dim"] - 1, m).round().astype(int)
 37    d = dict(d)
 38    d["ytr"] = d["ytr"][:, idx].copy(); d["yte"] = d["yte"][:, idx].copy()
 39    d["out_dim"] = m; d["n_grid"] = m
 40    return d
 41
 42
 43def run_one(seed, cfg, m=M_FINE):
 44    seed_all(seed)
 45    d = dataset(seed, int(cfg["n_train"]), m=m)
 46    for key in ("xtr", "ytr", "xte", "yte"):
 47        d[key] = torch.from_numpy(np.asarray(d[key], dtype=np.float32))
 48    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
 49    t0 = time.perf_counter()
 50    net, metric, hist = train_model(net, d, epochs=EPOCHS, lr=float(cfg["lr"]), batch=BATCH, log=lambda *_: None)
 51    wall = time.perf_counter() - t0
 52    if net is None: raise RuntimeError("canonical training failed")
 53    # Return trained behavior for the mechanism signature.
 54    with torch.no_grad():
 55        try:
 56            device = next(net.parameters()).device
 57            pred = net(d["xte"].to(device)).detach().cpu().numpy()
 58        except Exception:
 59            pred = np.zeros_like(d["yte"])
 60    return float(metric), {"model": net, "pred": pred, "truth": d["yte"].cpu().numpy(), "wall": wall,
 61                           "history": hist, "n_train": int(cfg["n_train"]), "m": m}
 62
 63
 64def factory(cfg):
 65    return lambda seed: run_one(seed, cfg)[0]
 66
 67
 68def main():
 69    t0 = time.perf_counter()
 70    base = sweep_baseline(factory, BASE_GRID, seeds=SWEEP_SEEDS)
 71    best_base_cfg = base["best_cfg"]
 72    # Explicitly run idea at baseline's lr and two nearby learning-rate settings;
 73    # each setting uses one of the predeclared resolution-aware budgets.
 74    idea_cfgs = [
 75        {"lr": best_base_cfg["lr"], "n_train": IDEA_N[1.0], "kappa": 1.0},
 76        {"lr": best_base_cfg["lr"], "n_train": IDEA_N[1.5], "kappa": 1.5},
 77        {"lr": 0.006 if best_base_cfg["lr"] != 0.006 else 0.0015,
 78         "n_train": IDEA_N[1.5], "kappa": 1.5},
 79    ]
 80    idea_runs = []
 81    for cfg in idea_cfgs:
 82        r = evaluate(factory(cfg), seeds=SEEDS)
 83        idea_runs.append({"cfg": cfg, "result": r})
 84    best_idea = min(idea_runs, key=lambda z: z["result"]["mean"])
 85    report = make_report("poisson_jfb_short_trace", "mlp_tiny", base,
 86                         best_idea["result"], extra={
 87        "mechanism_signature": mechanism_signature(best_base_cfg, best_idea["cfg"]),
 88        "audit": {"idea_sweep": idea_runs, "epochs": EPOCHS, "batch": BATCH,
 89                  "resolution_scaling": resolution_scaling(best_base_cfg)},
 90        "custom_track": {"name": "poisson_jfb_short_trace",
 91                         "file": "/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_jfb_short_trace.py",
 92                         "domain": "pde"}
 93    })
 94    report["elapsed_seconds"] = time.perf_counter() - t0
 95    Path("bench_report.json").write_text(json.dumps(report, indent=2))
 96    print(json.dumps(report, indent=2))
 97
 98
 99def mechanism_signature(base_cfg, idea_cfg):
100    # Re-test the claimed mechanism on trained models: increasing N should
101    # improve prediction error and input coverage, not merely an analytic proxy.
102    vals = []
103    for n in (IDEA_N[0.5], IDEA_N[1.0], IDEA_N[1.5]):
104        metric, aux = run_one(0, {"lr": idea_cfg["lr"], "n_train": n})
105        err = float(np.mean((aux["pred"] - aux["truth"]) ** 2))
106        vals.append({"n_train": n, "observed_test_mse": metric,
107                     "observed_prediction_mse_recomputed": err})
108    confirmed = vals[-1]["observed_test_mse"] < vals[0]["observed_test_mse"]
109    return {"claim": "more operator pairs at fixed output resolution reduce learned error",
110            "trained_model_observations": vals, "confirmed": bool(confirmed)}
111
112
113def resolution_scaling(base_cfg):
114    out = []
115    for m in (10, 20, 30):
116        n = int(math.ceil(m ** 1.0))
117        metric, aux = run_one(0, {"lr": base_cfg["lr"], "n_train": n}, m=m)
118        out.append({"m": m, "n_train": n, "test_mse": metric, "wall_seconds": aux["wall"]})
119    return out
120
121if __name__ == "__main__": main()