Contact-Splitting Momentum Optimizer / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1729
  6np.random.seed(SEED)
  7random.seed(SEED)
  8
  9# Contact splitting for f(x)=lambda*x^2/2, M=1.
 10def split_step(x, p, s, h, lam, gamma):
 11    # K(h/2)
 12    s += 0.5*h*0.5*p*p
 13    x += 0.5*h*p
 14    # V(h/2)
 15    f = 0.5*lam*x*x
 16    p -= 0.5*h*lam*x
 17    s -= 0.5*h*f
 18    # D(h)
 19    r = math.exp(-gamma*h)
 20    p *= r
 21    s *= r
 22    # V(h/2)
 23    f = 0.5*lam*x*x
 24    p -= 0.5*h*lam*x
 25    s -= 0.5*h*f
 26    # K(h/2)
 27    s += 0.5*h*0.5*p*p
 28    x += 0.5*h*p
 29    return x, p, s
 30
 31def H(x, p, s, lam, gamma):
 32    return 0.5*p*p + 0.5*lam*x*x + gamma*s
 33
 34def run_split(h, lam=1.0, gamma=0.2, n=2000, x=1.0, p=0.0, s=0.0):
 35    hs, xs = [], []
 36    for _ in range(n + 1):
 37        hs.append(H(x, p, s, lam, gamma))
 38        xs.append(x)
 39        x, p, s = split_step(x, p, s, h, lam, gamma)
 40    return np.asarray(hs), np.asarray(xs)
 41
 42def map_matrix(h, lam, gamma):
 43    r = math.exp(-gamma*h)
 44    b = h/2
 45    a = h*lam/2
 46    def mp(x, p):
 47        x1 = x + b*p
 48        p1 = p - a*x1
 49        p1 = r*p1
 50        p1 = p1 - a*x1
 51        return np.array([x1 + b*p1, p1])
 52    return np.column_stack([mp(1, 0), mp(0, 1)])
 53
 54def slope_log_abs(y, h, start=100):
 55    z = np.abs(y[start:])
 56    ok = z > 1e-12
 57    if ok.sum() < 20:
 58        return float('nan')
 59    t = np.arange(len(z))[ok] * h
 60    return float(np.polyfit(t, np.log(z[ok]), 1)[0])
 61
 62def split_opt_step(x, p, s, h, lam, gamma):
 63    # Same contact splitting, interpreted as an optimizer update.
 64    return split_step(x, p, s, h, lam, gamma)
 65
 66def momentum_step(x, v, lr, beta, lam):
 67    v = beta*v - lr*lam*x
 68    return x + v, v
 69
 70def main():
 71    lam, gamma = 1.0, 0.2
 72    results = {"seed": SEED, "predictions": {}, "comparison": {}}
 73
 74    # Prediction 1: exact contact trajectory has log-energy slope -gamma.
 75    hs, _ = run_split(1e-4, lam, gamma, n=1000)
 76    slope = slope_log_abs(hs, 1e-4, start=100)
 77    results["predictions"]["exact_contact_rate"] = {
 78        "predicted": -gamma, "observed": slope,
 79        "abs_error": abs(slope + gamma)
 80    }
 81
 82    # Prediction 2: symmetric splitting global error is O(h^2): halving h
 83    # should reduce fixed-time state error by about 4.
 84    T = 8.0
 85    ref_h = 1e-4
 86    xr = run_split(ref_h, lam, gamma, n=round(T/ref_h))[1][-1]
 87    errs = []
 88    for h in [0.08, 0.04, 0.02, 0.01]:
 89        xh = run_split(h, lam, gamma, n=round(T/h))[1][-1]
 90        errs.append((h, abs(xh-xr)))
 91    ratios = [errs[i][1]/errs[i+1][1] for i in range(len(errs)-1)]
 92    results["predictions"]["second_order_state_error"] = {
 93        "predicted_halving_ratio": 4.0,
 94        "observed_errors": errs,
 95        "observed_halving_ratios": ratios,
 96        "median_ratio": float(np.median(ratios))
 97    }
 98
 99    # Prediction 3: stability boundary is determined by spectral radius=1,
