1import os, sys, json, math, 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, sweep_baseline, evaluate, make_report
  9
 10OUT = Path(__file__).with_name("bench_report.json")
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = (0, 1, 2, 3)
 13EPOCHS = 12
 14NTRAIN, NTEST = 1200, 400
 15BATCH = 128
 16# The union is deliberately shared by both sides.
 17LRS = (1e-3, 3e-3, 1e-2)
 18IDEA_GRID = ({"lr": 1e-3, "k": 0.5, "delta": 0.02},
 19             {"lr": 3e-3, "k": 0.5, "delta": 0.05},
 20             {"lr": 1e-2, "k": 0.5, "delta": 0.10})
 21BASE_GRID = [{"lr": x, "weight_decay": wd} for x in LRS for wd in (0.0, 1e-4)]
 22
 23class MatchedGRU(nn.Module):
 24    """GRUCell(3,64)+linear, optionally followed by phase control."""
 25    def __init__(self, out_dim=1, mode="baseline", k=0.5, delta=0.05):
 26        super().__init__()
 27        self.cell = nn.GRUCell(3, 64)
 28        self.head = nn.Linear(64, out_dim)
 29        self.mode, self.k, self.delta = mode, k, delta
 30        self.last_stats = {}
 31
 32    @staticmethod
 33    def control(h, k):
 34        p = h.view(h.shape[0], 32, 2)
 35        z = torch.complex(p[..., 0], p[..., 1])
 36        z = z / (torch.abs(z) + 1e-6)
 37        r = z.mean(1, keepdim=True)
 38        return (2*k/32) * torch.imag(z * torch.conj(r))
 39
 40    @staticmethod
 41    def rotate(h, u, dt=1.0):
 42        p = h.view(h.shape[0], 32, 2)
 43        a = dt*u
 44        c, s = torch.cos(a), torch.sin(a)
 45        x, y = p[..., 0], p[..., 1]
 46        return torch.stack((x*c-y*s, x*s+y*c), -1).reshape_as(h)
 47
 48    def forward(self, x, collect=False):
 49        seq = x.view(x.shape[0], -1, 3)
 50        h = torch.zeros(x.shape[0], 64, device=x.device, dtype=x.dtype)
 51        held = torch.zeros(x.shape[0], 32, device=x.device, dtype=x.dtype)
 52        events = 0
 53        vs, exacts, holds, dv_exact, dv_held = [], [], [], [], []
 54        for t in range(seq.shape[1]):
 55            h = self.cell(seq[:, t], h)
 56            if self.mode != "baseline":
 57                ue = self.control(h, self.k)
 58                # Event logic is a controller implementation detail; no STE.
 59                trigger = torch.linalg.vector_norm(ue - held, dim=1) >= self.delta
 60                held = torch.where(trigger[:, None], ue, held)
 61                events += int(trigger.sum().item())
 62                if collect:
 63                    ph = h.view(h.shape[0], 32, 2)
 64                    z0 = torch.complex(ph[...,0], ph[...,1]); z0 = z0/(torch.abs(z0)+1e-6)
 65                    v0 = torch.abs(z0.mean(1))**2
 66                    he = self.rotate(h, ue, 1.0); hh = self.rotate(h, held, 1.0)
 67                    def vv(q):
 68                        qp=q.view(q.shape[0],32,2); zz=torch.complex(qp[...,0],qp[...,1]); zz=zz/(torch.abs(zz)+1e-6); return torch.abs(zz.mean(1))**2
 69                    dv_exact.append((vv(he)-v0).detach()); dv_held.append((vv(hh)-v0).detach())
 70                    z = torch.complex(h.view(h.shape[0],32,2)[...,0], h.view(h.shape[0],32,2)[...,1])
 71                    z = z/(torch.abs(z)+1e-6)
 72                    vs.append(torch.abs(z.mean(1))**2)
 73                    exacts.append(ue.detach()); holds.append(held.detach())
 74                h = self.rotate(h, held, 1.0)
 75        self.last_stats = {"events": events, "steps": x.shape[0]*seq.shape[1],
 76                           "v": vs, "u": exacts, "held": holds, "dv_exact": dv_exact, "dv_held": dv_held}
 77        return self.head(h)
 78
 79def seed_all(seed):
 80    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 81
 82def run_one(seed, cfg, mode, collect=False):
 83    seed_all(seed)
 84    d = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 85    model = MatchedGRU(mode=mode, k=cfg.get("k",0.0), delta=cfg.get("delta",1e9))
 86    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 87    try:
 88        model.to(dev); xtr,ytr=d["xtr"].to(dev),d["ytr"].to(dev)
 89        xte,yte=d["xte"].to(dev),d["yte"].to(dev)
 90        opt=torch.optim.Adam(model.parameters(), lr=cfg["lr"], weight_decay=cfg.get("weight_decay",0.0))
 91        for ep in range(EPOCHS):
 92            model.train(); perm=torch.randperm(len(xtr),device=dev)
 93            for i in range(0,len(xtr),BATCH):
 94                ix=perm[i:i+BATCH]; loss=((model(xtr[ix])-ytr[ix])**2).mean()
 95                opt.zero_grad(); loss.backward(); opt.step()
 96        model.eval()
 97        with torch.no_grad():
 98            pred=model(xte, collect=collect); metric=((pred-yte)**2).mean().item()
 99        stats=model.last_stats
