1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  9
 10TRACK = "dynamics"
 11MODEL = "rnn_small"
 12DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 13
 14class Potential(nn.Module):
 15    def __init__(self, dim):
 16        super().__init__()
 17        self.net = nn.Sequential(nn.Linear(dim, 32), nn.Tanh(), nn.Linear(32, 1))
 18    def forward(self, x):
 19        z = self.net(x)
 20        return z - z.mean()                 # minibatch gauge fixing
 21
 22def seed_all(seed):
 23    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 24    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 25
 26def make_next(x, pred):
 27    # A learned one-coordinate state transition: retain the observed window and
 28    # replace its final angle by the model's next-angle prediction.
 29    z = x.clone()
 30    z[:, -3] = pred[:, 0]
 31    return z
 32
 33def train_one(seed, cfg, mode, capture=False):
 34    global DEVICE
 35    seed_all(seed)
 36    ds = get_dataset(TRACK, seed, n_train=400, n_test=200)
 37    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 38    pot = Potential(ds["xtr"].shape[1]) if mode == "coh" else None
 39    c = nn.Parameter(torch.tensor(0.0)) if mode == "coh" else None
 40    params = list(net.parameters()) + ([] if pot is None else list(pot.parameters()) + [c])
 41    opt = torch.optim.Adam(params, lr=cfg["lr"])
 42    dev = DEVICE
 43    try:
 44        net.to(dev); dsx, dsy = ds["xtr"].to(dev), ds["ytr"].to(dev)
 45        if pot is not None: pot.to(dev); c.data = c.data.to(dev)
 46        for ep in range(cfg["epochs"]):
 47            net.train()
 48            perm = torch.randperm(len(dsx), device=dev)
 49            for ii in range(0, len(dsx), 64):
 50                x = dsx[perm[ii:ii+64]].detach().requires_grad_(True)
 51                pred = net(x)
 52                task = ((pred - dsy[perm[ii:ii+64]]) ** 2).mean()
 53                # selected 1D unstable-coordinate log Jacobian proxy
 54                g = torch.autograd.grad(pred[:, 0].sum(), x, create_graph=True)[0]
 55                ell = torch.log(g[:, -3].abs() + 1e-3)
 56                if mode == "none":
 57                    reg = 0.0
 58                elif mode == "point":
 59                    reg = ((ell - c0(ell)) ** 2).mean()
 60                else:
 61                    nxt = make_next(x, pred)
 62                    reg = (ell - pot(nxt)[:, 0] + pot(x)[:, 0] - c) .pow(2).mean()
 63                loss = task + cfg["lam"] * reg
 64                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 5.0); opt.step()
 65        net.eval()
 66        with torch.no_grad():
 67            metric = float(((net(ds["xte"].to(dev)) - ds["yte"].to(dev))**2).mean())
 68        if capture:
 69            return metric, net, pot, c, ds
 70        return metric
 71    except RuntimeError:
 72        # Explicit CPU fallback for constrained CUDA/cuDNN environments.
 73        if dev != "cpu":
 74            DEVICE = "cpu"
 75            return train_one(seed, cfg, mode, capture)
 76        raise
 77
 78def c0(ell):
 79    return ell.mean().detach()  # pointwise constant baseline, no learned potential
 80
 81def signature(net, pot, c, ds, seed=0):
 82    if pot is None: return {"confirmed": False, "reason": "no potential"}
 83    net.eval(); pot.eval(); dev = next(net.parameters()).device
 84    x = ds["xte"][:48].to(dev)
 85    rows = []
 86    for k in (1, 4, 8):
 87        z = x.clone(); s = torch.zeros(len(z), device=dev)
 88        for _ in range(k):
 89            z.requires_grad_(True); p = net(z)
 90            gg = torch.autograd.grad(p[:,0].sum(), z, create_graph=False)[0]
 91            ell = torch.log(gg[:, -3].abs() + 1e-3)
 92            s += ell
 93            z = make_next(z.detach(), p.detach())
 94        with torch.no_grad():
 95            r = s - (pot(z)[:,0] - pot(x)[:,0] + k*c)
 96            rows.append({"k": k, "observed_std_R_over_k": float((r/k).std()),
 97                         "observed_std_ell_minus_c": float((s/k-c).std())})
 98    vals = np.array([q["observed_std_R_over_k"] for q in rows])
 99    # Stage-1 prediction: coboundary residual should not grow linearly with horizon.
100    confirmed = bool(vals[-1] <= 1.25 * vals[0] + 1e-8)
101    return {"prediction": "finite-horizon residual per step remains bounded rather than accumulating linearly",
102            "observed": rows, "confirmed": confirmed}
103
104def main():
105    # Same union is used by baseline sweep and idea settings (parity).
106    grid = [{"lr": lr, "lam": lam, "epochs": 10} for lr in (1e-3, 3e-3, 1e-2) for lam in (1e-3, 1e-2)]
107    cache = {}
108    def maker(mode):
109        def f(cfg):
110            return lambda seed: train_one(seed, cfg, mode)
111        return f
112    base = sweep_baseline(maker("point"), grid)
113    best = base["best_cfg"]
114    # Best baseline config plus two nearby settings are all represented in grid.
115    idea_grid = [best, {"lr": 1e-3, "lam": best["lam"], "epochs": 10},
116                 {"lr": 1e-2, "lam": best["lam"], "epochs": 10}]
117    idea_runs = []
118    chosen = idea_grid[0]
119    for cfg in idea_grid:
120        r = evaluate(maker("coh")(cfg))
121        idea_runs.append({"cfg": cfg, "result": r})
122    chosen_run = min(idea_runs, key=lambda q: q["result"]["mean"])
123    idea = chosen_run["result"]
124    # Refit one paired seed for a mechanism signature from trained networks.
125    m, n, p, cc, ds = train_one(0, chosen_run["cfg"], "coh", True)
126    sig = signature(n, p, cc, ds)
127    report = make_report(TRACK, MODEL, base, idea, {"mechanism_signature": sig,
128        "idea_config_sweep": idea_runs, "device": DEVICE,
129        "primary_metric": "test MSE"})
130    report["baseline"]["method"] = "pointwise constant selected-coordinate log-Jacobian"
131    report["idea"]["method"] = "cohomological potential residual"
132    report["protocol_notes"] = "8 paired seeds; baseline sweep uses 4 seeds then full reevaluation; shared lr/lambda union"
133    Path("bench_report.json").write_text(json.dumps(report, indent=2))
134    print(json.dumps(report, indent=2))
135
136if __name__ == "__main__": main()