1import json
  2import math
  3import random
  4from dataclasses import dataclass
  5import numpy as np
  6
  7SEED = 2753
  8np.random.seed(SEED)
  9random.seed(SEED)
 10
 11@dataclass
 12class AdaptiveDisturbance:
 13    c: np.ndarray
 14    q: np.ndarray
 15    alpha: float = 0.25
 16    beta: float = 0.10
 17    eps: float = 0.005
 18
 19    def update(self, residuals):
 20        residuals = np.asarray(residuals, dtype=float).reshape(-1, self.c.size)
 21        med = np.median(residuals, axis=0)
 22        # The specified update uses the newly updated center in the radius.
 23        self.c = (1.0 - self.alpha) * self.c + self.alpha * med
 24        observed = np.max(np.abs(residuals - self.c), axis=0) + self.eps
 25        self.q = np.maximum((1.0 - self.beta) * self.q, observed)
 26        return self.c.copy(), self.q.copy()
 27
 28
 29def propagate_zonotope(A, B, x_center, x_generators, u, d, w_center, w_generators):
 30    """Exact affine zonotope propagation in generator form."""
 31    c = A @ x_center + B @ u + d + w_center
 32    G = np.concatenate((A @ x_generators, w_generators), axis=1)
 33    return c, G
 34
 35
 36def box_contains(c, G, lower, upper):
 37    radius = np.sum(np.abs(G), axis=1)
 38    return bool(np.all(c - radius >= lower - 1e-12) and np.all(c + radius <= upper + 1e-12))
 39
 40
 41def interval_prediction_check():
 42    # For an affine zonotope and axis-aligned box, support in coordinate j is
 43    # c_j +/- sum_k |G_jk|; this is both an over-approximation and exact.
 44    rng = np.random.default_rng(SEED)
 45    errors = []
 46    disagreements = 0
 47    for _ in range(3000):
 48        n, p = 3, 7
 49        c = rng.normal(size=n)
 50        G = rng.normal(size=(n, p))
 51        lo = c - np.sum(np.abs(G), axis=1) - rng.uniform(0, 1, n)
 52        hi = c + np.sum(np.abs(G), axis=1) + rng.uniform(0, 1, n)
 53        predicted = np.maximum(np.abs(G), 0).sum(axis=1)
 54        # Enumerating vertices is exact for each coordinate support; sample
 55        # random signs as an independent numerical check.
 56        signs = rng.choice([-1.0, 1.0], size=(1000, p))
 57        sampled = np.max(c[None, :] + signs @ G.T, axis=0)
 58        sampled_min = np.min(c[None, :] + signs @ G.T, axis=0)
 59        errors.append(float(np.max(np.maximum(sampled - (c + predicted),
 60                                              (c - predicted) - sampled_min))))
 61        a = box_contains(c, G, lo, hi)
 62        b = bool(np.all(c - predicted >= lo - 1e-12) and np.all(c + predicted <= hi + 1e-12))
 63        disagreements += int(a != b)
 64    return {"max_sample_support_gap": max(errors), "containment_disagreements": disagreements,
 65            "cases": 3000}
 66
 67
 68def stationary_scaling_sweep():
 69    eps = 0.005
 70    amp_list = [0.02, 0.05, 0.10, 0.20, 0.40]
 71    rows = []
 72    for amp in amp_list:
 73        est = AdaptiveDisturbance(np.zeros(1), np.zeros(1), alpha=0.25, beta=0.10, eps=eps)
 74        rng = np.random.default_rng(SEED + int(1000 * amp))
 75        # Batch maxima make the intended q update visible without assuming
 76        # an unobserved true bound.
 77        for _ in range(40):
 78            residuals = rng.uniform(-amp, amp, size=(64, 1))
 79            est.update(residuals)
 80        heldout = rng.uniform(-amp, amp, size=(10000, 1))
 81        covered = np.abs(heldout[:, 0] - est.c[0]) <= est.q[0]
 82        predicted = amp + eps
 83        rows.append({"amplitude": amp, "predicted_q": predicted,
 84                     "observed_q": float(est.q[0]),
 85                     "q_error": float(est.q[0] - predicted),
 86                     "coverage": float(np.mean(covered))})
 87    max_error = max(abs(r["q_error"]) for r in rows)
 88    # With finite batches q can be below amp+eps, but should remain close and
 89    # coverage should be high; report rather than silently changing the rule.
 90    return {"rows": rows, "max_abs_q_prediction_error": max_error,
 91            "minimum_coverage": min(r["coverage"] for r in rows)}
 92
 93
 94def jump_recovery_check():
 95    eps = 0.005
 96    est = AdaptiveDisturbance(np.zeros(1), np.zeros(1), alpha=0.25, beta=0.10, eps=eps)
 97    rng = np.random.default_rng(SEED + 99)
 98    for _ in range(12):
 99        est.update(rng.uniform(-0.10, 0.10, size=(64, 1)))
