Averaged Contractive State-Space Network / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2import time
  3import numpy as np
  4from scipy.integrate import solve_ivp
  5
  6SEED = 2040
  7np.random.seed(SEED)
  8
  9# Periodic contractive linear ODE:
 10#   h' = A(s) h + b, A(s) = -alpha I + q sin(2 pi s) K,
 11# with K=diag(1,-1).  Its phase average is Abar=-alpha I.
 12# Since K is symmetric, mu_2(A(s)) = -alpha + q sin(2 pi s)*lambda_max(K)
 13# and therefore mu_2 <= -(alpha-q), for q < alpha.
 14ALPHA = 1.2
 15K = np.diag([1.0, -1.0])
 16B = np.array([0.7, -0.35])
 17H0 = np.array([0.4, -0.8])
 18T_HORIZON = 4.0
 19
 20
 21def A_of_phase(s, q):
 22    return -ALPHA * np.eye(2) + q * np.sin(2.0 * np.pi * s) * K
 23
 24
 25def rhs_fast(t, h, eps, q):
 26    return A_of_phase(t / eps, q).dot(h) + B
 27
 28
 29def rhs_avg(t, h, q):
 30    # Exact phase average: average(sin)=0. q is retained only for API symmetry.
 31    return -ALPHA * h + B
 32
 33
 34def solve_fast(eps, q, h0=H0, horizon=T_HORIZON, dense=False):
 35    # Resolving each oscillation with at least 20 adaptive solver steps keeps
 36    # integration error well below the averaging error in the sweep.
 37    sol = solve_ivp(lambda t, h: rhs_fast(t, h, eps, q), (0, horizon), h0,
 38                    rtol=2e-10, atol=2e-12, max_step=eps / 25.0,
 39                    dense_output=dense)
 40    if not sol.success:
 41        raise RuntimeError(sol.message)
 42    return sol
 43
 44
 45def solve_avg(h0=H0, horizon=T_HORIZON, dense=False):
 46    sol = solve_ivp(lambda t, h: rhs_avg(t, h, 0.0), (0, horizon), h0,
 47                    rtol=2e-11, atol=2e-13, max_step=0.03, dense_output=dense)
 48    if not sol.success:
 49        raise RuntimeError(sol.message)
 50    return sol
 51
 52
 53def trajectory_error(eps, q, n=2001):
 54    ts = np.linspace(0, T_HORIZON, n)
 55    fast = solve_fast(eps, q, dense=True).sol(ts).T
 56    avg = solve_avg(dense=True).sol(ts).T
 57    err = np.linalg.norm(fast - avg, axis=1)
 58    return float(np.max(err)), float(np.sqrt(np.mean(err * err)))
 59
 60
 61def contraction_sweep(qs):
 62    rows = []
 63    for q in qs:
 64        # Direct grid verification of the matrix-measure certificate.
 65        phases = np.linspace(0, 1, 10001)
 66        mus = np.array([np.linalg.eigvalsh((A_of_phase(s, q) + A_of_phase(s, q).T) / 2).max()
 67                        for s in phases])
 68        predicted = -ALPHA + abs(q)
 69        rows.append({"q": float(q), "predicted_max_mu": float(predicted),
 70                     "observed_max_mu": float(mus.max()),
 71                     "certificate_holds": bool(mus.max() < 0)})
 72    return rows
 73
 74
 75def perturbation_decay(q, horizon=12.0, n=2401):
 76    # Same input, two initial states: their difference has an exact contraction
 77    # envelope. Measure empirical slope after transients, and test the pointwise
 78    # certificate log ratio <= -(alpha-q)t.
 79    d0 = np.array([1.0, -0.7])
 80    sol1 = solve_fast(0.08, q, h0=np.zeros(2), horizon=horizon, dense=True)
 81    sol2 = solve_fast(0.08, q, h0=d0, horizon=horizon, dense=True)
 82    ts = np.linspace(0, horizon, n)
 83    d = np.linalg.norm(sol2.sol(ts).T - sol1.sol(ts).T, axis=1)
 84    log_ratio = np.log(np.maximum(d, 1e-300) / np.linalg.norm(d0))
 85    # Fit the late-time slope, where periodic modulation averages out.
 86    mask = (ts >= 4.0) & (ts <= horizon)
 87    slope = float(np.polyfit(ts[mask], log_ratio[mask], 1)[0])
 88    cert_rate = ALPHA - abs(q)
 89    envelope_margin = float(np.max(log_ratio + cert_rate * ts))
 90    return {"q": float(q), "predicted_long_run_slope": -ALPHA,
 91            "observed_late_log_slope": slope,
 92            "predicted_certificate_rate": cert_rate,
 93            "max_log_envelope_margin": envelope_margin}
 94
 95
 96def rollout_speed(eps=0.0625, q=0.8, horizon=4.0, repeats=3):
 97    # Compare expensive phase-resolved solve with the averaged rollout. Both
 98    # solve the same initial-value problem and are evaluated at 2001 points.
 99    ts = np.linspace(0, horizon, 2001)
100    fast_times, avg_times = [], []
101    for _ in range(repeats):
102        t0 = time.perf_counter()
103        sf = solve_fast(eps, q, horizon=horizon, dense=True)
104        _ = sf.sol(ts)
105        fast_times.append(time.perf_counter() - t0)
106        t0 = time.perf_counter()
107        sa = solve_avg(horizon=horizon, dense=True)
108        _ = sa.sol(ts)
109        avg_times.append(time.perf_counter() - t0)
110    return {"eps": eps, "fast_seconds_median": float(np.median(fast_times)),
111            "averaged_seconds_median": float(np.median(avg_times)),
112            "speedup": float(np.median(fast_times) / np.median(avg_times))}
113
114
115def main():
116    q = 0.8
117    eps_list = [0.5, 0.25, 0.125, 0.0625, 0.03125]
118    errors = []
119    for eps in eps_list:
120        mx, rms = trajectory_error(eps, q)
121        errors.append({"eps": eps, "max_error": mx, "rms_error": rms})
122    loge = np.log([r["max_error"] for r in errors])
123    logeps = np.log(eps_list)
124    scaling = float(np.polyfit(logeps, loge, 1)[0])
125    result = {
126        "seed": SEED,
127        "system": {"alpha": ALPHA, "q": q, "K": K.tolist(), "B": B.tolist()},
128        "predictions": {
129            "averaging_error_power": "O(eps), predicted log-log slope 1",
130            "contraction_certificate": "max_s mu_2(A(s)) = -alpha+q; boundary q=alpha",
131            "perturbation_decay": "late log-distance slope equals -alpha, and is bounded by -(alpha-q)"
132        },
133        "averaging_sweep": {"rows": errors, "observed_loglog_slope": scaling},
134        "contraction_boundary_sweep": contraction_sweep([0.0, 0.6, 1.19, 1.2, 1.4]),
135        "perturbation_decay": [perturbation_decay(v) for v in [0.0, 0.4, 0.8, 1.1]],
136        "rollout_speed": rollout_speed()
137    }
138    with open("results.json", "w") as f:
139        json.dump(result, f, indent=2)
140    print(json.dumps(result, indent=2))
141
142
143if __name__ == "__main__":
144    main()