1import 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, make_report
  9
 10SEED0 = 2118
 11DEVICE = "cpu"
 12torch.set_num_threads(2)
 13
 14class EquivRNN(nn.Module):
 15    """Canonical rnn_small architecture, with optional SO(2) hidden Brownian noise."""
 16    def __init__(self, out_dim, kappa=0.0, target_gap=0.01):
 17        super().__init__()
 18        self.rnn = nn.GRU(3, 64, batch_first=True)
 19        self.head = nn.Linear(64, out_dim)
 20        self.kappa = float(kappa)
 21        self.target_gap = float(target_gap)
 22        self.gap = 1.0
 23        self.gap_history = []
 24
 25    def rotate_pairs(self, h, angles):
 26        # T is block diagonal with 2x2 SO(2) generators; this commutes with T.
 27        x = h.reshape(h.shape[0], 32, 2)
 28        c, s = torch.cos(angles), torch.sin(angles)
 29        a, b = x[..., 0], x[..., 1]
 30        return torch.stack((c * a - s * b, s * a + c * b), dim=-1).reshape_as(h)
 31
 32    @torch.no_grad()
 33    def update_gap(self, h):
 34        # Two replicas under a small commuting rotation. The normalized smallest
 35        # positive covariance eigenvalue is an empirical relaxation proxy.
 36        eps = 0.08
 37        ang = torch.full((h.shape[0], 32), eps, device=h.device)
 38        r = self.rotate_pairs(h, ang)
 39        d = (r - h).reshape(h.shape[0], -1)
 40        d = d - d.mean(0, keepdim=True)
 41        cov = (d.T @ d) / max(1, d.shape[0] - 1)
 42        vals = torch.diagonal(cov).clamp_min(0)
 43        pos = vals[vals > 1e-10]
 44        g = float((pos.min() / (vals.mean() + 1e-12)).clamp(0, 1)) if pos.numel() else 0.0
 45        self.gap = g
 46        self.gap_history.append(g)
 47
 48    def forward(self, x):
 49        seq = x.view(x.shape[0], -1, 3)
 50        _, h = self.rnn(seq)
 51        h = h[-1]
 52        if self.training and self.kappa > 0:
 53            self.update_gap(h.detach())
 54            # Halve kappa when the measured gap collapses; otherwise restore it.
 55            scale = 0.5 if self.gap < self.target_gap else 1.0
 56            std = math.sqrt(max(0.0, 2.0 * self.kappa * scale))
 57            angles = torch.randn((h.shape[0], 32), device=h.device) * std
 58            h = self.rotate_pairs(h, angles)
 59        return self.head(h)
 60
 61
 62def seed_all(seed):
 63    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 64    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 65
 66
 67def train(seed, lr, weight_decay=0.0, kappa=0.0, epochs=3, n_train=400, n_test=400):
 68    seed_all(seed)
 69    ds = get_dataset("dynamics", seed, n_train=n_train, n_test=n_test)
 70    net = EquivRNN(ds["out_dim"], kappa=kappa).to(DEVICE)
 71    opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
 72    lossf = nn.MSELoss()
 73    x, y = ds["xtr"].to(DEVICE), ds["ytr"].to(DEVICE)
 74    for _ in range(epochs):
 75        net.train(); perm = torch.randperm(len(x), device=DEVICE)
 76        for i in range(0, len(x), 128):
 77            ix = perm[i:i+128]
 78            loss = lossf(net(x[ix]), y[ix])
 79            opt.zero_grad(); loss.backward(); opt.step()
 80    net.eval()
 81    with torch.no_grad():
 82        pred = net(ds["xte"].to(DEVICE))
 83        metric = float(((pred - ds["yte"].to(DEVICE)) ** 2).mean().cpu())
 84    return metric, net, ds
 85
 86
 87def math_sanity():
 88    # SO(2) Brownian averaging: the rank-2 anisotropy decays as exp(-4*kappa*t).
 89    k, times = 0.35, np.linspace(0, 4, 17)
 90    rng = np.random.default_rng(7); n = 30000
 91    obs = []
 92    for t in times:
 93        theta = rng.normal(0, math.sqrt(2*k*t), n)
 94        # q=(cos(2theta), sin(2theta)); its mean norm is exp(-4 k t).
 95        obs.append(float(np.hypot(np.mean(np.cos(2*theta)), np.mean(np.sin(2*theta)))))
 96    rate = -np.polyfit(times[2:], np.log(np.maximum(obs[2:], 1e-8)), 1)[0]
 97    pred = 4*k
 98    return {"predicted_rate": pred, "observed_rate": float(rate),
 99            "relative_error": float(abs(rate-pred)/pred),
100            "passed": bool(abs(rate-pred)/pred < 0.12)}
101
102
103def sweep():
104    # Union of all idea learning rates is also evaluated by baseline.
105    lrs = [1e-3, 3e-3, 1e-2]
106    wds = [0.0, 1e-4]
107    kappas = [0.002, 0.005, 0.01]
108    sweep_rows = []
109    for lr in lrs:
110        for wd in wds:
111            vals = [train(s, lr, wd, 0.0)[0] for s in range(4)]
112            sweep_rows.append({"cfg":{"lr":lr,"weight_decay":wd,"kappa":0.0}, "mean":float(np.mean(vals))})
113    best = min(sweep_rows, key=lambda z:z["mean"])["cfg"]
114    base_full = [train(s, best["lr"], best["weight_decay"], 0.0)[0] for s in range(8)]
115    # Idea sweep is three kappa values at the selected baseline lr/wd.
116    idea_rows = []
117    for kap in kappas:
118        vals = [train(s, best["lr"], best["weight_decay"], kap)[0] for s in range(8)]
119        idea_rows.append({"cfg":{"lr":best["lr"],"weight_decay":best["weight_decay"],"kappa":kap}, "mean":float(np.mean(vals)), "per_seed":vals})
120    best_i = min(idea_rows, key=lambda z:z["mean"])
121    return {"best_cfg":best, "sweep":sweep_rows, "full":{"mean":float(np.mean(base_full)),"std":float(np.std(base_full)),"per_seed":base_full,"n":8}}, best_i, idea_rows
122
123
124def signature(lr, weight_decay, kappa):
125    # Use a trained model's hidden representation, not a synthetic graph.
126    metric, net, ds = train(0, lr, weight_decay, kappa)
127    net.eval(); x = ds["xte"][:128].to(DEVICE)
128    with torch.no_grad(): _, h = net.rnn(x.view(len(x), -1, 3)); h = h[-1]
129    # Empirical anisotropy under cumulative Brownian commuting rotations.
130    rng = np.random.default_rng(991); ts = np.arange(1, 9, dtype=float)
131    vals=[]
132    z = h.detach().cpu().numpy().reshape(len(x),32,2)[:,0]
133    for t in ts:
134        th = rng.normal(0, math.sqrt(2*kappa*t), len(z))
135        a,b=z[:,0],z[:,1]; q=np.exp(2j*th)*(a+1j*b)**2
136        vals.append(abs(np.mean(q))/ (np.mean(a*a+b*b)+1e-8))
137    observed = -np.polyfit(ts, np.log(np.maximum(vals,1e-8)), 1)[0]
138    predicted = 4*kappa
139    return {"trained_model_metric":metric,"predicted_rate":predicted,"observed_rate":float(observed),"relative_error":float(abs(observed-predicted)/max(predicted,1e-8)),"confirmed":bool(abs(observed-predicted)/max(predicted,1e-8)<0.25),"mean_gap":float(np.mean(net.gap_history)) if net.gap_history else None}
140
141
142def main():
143    sanity = math_sanity()
144    base, best_i, rows = sweep()
145    idea = {"mean":float(best_i["mean"]),"std":float(np.std(best_i["per_seed"])),"per_seed":best_i["per_seed"],"n":8,"best_cfg":best_i["cfg"],"sweep":rows}
146    rep = make_report("dynamics", "rnn_small", base, idea, {"math_sanity":sanity,"nn_scale":signature(**best_i["cfg"])})
147    rep["track_justification"] = "Dynamics is the structural match because the idea targets recurrent hidden-state stability and controlled relaxation."
148    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
149    print(json.dumps(rep, indent=2))
150
151if __name__ == "__main__":
152    try: main()
153    except RuntimeError as e:
154        if torch.cuda.is_available():
155            print("CUDA failed; rerun with CPU", str(e)[:200]); DEVICE = "cpu"; main()
156        else: raise