100    before = float(est.q[0])
101    new_amp = 0.30
102    residuals = rng.uniform(-new_amp, new_amp, size=(64, 1))
103    est.update(residuals)
104    after = float(est.q[0])
105    # The max rule predicts expansion in the first batch containing the jump.
106    coverage = float(np.mean(np.abs(residuals[:, 0] - est.c[0]) <= est.q[0]))
107    return {"old_amplitude": 0.10, "new_amplitude": new_amp,
108            "q_before_jump": before, "q_after_one_batch": after,
109            "predicted_min_after_batch": new_amp + eps - 0.03,
110            "jump_batch_coverage": coverage,
111            "expanded_in_one_batch": after > before}
112
113
114def point_mass_comparison(episodes=250, horizon=40):
115    """Tiny shield comparison; rejected proposals use zero acceleration backup."""
116    rng = np.random.default_rng(SEED + 7)
117    dt = 1.0
118    pos_limit, vel_limit, u_limit = 1.0, 1.5, 0.45
119    true_amp = 0.12
120    fixed_q = np.array([0.0, true_amp + 0.005])
121    stats = {"fixed": [0, 0], "adaptive": [0, 0]} # rejects, violations
122    for method in stats:
123        for ep in range(episodes):
124            x = np.array([0.0, 0.0])
125            est = AdaptiveDisturbance(np.zeros(2), np.array([0.0, 0.35]), alpha=.25, beta=.10, eps=.005)
126            history = []
127            for t in range(horizon):
128                # Mildly aggressive policy, making the shield relevant.
129                proposed = np.clip(-0.55*x[0] - 0.25*x[1] + rng.normal(0, .08), -u_limit, u_limit)
130                A = np.array([[1., dt], [0., 1.]])
131                B = np.array([[0.], [dt]])
132                # Current point state, disturbance zonotope in acceleration.
133                q = fixed_q if method == "fixed" else est.q
134                wc = np.zeros(2) if method == "fixed" else est.c
135                W = np.diag(np.asarray(q, dtype=float).reshape(-1))
136                c, G = propagate_zonotope(A, B, x, np.zeros((2, 0)), np.array([proposed]),
137                                           np.zeros(2), wc, W)
138                safe = box_contains(c, G, np.array([-pos_limit, -vel_limit]),
139                                    np.array([pos_limit, vel_limit]))
140                u = proposed if safe else 0.0
141                stats[method][0] += int(not safe)
142                w = np.array([0., rng.uniform(-true_amp, true_amp)])
143                xn = (A @ np.asarray(x).reshape(2) + B[:, 0] * float(u) + w).reshape(2)
144                stats[method][1] += int(abs(xn[0]) > pos_limit or abs(xn[1]) > vel_limit)
145                if method == "adaptive":
146                    residual = xn - (A @ x + B * u)
147                    history.append(residual)
148                    if len(history) >= 8:
149                        est.update(np.array(history[-8:]))
150                x = xn
151    denom = episodes * horizon
152    return {m: {"rejection_rate": v[0] / denom, "violation_rate": v[1] / denom}
153            for m, v in stats.items()}
154
155
156def main():
157    result = {"seed": SEED,
158              "math_check": interval_prediction_check(),
159              "stationary_scaling": stationary_scaling_sweep(),
160              "jump_recovery": jump_recovery_check(),
161              "shield_comparison": point_mass_comparison()}
162    print(json.dumps(result, indent=2))
163
164if __name__ == "__main__":
165    main()