1import sys, json, math, time
  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, train_model, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11# Union is shared by baseline and idea; baseline sweep covers every idea LR.
 12GRID = [
 13    {"lr": 1e-3, "epochs": 12},
 14    {"lr": 3e-3, "epochs": 12},
 15    {"lr": 1e-2, "epochs": 12},
 16]
 17HIDDEN = 32
 18SWEEPS = 3
 19CLUSTERS = 4
 20ETA = 0.55
 21R_BINS = 4
 22A_MAX = 8.0
 23
 24class IterativeGraphNet(nn.Module):
 25    """Small recurrent message-passing solver over the 8-step pendulum window."""
 26    def __init__(self, mode="full", eta=ETA, clusters=CLUSTERS, sweeps=SWEEPS):
 27        super().__init__()
 28        self.mode, self.eta = mode, eta
 29        self.clusters, self.sweeps = clusters, sweeps
 30        self.inp = nn.Linear(3, HIDDEN)
 31        self.msg = nn.Linear(HIDDEN, HIDDEN, bias=False)
 32        self.self_proj = nn.Linear(HIDDEN, HIDDEN)
 33        self.head = nn.Linear(HIDDEN, 1)
 34
 35    def one_update(self, h, x, ids):
 36        # Gather all neighbors from the pre-update tensor, then scatter together.
 37        old = h
 38        left = torch.roll(old, 1, dims=1)
 39        right = torch.roll(old, -1, dims=1)
 40        neigh = 0.5 * (left + right)
 41        z = torch.tanh(self.self_proj(old) + self.msg(neigh) + self.inp(x))
 42        out = old.clone()
 43        out[:, ids, :] = (1.0 - self.eta) * old[:, ids, :] + self.eta * z[:, ids, :]
 44        return out
 45
 46    def histograms(self, h):
 47        # Observable local mismatch: disagreement with ring-neighbor states.
 48        d = 0.5 * (torch.roll(h, 1, 1) - h).norm(dim=-1) + \
 49            0.5 * (torch.roll(h, -1, 1) - h).norm(dim=-1)
 50        omega = torch.clamp((d / (d.detach().mean(dim=1, keepdim=True) + 1e-6) * 4).long(), 0, int(A_MAX))
 51        q = torch.clamp((R_BINS * omega / A_MAX).long(), 0, R_BINS - 1)
 52        hs = []
 53        for a in range(self.clusters):
 54            ids = list(range(a * 8 // self.clusters, (a + 1) * 8 // self.clusters))
 55            hs.append(torch.stack([(q[:, ids] == r).float().mean(1) for r in range(R_BINS)], 1))
 56        return torch.stack(hs, 1), d
 57
 58    def forward(self, x):
 59        # x is [batch, 24], interpreted as 8 nodes with 3 features.
 60        x = x.view(x.shape[0], 8, 3)
 61        h = torch.tanh(self.inp(x))
 62        if self.mode == "full":
 63            for _ in range(self.sweeps):
 64                h = self.one_update(h, x, list(range(8)))
 65        else:
 66            # Equal node-update budget: clusters*sweeps decisions.
 67            for _ in range(self.sweeps * self.clusters):
 68                hist, _ = self.histograms(h.detach())
 69                action = hist[:, :, -1].argmax(1)
 70                # A batch has one action per example; grouping preserves synchronous updates.
 71                nxt = h.clone()
 72                for a in range(self.clusters):
 73                    mask = action == a
 74                    if mask.any():
 75                        ids = list(range(a * 8 // self.clusters, (a + 1) * 8 // self.clusters))
 76                        nxt[mask] = self.one_update(h[mask], x[mask], ids)
 77                h = nxt
 78        return self.head(h[:, -1, :])
 79
 80def train_idea(seed, cfg):
 81    torch.manual_seed(seed); np.random.seed(seed)
 82    ds = get_dataset("dynamics", seed, n_train=400, n_test=100)
 83    model = IterativeGraphNet("cluster")
 84    device = "cuda" if torch.cuda.is_available() else "cpu"
 85    try:
 86        model = model.to(device)
 87        opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
 88        lossf = nn.MSELoss(); x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 89        for _ in range(cfg["epochs"]):
 90            model.train(); p = torch.randperm(len(x), device=device)
 91            for i in range(0, len(x), 128):
 92                ix = p[i:i+128]; loss = lossf(model(x[ix]), y[ix])
 93                opt.zero_grad(); loss.backward(); opt.step()
 94        model.eval()
 95        with torch.no_grad():
 96            metric = float(lossf(model(ds["xte"].to(device)), ds["yte"].to(device)))
 97        return metric, model.cpu(), ds
 98    except (RuntimeError, torch.cuda.CudaError):
 99        # Required robust fallback; recreate model to avoid partial CUDA state.
100        model = IterativeGraphNet("cluster")
101        opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
102        lossf = nn.MSELoss(); x, y = ds["xtr"], ds["ytr"]
103        for _ in range(cfg["epochs"]):
104            p = torch.randperm(len(x))
105            for i in range(0, len(x), 128):
106                ix = p[i:i+128]; loss = lossf(model(x[ix]), y[ix])
107                opt.zero_grad(); loss.backward(); opt.step()
108        with torch.no_grad(): metric = float(lossf(model(ds["xte"]), ds["yte"]))
109        return metric, model, ds
110
111def baseline_factory(cfg):
112    def run(seed):
113        torch.manual_seed(seed); np.random.seed(seed)
114        ds = get_dataset("dynamics", seed, n_train=400, n_test=100)
115        model = IterativeGraphNet("full")
116        _, metric, _ = train_model(model, ds, epochs=cfg["epochs"], lr=cfg["lr"], batch=128, log=lambda *_: None)
117        return metric
118    return run
119
120def idea_eval(cfg):
121    vals = []
122    for s in SEEDS:
123        v, _, _ = train_idea(s, cfg); vals.append(v)
124    return {"mean": float(np.mean(vals)), "std": float(np.std(vals)), "per_seed": vals, "n": len(vals)}
125
126def mechanism_signature(base_cfg, idea_cfg):
127    # Re-test eta^2 ordering prediction on trained models, not a toy identity.
128    b, bm, ds = train_idea(0, idea_cfg)
129    bm.eval(); x = ds["xte"][:32].view(32, 8, 3)
130    rows = []
131    with torch.no_grad():
132        for eta in (0.2, 0.4, 0.6, 0.8):
133            bm.eta = eta
134            h = torch.tanh(bm.inp(x)); old = h.clone()
135            ids = [0, 1]
136            # synchronous cluster update
137            sync = bm.one_update(old, x, ids)
138            # deliberately sequential ordering on the same trained update rule
139            seq = old.clone()
140            for i in ids: seq = bm.one_update(seq, x, [i])
141            err = float((sync[:, ids] - seq[:, ids]).abs().max())
142            rows.append({"eta": eta, "observed_linf": err, "eta2": eta * eta})
143    ratios = [r["observed_linf"] / max(r["eta2"], 1e-9) for r in rows]
144    # Prediction is quantitative scaling, assessed by normalized variation.
145    confirmed = bool(max(ratios) < 1.5 * max(min(ratios), 1e-9))
146    return {"prediction": "trained-model synchronous-vs-sequential discrepancy scales approximately eta^2", "rows": rows, "eta2_ratio_range": [float(min(ratios)), float(max(ratios))], "confirmed": confirmed}
147
148def main():
149    t0 = time.time()
150    base = sweep_baseline(baseline_factory, GRID)
151    idea_trials = [(cfg, idea_eval(cfg)) for cfg in GRID]
152    idea_cfg, idea = min(idea_trials, key=lambda z: z[1]["mean"])
153    rep = make_report("dynamics", "rnn_small", base, idea, {
154        "track_structure": "controlled pendulum multi-step dynamics; recurrent state updates",
155        "scheduler": "largest high-bin hidden-state disagreement histogram",
156        "clusters": CLUSTERS, "histogram_bins": R_BINS,
157        "signature": mechanism_signature(base["best_cfg"], base["best_cfg"])
158    })
159    rep["runtime_sec"] = time.time() - t0
160    rep["protocol_notes"] = {"paired_seeds": list(SEEDS), "dataset_sizes": [400, 100], "idea_grid": [{"cfg": c, "mean": r["mean"]} for c, r in idea_trials], "selected_idea_cfg": idea_cfg, "baseline_grid_union": GRID}
161    Path("bench_report.json").write_text(json.dumps(rep, indent=2))
162    print(json.dumps(rep, indent=2))
163
164if __name__ == "__main__": main()