1"""Toy verification of conformal early rejection for synthetic diffusion/NAS trials.
  2
  3The simulator deliberately keeps final evaluation cheap (a latent table lookup), while
  4counting it as the expensive operation.  The monitor is calibrated only on a split of
  5completed trajectories and is never used as the authoritative final label.
  6"""
  7import json
  8from pathlib import Path
  9import numpy as np
 10
 11SEED = 2762
 12CHECKPOINTS = (0.75, 0.50, 0.25)
 13ALPHAS = (0.05, 0.10)
 14
 15
 16def sigmoid(x):
 17    return 1.0 / (1.0 + np.exp(-np.clip(x, -40, 40)))
 18
 19
 20def draw_trajectories(n, rng, beta=5.0, shift=0.0, noise=0.45):
 21    # x is architecture quality; failures are the hard external evaluation label.
 22    x = rng.normal(loc=shift, scale=1.0, size=n)
 23    failure = x < 0.0
 24    scores = {}
 25    for j, t in enumerate(CHECKPOINTS):
 26        # Later denoising checkpoints have more informative hidden states.
 27        # The same latent x is observed with independent checkpoint noise.
 28        s = sigmoid(beta * (-x + rng.normal(0, noise * (1.15 - 0.15*j), n)))
 29        scores[t] = s
 30    # authoritative accuracy/utility, only available after full evaluation
 31    accuracy = 0.70 + 0.20 * sigmoid(1.6*x) + rng.normal(0, .004, n)
 32    feasible = (~failure) & (accuracy >= 0.80)
 33    return x, failure, feasible, accuracy, scores
 34
 35
 36def conformal_threshold(values, alpha):
 37    # Split-conformal order statistic: k=ceil((n+1)(1-alpha)), one-indexed.
 38    n = len(values)
 39    k = int(np.ceil((n + 1) * (1-alpha)))
 40    k = max(1, min(n, k))
 41    return np.sort(values)[k-1]
 42
 43
 44def thresholds(cal_scores, alpha):
 45    return {t: conformal_threshold(cal_scores[t], alpha) for t in CHECKPOINTS}
 46
 47
 48def run_controller(scores, failure, feasible, accuracy, tau=None, mode='conformal'):
 49    n = len(failure)
 50    if mode == 'none':
 51        rejected = np.zeros(n, dtype=bool)
 52    elif mode == 'uncalibrated':
 53        rejected = np.zeros(n, dtype=bool)
 54        for t in CHECKPOINTS:
 55            rejected |= scores[t] > 0.5
 56    else:
 57        rejected = np.zeros(n, dtype=bool)
 58        for t in CHECKPOINTS:
 59            rejected |= scores[t] > tau[t]
 60    evaluated = ~rejected
 61    accepted_failure = float(failure[evaluated].mean()) if evaluated.any() else 1.0
 62    # best authoritative result among evaluated candidates, with infeasible candidates
 63    # excluded from the utility metric.
 64    good = evaluated & feasible
 65    best = float(accuracy[good].max()) if good.any() else float('nan')
 66    return dict(rejected=int(rejected.sum()), evaluated=int(evaluated.sum()),
 67                reject_rate=float(rejected.mean()), accepted_failure=accepted_failure,
 68                best_accuracy=best)
 69
 70
 71def iid_trial(ncal, ntest, beta, alpha, seed, shift_cal=0., shift_test=0.):
 72    rng = np.random.default_rng(seed)
 73    _, fc, _, _, sc = draw_trajectories(ncal, rng, beta=beta, shift=shift_cal)
 74    _, ft, fe, acc, st = draw_trajectories(ntest, rng, beta=beta, shift=shift_test)
 75    tau = thresholds(sc, alpha)
 76    # marginal calibration property: probability that a fresh score exceeds its
 77    # split-conformal threshold, checked checkpoint-wise.
 78    exceed = {str(t): float((st[t] > tau[t]).mean()) for t in CHECKPOINTS}
 79    return tau, exceed, run_controller(st, ft, fe, acc, tau, 'conformal'), run_controller(st, ft, fe, acc, mode='uncalibrated'), run_controller(st, ft, fe, acc, mode='none')
 80
 81
 82def aggregate(rows):
 83    keys = ['rejected','evaluated','reject_rate','accepted_failure','best_accuracy']
 84    return {k: float(np.nanmean([r[k] for r in rows])) for k in keys}
 85
 86
 87def main():
 88    out = Path('results.json')
 89    # Math sanity check: order-statistic exceedance and its finite-sample prediction.
 90    coverage = []
 91    for ncal in (20, 50, 100, 500):
 92        vals = []
 93        for rep in range(300):
 94            _, ex, *_ = iid_trial(ncal, 2000, 5.0, .10, 10000+ncal*1000+rep)
 95            vals.append(np.mean(list(ex.values())))
 96        observed = float(np.mean(vals))
 97        # For continuous iid scores, P(fresh > kth order statistic) is approximately
 98        # (n+1-k)/(n+1), hence <= alpha with this conservative order statistic.
 99        k = int(np.ceil((ncal+1)*.90))
100        predicted = float((ncal-k+1)/(ncal+1))
101        coverage.append({'ncal': ncal, 'predicted_exceedance': predicted, 'observed_exceedance': observed})
102
103    # Prediction 1/2: stronger score separability gives lower accepted failure and
104    # more useful rejection; prediction 3: exchangeability shift breaks calibration.
105    separability = []
106    for beta in (0.5, 1., 2., 5., 10.):
107        rows = []
108        for rep in range(80):
109            _, _, c, u, no = iid_trial(400, 1000, beta, .10, 20000+rep+int(beta*100))
110            rows.append(c)
111        separability.append({'beta': beta, **aggregate(rows)})
112
113    shift = []
114    for d in (0., -0.25, -0.5, -1.0):
115        rows = []
116        for rep in range(100):
117            _, _, c, u, no = iid_trial(400, 1000, 5., .10, 30000+rep+int(abs(d)*100), shift_test=d)
118            rows.append(c)
119        a = aggregate(rows)
120        shift.append({'test_shift': d, **a})
121
122    # Fixed proposal stream comparison at two alpha values.
123    controllers = {}
124    for alpha in ALPHAS:
125        allc, allu, alln = [], [], []
126        for rep in range(100):
127            _, _, c, u, no = iid_trial(400, 1000, 5., alpha, 40000+rep)
128            allc.append(c); allu.append(u); alln.append(no)
129        controllers[str(alpha)] = {'conformal': aggregate(allc), 'uncalibrated': aggregate(allu), 'none': aggregate(alln)}
130
131    result = {'seed': SEED, 'checkpoints': CHECKPOINTS, 'coverage_sweep': coverage,
132              'separability_sweep': separability, 'distribution_shift_sweep': shift,
133              'controller_comparison': controllers}
134    out.write_text(json.dumps(result, indent=2))
135    print(json.dumps(result, indent=2))
136
137if __name__ == '__main__':
138    main()