1import json
  2import sys
  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, evaluate, sweep_baseline, make_report
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12EPOCHS = 15
 13N_TRAIN, N_TEST = 800, 200
 14BATCH = 128
 15
 16
 17def seed_all(seed):
 18    np.random.seed(seed)
 19    torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        torch.cuda.manual_seed_all(seed)
 22
 23
 24class GRUBaseline(nn.Module):
 25    def __init__(self, hidden=64):
 26        super().__init__()
 27        self.rnn = nn.GRU(3, hidden, batch_first=True)
 28        self.head = nn.Linear(hidden, 1)
 29
 30    def forward(self, x):
 31        _, h = self.rnn(x.view(x.shape[0], -1, 3))
 32        return self.head(h[-1])
 33
 34
 35class TwoColumnCore(nn.Module):
 36    """Two transfer channels: single-site A and adjacent-pair B."""
 37    def __init__(self, alpha=0.9, hidden=64, p=8, chi=32):
 38        super().__init__()
 39        self.alpha = float(alpha)
 40        self.p, self.chi = p, chi
 41        self.embed = nn.Linear(3, p)
 42        self.A_raw = nn.Parameter(torch.randn(p, chi, chi) * 0.08)
 43        self.B_raw = nn.Parameter(torch.randn(2 * p, chi, chi) * 0.08)
 44        self.inject1 = nn.Linear(p, chi, bias=False)
 45        self.inject2 = nn.Linear(2 * p, chi, bias=False)
 46        self.head = nn.Sequential(nn.Linear(p + 2 * chi, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 47
 48    def _bounded(self, mats):
 49        # Normalize each feature slice; differentiable and directly measured in signature.
 50        flat = mats.reshape(mats.shape[0], -1)
 51        norms = torch.linalg.matrix_norm(flat, ord=2).clamp_min(1e-6)
 52        return mats * (self.alpha / norms).reshape(-1, 1, 1)
 53
 54    def forward_features(self, x, perturb=0.0):
 55        e = torch.tanh(self.embed(x.view(x.shape[0], -1, 3)))
 56        if perturb:
 57            e = e + perturb * torch.randn_like(e)
 58        h1 = torch.zeros(e.shape[0], self.chi, device=e.device)
 59        h2 = torch.zeros_like(h1)
 60        prev = torch.zeros_like(e[:, 0])
 61        for t in range(e.shape[1]):
 62            A = torch.einsum("bp,pij->bij", e[:, t], self._bounded(self.A_raw))
 63            z = torch.cat((prev, e[:, t]), dim=-1)
 64            B = torch.einsum("bp,pij->bij", z, self._bounded(self.B_raw))
 65            h1 = torch.bmm(A, h1.unsqueeze(-1)).squeeze(-1) + self.inject1(e[:, t])
 66            h2 = torch.bmm(B, h2.unsqueeze(-1)).squeeze(-1) + self.inject2(z)
 67            h1 = torch.tanh(h1)
 68            h2 = torch.tanh(h2)
 69            prev = e[:, t]
 70        return torch.cat((prev, h1, h2), dim=-1)
 71
 72    def forward(self, x):
 73        return self.head(self.forward_features(x))
 74
 75
 76def train_one(kind, cfg, seed, return_model=False):
 77    seed_all(seed)
 78    ds = get_dataset("dynamics", seed, n_train=N_TRAIN, n_test=N_TEST)
 79    if kind == "baseline":
 80        net = GRUBaseline()
 81    else:
 82        net = TwoColumnCore(alpha=cfg["alpha"])
 83    net, metric, history = train_model(net, ds, epochs=EPOCHS, lr=cfg["lr"],
 84                                       batch=BATCH, weight_decay=cfg.get("weight_decay", 0.0),
 85                                       log=lambda *_: None)
 86    if net is None or metric is None:
 87        return float("nan") if not return_model else (float("nan"), None, ds)
 88    if return_model:
 89        return float(metric), net, ds
 90    return float(metric)
 91
 92
 93def baseline_factory(cfg):
 94    return lambda seed: train_one("baseline", cfg, seed)
 95
 96
 97def idea_factory(cfg):
 98    return lambda seed: train_one("idea", cfg, seed)
 99
100
101def mechanism_signature(cfg, seeds):
102    predicted = []
103    observed = []
104    for seed in seeds:
105        metric, net, ds = train_one("idea", cfg, seed, return_model=True)
106        if net is None:
107            continue
108        net.eval()
109        device = next(net.parameters()).device
110        x = ds["xte"][:64].to(device)
111        with torch.no_grad():
112            a = net.forward_features(x)
113            b = net.forward_features(x, perturb=1e-3)
114            ratio = torch.linalg.vector_norm(b - a, dim=1).mean().item() / 1e-3
115        observed.append(float(ratio))
116        predicted.append(float(cfg["alpha"]))
117    # State nonlinearities and input injection mean this is an approximate prediction;
118    # confirmation requires observed attenuation to be near the claimed alpha bound.
119    obs = float(np.mean(observed)) if observed else float("nan")
120    pred = float(np.mean(predicted)) if predicted else float("nan")
121    return {"quantity": "finite perturbation amplification of trained memory features",
122            "predicted_alpha_bound": pred, "observed_amplification": obs,
123            "n_models": len(observed), "tolerance": "observed <= alpha + 0.15",
124            "confirmed": bool(np.isfinite(obs) and obs <= pred + 0.15)}
125
126
127def main():
128    # Union of all tried learning rates is shared by both systems.
129    lrs = [1e-3, 3e-3, 1e-2]
130    base_grid = [{"lr": lr} for lr in lrs]
131    base = sweep_baseline(baseline_factory, base_grid, seeds=SWEEP_SEEDS)
132    best_lr = base["best_cfg"]["lr"]
133    idea_grid = [{"lr": lr, "alpha": a} for lr, a in
134                 [(best_lr, 0.8), (best_lr, 0.9), (best_lr, 1.0)]]
135    # Same lr union on baseline side is already evaluated; idea uses best alpha at best lr.
136    idea_scores = []
137    for cfg in idea_grid:
138        r = evaluate(idea_factory(cfg), seeds=SEEDS)
139        idea_scores.append((r["mean"], cfg, r))
140    _, best_cfg, idea = min(idea_scores, key=lambda q: q[0])
141    report = make_report("dynamics", "rnn_small", base, idea,
142                         extra=mechanism_signature(best_cfg, SEEDS))
143    report["idea_sweep"] = [{"cfg": c, "mean": m} for m, c, _ in idea_scores]
144    report["protocol_notes"] = {"epochs": EPOCHS, "n_train": N_TRAIN, "n_test": N_TEST,
145                                "matched_task": "controlled damped pendulum horizon-8",
146                                "architecture_difference": "GRU recurrence versus two-column A/B memory"}
147    with open("bench_report.json", "w") as f:
148        json.dump(report, f, indent=2)
149    print(json.dumps(report, indent=2))
150
151
152if __name__ == "__main__":
153    main()