1"""Stage-2 bench for Maslov phase budget on the matched dynamics RNN.
  2Run from this experiment directory.
  3"""
  4import os, sys, json, math, random
  5import numpy as np
  6import torch
  7import torch.nn as nn
  8
  9sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 10from bench import get_dataset, sweep_baseline, evaluate, make_report
 11
 12TRACK = "dynamics"
 13EPOCHS = 12
 14BATCH = 128
 15NTRAIN, NTEST = 400, 200
 16DT = 1.0
 17
 18
 19def seed_all(seed):
 20    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 21    if torch.cuda.is_available():
 22        torch.cuda.manual_seed_all(seed)
 23
 24
 25def phase_math_check():
 26    # K(exp(-i theta)) convention: q_fd equals sum(theta velocity).
 27    th = np.array([.23, -.41]); vel = np.array([.37, -.22]); eps = 1e-6
 28    c, s = np.cos(th), np.sin(th)
 29    cp, sp = np.cos(th + eps*vel), np.sin(th + eps*vel)
 30    u = np.zeros((4,4)); up = np.zeros((4,4))
 31    u[:2,:2] = np.diag(c); u[:2,2:] = np.diag(s)
 32    u[2:,:2] = np.diag(-s); u[2:,2:] = np.diag(c)
 33    up[:2,:2] = np.diag(cp); up[:2,2:] = np.diag(sp)
 34    up[2:,:2] = np.diag(-sp); up[2:,2:] = np.diag(cp)
 35    J = np.block([[np.zeros((2,2)), np.eye(2)],[-np.eye(2),np.zeros((2,2))]])
 36    right = np.linalg.solve(u.T, (up-u).T).T
 37    q = np.trace(J @ right) / (2*eps)
 38    expected = -float(vel.sum())  # this K convention embeds exp(-i theta)
 39    return {"expected_sum_theta_dot": expected, "finite_difference_q": float(q),
 40            "absolute_error": float(abs(q-expected)), "passed": bool(abs(q-expected) < 1e-7)}
 41
 42
 43class PhaseGRU(nn.Module):
 44    """Same 64-unit GRU as bench rnn_small, exposing hidden trajectory for loss."""
 45    def __init__(self, out_dim=1, hidden=64):
 46        super().__init__()
 47        self.rnn = nn.GRU(3, hidden, batch_first=True)
 48        self.head = nn.Linear(hidden, out_dim)
 49        self.hidden = hidden
 50
 51    def forward(self, x, return_phase=False):
 52        seq = x.view(x.shape[0], -1, 3)
 53        hs, h = self.rnn(seq)
 54        out = self.head(h[-1])
 55        if not return_phase:
 56            return out
 57        # Pair coordinates into n planar planes. This is the observable
 58        # compact/unitary phase proxy of the recurrent hidden trajectory.
 59        a, b = hs[..., 0::2], hs[..., 1::2]
 60        cross = a[..., 1:] * b[..., :-1] - b[..., 1:] * a[..., :-1]
 61        dot = a[..., 1:] * a[..., :-1] + b[..., 1:] * b[..., :-1]
 62        # Consecutive hidden vectors define finite planar rotations. atan2 is
 63        # stable here and remains differentiable; aggregate over planes.
 64        q = torch.atan2(cross, dot).mean(dim=-1) / DT
 65        return out, q
 66
 67
 68def train_one(seed, lr, phase_lambda=0.0, tv_lambda=0.0, weight_decay=0.0,
 69              collect=False):
 70    seed_all(seed)
 71    ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST)
 72    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 73    model = PhaseGRU().to(dev)
 74    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
 75    lossf = nn.MSELoss()
 76    x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
 77    try:
 78        for _ in range(EPOCHS):
 79            model.train(); perm = torch.randperm(len(x), device=dev)
 80            for i in range(0, len(x), BATCH):
 81                idx = perm[i:i+BATCH]
 82                pred, q = model(x[idx], return_phase=True)
 83                task = lossf(pred, y[idx])
 84                phase = phase_lambda * (q*q).mean()
 85                tv = tv_lambda * ((q[:,1:] - q[:,:-1])**2).mean()
 86                loss = task + phase + tv
 87                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 88        model.eval()
 89        with torch.no_grad():
 90            pred, q = model(ds["xte"].to(dev), return_phase=True)
 91            metric = float(lossf(pred, ds["yte"].to(dev)).cpu())
 92            qn = q.detach().cpu().numpy()
 93            stats = {"q_variance": float(qn.var()), "q_abs_mean": float(np.abs(qn).mean()),
 94                     "q_tv": float(np.diff(qn, axis=1).var())}
 95        return (metric, stats) if collect else metric
 96    except RuntimeError:
 97        # Explicit CPU fallback for shared/fragile CUDA environments.
 98        if dev.type == "cuda":
 99            torch.cuda.empty_cache()
