1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6# Adaptive reset neural-ODE MVP.  The scalar flow is deliberately chosen so that
  7# its exact solution makes the paper's stopping rule quantitatively testable.
  8SEED = 2093
  9rng = np.random.default_rng(SEED)
 10
 11
 12def q95(x):
 13    return float(np.quantile(np.asarray(x), 0.95))
 14
 15
 16def constant_flow_error(s, lam, lam_hat, y0):
 17    return np.abs(np.exp(lam * s) - np.exp(lam_hat * s)) * np.abs(y0)
 18
 19
 20def analytic_crossing(eps, lam, lam_hat, yq, horizon=10.0):
 21    # First crossing of the exact supremum/quantile error.  Bisection is used
 22    # only for the analytic verification; the implementation uses a grid.
 23    def f(s):
 24        return abs(math.exp(lam*s) - math.exp(lam_hat*s)) * yq - eps
 25    if f(horizon) <= 0:
 26        return horizon
 27    lo, hi = 0.0, horizon
 28    for _ in range(80):
 29        mid = (lo + hi) / 2
 30        if f(mid) > 0:
 31            hi = mid
 32        else:
 33            lo = mid
 34    return (lo + hi) / 2
 35
 36
 37def grid_crossing(eps, lam, lam_hat, y0, dt=0.002, horizon=10.0):
 38    ts = np.arange(dt, horizon + 0.5 * dt, dt)
 39    e = np.array([q95(constant_flow_error(t, lam, lam_hat, y0)) for t in ts])
 40    hit = np.flatnonzero(e > eps)
 41    return float(ts[hit[0]]) if len(hit) else horizon
 42
 43
 44def verify_math():
 45    y0 = np.linspace(0.8, 1.2, 101)
 46    yq = np.quantile(y0, .95)
 47    rows = []
 48    # Prediction 1: grid boundary converges to the analytic first crossing.
 49    for eps in (.01, .03, .08):
 50        exact = analytic_crossing(eps, .5, .35, yq)
 51        observed = grid_crossing(eps, .5, .35, y0)
 52        rows.append({"kind": "tolerance", "epsilon": eps, "predicted": exact,
 53                     "observed": observed, "abs_error": abs(exact-observed)})
 54    # Prediction 2: increasing mismatch causes an earlier boundary.
 55    mismatch_rows = []
 56    for mh in (.49, .45, .40, .30):
 57        exact = analytic_crossing(.03, .5, mh, yq)
 58        observed = grid_crossing(.03, .5, mh, y0)
 59        mismatch_rows.append({"lambda_hat": mh, "predicted": exact, "observed": observed})
 60    # Prediction 3: at fixed mismatch, tolerance monotonically increases window.
 61    tol_exact = [analytic_crossing(e, .5, .35, yq) for e in (.005,.01,.02,.04,.08)]
 62    tol_obs = [grid_crossing(e, .5, .35, y0) for e in (.005,.01,.02,.04,.08)]
 63    # Strict checks are against the discrete-grid resolution, not an arbitrary win.
 64    max_grid_error = max(r["abs_error"] for r in rows)
 65    monotone_mismatch = all(mismatch_rows[i]["observed"] > mismatch_rows[i+1]["observed"]
 66                             for i in range(len(mismatch_rows)-1))
 67    monotone_tol = all(tol_obs[i] < tol_obs[i+1] for i in range(len(tol_obs)-1))
 68    passed = max_grid_error <= .00201 and monotone_mismatch and monotone_tol
 69    return {"passed": passed, "y95": float(yq), "tolerance_sweep": rows,
 70            "mismatch_sweep": mismatch_rows,
 71            "tolerance_monotonicity": {"predicted": tol_exact, "observed": tol_obs},
 72            "max_grid_abs_error": max_grid_error}
 73
 74
 75def teacher_a(t, difficulty=1.0):
 76    # A slowly changing vector-field coefficient: the local model is theta*y.
 77    return .35 + difficulty * (.30*np.sin(.9*t) + .22*np.sin(2.1*t + .4))
 78
 79
 80def teacher_solution(times, y0, difficulty=1.0):
 81    a = teacher_a(times, difficulty)
 82    integ = np.zeros_like(times)
 83    integ[1:] = np.cumsum(.5*(a[1:]+a[:-1])*np.diff(times))
 84    return y0[:, None] * np.exp(integ[None, :])
 85
 86
 87def fit_local_theta(times, ys, difficulty=1.0):
 88    # Teacher-forced derivative regression for f_theta(y)=theta*y.
 89    a = teacher_a(times, difficulty)
 90    y = ys[:, :-1]
 91    dy = a[:-1][None, :] * y
 92    return float(np.sum(y*dy) / (np.sum(y*y) + 1e-12))
 93
 94
 95def adaptive_windows(times, truth, y0, eps=.035, min_steps=4, cap_steps=35,
 96                     difficulty=1.0):
 97    windows, thetas = [], []
 98    start = 0
 99    n = len(times)-1
