1import sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import get_dataset, sweep_baseline, evaluate, make_report
  8
  9NMOD = 4
 10EPOCHS = 10
 11BATCH = 128
 12SEEDS = tuple(range(8))
 13SWEEP_SEEDS = tuple(range(4))
 14
 15
 16def graph_laplacian():
 17    A = np.array([[0,.8,.2,.1],[.8,0,.5,.15],[.2,.5,0,.7],[.1,.15,.7,0]], dtype=float)
 18    return A, np.diag(A.sum(1)) - A
 19
 20
 21def grounded_gap(L, pins, strength):
 22    p = np.zeros(len(L)); p[list(pins)] = strength
 23    return float(np.linalg.eigvalsh(L + np.diag(p))[0])
 24
 25
 26def greedy_pins(L, m=2, strength=.5):
 27    chosen = []
 28    for _ in range(m):
 29        candidates = [(grounded_gap(L, chosen + [i], strength), i)
 30                      for i in range(len(L)) if i not in chosen]
 31        chosen.append(max(candidates, key=lambda z: (z[0], -z[1]))[1])
 32    return chosen, grounded_gap(L, chosen, strength)
 33
 34
 35class Module(nn.Module):
 36    def __init__(self):
 37        super().__init__()
 38        self.rnn = nn.GRU(3, 32, batch_first=True)
 39        self.head = nn.Linear(32, 1)
 40
 41    def forward(self, x):
 42        _, h = self.rnn(x.view(x.shape[0], 8, 3))
 43        return self.head(h[-1]), h[-1]
 44
 45
 46def train_system(seed, lr, consensus, pin_strength=0.0, return_signature=False):
 47    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 48    try:
 49        device = "cuda" if torch.cuda.is_available() else "cpu"
 50        if device == "cuda":
 51            try:
 52                torch.cuda.set_device(0)
 53                torch.backends.cudnn.enabled = False
 54            except Exception:
 55                device = "cpu"
 56    except Exception:
 57        device = "cpu"
 58    ds = get_dataset("dynamics", seed, n_train=400, n_test=200)
 59    xtr, ytr = ds["xtr"].to(device), ds["ytr"].to(device).reshape(-1, 1)
 60    xte, yte = ds["xte"].to(device), ds["yte"].to(device).reshape(-1, 1)
 61    nets = nn.ModuleList([Module() for _ in range(NMOD)]).to(device)
 62    opt = torch.optim.Adam(nets.parameters(), lr=lr)
 63    _, L = graph_laplacian()
 64    pins, gap = greedy_pins(L, 2, pin_strength if pin_strength else .5)
 65    pinmask = torch.zeros(NMOD, device=device)
 66    if pin_strength > 0: pinmask[pins] = 1.
 67    history = []
 68    for ep in range(EPOCHS):
 69        perm = torch.randperm(len(xtr), device=device)
 70        for start in range(0, len(xtr), BATCH):
 71            idx = perm[start:start+BATCH]
 72            preds, hs = [], []
 73            for net in nets:
 74                q, h = net(xtr[idx]); preds.append(q); hs.append(h)
 75            pred = torch.stack(preds)
 76            h = torch.stack(hs)
 77            loss = ((pred - ytr[idx].unsqueeze(0)) ** 2).mean()
 78            # L2 graph-consensus regularization, plus stronger corrective
 79            # teacher/anchor loss on spectral-greedy pinned modules.
 80            diff = h[:, None] - h[None, :]
 81            loss = loss + consensus * sum(L[i,j] * (h[i]-h[j]).pow(2).mean()
 82                                          for i in range(NMOD) for j in range(NMOD)) / 2
 83            if pin_strength > 0:
 84                mean = h.mean(0, keepdim=True).detach()
 85                loss = loss + pin_strength * ((h - mean).pow(2).mean(1) * pinmask[:, None]).mean()
 86            opt.zero_grad(); loss.backward(); opt.step()
 87        with torch.no_grad():
 88            ev = torch.stack([net(xte)[0] for net in nets])
 89            history.append(float(((ev-yte.unsqueeze(0))**2).mean()))
 90    with torch.no_grad():
 91        out = torch.stack([net(xte)[0] for net in nets])
 92        mse = float(((out-yte.unsqueeze(0))**2).mean())
 93        module_mean = out.mean(0, keepdim=True)
 94        disagreement = float(.5 * ((out-module_mean)**2).mean())
 95        # Measured trained-model behavior, not an analytical identity.
 96        train_h = torch.stack([net(xte)[1] for net in nets])
 97        measured_dis = float(.5 * ((train_h-train_h.mean(0,keepdim=True))**2).mean())
 98    if return_signature:
 99        return mse, {"pins": pins, "lambda_min_grounded": gap,
100                      "test_representation_disagreement": measured_dis,
101                      "task_prediction_disagreement": disagreement}
102    return mse
103
104
105def make_fn(cfg):
106    return lambda seed: train_system(seed, cfg["lr"], cfg["consensus"], 0.0)
107
108
109def main():
110    # Search-space parity: every idea lr/consensus pair is also evaluated by
111    # the baseline sweep. The baseline method's central consensus knob is swept.
112    grid = [{"lr": lr, "consensus": c} for lr in [1e-3, 3e-3, 1e-2]
113            for c in [0.0, 0.01, 0.05]]
114    base = sweep_baseline(make_fn, grid, seeds=SWEEP_SEEDS)
115    best = base["best_cfg"]
116    idea_cfgs = [best,
117                 {"lr": best["lr"], "consensus": best["consensus"] + .01},
118                 {"lr": best["lr"], "consensus": max(0.0, best["consensus"] - .01)}]
119    # Keep the idea's three settings on the same union grid where possible.
120    all_grid = list(grid)
121    for cfg in idea_cfgs:
122        if cfg not in all_grid: all_grid.append(cfg)
123    base = sweep_baseline(make_fn, all_grid, seeds=SWEEP_SEEDS)
124    idea_trials = []
125    for cfg in idea_cfgs:
126        r = evaluate(lambda s, cfg=cfg: train_system(s, cfg["lr"], cfg["consensus"], .5), seeds=SEEDS)
127        idea_trials.append({"cfg": cfg, "result": r})
128    chosen = min(idea_trials, key=lambda z: z["result"]["mean"])
129    idea = chosen["result"]
130    best_idea_cfg = chosen["cfg"]
131    sig_mse, sig = train_system(0, best_idea_cfg["lr"], best_idea_cfg["consensus"], .5, True)
132    _, L = graph_laplacian()
133    random_gap = float(np.mean([grounded_gap(L, p, .5) for p in ([0,1],[1,2],[2,3],[0,3])]))
134    sig.update({"baseline_random_gap": random_gap,
135                "observed_vs_predicted": "task-independent disagreement is measured on trained representations; no exact decay slope is identifiable from final snapshots",
136                "confirmed": False})
137    report = make_report("dynamics", "rnn_small", base, idea, sig)
138    report["idea_sweep"] = idea_trials
139    report["baseline_sweep_union"] = all_grid
140    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
141    print(json.dumps(report, indent=2))
142
143if __name__ == "__main__": main()