"""Stage-2 bench for Maslov phase budget on the matched dynamics RNN. Run from this experiment directory. """ import os, sys, json, math, random import numpy as np import torch import torch.nn as nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, sweep_baseline, evaluate, make_report TRACK = "dynamics" EPOCHS = 12 BATCH = 128 NTRAIN, NTEST = 400, 200 DT = 1.0 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def phase_math_check(): # K(exp(-i theta)) convention: q_fd equals sum(theta velocity). th = np.array([.23, -.41]); vel = np.array([.37, -.22]); eps = 1e-6 c, s = np.cos(th), np.sin(th) cp, sp = np.cos(th + eps*vel), np.sin(th + eps*vel) u = np.zeros((4,4)); up = np.zeros((4,4)) u[:2,:2] = np.diag(c); u[:2,2:] = np.diag(s) u[2:,:2] = np.diag(-s); u[2:,2:] = np.diag(c) up[:2,:2] = np.diag(cp); up[:2,2:] = np.diag(sp) up[2:,:2] = np.diag(-sp); up[2:,2:] = np.diag(cp) J = np.block([[np.zeros((2,2)), np.eye(2)],[-np.eye(2),np.zeros((2,2))]]) right = np.linalg.solve(u.T, (up-u).T).T q = np.trace(J @ right) / (2*eps) expected = -float(vel.sum()) # this K convention embeds exp(-i theta) return {"expected_sum_theta_dot": expected, "finite_difference_q": float(q), "absolute_error": float(abs(q-expected)), "passed": bool(abs(q-expected) < 1e-7)} class PhaseGRU(nn.Module): """Same 64-unit GRU as bench rnn_small, exposing hidden trajectory for loss.""" def __init__(self, out_dim=1, hidden=64): super().__init__() self.rnn = nn.GRU(3, hidden, batch_first=True) self.head = nn.Linear(hidden, out_dim) self.hidden = hidden def forward(self, x, return_phase=False): seq = x.view(x.shape[0], -1, 3) hs, h = self.rnn(seq) out = self.head(h[-1]) if not return_phase: return out # Pair coordinates into n planar planes. This is the observable # compact/unitary phase proxy of the recurrent hidden trajectory. a, b = hs[..., 0::2], hs[..., 1::2] cross = a[..., 1:] * b[..., :-1] - b[..., 1:] * a[..., :-1] dot = a[..., 1:] * a[..., :-1] + b[..., 1:] * b[..., :-1] # Consecutive hidden vectors define finite planar rotations. atan2 is # stable here and remains differentiable; aggregate over planes. q = torch.atan2(cross, dot).mean(dim=-1) / DT return out, q def train_one(seed, lr, phase_lambda=0.0, tv_lambda=0.0, weight_decay=0.0, collect=False): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST) dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = PhaseGRU().to(dev) opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) lossf = nn.MSELoss() x, y = ds["xtr"].to(dev), ds["ytr"].to(dev) try: for _ in range(EPOCHS): model.train(); perm = torch.randperm(len(x), device=dev) for i in range(0, len(x), BATCH): idx = perm[i:i+BATCH] pred, q = model(x[idx], return_phase=True) task = lossf(pred, y[idx]) phase = phase_lambda * (q*q).mean() tv = tv_lambda * ((q[:,1:] - q[:,:-1])**2).mean() loss = task + phase + tv opt.zero_grad(set_to_none=True); loss.backward(); opt.step() model.eval() with torch.no_grad(): pred, q = model(ds["xte"].to(dev), return_phase=True) metric = float(lossf(pred, ds["yte"].to(dev)).cpu()) qn = q.detach().cpu().numpy() stats = {"q_variance": float(qn.var()), "q_abs_mean": float(np.abs(qn).mean()), "q_tv": float(np.diff(qn, axis=1).var())} return (metric, stats) if collect else metric except RuntimeError: # Explicit CPU fallback for shared/fragile CUDA environments. if dev.type == "cuda": torch.cuda.empty_cache() old = torch.cuda.is_available torch.cuda.is_available = lambda: False try: return train_one(seed, lr, phase_lambda, tv_lambda, weight_decay, collect) finally: torch.cuda.is_available = old raise def make_train(cfg, idea=False, collect=False): def f(seed): return train_one(seed, cfg["lr"], cfg.get("phase", 0.0) if idea else 0.0, cfg.get("tv", 0.0) if idea else 0.0, cfg.get("wd", 0.0), collect) return f def main(): # Union of learning rates is shared by both sides; baseline also sweeps wd. baseline_grid = [{"lr": lr, "wd": wd} for lr in (1e-3, 3e-3, 6e-3) for wd in (0.0, 1e-4)] base = sweep_baseline(make_train, baseline_grid) best = base["best_cfg"] # Three idea settings: baseline-best and two nearby phase budgets. idea_grid = [dict(best, phase=0.0, tv=0.0), dict(best, phase=0.002, tv=0.01), dict(best, phase=0.008, tv=0.04)] idea_sweep = [] for cfg in idea_grid: r = evaluate(make_train(cfg, idea=True), seeds=(0,1,2,3)) idea_sweep.append({"cfg": cfg, "mean": r["mean"]}) best_idea_cfg = min(idea_grid, key=lambda c: next(z["mean"] for z in idea_sweep if z["cfg"] == c)) idea_res = evaluate(make_train(best_idea_cfg, idea=True), seeds=tuple(range(8))) # Re-test behavior on the actual final trained systems, not a toy graph. bstats, istats = [], [] for s in range(8): _, sb = train_one(s, best["lr"], 0.0, 0.0, best.get("wd",0.0), True) _, si = train_one(s, best_idea_cfg["lr"], best_idea_cfg["phase"], best_idea_cfg["tv"], best_idea_cfg.get("wd",0.0), True) bstats.append(sb); istats.append(si) bvar = float(np.mean([z["q_variance"] for z in bstats])); ivar = float(np.mean([z["q_variance"] for z in istats])) sig = {"quantity": "test hidden angular-velocity variance q", "prediction": "phase budget lowers q variance", "baseline_mean": bvar, "idea_mean": ivar, "ratio_idea_over_baseline": ivar/max(bvar,1e-12), "relative_reduction": 1.0-ivar/max(bvar,1e-12), "confirmed": bool(ivar < 0.9*bvar)} extra = {"track_justification": "dynamics is the matched actuated-pendulum rollout track for stability/control ideas", "idea_sweep": idea_sweep, "mechanism_signature": sig, "math_check": phase_math_check(), "custom_track": None} rep = make_report(TRACK, "rnn_small", base, idea_res, extra) with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()