1import os, sys, json
  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, train_model, evaluate, sweep_baseline, make_report
  8
  9TRACK = "dynamics"
 10EPOCHS = 18
 11LRS = [1e-3, 3e-3, 1e-2]
 12D = 4
 13
 14
 15def fixed_ops():
 16    # Hermitian Pauli interactions on two latent qubits; X and Z terms do not commute.
 17    I = torch.eye(2, dtype=torch.float32)
 18    X = torch.tensor([[0., 1.], [1., 0.]])
 19    Y = torch.tensor([[0., -1.], [1., 0.]])  # real representation of iY
 20    Z = torch.tensor([[1., 0.], [0., -1.]])
 21    kron = torch.kron
 22    return torch.stack([kron(X, I), kron(Z, X), kron(Y, Y), kron(Z, Z)])
 23
 24
 25class MatchedRNN(nn.Module):
 26    def __init__(self, idea, seed):
 27        super().__init__()
 28        # This is the bench rnn_small backbone, kept identical between systems.
 29        self.rnn = nn.GRU(3, 64, batch_first=True)
 30        self.proj = nn.Linear(64, D, bias=False)
 31        g = torch.Generator().manual_seed(10000 + int(seed))
 32        with torch.no_grad():
 33            self.proj.weight.copy_(torch.randn(D, 64, generator=g) / np.sqrt(64))
 34        for p in self.proj.parameters():
 35            p.requires_grad_(False)
 36        self.idea = idea
 37        if idea:
 38            self.theta = nn.Parameter(torch.randn(D, generator=g) * 0.15)
 39            self.register_buffer("ops", fixed_ops())
 40        else:
 41            self.readout = nn.Linear(D, 1, bias=False)
 42
 43    def forward(self, x):
 44        seq = x.view(x.shape[0], -1, 3)
 45        _, h = self.rnn(seq)
 46        z = self.proj(h[-1])
 47        if not self.idea:
 48            return self.readout(z)
 49        # rho is a pure-state density matrix made from the shared latent feature.
 50        q = z / (torch.linalg.vector_norm(z, dim=1, keepdim=True) + 1e-6)
 51        rho = q.unsqueeze(2) * q.unsqueeze(1)
 52        H = torch.einsum("j,jab->ab", self.theta, self.ops)
 53        lam, U = torch.linalg.eigh(H)
 54        A = (U * torch.tanh(lam).unsqueeze(0)) @ U.T
 55        pred = torch.einsum("bij,ji->b", rho, A)
 56        return pred.unsqueeze(1)
 57
 58
 59def run_one(idea, lr, seed, keep_model=False):
 60    torch.manual_seed(7000 + int(seed))
 61    np.random.seed(7000 + int(seed))
 62    ds = get_dataset(TRACK, int(seed), n_train=400, n_test=400)
 63    model = MatchedRNN(idea, seed)
 64    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 65    if net is None:
 66        raise RuntimeError("bench training failed")
 67    return (float(metric), net, ds) if keep_model else float(metric)
 68
 69
 70def make_train(idea, cfg=None):
 71    cfg = cfg or {"lr": 3e-3, "epochs": EPOCHS}
 72    return lambda seed: run_one(idea, float(cfg["lr"]), seed)
 73
 74
 75def behavior_signature(idea_res):
 76    # Re-test the claimed nonlinear expressivity on predictions of trained systems.
 77    vals = []
 78    nonlinear = []
 79    comm = fixed_ops()
 80    comm_norm = float(torch.linalg.matrix_norm(comm[0] @ comm[1] - comm[1] @ comm[0]))
 81    for seed in range(8):
 82        metric, net, ds = run_one(True, idea_res["cfg"]["lr"], seed, True)
 83        net = net.to("cpu")
 84        with torch.no_grad():
 85            p = net(ds["xte"]).squeeze().numpy()
 86            # Shared latent coordinates are reconstructed from the trained backbone.
 87            seq = ds["xte"].view(len(ds["xte"]), -1, 3)
 88            _, h = net.rnn(seq)
 89            z = net.proj(h[-1]).numpy()
 90        X = np.column_stack([np.ones(len(z)), z])
 91        fit = X @ np.linalg.lstsq(X, p, rcond=None)[0]
 92        nl = float(np.sqrt(np.mean((p - fit) ** 2)))
 93        vals.append(float(metric)); nonlinear.append(nl)
 94    return {
 95        "prediction": "noncommuting Hamiltonian terms should create a measurable nonlinear readout",
 96        "commutator_frobenius": comm_norm,
 97        "trained_idea_test_mse_mean": float(np.mean(vals)),
 98        "trained_idea_nonlinear_residual_mean": float(np.mean(nonlinear)),
 99        "confirmed": bool(comm_norm > 1e-6 and np.mean(nonlinear) > 1e-4),
100    }
101
102
103def main():
104    # Baseline and idea use the same union of step sizes, satisfying search parity.
105    grid = [{"lr": lr, "epochs": EPOCHS} for lr in LRS]
106    base = sweep_baseline(lambda cfg: make_train(False, cfg), grid)
107    idea_trials = []
108    best = None
109    for cfg in grid:
110        r = evaluate(make_train(True, cfg))
111        idea_trials.append({"cfg": cfg, **r})
112        if best is None or r["mean"] < best["mean"]:
113            best = {"cfg": cfg, **r}
114    sig = behavior_signature(best)
115    report = make_report(TRACK, "rnn_small_matched_spectral_head", base, best,
116                         {"mechanism_signature": sig,
117                          "idea_sweep": idea_trials,
118                          "structural_match": "dynamics: controlled pendulum rollout"})
119    report["budget"] = {"epochs": EPOCHS, "n_train": 400, "n_test": 400, "seeds": 8}
120    with open("bench_report.json", "w") as f:
121        json.dump(report, f, indent=2)
122    print(json.dumps(report, indent=2))
123
124
125if __name__ == "__main__":
126    main()