1"""MVP for safe receding-horizon topology switching.
  2
  3The toy has correlated features x2=x1+noise.  The deployed model uses both
  4features with weights (1,1), while the target topology drops feature 2 and
  5uses target weight (2,0).  The endpoint is accurate, but immediately after
  6changing the mask, retaining the old active weight gives a large error.
  7"""
  8import json
  9import math
 10from dataclasses import dataclass
 11import numpy as np
 12
 13SEED = 217
 14
 15@dataclass
 16class Toy:
 17    x1: np.ndarray
 18    x2: np.ndarray
 19    y: np.ndarray
 20
 21    @property
 22    def a(self):
 23        # residual after topology switch and interpolation alpha:
 24        # (1+alpha)x1 - (x1+x2) = alpha*x1-x2
 25        return float(np.mean(self.x1 * self.x1))
 26    @property
 27    def b(self):
 28        return float(np.mean(self.x1 * self.x2))
 29    @property
 30    def c(self):
 31        return float(np.mean(self.x2 * self.x2))
 32    def loss_alpha(self, alpha):
 33        return self.a * alpha * alpha - 2*self.b*alpha + self.c
 34    def old_loss(self):
 35        return float(np.mean((self.x1 + self.x2 - self.y)**2))
 36    def target_loss(self):
 37        return self.loss_alpha(1.0)
 38    def predicted_root(self, eps):
 39        # smallest alpha in [0,1] with loss_alpha <= eps
 40        # positive root of a alpha^2 - 2b alpha + c-eps=0
 41        d = self.b*self.b - self.a*(self.c-eps)
 42        if d < 0: return None
 43        roots = [(self.b-math.sqrt(max(0,d)))/self.a,
 44                 (self.b+math.sqrt(max(0,d)))/self.a]
 45        feasible = [r for r in roots if 0 <= r <= 1]
 46        if self.loss_alpha(0) <= eps: return 0.0
 47        return min(feasible) if feasible else float('inf')
 48
 49def make_toy(noise, n=20000):
 50    rng = np.random.default_rng(SEED + int(round(noise*1000)))
 51    x1 = rng.normal(size=n)
 52    x2 = x1 + noise*rng.normal(size=n)
 53    return Toy(x1, x2, x1+x2)
 54
 55def certify(toy, alpha, eps):
 56    return toy.loss_alpha(alpha) - toy.old_loss() <= eps + 1e-12
 57
 58def filtered_switch(toy, eps, candidates):
 59    rejected = []
 60    for alpha in candidates:
 61        if certify(toy, alpha, eps):
 62            return alpha, rejected
 63        rejected.append(alpha)
 64    return None, rejected
 65
 66def run():
 67    # Core math verification: sweep noise and tolerance, compare quadratic root
 68    # to the observed first accepted planner candidate.
 69    rows = []
 70    candidate_grid = np.linspace(0.0, 1.0, 101)
 71    for noise in [0.05, 0.10, 0.20, 0.35]:
 72        toy = make_toy(noise)
 73        for eps in [0.10, 0.25, 0.50]:
 74            root = toy.predicted_root(eps)
 75            observed, rej = filtered_switch(toy, eps, candidate_grid)
 76            rows.append({
 77                'noise': noise, 'epsilon': eps,
 78                'predicted_alpha_min': root,
 79                'observed_alpha_min_grid': observed,
 80                'grid_error': None if observed is None or root is None else observed-root,
 81                'rejected_candidates': len(rej),
 82                'endpoint_loss': toy.target_loss(),
 83                'initial_new_mask_loss': toy.loss_alpha(0),
 84            })
 85
 86    # Prediction 1: residual follows the derived quadratic to numerical precision.
 87    toy = make_toy(0.20)
 88    alphas = np.linspace(0,1,51)
 89    empirical = np.array([toy.loss_alpha(a) for a in alphas])
 90    formula = toy.a*alphas**2 - 2*toy.b*alphas + toy.c
 91    quadratic_max_abs_error = float(np.max(np.abs(empirical-formula)))
 92
 93    # Prediction 2: larger tolerance monotonically lowers the required transition.
 94    roots = [toy.predicted_root(e) for e in [0.10,0.25,0.50]]
 95    tolerance_monotone = bool(all(roots[i] >= roots[i+1] for i in range(2)))
 96
 97    # Baselines at the same candidate schedule. One-shot has an unsafe
 98    # intermediate (alpha=0 under new mask); unfiltered gradual schedule also
 99    # violates until its first safe alpha. The filter executes only safe alpha.
100    eps = 0.25
101    one_shot = make_toy(0.20)
102    gradual = np.linspace(0,1,11)
103    unfiltered_violations = int(sum(not certify(one_shot,a,eps) for a in gradual))
104    accepted, rejected = filtered_switch(one_shot, eps, gradual)
105    filtered_violations = 0 if accepted is not None and certify(one_shot, accepted, eps) else 1
106
107    # Cut reuse: without cuts, each replan retries all unsafe candidates;
108    # with cuts, rejected alpha values are remembered and not evaluated again.
109    unsafe = [float(a) for a in gradual if not certify(one_shot,a,eps)]
110    replans = 5
111    no_cut_evals = replans * len(unsafe)
112    with_cut_evals = len(unsafe)
113
114    out = {
115        'seed': SEED,
116        'quadratic_max_abs_error': quadratic_max_abs_error,
117        'predicted_vs_observed': rows,
118        'prediction_checks': {
119            'quadratic_formula_exact': quadratic_max_abs_error < 1e-12,
120            'tolerance_lowers_required_alpha': tolerance_monotone,
121            'cut_reuse_reduces_rejected_evaluations': with_cut_evals < no_cut_evals,
122        },
123        'comparison': {
124            'epsilon': eps,
125            'one_shot_endpoint_loss': one_shot.target_loss(),
126            'one_shot_intermediate_violations': 1 if not certify(one_shot,0,eps) else 0,
127            'unfiltered_gradual_violations_of_11': unfiltered_violations,
128            'filtered_first_accepted_alpha': accepted,
129            'filtered_rejected_before_accept': len(rejected),
130            'filtered_accepted_step_violations': filtered_violations,
131            'cut_reuse_evaluations_over_5_replans': {'without_cuts': no_cut_evals, 'with_cuts': with_cut_evals},
132        },
133    }
134    print(json.dumps(out, indent=2, sort_keys=True))
135
136if __name__ == '__main__':
137    run()