1import json
  2import math
  3from pathlib import Path
  4
  5import numpy as np
  6from scipy.special import lambertw
  7
  8SEED = 2900
  9rng = np.random.default_rng(SEED)
 10
 11
 12def dominant_root(k, tau):
 13    """Dominant characteristic root of lambda + k exp(-lambda tau)=0."""
 14    # lambda*tau = W_j(-k*tau); enumerate branches near the dominant pair.
 15    roots = [lambertw(-k * tau, j) / tau for j in range(-8, 9)]
 16    return max(roots, key=lambda z: z.real)
 17
 18
 19def simulate_linear(k, tau, dt=0.001, duration=80.0, z0=1e-6):
 20    """Euler integration of dz/dt=-k z(t-tau), with constant prehistory."""
 21    n = int(duration / dt)
 22    m = max(1, int(round(tau / dt)))
 23    z = np.full(n + m + 1, z0, dtype=float)
 24    for i in range(m, n + m):
 25        z[i + 1] = z[i] - dt * k * z[i - m]
 26        if abs(z[i + 1]) > 1e100:
 27            z[i + 1:] = z[i + 1]
 28            break
 29    t = np.arange(n + 1) * dt
 30    return t, z[m : m + n + 1]
 31
 32
 33def measured_growth(t, z, unstable):
 34    """Fit log magnitude after a short transient; ignore numerical zeros."""
 35    if not unstable:
 36        return float(np.polyfit(t[-20_000:], np.log(np.maximum(np.abs(z[-20_000:]), 1e-300)), 1)[0])
 37    # A broad fit before saturation/overflow gives the envelope rate.
 38    mask = (np.abs(z) > 1e-8) & (np.abs(z) < 1e50) & (t > 5.0)
 39    if mask.sum() < 20:
 40        return float("nan")
 41    return float(np.polyfit(t[mask], np.log(np.abs(z[mask])), 1)[0])
 42
 43
 44def escape_times(k, tau, Ds, R=0.1, dt=0.001, duration=220.0):
 45    """Direct deterministic delayed evolution from fluctuation amplitude sqrt(D)."""
 46    r = max(0.0, dominant_root(k, tau).real)
 47    out = []
 48    for D in Ds:
 49        t, z = simulate_linear(k, tau, dt=dt, duration=duration, z0=math.sqrt(D))
 50        hit = np.flatnonzero(np.abs(z) >= R)
 51        out.append(float(t[hit[0]]) if len(hit) else float("inf"))
 52    return r, out
 53
 54
 55def double_well_grad(x):
 56    return x * (x * x - 1.0)
 57
 58
 59def run_double_well(eta, steps, delayed_m=0, noise=0.018, seed=0):
 60    """Same noisy SGD setup; delayed_m=0 is current-gradient SGD."""
 61    rg = np.random.default_rng(seed)
 62    x = 0.95
 63    history = [x]
 64    gradients = []
 65    for _ in range(steps):
 66        gradients.append(double_well_grad(x))
 67        g = gradients[-1] if delayed_m == 0 or len(gradients) <= delayed_m else gradients[-1 - delayed_m]
 68        x -= eta * g + math.sqrt(eta) * noise * rg.normal()
 69        history.append(x)
 70    h = np.asarray(history)
 71    # A crossing into the opposite basin is a clear toy escape event.
 72    crossings = np.flatnonzero(h < -0.5)
 73    return int(crossings[0]) if len(crossings) else None, float(np.mean((h[-1000:] ** 2 - 1.0) ** 2) / 4.0)
 74
 75
 76def main():
 77    k = 1.0
 78    tau_c = math.pi / (2 * k)
 79    multiples = [0.5, 0.9, 1.05, 1.3, 2.0]
 80    boundary_rows = []
 81    for mult in multiples:
 82        tau = mult * tau_c
 83        root = dominant_root(k, tau)
 84        unstable = root.real > 1e-10
 85        t, z = simulate_linear(k, tau)
 86        fit = measured_growth(t, z, unstable)
 87        boundary_rows.append({
 88            "multiple_tau_c": mult, "tau": tau,
 89            "predicted_real_lambda": float(root.real),
 90            "predicted_unstable": bool(unstable),
 91            "observed_growth_fit": fit,
 92            "observed_unstable": bool(fit > 1e-4),
 93        })
 94
 95    # Prediction 1: transition occurs at k*tau=pi/2.
 96    # Prediction 2: above threshold, measured growth follows Re(lambda+).
 97    growth_rows = [r for r in boundary_rows if r["predicted_unstable"]]
 98    growth_abs_errors = [abs(r["observed_growth_fit"] - r["predicted_real_lambda"]) for r in growth_rows]
 99
100    # Prediction 3: escape time is affine in log(R/sqrt(D)), with slope 1/r.
101    tau_escape = 1.3 * tau_c
102    Ds = np.logspace(-12, -5, 8)
103    r, times = escape_times(k, tau_escape, Ds)
104    x = np.log(0.1 / np.sqrt(Ds))
105    slope, intercept = np.polyfit(x, np.asarray(times), 1)
106    escape_rows = [{"D": float(D), "predicted_T": float(T)} for D, T in zip(Ds, times)]
107
108    # Secondary mini-comparison: fixed stale gradients versus calibrated burst delay.
109    # At the well curvature k=2, tau_c=pi/4; m=ceil(1.1*tau_c/eta).
110    eta = 0.02
111    burst_m = int(math.ceil(1.1 * (math.pi / 4) / eta))
112    baseline = [run_double_well(eta, 30000, delayed_m=0, seed=s) for s in range(12)]
113    fixed_stale = [run_double_well(eta, 30000, delayed_m=burst_m, seed=s) for s in range(12)]
114    def summarize(a):
115        events = [x[0] for x in a if x[0] is not None]
116        return {"escape_fraction": len(events) / len(a), "median_escape_step": float(np.median(events)) if events else None,
117                "final_loss": float(np.mean([x[1] for x in a]))}
118
119    result = {
120        "seed": SEED, "tau_c_k1": tau_c,
121        "boundary_sweep": boundary_rows,
122        "boundary_prediction": {"predicted_multiple": 1.0,
123                                 "observed_first_unstable_multiple": next((r["multiple_tau_c"] for r in boundary_rows if r["observed_unstable"]), None)},
124        "growth_prediction": {"mean_absolute_rate_error": float(np.mean(growth_abs_errors)) if growth_abs_errors else None,
125                               "predicted_rates": [r["predicted_real_lambda"] for r in growth_rows],
126                               "observed_rates": [r["observed_growth_fit"] for r in growth_rows]},
127        "escape_prediction": {"tau": tau_escape, "r": r, "predicted_slope_1_over_r": 1.0 / r,
128                               "observed_slope": float(slope), "slope_relative_error": float(abs(slope - 1.0 / r) / (1.0 / r)),
129                               "intercept": float(intercept), "points": escape_rows},
130        "double_well": {"eta": eta, "calibrated_delay_steps": burst_m,
131                         "current_gradient_sgd": summarize(baseline),
132                         "fixed_delayed_gradient_sgd": summarize(fixed_stale)},
133    }
134    Path("results.json").write_text(json.dumps(result, indent=2))
135    print(json.dumps(result, indent=2))
136
137
138if __name__ == "__main__":
139    main()