100    while start < n:
101        stop_cap = min(n, start + cap_steps)
102        # Warm-start analogue: use teacher derivative data in the candidate cap.
103        theta = fit_local_theta(times[start:stop_cap+1],
104                                truth[:, start:stop_cap+1], difficulty)
105        offsets = np.arange(min_steps, stop_cap-start+1)
106        chosen = stop_cap
107        for off in offsets:
108            s = times[start+off] - times[start]
109            pred = truth[:, start:start+1] * np.exp(theta*s)
110            err = q95(np.abs(truth[:, start+off] - pred[:, 0]))
111            if err > eps:
112                chosen = start + int(off)
113                break
114        # enforce min window even if the first grid point crosses
115        chosen = max(chosen, min(n, start + min_steps))
116        windows.append((start, chosen))
117        thetas.append(theta)
118        start = chosen
119    return windows, thetas
120
121
122def rollout(times, y0, windows, thetas):
123    out = np.empty((len(y0), len(times)))
124    out[:, 0] = y0
125    for (start, stop), theta in zip(windows, thetas):
126        # At deployment the reset is the model's terminal state, not teacher data.
127        for j in range(start+1, stop+1):
128            out[:, j] = out[:, start] * np.exp(theta*(times[j]-times[start]))
129    return out
130
131
132def mini_experiment():
133    y0 = np.linspace(.8, 1.2, 128)
134    times = np.linspace(0, 12, 1201)
135    truth = teacher_solution(times, y0)
136    # Shared field is the standard one-field teacher-forced derivative fit.
137    theta = fit_local_theta(times, truth)
138    shared = y0[:, None] * np.exp(theta*times[None, :])
139    windows, thetas = adaptive_windows(times, truth, y0)
140    adaptive = rollout(times, y0, windows, thetas)
141    shared_rmse = float(np.sqrt(np.mean((shared-truth)**2)))
142    adaptive_rmse = float(np.sqrt(np.mean((adaptive-truth)**2)))
143    shared_final = float(np.sqrt(np.mean((shared[:,-1]-truth[:,-1])**2)))
144    adaptive_final = float(np.sqrt(np.mean((adaptive[:,-1]-truth[:,-1])**2)))
145    boundary_errors = [q95(np.abs(truth[:,s]-
146                         (truth[:,s-1] if False else adaptive[:,s]))) for s,_ in windows[1:]]
147    return {"shared_theta": theta, "windows": len(windows),
148            "window_lengths": [b-a for a,b in windows],
149            "shared_rmse": shared_rmse, "adaptive_rmse": adaptive_rmse,
150            "shared_final_rmse": shared_final, "adaptive_final_rmse": adaptive_final,
151            "adaptive_boundary_count": len(boundary_errors),
152            "adaptive_boundary_terminal_error_q95": boundary_errors[:8]}
153
154
155def main():
156    result = {"seed": SEED, "math_verification": verify_math(),
157              "mini_experiment": mini_experiment()}
158    Path("results.json").write_text(json.dumps(result, indent=2))
159    print(json.dumps(result, indent=2))
160
161if __name__ == "__main__":
162    main()