100    # and for gamma=0 the undamped oscillator has the Verlet boundary h*sqrt(lam)=2.
101    boundary_rows = []
102    for g in [0.0, 0.2, 1.0]:
103        grid = np.linspace(0.01, 4.0, 4000)
104        rho = np.array([max(abs(np.linalg.eigvals(map_matrix(h, lam, g)))) for h in grid])
105        stable = np.where(rho <= 1.000001)[0]
106        observed = float(grid[stable[-1]]) if len(stable) else 0.0
107        boundary_rows.append({"gamma": g, "predicted_gamma0_boundary": 2.0/math.sqrt(lam), "observed_boundary": observed})
108    results["predictions"]["stability_boundary"] = boundary_rows
109
110    # Prediction 2b: the symmetric map's local certificate-rate defect is O(h^2).
111    # Ignore the final transient where H can approach floating-point zero.
112    residual_rows = []
113    for h in [0.08, 0.04, 0.02, 0.01]:
114        eh, _ = run_split(h, lam, gamma, n=round(4.0/h), x=1.0, p=0.0, s=0.3)
115        valid = (np.abs(eh[:-1]) > 1e-8) & (np.abs(eh[1:]) > 1e-8)
116        q = np.log(np.abs(eh[1:][valid])/np.abs(eh[:-1][valid]))/h + gamma
117        residual_rows.append((h, float(np.median(np.abs(q)))))
118    residual_ratios = [residual_rows[i][1]/residual_rows[i+1][1] for i in range(3)]
119    results["predictions"]["certificate_rate_residual"] = {
120        "predicted_halving_ratio": 4.0,
121        "observed_median_abs_residuals": residual_rows,
122        "observed_halving_ratios": residual_ratios,
123        "median_ratio": float(np.median(residual_ratios))
124    }
125
126    # Small optimization comparison: quadratic with curvature spectrum, equal
127    # number of gradient evaluations (split uses two; momentum uses two updates).
128    curv = np.array([0.2, 1.0, 4.0, 12.0])
129    x0 = np.ones_like(curv)
130    def f(x): return float(0.5*np.sum(curv*x*x))
131    h = 0.28; gamma2 = 0.2; beta = math.exp(-gamma2*h)
132    x, p, s = x0.copy(), np.zeros_like(x0), 0.0
133    xb, v = x0.copy(), np.zeros_like(x0)
134    idea_losses, base_losses = [], []
135    for k in range(120):
136        idea_losses.append(f(x)); base_losses.append(f(xb))
137        # vectorized separable contact split
138        s += 0.5*h*0.5*np.sum(p*p)
139        x += 0.5*h*p
140        p -= 0.5*h*curv*x; s -= 0.5*h*f(x)
141        p *= math.exp(-gamma2*h); s *= math.exp(-gamma2*h)
142        p -= 0.5*h*curv*x; s -= 0.5*h*f(x)
143        s += 0.5*h*0.5*np.sum(p*p); x += 0.5*h*p
144        # two standard momentum updates for equal gradient evaluations
145        for _ in range(2): xb, v = momentum_step(xb, v, h/2, beta, curv)
146    results["comparison"] = {
147        "h": h, "steps": 120, "idea_final_loss": idea_losses[-1],
148        "baseline_final_loss": base_losses[-1],
149        "idea_min_loss": min(idea_losses), "baseline_min_loss": min(base_losses),
150        "idea_losses_first_last": [idea_losses[0], idea_losses[-1]],
151        "baseline_losses_first_last": [base_losses[0], base_losses[-1]]
152    }
153    Path("results.json").write_text(json.dumps(results, indent=2))
154    print(json.dumps(results, indent=2))
155
156if __name__ == "__main__":
157    main()