1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6import sys
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10SEED = 2841
 11EPOCHS = 12
 12NTRAIN, NTEST = 1200, 400
 13LR_GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 14
 15class OrbitalRNN(nn.Module):
 16    """Phase-only recurrent replacement for bench rnn_small.
 17
 18    The directed ring plus two-step skip graph is fixed and positive. Input
 19    drive, common frequency, and readout are learned end-to-end; amplitudes
 20    are fixed at one, as in the MVP implementation plan.
 21    """
 22    def __init__(self, hidden=64, dt=0.15, kappa=0.8):
 23        super().__init__()
 24        self.hidden, self.dt, self.kappa = hidden, dt, kappa
 25        self.inp = nn.Linear(3, hidden, bias=False)
 26        self.omega = nn.Parameter(torch.randn(hidden) * 0.05 + 0.12)
 27        A = torch.zeros(hidden, hidden)
 28        for i in range(hidden):
 29            A[i, (i + 1) % hidden] = 0.35
 30            A[i, (i + 2) % hidden] = 0.15
 31        self.register_buffer("A", A)
 32        self.head = nn.Linear(2 * hidden, 1)
 33
 34    def phase_step(self, theta, drive):
 35        # d_ij = theta_j - theta_i; A[i,j] is directed i <- j.
 36        diff = theta[:, None, :] - theta[:, :, None]
 37        coupling = torch.einsum("bij,ij->bi", torch.sin(diff), self.A)
 38        return theta + self.dt * (self.omega + 0.08 * drive + self.kappa * coupling)
 39
 40    def forward_with_theta(self, x, theta0=None, return_states=False):
 41        b = x.shape[0]
 42        seq = x.view(b, -1, 3)
 43        theta = torch.zeros(b, self.hidden, device=x.device) if theta0 is None else theta0
 44        states = []
 45        for k in range(seq.shape[1]):
 46            theta = self.phase_step(theta, self.inp(seq[:, k]))
 47            states.append(torch.cat((torch.cos(theta), torch.sin(theta)), dim=1))
 48        out = self.head(states[-1])
 49        return (out, torch.stack(states, dim=1), theta) if return_states else out
 50
 51    def forward(self, x):
 52        return self.forward_with_theta(x)
 53
 54
 55def seed_all(seed):
 56    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 57    if torch.cuda.is_available():
 58        torch.cuda.manual_seed_all(seed)
 59
 60
 61def train_one(kind, lr, seed, keep_model=False):
 62    seed_all(seed)
 63    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 64    model = make_model("rnn_small", ds["input_shape"], ds["out_dim"]) if kind == "baseline" else OrbitalRNN()
 65    try:
 66        _, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 67    except Exception:
 68        # train_model already has a CUDA fallback; this protects unusual driver errors.
 69        model = model.cpu(); ds = {k: (v.cpu() if torch.is_tensor(v) else v) for k, v in ds.items()}
 70        _, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 71    if metric is None: raise RuntimeError("benchmark training failed")
 72    return (float(metric), model, ds) if keep_model else float(metric)
 73
 74
 75def baseline_factory(cfg):
 76    return lambda seed: train_one("baseline", float(cfg["lr"]), seed)
 77
 78def idea_factory(cfg):
 79    return lambda seed: train_one("idea", float(cfg["lr"]), seed)
 80
 81
 82def mechanism_signature(model, ds):
 83    """Retest orbital prediction on trajectories produced by the trained model.
 84
 85    At each observed phase state, compute the transverse Jacobian spectral
 86    abscissa and compare it with finite-perturbation decay of the trained cell.
 87    """
 88    device = next(model.parameters()).device
 89    x = ds["xte"][:64].to(device)
 90    model.eval()
 91    with torch.no_grad():
 92        _, _, theta = model.forward_with_theta(x, return_states=True)
 93    th = theta[0].detach().cpu().numpy()
 94    A = model.A.detach().cpu().numpy(); k = model.kappa
 95    J = np.zeros((model.hidden, model.hidden))
 96    for i in range(model.hidden):
 97        for j in range(model.hidden):
 98            if i != j:
 99                J[i, j] = k * A[i, j] * math.cos(float(th[j] - th[i]))
