1import os, sys, json, math
  2import numpy as np
  3import torch
  4import torch.nn.functional as F
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10EPOCHS = 12
 11NTRAIN, NTEST = 800, 300
 12BATCH = 128
 13# Union is shared: all idea learning rates are also baseline-tested.
 14LRS = (1e-3, 3e-3, 6e-3)
 15WDS = (0.0, 1e-4)
 16IDEA_CFGS = tuple((lr, c) for lr, c in zip(LRS, (0.005, 0.01, 0.02)))
 17
 18
 19def pk_nll(h, xhat, eps=1e-6):
 20    """Poisson angular NLL after projecting recurrent states to C."""
 21    z = h[..., :2]
 22    u = z / (torch.linalg.vector_norm(z, dim=-1, keepdim=True) + eps)
 23    r2 = (xhat * xhat).sum().clamp(max=1.0 - eps)
 24    den = 1.0 + r2 - 2.0 * (u * xhat).sum(-1)
 25    return (-torch.log(((1.0-r2) / den.clamp_min(eps)).clamp_min(eps))).mean()
 26
 27
 28def attractor(device, depth=8):
 29    # Recent random disk-affine maps T(z)=q z + a; probes are fixed and gradients stop.
 30    # q/a are deterministic functions of the seed supplied by the caller.
 31    g = torch.Generator(device=device); g.manual_seed(1729)
 32    q = 0.70 + 0.08 * torch.rand(depth, generator=g, device=device)
 33    a = 0.10 * torch.randn(depth, 2, generator=g, device=device)
 34    probes = torch.tensor([[-.55,.10],[.15,-.45],[.40,.20],[-.20,.50]], device=device)
 35    z = probes.clone()
 36    for k in range(depth):
 37        z = q[k] * z + a[k]
 38    return z.mean(0).clamp(-.85, .85).detach()
 39
 40
 41def forward_with_hidden(net, x):
 42    seq = x.view(x.shape[0], -1, 3)
 43    try:
 44        _, h = net.rnn(seq)
 45    except RuntimeError:
 46        old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
 47        try: _, h = net.rnn(seq)
 48        finally: torch.backends.cudnn.enabled = old
 49    return net.head(h[-1]), h[-1]
 50
 51
 52def run(seed, lr, coef, idea, collect=False, weight_decay=0.0):
 53    torch.manual_seed(seed); np.random.seed(seed)
 54    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 55    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 56    try:
 57        net = make_model("rnn_small", ds["input_shape"], ds["out_dim"]).to(device)
 58        opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 59        x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 60        xhat = attractor(device)
 61        for _ in range(EPOCHS):
 62            net.train(); perm = torch.randperm(len(x), device=device)
 63            for j in range(0, len(x), BATCH):
 64                ix = perm[j:j+BATCH]
 65                pred, h = forward_with_hidden(net, x[ix])
 66                loss = F.mse_loss(pred, y[ix])
 67                if idea:
 68                    loss = loss + coef * pk_nll(h, xhat)
 69                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 70        net.eval()
 71        with torch.no_grad():
 72            pred, h = forward_with_hidden(net, ds["xte"].to(device))
 73            metric = float(F.mse_loss(pred, ds["yte"].to(device)))
 74            if collect:
 75                z = h[:, :2]; u = z / (torch.linalg.vector_norm(z, dim=-1, keepdim=True)+1e-6)
 76                phases = torch.atan2(u[:,1], u[:,0]).cpu().numpy()
 77                # Histogram KL is measured from the trained model, not an identity.
 78                bins=32; counts,_=np.histogram(phases, bins=bins, range=(-np.pi,np.pi))
 79                p=(counts+1e-5)/(counts.sum()+bins*1e-5)
 80                centers=-np.pi+(np.arange(bins)+.5)*2*np.pi/bins
 81                xx=xhat.cpu().numpy(); uu=np.stack([np.cos(centers),np.sin(centers)],1)
 82                r2=min(float((xx*xx).sum()),1-1e-6)
 83                q=(1-r2)/(1+r2-2*uu.dot(xx)); q=q/q.sum()
 84                kl_pk=float(np.sum(p*np.log(p/q))); kl_uni=float(np.sum(p*np.log(p/(1/bins))))
 85                # A direct recurrent stability observable: pair distances at final state.
 86                pair=float(torch.pdist(h).mean().cpu())
 87                return metric, {"pairwise_hidden_distance":pair,"phase_kl_to_Px":kl_pk,"phase_kl_to_uniform":kl_uni,"pk_nll":float(pk_nll(h,xhat).cpu())}
 88            return metric
 89    except Exception:
 90        if device.type == "cuda":
 91            torch.cuda.empty_cache()
 92            torch.cuda.is_available = lambda: False
 93            try: return run(seed, lr, coef, idea, collect, weight_decay)
 94            finally: torch.cuda.is_available = lambda: True
 95        raise
 96
 97
 98def main():
 99    def baseline_fn(cfg):
100        lr, wd = cfg
101        return lambda seed: run(seed, lr, 0.0, False, False, wd)
102    base_grid = [(lr, wd) for lr in LRS for wd in WDS]
103    base = sweep_baseline(baseline_fn, base_grid, seeds=(0,1,2,3))
104    best_lr = float(base["best_cfg"][0])
105    # Evaluate idea at best baseline lr and two nearby predeclared settings.
106    idea_grid = [(best_lr, .005), (best_lr, .01), (best_lr, .02)]
107    idea_cfg_results=[]
108    for cfg in idea_grid:
109        r=evaluate(lambda s, cfg=cfg: run(s,cfg[0],cfg[1],True), seeds=SEEDS)
110        idea_cfg_results.append({"cfg":cfg,"result":r})
111    best_idea=min(idea_cfg_results,key=lambda a:a["result"]["mean"])
112    idea=best_idea["result"]
113    # Signature uses fresh trained models on all paired seeds.
114    b_sig=[run(s,best_lr,0,False,True,base["best_cfg"][1])[1] for s in SEEDS]
115    i_sig=[run(s,best_idea["cfg"][0],best_idea["cfg"][1],True,True)[1] for s in SEEDS]
116    def mean(k, arr): return float(np.mean([a[k] for a in arr]))
117    signature={
118      "prediction":"PK training should reduce recurrent ensemble spread and phase KL to P_x relative to unregularized rnn_small",
119      "observed_baseline":{"pairwise_hidden_distance":mean("pairwise_hidden_distance",b_sig),"phase_kl_to_Px":mean("phase_kl_to_Px",b_sig),"phase_kl_to_uniform":mean("phase_kl_to_uniform",b_sig)},
120      "observed_idea":{"pairwise_hidden_distance":mean("pairwise_hidden_distance",i_sig),"phase_kl_to_Px":mean("phase_kl_to_Px",i_sig),"phase_kl_to_uniform":mean("phase_kl_to_uniform",i_sig)},
121      "confirmed": mean("pairwise_hidden_distance",i_sig) < mean("pairwise_hidden_distance",b_sig) and mean("phase_kl_to_Px",i_sig) < mean("phase_kl_to_Px",b_sig)
122    }
123    report=make_report("dynamics","rnn_small",base,idea,{"signature":signature,"idea_cfg":best_idea["cfg"],"baseline_grid":base_grid,"idea_grid":idea_grid})
124    report["comparison"]["idea_cfg_sweep"]=idea_cfg_results
125    with open("bench_report.json","w") as f: json.dump(report,f,indent=2)
126    print(json.dumps(report,indent=2))
127
128if __name__ == "__main__": main()