Dissipation-Budgeted Nonreversible Sampling / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5from scipy.stats import linregress
  6
  7# Dissipation-Budgeted Nonreversible Sampling: minimal first-passage MVP.
  8# Dynamics: dX = [-V'(X) + u] dt + sqrt(2D) dW,
  9# V(x)=(x^2-1)^2/4, target is the right well x >= 0.8.
 10
 11SEED = 3131
 12DT = 0.002
 13TMAX = 12.0
 14D = 0.08
 15X0 = -1.0
 16TARGET = 0.8
 17N = 18000
 18SCALES = [0.0, 0.15, 0.30, 0.50, 0.80]
 19BUDGETS = [0.02, 0.05, 0.15, 0.50]
 20
 21
 22def force(x):
 23    # -V'(x) = x - x^3
 24    return x - x**3
 25
 26
 27def simulate(u, budget=np.inf, seed=0):
 28    rng = np.random.default_rng(seed)
 29    nstep = int(round(TMAX / DT))
 30    x = np.full(N, X0, dtype=np.float64)
 31    alive = np.ones(N, dtype=bool)
 32    hit = np.full(N, TMAX + DT, dtype=np.float64)
 33    q = np.zeros(N, dtype=np.float64)
 34    # For a constant applied drift, q increment is exactly u^2/(2D) dt.
 35    qrate = u * u / (2.0 * D)
 36    # Store survival at regular reporting times for an independent hazard check.
 37    report_every = 50
 38    survival = []
 39    times = []
 40    for k in range(nstep):
 41        t = k * DT
 42        active = alive
 43        if np.any(active):
 44            # A cumulative budget is enforced before each step. Partial final
 45            # drift is represented by a multiplicative control gate.
 46            remaining = np.maximum(budget - q[active], 0.0)
 47            gate = np.minimum(1.0, remaining / (qrate * DT)) if qrate > 0 else np.ones(np.sum(active))
 48            drift_u = u * gate
 49            q[active] += (drift_u * drift_u / (2.0 * D)) * DT
 50            noise = np.sqrt(2.0 * D * DT) * rng.standard_normal(np.sum(active))
 51            xa = x[active]
 52            xnew = xa + (force(xa) + drift_u) * DT + noise
 53            x[active] = xnew
 54            entered = xnew >= TARGET
 55            inds = np.flatnonzero(active)
 56            if np.any(entered):
 57                ii = inds[entered]
 58                # Linear interpolation makes first-passage estimates less dt-dependent.
 59                prev = xa[entered]
 60                frac = np.clip((TARGET - prev) / (xnew[entered] - prev + 1e-30), 0, 1)
 61                hit[ii] = t + frac * DT
 62                alive[ii] = False
 63        if (k + 1) % report_every == 0:
 64            times.append((k + 1) * DT)
 65            survival.append(np.mean(alive))
 66    return {"u": u, "budget": budget, "hit": hit, "q": q,
 67            "times": np.asarray(times), "survival": np.asarray(survival)}
 68
 69
 70def summarize(r):
 71    hit = r["hit"] <= TMAX
 72    ht = r["hit"][hit]
 73    # Censored mean is not used as the primary speed metric. Estimate a late
 74    # approximately constant hazard from log survival, with a small floor.
 75    t = r["times"]
 76    s = np.maximum(r["survival"], 1.0 / N)
 77    use = (t >= 3.0) & (s > 0.03) & (s < 0.98)
 78    if np.sum(use) >= 3:
 79        fit = linregress(t[use], np.log(s[use]))
 80        rate = max(0.0, -fit.slope)
 81        r2 = fit.rvalue**2
 82    else:
 83        rate, r2 = float("nan"), float("nan")
 84    # Empirical all-time cumulative hazard identity: H=-log S.
 85    s0 = max(r["survival"][0], 1.0 / N)
 86    send = max(r["survival"][-1], 1.0 / N)
 87    H = -math.log(send / s0)
 88    # q is pathwise integrated only until first hit, as required by stopping.
 89    mean_q_hit = float(np.mean(r["q"][hit])) if np.any(hit) else float("nan")
 90    return {"u": r["u"], "budget": r["budget"], "hit_fraction": float(np.mean(hit)),
 91            "mean_hit_time": float(np.mean(ht)) if np.any(hit) else None,
 92            "median_hit_time": float(np.median(ht)) if np.any(hit) else None,
 93            "hazard_tail": float(rate), "hazard_tail_r2": float(r2),
 94            "mean_q_hit": mean_q_hit, "cum_hazard": H,
 95            "log_survival_ratio": float(-math.log(send / s0))}
 96
 97
 98def main():
 99    out = Path("results.json")
100    # Fixed seeds make every control reproducible. Independent seeds avoid
101    # pretending that noisy first-passage paths are exactly paired.
102    runs = []
103    for j, u in enumerate(SCALES):
104        runs.append(simulate(u, np.inf, SEED + j))
105    for j, b in enumerate(BUDGETS):
106        runs.append(simulate(0.5, b, SEED + 100 + j))
107    summaries = [summarize(r) for r in runs]
108
109    # Verify S(t)=exp(-integral hazard) in the discrete empirical sense using
110    # interval hazard h_k=-log(S_{k+1}/S_k)/dt.
111    base = runs[0]
112    ss = np.maximum(base["survival"], 1.0 / N)
113    h = -np.diff(np.log(ss)) / np.diff(base["times"])
114    H_from_h = float(np.sum(h * np.diff(base["times"])))
115    H_direct = float(-np.log(ss[-1] / ss[0]))
116    survival_identity_error = abs(H_from_h - H_direct)
117
118    # Check Q formula on an uncapped constant-drift run: q should equal
119    # u^2*T_hit/(2D) up to the final discretization/interpolation step.
120    unc = runs[SCALES.index(0.5)]
121    valid = unc["hit"] <= TMAX
122    expected_q = 0.5**2 * unc["hit"][valid] / (2 * D)
123    q_error = float(np.max(np.abs(unc["q"][valid] - expected_q))) if np.any(valid) else None
124
125    # Falsifiable functional-form check: log hazard improvement against Q.
126    sweep = [x for x in summaries[:len(SCALES)] if x["u"] > 0 and np.isfinite(x["hazard_tail"])]
127    cal = [x for x in sweep if x["hazard_tail"] > 0 and x["mean_q_hit"] > 0]
128    if len(cal) >= 3:
129        xs = np.array([x["mean_q_hit"] for x in cal])
130        ys = np.log(np.array([x["hazard_tail"] for x in cal]) /
131                    max(summaries[0]["hazard_tail"], 1e-12))
132        lr = linregress(xs, ys)
133        q_fit = {"slope_C": float(lr.slope), "intercept": float(lr.intercept),
134                 "r2": float(lr.rvalue**2), "n": len(cal)}
135    else:
136        q_fit = {"slope_C": None, "intercept": None, "r2": None, "n": len(cal)}
137
138    result = {"config": {"N": N, "dt": DT, "Tmax": TMAX, "D": D,
139                          "target": TARGET, "seed": SEED},
140              "summaries": summaries,
141              "math_checks": {"survival_identity_abs_error": survival_identity_error,
142                              "q_constant_drift_max_abs_error": q_error,
143                              "log_hazard_vs_q_fit": q_fit},
144              "interpretation": {
145                  "baseline": summaries[0],
146                  "uncapped_u_0.5": summaries[SCALES.index(0.5)],
147                  "budgeted": summaries[len(SCALES):]}}
148    out.write_text(json.dumps(result, indent=2))
149    print(json.dumps(result, indent=2))
150
151
152if __name__ == "__main__":
153    main()