1import sys, json, random
  2import numpy as np
  3sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  4import torch
  5import torch.nn as nn
  6from bench import get_dataset, evaluate, sweep_baseline, permutation_pvalue, make_report
  7
  8SEEDS = tuple(range(8))
  9GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 10EPOCHS = 18
 11NTR, NTE = 1200, 300
 12
 13
 14def seed_all(seed):
 15    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 16    if torch.cuda.is_available():
 17        torch.cuda.manual_seed_all(seed)
 18
 19
 20class SharedEncoder(nn.Module):
 21    def __init__(self):
 22        super().__init__()
 23        self.rnn = nn.GRU(3, 64, batch_first=True)
 24        self.phi = nn.Linear(64, 1)
 25        self.readout = nn.Linear(3, 1)
 26        self.no_cudnn = False
 27
 28    def encode(self, x):
 29        seq = x.view(x.shape[0], -1, 3)
 30        try:
 31            _, h = self.rnn(seq)
 32        except RuntimeError:
 33            self.no_cudnn = True
 34        if self.no_cudnn:
 35            old = torch.backends.cudnn.enabled
 36            torch.backends.cudnn.enabled = False
 37            try:
 38                _, h = self.rnn(seq)
 39            finally:
 40                torch.backends.cudnn.enabled = old
 41        return self.phi(h[-1]).squeeze(-1)
 42
 43    def q(self, x):
 44        z = self.encode(x)
 45        # Controlled Koopman regressor: [z_t, terminal control, 1].
 46        u = x.view(x.shape[0], -1, 3)[:, -1, 2]
 47        return torch.stack((z, u, torch.ones_like(z)), dim=1)
 48
 49    def forward(self, x):
 50        return self.readout(self.q(x)).squeeze(-1)
 51
 52
 53def fit_baseline(seed, lr):
 54    seed_all(seed)
 55    d = get_dataset("dynamics", seed, NTR, NTE)
 56    net = SharedEncoder()
 57    opt = torch.optim.Adam(net.parameters(), lr=lr)
 58    x, y = d["xtr"], d["ytr"].squeeze(1)
 59    net.train()
 60    for _ in range(EPOCHS):
 61        opt.zero_grad()
 62        loss = ((net(x) - y) ** 2).mean()
 63        loss.backward(); opt.step()
 64    net.eval()
 65    with torch.no_grad():
 66        return float(((net(d["xte"]) - d["yte"].squeeze(1)) ** 2).mean())
 67
 68
 69def mn_posterior(q, y, prior_scale):
 70    # d=1, r=3 Matrix Normal-Inverse-Wishart update.
 71    dtype = q.dtype; dev = q.device
 72    K0 = prior_scale * torch.eye(3, dtype=dtype, device=dev)
 73    M0 = torch.zeros((1, 3), dtype=dtype, device=dev)
 74    S0 = torch.ones((1, 1), dtype=dtype, device=dev) * 0.02
 75    nu0 = 4.0
 76    K = K0 + q.T @ q
 77    B = M0 @ K0 + y.reshape(1, -1) @ q
 78    M = torch.linalg.solve(K, B.T).T
 79    nu = nu0 + q.shape[0]
 80    S = S0 + y.reshape(1, -1) @ y.reshape(-1, 1) + M0 @ K0 @ M0.T - M @ K @ M.T
 81    S = (S + S.T) / 2 + 1e-5 * torch.eye(1, dtype=dtype, device=dev)
 82    return K, M, S, nu
 83
 84
 85def fit_idea(seed, lr, return_signature=False):
 86    seed_all(seed)
 87    d = get_dataset("dynamics", seed, NTR, NTE)
 88    net = SharedEncoder()
 89    # Meta-like prior fitting: train encoder using differentiable closed-form posterior.
 90    opt = torch.optim.Adam(net.parameters(), lr=lr)
 91    x, y = d["xtr"], d["ytr"].squeeze(1)
 92    net.train()
 93    for _ in range(EPOCHS):
 94        opt.zero_grad()
 95        q = net.q(x)
 96        K, M, S, nu = mn_posterior(q, y, 0.7)
 97        pred = (q @ M.T).squeeze(1)
 98        h = torch.diagonal(q @ torch.linalg.solve(K, q.T))
 99        # Gaussian training objective with posterior predictive variance;
100        # this is the scalar Student-t quadratic/NLL surrogate.
101        var = ((1.0 + h) * S.squeeze() / (nu - 1.0)).clamp_min(1e-5)
102        loss = (0.5 * ((y - pred) ** 2 / var + torch.log(var))).mean()
103        loss.backward(); opt.step()
104    net.eval()
105    with torch.no_grad():
106        qc = net.q(x); K, M, S, nu = mn_posterior(qc, y, 0.7)
107        qt = net.q(d["xte"])
108        pred = (qt @ M.T).squeeze(1)
109        h = torch.diagonal(qt @ torch.linalg.solve(K, qt.T))
110        var = ((1.0 + h) * S.squeeze() / (nu - 1.0)).clamp_min(1e-5)
111        yt = d["yte"].squeeze(1)
112        mse = float(((pred - yt) ** 2).mean())
113        if return_signature:
114            # Measured trained-model behavior: compare low/high leverage subsets.
115            med = torch.median(h)
116            near = var[h <= med].mean().item(); far = var[h > med].mean().item()
117            empirical_ratio = far / max(near, 1e-12)
118            expected_ratio = ((1 + h[h > med]).mean() / (1 + h[h <= med]).mean()).item()
119            coverage = float(((yt - pred).abs() <= 1.645 * torch.sqrt(var)).float().mean())
120            return mse, {"near_predictive_variance": near, "far_predictive_variance": far,
121                         "observed_far_near_ratio": empirical_ratio,
122                         "predicted_far_near_ratio": expected_ratio,
123                         "interval_90pct_coverage": coverage,
124                         "confirmed": bool(empirical_ratio > 1.0 and abs(empirical_ratio-expected_ratio) / max(expected_ratio,1e-9) < 0.15)}
125        return mse
126
127
128def main():
129    base = sweep_baseline(lambda cfg: lambda seed: fit_baseline(seed, cfg["lr"]), GRID,
130                          seeds=(0, 1, 2, 3))
131    # Same union of learning rates is evaluated for the idea; report best full-seed result.
132    idea_sweep = []
133    for cfg in GRID:
134        r = evaluate(lambda seed, lr=cfg["lr"]: fit_idea(seed, lr), SEEDS)
135        idea_sweep.append({"cfg": cfg, "mean": r["mean"], "full": r})
136    best = min(idea_sweep, key=lambda z: z["mean"])
137    idea = best["full"]
138    # Explicit paired deltas at selected best settings.
139    bvals = [fit_baseline(s, base["best_cfg"]["lr"]) for s in SEEDS]
140    ivals = [fit_idea(s, best["cfg"]["lr"]) for s in SEEDS]
141    diffs = [i-b for i,b in zip(ivals,bvals)]
142    idea["per_seed"] = ivals; idea["mean"] = float(np.mean(ivals)); idea["std"] = float(np.std(ivals))
143    report = make_report("dynamics", "shared_gru_latent_linear_head", base, idea,
144                         {"idea_sweep": idea_sweep,
145                          "paired_delta": {"per_seed": diffs, "mean": float(np.mean(diffs)),
146                                           "permutation_p": permutation_pvalue(diffs)},
147                          "mechanism_signature": fit_idea(0, best["cfg"]["lr"], True)[1],
148                          "structural_match": "controlled damped pendulum multi-step dynamics"})
149    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
150    print(json.dumps(report, indent=2))
151
152if __name__ == "__main__": main()