1import json, math
  2import numpy as np
  3
  4# OU Langevin certificate and a certificate-controlled noise schedule.
  5# dX=-m X dt+sqrt(2T)dW, with stationary pi=N(0,T/m).
  6SEED = 551
  7m, T, eta = 1.0, 0.5, 0.02
  8x0_mean = 3.0
  9x0_var = T / m                 # shifted stationary Gaussian initial law
 10H = 1.5                         # unsafe region A={x>H}
 11steps = 500
 12n_particles = 20000
 13delta = 0.10
 14
 15
 16def normal_tail(z):
 17    return 0.5 * math.erfc(z / math.sqrt(2.0))
 18
 19
 20def stationary_risk(h=H):
 21    return normal_tail(h / math.sqrt(T / m))
 22
 23
 24def exact_risk(t, h=H):
 25    mu = x0_mean * math.exp(-m * t)
 26    var = (T / m) + (x0_var - T / m) * math.exp(-2 * m * t)
 27    return normal_tail((h - mu) / math.sqrt(var))
 28
 29
 30def exact_chi2():
 31    # chi^2(N(mu,s2)||N(0,s2)) = exp(mu^2/s2)-1.
 32    return math.exp(x0_mean ** 2 / x0_var) - 1.0
 33
 34
 35def certificate(t):
 36    p = stationary_risk()
 37    return p + math.sqrt(p * exact_chi2()) * math.exp(-m * t)
 38
 39
 40def sufficient_burn_in():
 41    p = stationary_risk()
 42    numerator = math.sqrt(p * exact_chi2())
 43    if delta <= p:
 44        return math.inf
 45    return max(0.0, math.log(numerator / (delta - p)) / m)
 46
 47
 48def verify_envelope():
 49    ts = np.linspace(0, 12, 241)
 50    risks = np.array([exact_risk(float(t)) for t in ts])
 51    bounds = np.array([certificate(float(t)) for t in ts])
 52    return {
 53        "max_exact_minus_bound": float(np.max(risks - bounds)),
 54        "min_bound_minus_exact": float(np.min(bounds - risks)),
 55        "exact_risk_t0": exact_risk(0.0),
 56        "exact_risk_t1": exact_risk(1.0),
 57        "exact_risk_t4": exact_risk(4.0),
 58        "stationary_pi_A": stationary_risk(),
 59        "chi0_squared": exact_chi2(),
 60        "bound_t0": certificate(0.0),
 61        "bound_t1": certificate(1.0),
 62        "bound_t4": certificate(4.0),
 63        "analytic_burn_in_for_delta": sufficient_burn_in(),
 64    }
 65
 66
 67def run_schedule(certificate_stop=False):
 68    # Same initial law and same random seed for both schedules. The controlled
 69    # schedule consumes no further noise after the certificate crossing.
 70    r = np.random.default_rng(SEED)
 71    x = r.normal(x0_mean, math.sqrt(x0_var), size=n_particles)
 72    unsafe_fractions, losses = [], []
 73    stop_step = None
 74    for k in range(steps):
 75        t = k * eta
 76        if certificate_stop and stop_step is None and certificate(t) <= delta:
 77            stop_step = k
 78        noise = 0.0 if (certificate_stop and stop_step is not None) else math.sqrt(2 * T * eta) * r.normal(size=n_particles)
 79        x = x - eta * m * x + noise
 80        unsafe_fractions.append(float(np.mean(x > H)))
 81        losses.append(float(np.mean(0.5 * m * x * x)))
 82    post = unsafe_fractions[stop_step:] if stop_step is not None else unsafe_fractions
 83    return {
 84        "max_unsafe_fraction": float(np.max(unsafe_fractions)),
 85        "mean_unsafe_fraction": float(np.mean(unsafe_fractions)),
 86        "final_unsafe_fraction": unsafe_fractions[-1],
 87        "mean_post_stop_unsafe_fraction": float(np.mean(post)),
 88        "final_loss": losses[-1],
 89        "mean_last_50_loss": float(np.mean(losses[-50:])),
 90        "stop_step": stop_step,
 91        "stop_time": None if stop_step is None else stop_step * eta,
 92    }
 93
 94
 95def main():
 96    out = {
 97        "setup": {"m": m, "T": T, "eta": eta, "H": H, "steps": steps,
 98                  "particles": n_particles, "delta": delta, "seed": SEED},
 99        "verification": verify_envelope(),
100        "baseline": run_schedule(False),
101        "idea": run_schedule(True),
102    }
103    with open("results.json", "w") as f:
104        json.dump(out, f, indent=2)
105    print(json.dumps(out, indent=2))
106
107
108if __name__ == "__main__":
109    main()