1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 535
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7try:
  8    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  9except Exception:
 10    device = torch.device('cpu')
 11
 12# Scalar box-constrained output: y in [-1, 1]. The feature map is frozen;
 13# only the affine output head is trained, so the nominal objective is convex.
 14def make_data(n, noise, rng):
 15    x = rng.uniform(-1.0, 1.0, size=(n, 1)).astype(np.float32)
 16    y_clean = (0.82*x[:, 0] + 0.03*np.sin(3*x[:, 0])).astype(np.float32)
 17    y_obs = y_clean + rng.normal(0, noise, size=n).astype(np.float32)
 18    return x, y_obs, y_clean
 19
 20def fit_nominal(x, y, steps=1000, lr=0.04):
 21    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
 22    Y = torch.tensor(y[:, None], device=device)
 23    w = torch.zeros(2, 1, device=device, requires_grad=True)
 24    opt = torch.optim.Adam([w], lr=lr)
 25    for _ in range(steps):
 26        opt.zero_grad(); pred = X @ w
 27        loss = ((pred-Y)**2).mean() + 1e-4*(w[0]**2)
 28        loss.backward(); opt.step()
 29    return w.detach()
 30
 31def fit_scenario(x, y, residuals, mu=3.0, scen=48, steps=1200, lr=0.025):
 32    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
 33    Y = torch.tensor(y[:, None], device=device)
 34    R = torch.tensor(residuals[:, None], device=device)
 35    w = torch.zeros(2, 1, device=device, requires_grad=True)
 36    opt = torch.optim.Adam([w], lr=lr)
 37    g = torch.Generator(device=device); g.manual_seed(SEED+7)
 38    for _ in range(steps):
 39        opt.zero_grad(); pred = X @ w
 40        # Bootstrap residual scenarios independently for each minibatch item.
 41        idx = torch.randint(R.shape[0], (scen, X.shape[0]), generator=g, device=device)
 42        s = R[idx].squeeze(-1).T  # [batch, scenarios]
 43        ys = pred + s
 44        h = torch.relu(-1.0-ys) + torch.relu(ys-1.0)
 45        loss = ((pred-Y)**2).mean() + mu*h.mean() + 1e-4*(w[0]**2)
 46        loss.backward(); opt.step()
 47    return w.detach()
 48
 49def predict(w, x):
 50    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
 51    return (X @ w).detach().cpu().numpy()[:, 0]
 52
 53def violation_metrics(pred, clean, noise, rng):
 54    # Independent held-out disturbance distribution, with repeated draws to
 55    # estimate rare violations without conflating them with training labels.
 56    draws = clean[:, None] + rng.normal(0, noise, size=(len(clean), 80))
 57    # Predictor is nominal; each actual output is prediction error scenario.
 58    errors = draws - clean[:, None]
 59    actual = pred[:, None] + errors
 60    viol = np.maximum(0, -1-actual) + np.maximum(0, actual-1)
 61    return {
 62        'nominal_mse_to_clean': float(np.mean((pred-clean)**2)),
 63        'p95_abs_error': float(np.quantile(np.abs(pred[:, None]-draws), .95)),
 64        'trajectory_sample_violation_rate': float(np.mean(viol > 0)),
 65        'mean_slack': float(np.mean(viol)),
 66        'p99_slack': float(np.quantile(viol, .99)),
 67    }
 68
 69def exact_slack_check():
 70    # For each scenario, min_{h>=0} mu*h subject to y+s <= 1+h and
 71    # -1-h <= y+s is attained at h=max(0, |y+s|-1).
 72    vals = np.array([-1.4, -1.0, -0.2, 0.7, 1.0, 1.25])
 73    mu = 2.7
 74    eliminated = mu*np.maximum(0, np.abs(vals)-1)
 75    # brute force grid minimization of the defining one-dimensional problem
 76    brute = []
 77    for v in vals:
 78        hs = np.linspace(0, 1.5, 30001)
 79        feasible = (v <= 1+hs) & (v >= -1-hs)
 80        brute.append(np.min(np.where(feasible, mu*hs, np.inf)))
 81    brute = np.array(brute)
 82    return {'max_abs_error': float(np.max(np.abs(brute-eliminated))),
 83            'slacks': eliminated.tolist(), 'passed': bool(np.max(np.abs(brute-eliminated)) < 1e-4)}
 84
 85def main():
 86    rng = np.random.default_rng(SEED)
 87    xtr, ytr, cleantr = make_data(700, .10, rng)
 88    xte, yte, cleante = make_data(5000, .10, rng)
 89    # Residual buffer from observed prediction errors of a nominal fitted head.
 90    w0 = fit_nominal(xtr, ytr)
 91    residual_buffer = ytr - predict(w0, xtr)
 92    wb = fit_nominal(xtr, ytr)
 93    ws = fit_scenario(xtr, ytr, residual_buffer)
 94    result = {
 95      'device': str(device), 'exact_slack_check': exact_slack_check(),
 96      'baseline': violation_metrics(predict(wb,xte), cleante, .10, np.random.default_rng(SEED+2)),
 97      'residual_scenario': violation_metrics(predict(ws,xte), cleante, .10, np.random.default_rng(SEED+2)),
 98      'weights': {'baseline': wb.cpu().numpy().ravel().tolist(), 'scenario': ws.cpu().numpy().ravel().tolist()},
 99      'buffer': {'n': int(len(residual_buffer)), 'mean': float(residual_buffer.mean()), 'std': float(residual_buffer.std())}
100    }
101    print(json.dumps(result, indent=2))
102
103if __name__ == '__main__': main()