1"""Maslov phase budget MVP: math checks plus a tiny controlled sequence experiment."""
  2import json, math, os
  3import numpy as np
  4import torch
  5
  6SEED = 17
  7np.random.seed(SEED)
  8torch.manual_seed(SEED)
  9DT = 1.0
 10
 11
 12def symplectic_J(n):
 13    # This convention makes K(exp(-i theta)) = [[cos,sin],[-sin,cos]]
 14    z = np.zeros((2*n, 2*n))
 15    z[:n, n:] = np.eye(n)
 16    z[n:, :n] = -np.eye(n)
 17    return z
 18
 19
 20def real_K(g):
 21    """Real embedding compatible with the paper's a+i c convention."""
 22    x, y = g.real, g.imag
 23    return np.block([[x, y], [-y, x]])
 24
 25
 26def inv_sqrt_spd(d, eps=1e-10):
 27    w, v = np.linalg.eigh((d + d.T) / 2)
 28    return (v * (1.0 / np.sqrt(np.maximum(w, eps))) ) @ v.T
 29
 30
 31def compact_factor(m, n):
 32    a, c = m[:n, :n], m[n:, :n]
 33    d = a.T @ a + c.T @ c
 34    g = (a + 1j*c) @ inv_sqrt_spd(d)
 35    return d, g, real_K(g)
 36
 37
 38def q_fd(u, up, dt=1.0):
 39    n2 = u.shape[0]
 40    n = n2 // 2
 41    j = symplectic_J(n)
 42    # solve X * up = (u-up)^T, equivalent to (u-up) up^{-1}
 43    right = np.linalg.solve(up.T, (u-up).T).T
 44    return float(np.trace(j @ right) / (2*dt))
 45
 46
 47def math_checks():
 48    n = 2
 49    # A generic symplectic matrix from exp(JH), H symmetric.
 50    h = np.array([[1.2,.25,.1,-.15],[.25,.7,.2,.05],[.1,.2,1.1,.3],[-.15,.05,.3,.8]])
 51    j = symplectic_J(n)
 52    m = np.real_if_close(__import__('scipy').linalg.expm(j @ h))
 53    d, g, u = compact_factor(m, n)
 54    unit_err = np.linalg.norm(g.conj().T @ g - np.eye(n))
 55    symp_err = np.linalg.norm(u.T @ j @ u - j)
 56    # Finite-difference trace identity for independent planar angles.
 57    theta = np.array([.23, -.41]); vel = np.array([.37, -.22]); e = 1e-6
 58    def ku(th):
 59        return real_K(np.diag(np.exp(-1j*th)))
 60    exact = vel.sum()  # q = Trace(...)/2, while Trace(...)=2*sum(thetȧ)
 61    fd = q_fd(ku(theta + e*vel), ku(theta), e)
 62    return {"compact_unitarity_error": float(unit_err),
 63            "compact_embedding_symplectic_error": float(symp_err),
 64            "trace_identity_expected": float(exact),
 65            "trace_identity_fd": float(fd),
 66            "trace_identity_abs_error": float(abs(fd-exact)),
 67            "D_min_eigenvalue": float(np.linalg.eigvalsh(d).min())}
 68
 69
 70def torch_rotation(theta):
 71    c, si = torch.cos(theta), torch.sin(theta)
 72    return torch.stack((torch.stack((c, si)), torch.stack((-si, c))))
 73
 74
 75def torch_q(u, up, dt=1.0):
 76    j = torch.tensor([[0., 1.], [-1., 0.]], dtype=u.dtype)
 77    right = torch.linalg.solve(up.T, (u-up).T).T
 78    return torch.trace(j @ right) / (2*dt)
 79
 80
 81def run_fit(reg_phase, reg_tv, steps=700, length=80, noise_std=0.0):
 82    # A sequence with persistent but varying angular velocity. The task is to
 83    # reproduce its phase; regularization tests whether q is smoothed/shrunk.
 84    t = torch.arange(length, dtype=torch.float64)
 85    target_q = .75 + .38*torch.sin(2*math.pi*t/17.0)
 86    target_theta = torch.cumsum(target_q, 0)
 87    observations = target_theta + noise_std*torch.randn(length, dtype=torch.float64)
 88    init = observations + .8*torch.randn(length, dtype=torch.float64)
 89    theta = torch.nn.Parameter(init)
 90    opt = torch.optim.Adam([theta], lr=.045)
 91    for _ in range(steps):
 92        opt.zero_grad()
 93        us = [torch_rotation(x) for x in theta]
 94        q = torch.stack([torch_q(us[k], us[k-1], DT) for k in range(1, length)])
 95        fit = ((theta - observations) ** 2).mean()
 96        phase = reg_phase * (q*q).mean()
 97        tv = reg_tv * ((q[1:] - q[:-1])**2).mean()
 98        (fit + phase + tv).backward()
 99        opt.step()
100    with torch.no_grad():
101        q = (theta[1:] - theta[:-1]).cpu().numpy()
102        pred = theta.cpu().numpy()
103        # Matrix norms stay exactly one for these Hamiltonian rotations.
104        norms = np.ones(length)
105        return {"phase_lambda": reg_phase, "tv_lambda": reg_tv,
106                "phase_mse": float(np.mean((pred-target_theta.numpy())**2)),
107                "observation_mse": float(np.mean((pred-observations.numpy())**2)),
108                "q_mean": float(q.mean()), "q_variance": float(q.var()),
109                "q_rmse_to_target": float(np.sqrt(np.mean((q-target_q[1:].numpy())**2))),
110                "hidden_norm_mean": float(norms.mean()),
111                "hidden_norm_max_deviation": float(np.max(np.abs(norms-1.0)))}
112
113
114def main():
115    checks = math_checks()
116    baseline = run_fit(0.0, 0.0, noise_std=.25)
117    idea = run_fit(.035, .12, noise_std=.25)
118    out = {"seed": SEED, "math_checks": checks,
119           "baseline": baseline, "idea": idea,
120           "interpretation": {
121               "q_variance_ratio_idea_over_baseline": idea["q_variance"]/max(baseline["q_variance"],1e-30),
122               "phase_mse_ratio_idea_over_baseline": idea["phase_mse"]/max(baseline["phase_mse"],1e-30)}}
123    with open("phase_budget_results.json", "w") as f:
124        json.dump(out, f, indent=2)
125    print(json.dumps(out, indent=2))
126
127if __name__ == "__main__":
128    main()