1import json
  2import numpy as np
  3from scipy.special import ndtr
  4from reachable_risk import reachable_risk, fixed_false_warning_recall
  5
  6SEED = 1449
  7rng = np.random.default_rng(SEED)
  8
  9
 10def mechanism_sweeps():
 11    # Prediction 1: q=Phi(-m), hence q(0)=.5 and q falls monotonically with m.
 12    margins = np.array([-3., -2., -1., 0., 1., 2., 3.])
 13    q = ndtr(-margins)
 14    monotone = bool(np.all(np.diff(q) < 0))
 15    q_at_zero = float(q[margins == 0][0])
 16
 17    # Prediction 2: scalar covariance follows s_h=lambda^2 s_(h-1)+Q.
 18    # For lambda<1 it saturates at Q/(1-lambda^2); lambda>1 grows exponentially.
 19    Q, s0, H = .04, .01, 20
 20    growth = {}
 21    for lam in [.5, .9, 1.0, 1.1, 1.4]:
 22        s = s0
 23        vals = []
 24        for _ in range(H):
 25            s = lam * lam * s + Q
 26            vals.append(s)
 27        growth[str(lam)] = {
 28            "s_h": vals,
 29            "ratio_last_first": float(vals[-1] / vals[0]),
 30            "stable_limit_if_applicable": (Q / (1-lam*lam)) if abs(lam) < 1 else None,
 31        }
 32    stable_sat_error = abs(growth["0.5"]["s_h"][-1] - Q/(1-.5**2))
 33    unstable_ratio = growth["1.4"]["ratio_last_first"]
 34
 35    # Prediction 3: for a fixed positive per-step q, union risk is 1-(1-q)^H.
 36    q0 = .1
 37    horizons = np.arange(1, 11)
 38    union = 1 - (1-q0) ** horizons
 39    exact_h5 = float(union[4])
 40
 41    return {
 42        "margin_sweep": {"margins": margins.tolist(), "q": q.tolist(),
 43                          "monotone_decreasing": monotone, "q_at_m_zero": q_at_zero,
 44                          "predicted_q_at_m_zero": .5},
 45        "covariance_sweep": growth,
 46        "covariance_predictions": {
 47            "stable_lambda_0.5_limit": Q/(1-.5**2),
 48            "observed_lambda_0.5_final": growth["0.5"]["s_h"][-1],
 49            "stable_absolute_error": float(stable_sat_error),
 50            "unstable_lambda_1.4_ratio_last_first": unstable_ratio,
 51            "predicted_unstable_ratio_gt_1": True,
 52        },
 53        "horizon_sweep": {"q_per_step": q0, "horizons": horizons.tolist(),
 54                          "union_risk": union.tolist(), "risk_at_h5": exact_h5,
 55                          "predicted_h5": 1-(1-q0)**5},
 56    }
 57
 58
 59def rollout_trial(x0, H, noise_sd, threshold):
 60    # state=(position, velocity), x_{t+1}=(position+velocity, velocity)+epsilon
 61    x = np.array(x0, dtype=float)
 62    states = []
 63    for _ in range(H):
 64        x = np.array([x[0] + x[1], x[1]]) + rng.normal(0, noise_sd, 2)
 65        states.append(x.copy())
 66    return np.asarray(states), bool(np.any(np.asarray(states)[:, 0] >= threshold))
 67
 68
 69def mini_experiment(n=2500, H=8):
 70    # Every episode has a fresh initial state; warning is generated before rollout.
 71    J = np.array([[1., 1.], [0., 1.]])
 72    Q = np.diag([.025**2, .012**2])
 73    init_cov = np.diag([.03**2, .02**2])
 74    halfspaces = [(np.array([1., 0.]), 1.0)]
 75    noise_sd = np.array([.025, .012])
 76    risk_scores, mean_scores, labels = [], [], []
 77    for _ in range(n):
 78        x0 = np.array([rng.uniform(.05, .85), rng.uniform(.01, .13)])
 79        controls = [None] * H
 80        def dyn(mu, u):
 81            return J @ mu, J
 82        risk = reachable_risk(x0, init_cov, dyn, controls, halfspaces, Q)
 83        # Fair deterministic baseline: mean-only endpoint warning probability
 84        # with a fixed noise scale, without reachable-set covariance growth.
 85        mean_future = np.array([x0[0] + H*x0[1], x0[1]])
 86        mean_scores.append(float(ndtr((mean_future[0] - 1.0) / 0.05)))
 87        risk_scores.append(float(risk))
 88        _, violated = rollout_trial(x0, H, noise_sd, 1.0)
 89        labels.append(violated)
 90    labels = np.asarray(labels)
 91    risk_scores, mean_scores = np.asarray(risk_scores), np.asarray(mean_scores)
 92    rr, rt = fixed_false_warning_recall(risk_scores, labels, .05)
 93    mr, mt = fixed_false_warning_recall(mean_scores, labels, .05)
 94    # 10 equal-count bins give a simple empirical expected calibration error.
 95    def ece(scores):
 96        order = np.argsort(scores); chunks = np.array_split(order, 10); out = 0.
 97        for ix in chunks:
 98            out += len(ix)/len(scores) * abs(np.mean(labels[ix])-np.mean(scores[ix]))
 99        return float(out)
100    return {"n": n, "horizon": H, "violation_rate": float(np.mean(labels)),
101            "reachable_risk_recall_at_5pct_fpr": rr,
102            "mean_score_recall_at_5pct_fpr": mr,
103            "reachable_threshold": rt, "mean_threshold": mt,
104            "reachable_risk_ece": ece(risk_scores), "mean_score_ece": ece(mean_scores),
105            "baseline_definition": "fixed-sigma endpoint Gaussian warning probability (sigma=0.05)"}
106
107
108if __name__ == "__main__":
109    result = {"seed": SEED, "mechanism": mechanism_sweeps(), "mini_experiment": mini_experiment()}
110    with open("results.json", "w") as f:
111        json.dump(result, f, indent=2)
112    print(json.dumps(result, indent=2))