100        J[i, i] = -np.sum([k * A[i, j] * math.cos(float(th[j] - th[i])) for j in range(model.hidden) if j != i])
101    eig = np.linalg.eigvals(J)
102    transverse = [z.real for z in eig if abs(z) > 1e-6]
103    predicted = float(max(transverse)) if transverse else 0.0
104    # Empirical cell perturbation on the same trained model and observed input.
105    xone = x[:1]
106    theta0 = torch.zeros(1, model.hidden, device=device)
107    d = torch.linspace(-1, 1, model.hidden, device=device).unsqueeze(0)
108    d = d - d.mean(); d = 1e-5 * d / d.norm()
109    norms = []
110    with torch.no_grad():
111        for t in range(8):
112            theta0 = model.phase_step(theta0, model.inp(xone[:, t*3:(t+1)*3]))
113            # Compare a perturbed parallel rollout at this same input step.
114            # Recompute from initial perturbation for a direct finite response.
115            if t == 0: pert = d.clone()
116            pert = model.phase_step(pert, model.inp(xone[:, t*3:(t+1)*3]))
117            norms.append(float((pert - theta0).norm().cpu()))
118    # fit log decay after first point; positive slope means instability.
119    observed = float(np.polyfit(np.arange(len(norms))[1:], np.log(np.maximum(norms, 1e-30))[1:], 1)[0])
120    # Euler predicts log multiplier log(1 + dt*alpha)/dt.
121    predicted_euler = math.log(max(1e-8, 1.0 + model.dt * predicted)) / model.dt
122    confirmed = bool(abs(observed - predicted_euler) < 0.35 and predicted < 0)
123    return {"predicted_transverse_alpha": predicted, "predicted_euler_slope": predicted_euler,
124            "observed_perturbation_slope": observed, "max_edge_angle_rad": float(np.max(np.abs(((th[:,None]-th[None,:]+np.pi)%(2*np.pi))-np.pi))),
125            "confirmed": confirmed}
126
127
128def main():
129    torch.set_num_threads(4)
130    # Same union of method step sizes on both sides; baseline is selected on 4 seeds.
131    base = sweep_baseline(baseline_factory, LR_GRID)
132    best_lr = float(base["best_cfg"]["lr"])
133    idea_grid = LR_GRID
134    # Explicitly evaluate all three idea settings, then report the best setting.
135    idea_runs = []
136    for cfg in idea_grid:
137        r = evaluate(idea_factory(cfg))
138        idea_runs.append({"cfg": cfg, "result": r})
139    best_idea = min(idea_runs, key=lambda z: z["result"]["mean"])
140    # Retain the trained best model on seed 0 for behavior signature.
141    _, trained, ds = train_one("idea", float(best_idea["cfg"]["lr"]), 0, keep_model=True)
142    report = make_report("dynamics", "rnn_small", base, best_idea["result"], {
143        "prediction": "stable transverse perturbations should contract with slope near log(1+dt*alpha)/dt",
144        "trained_model": "orbital phase RNN, seed 0, best benchmark lr",
145        **mechanism_signature(trained, ds)
146    })
147    report["idea_sweep"] = idea_runs
148    report["protocol"] = {"epochs": EPOCHS, "n_train": NTRAIN, "n_test": NTEST,
149                           "paired_seeds": list(range(8)), "lr_union": [c["lr"] for c in LR_GRID],
150                           "baseline_best_lr": best_lr, "idea_best_lr": best_idea["cfg"]["lr"]}
151    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
152    print(json.dumps(report, indent=2))
153
154if __name__ == "__main__": main()