100            old = torch.cuda.is_available
101            torch.cuda.is_available = lambda: False
102            try:
103                return train_one(seed, lr, phase_lambda, tv_lambda, weight_decay, collect)
104            finally:
105                torch.cuda.is_available = old
106        raise
107
108
109def make_train(cfg, idea=False, collect=False):
110    def f(seed):
111        return train_one(seed, cfg["lr"], cfg.get("phase", 0.0) if idea else 0.0,
112                         cfg.get("tv", 0.0) if idea else 0.0, cfg.get("wd", 0.0), collect)
113    return f
114
115
116def main():
117    # Union of learning rates is shared by both sides; baseline also sweeps wd.
118    baseline_grid = [{"lr": lr, "wd": wd} for lr in (1e-3, 3e-3, 6e-3)
119                     for wd in (0.0, 1e-4)]
120    base = sweep_baseline(make_train, baseline_grid)
121    best = base["best_cfg"]
122    # Three idea settings: baseline-best and two nearby phase budgets.
123    idea_grid = [dict(best, phase=0.0, tv=0.0),
124                 dict(best, phase=0.002, tv=0.01),
125                 dict(best, phase=0.008, tv=0.04)]
126    idea_sweep = []
127    for cfg in idea_grid:
128        r = evaluate(make_train(cfg, idea=True), seeds=(0,1,2,3))
129        idea_sweep.append({"cfg": cfg, "mean": r["mean"]})
130    best_idea_cfg = min(idea_grid, key=lambda c: next(z["mean"] for z in idea_sweep if z["cfg"] == c))
131    idea_res = evaluate(make_train(best_idea_cfg, idea=True), seeds=tuple(range(8)))
132    # Re-test behavior on the actual final trained systems, not a toy graph.
133    bstats, istats = [], []
134    for s in range(8):
135        _, sb = train_one(s, best["lr"], 0.0, 0.0, best.get("wd",0.0), True)
136        _, si = train_one(s, best_idea_cfg["lr"], best_idea_cfg["phase"], best_idea_cfg["tv"], best_idea_cfg.get("wd",0.0), True)
137        bstats.append(sb); istats.append(si)
138    bvar = float(np.mean([z["q_variance"] for z in bstats])); ivar = float(np.mean([z["q_variance"] for z in istats]))
139    sig = {"quantity": "test hidden angular-velocity variance q", "prediction": "phase budget lowers q variance",
140           "baseline_mean": bvar, "idea_mean": ivar, "ratio_idea_over_baseline": ivar/max(bvar,1e-12),
141           "relative_reduction": 1.0-ivar/max(bvar,1e-12), "confirmed": bool(ivar < 0.9*bvar)}
142    extra = {"track_justification": "dynamics is the matched actuated-pendulum rollout track for stability/control ideas",
143             "idea_sweep": idea_sweep, "mechanism_signature": sig,
144             "math_check": phase_math_check(), "custom_track": None}
145    rep = make_report(TRACK, "rnn_small", base, idea_res, extra)
146    with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2)
147    print(json.dumps(rep, indent=2))
148
149if __name__ == "__main__": main()