100        result={"metric":metric,"events_per_sequence":stats["events"]/NTEST,
101                "model":model,"stats":stats,"device":str(dev)}
102        return result
103    except RuntimeError:
104        if dev.type == "cuda":
105            torch.cuda.empty_cache(); os.environ["CUDA_VISIBLE_DEVICES"]=""
106            return run_one(seed,cfg,mode,collect)
107        raise
108
109def metric_runner(mode, cfg):
110    return lambda seed: run_one(seed,cfg,mode)["metric"]
111
112def mechanism_signature(cfg):
113    # Both values are measured from hidden states produced by trained models.
114    pred, exact_obs, held_obs, errs, counts = [], [], [], [], []
115    for seed in SEEDS:
116        r = run_one(seed, cfg, "event", collect=True); st = r["stats"]
117        for u, de, dh in zip(st["u"], st["dv_exact"], st["dv_held"]):
118            # First-order prediction for dt=1 under the exact tangent field.
119            pred.append((-cfg["k"] * ((u/cfg["k"])**2).sum(1)).mean().item())
120            exact_obs.append(de.mean().item()); held_obs.append(dh.mean().item())
121        for u, h in zip(st["u"], st["held"]):
122            errs.append(torch.linalg.vector_norm(u-h, dim=1).max().item())
123        counts.append(st["events"])
124    pm, em, hm = map(float, (np.mean(pred), np.mean(exact_obs), np.mean(held_obs)))
125    ratio = em/pm if pm else float("nan")
126    return {"quantity":"one-step V change on trained hidden states",
127            "predicted_exact_delta_V":pm, "observed_exact_delta_V":em,
128            "observed_held_delta_V":hm, "exact_prediction_ratio":ratio,
129            "max_control_hold_error":float(max(errs)),
130            "mean_events_per_sequence":float(np.mean(counts)/NTEST),
131            "confirmed":bool(np.isfinite(ratio) and abs(ratio-1)<0.10)}
132
133def main():
134    # cheap core math check first: finite difference of V under the exact law
135    rng=np.random.default_rng(123); th=rng.normal(0,.25,32); z=np.exp(1j*th); r=z.mean();
136    u=2/32*np.imag(z*np.conj(r)); eps=1e-6
137    v0=abs(r)**2; v1=abs(np.exp(1j*(th+eps*u)).mean())**2
138    math_check={"observed_vdot":float((v1-v0)/eps),"predicted_vdot":float(-np.sum(u*u)),
139                "ratio":float(((v1-v0)/eps)/(-np.sum(u*u)))}
140    base=sweep_baseline(lambda c: metric_runner("baseline",c), BASE_GRID, seeds=SWEEP_SEEDS)
141    best_lr=base["best_cfg"]["lr"]
142    # Idea is evaluated at best baseline lr and two nearby union-grid settings.
143    idea_cfgs=[dict(c) for c in IDEA_GRID]
144    idea_cfgs[1]["lr"]=best_lr
145    idea_sweep=[]
146    for c in idea_cfgs:
147        rr=evaluate(metric_runner("event",c), seeds=SWEEP_SEEDS)
148        idea_sweep.append({"cfg":c,"mean":rr["mean"]})
149    best_idea=min(idea_sweep,key=lambda q:q["mean"])["cfg"]
150    idea=evaluate(metric_runner("event",best_idea), seeds=SEEDS)
151    sig=mechanism_signature(best_idea)
152    report=make_report("dynamics","rnn_small",base,idea,
153        {"mechanism_signature":sig,"math_check":math_check,
154         "idea_sweep":idea_sweep,"protocol":{"epochs":EPOCHS,"n_train":NTRAIN,"n_test":NTEST,"paired_seeds":list(SEEDS)}})
155    OUT.write_text(json.dumps(report,indent=2,allow_nan=False))
156    print(json.dumps(report,indent=2))
157if __name__=="__main__": main()