1import json
 2import numpy as np
 3from trajectory_tube import TrajectoryTube, conformal_quantile
 4
 5SEED = 2834
 6rng = np.random.default_rng(SEED)
 7alpha = 0.1
 8alpha_L = 0.1
 9T = np.linspace(0.0, 1.0, 1001)
10obs = np.array([0.1, 0.3, 0.5, 0.75])
11
12# A continuous piecewise-linear trajectory makes the Lipschitz proof exact.
13def ramp(t, slope):
14    return slope * np.maximum(np.asarray(t) - 0.75, 0.0)
15
16# Calibration residuals are zero because the predictor is zero and observations
17# stop at the kink. Dense calibration gives the separate slope quantile.
18cal_slopes = np.linspace(0.4, 0.95, 20)
19cal_x = [ramp(T, a) for a in cal_slopes]
20cal_obs = [ramp(obs, a) for a in cal_slopes]
21q = conformal_quantile([np.max(np.abs(x)) for x in cal_obs], alpha)
22Lhat = conformal_quantile(cal_slopes, alpha_L)
23
24# Prediction 1: radius must be affine in nearest-observation distance.
25gamma = Lhat
26r, delta = TrajectoryTube(q, gamma).radius(T, obs)
27fit_slope, fit_intercept = np.polyfit(delta, r, 1)
28linear_max_error = np.max(np.abs(r - (q + gamma * delta)))
29
30# Prediction 2: conditional transition occurs at multiplier one for a truth
31# whose slope equals Lhat. A deliberately sublinear multiplier must fail in
32# the largest gap (at the first post-kink point).
33test_boundary = ramp(T, Lhat)
34mults = np.array([0.0, 0.25, 0.5, 0.75, 0.9, 1.0, 1.1, 1.25, 1.5])
35cover_boundary = []
36for m in mults:
37    tube = TrajectoryTube(q, m * Lhat)
38    cover_boundary.append(np.mean(tube.contains(test_boundary, np.zeros_like(T), T, obs)))
39cover_boundary = np.asarray(cover_boundary)
40full = np.where(cover_boundary >= 1 - 1e-12)[0]
41first_full = float(mults[full[0]]) if len(full) else None
42sub_m = 0.75
43inside_sub = TrajectoryTube(q, sub_m * Lhat).contains(test_boundary, np.zeros_like(T), T, obs)
44first_failure = float(T[np.where(~inside_sub)[0][0]]) if np.any(~inside_sub) else None
45
46# Prediction 3: marginal coverage is controlled by the slope event. Test slopes
47# are iid from the same range as calibration; at multiplier 1 the expected
48# slope event is 1-alpha_L, while a safety multiplier covers all this range.
49test_slopes = rng.uniform(0.4, 0.95, 1000)
50cover_m1 = []
51cover_m11 = []
52for a in test_slopes:
53    truth = ramp(T, a)
54    cover_m1.append(np.mean(TrajectoryTube(q, Lhat).contains(truth, np.zeros_like(T), T, obs)))
55    cover_m11.append(np.mean(TrajectoryTube(q, 1.1 * Lhat).contains(truth, np.zeros_like(T), T, obs)))
56marginal_m1 = float(np.mean(cover_m1))
57marginal_m11 = float(np.mean(cover_m11))
58slope_event = float(np.mean(test_slopes <= Lhat))
59
60# Direct numerical triangle-inequality check for random linear truth/predictor.
61violations = []
62for _ in range(100):
63    tt = np.linspace(0, 1, 101)
64    lt, lp = rng.uniform(0.1, 2.0), rng.uniform(0.1, 1.0)
65    truth, pred = lt * tt, lp * tt
66    idx = np.arange(0, 101, 20)
67    d = np.min(np.abs(tt[:, None] - tt[idx][None, :]), axis=1)
68    q2 = np.max(np.abs(truth[idx] - pred[idx]))
69    violations.append(np.max(np.abs(truth - pred) - (q2 + (lt + lp) * d)))
70
71# Standard baseline: no between-observation inflation.
72baseline = float(np.mean(np.abs(test_boundary) <= q + 1e-12))
73out = {
74    "seed": SEED, "alpha": alpha, "alpha_L": alpha_L, "q": q, "Lhat": Lhat,
75    "predictions": {
76        "radius_slope_predicted": float(gamma),
77        "radius_slope_observed": float(fit_slope),
78        "radius_intercept_predicted": float(q),
79        "radius_intercept_observed": float(fit_intercept),
80        "max_affine_error": float(linear_max_error),
81        "conditional_transition_predicted_multiplier": 1.0,
82        "conditional_transition_observed_first_full_multiplier": first_full,
83        "sublinear_multiplier": sub_m,
84        "sublinear_first_failure_time": first_failure,
85        "slope_event_predicted": 1 - alpha_L,
86        "slope_event_observed": slope_event,
87        "marginal_coverage_at_multiplier_1": marginal_m1,
88        "marginal_coverage_at_multiplier_1.1": marginal_m11,
89        "triangle_max_violation": float(max(violations))
90    },
91    "sweep_boundary": [{"multiplier": float(m), "coverage": float(c)} for m, c in zip(mults, cover_boundary)],
92    "comparison": {"baseline_constant_q_coverage": baseline,
93                   "idea_boundary_coverage_at_Lhat": float(cover_boundary[mults == 1][0]),
94                   "idea_marginal_coverage_at_Lhat": marginal_m1}
95}
96print(json.dumps(out, indent=2))