1"""Toy MVP for certified ambiguity gating.
  2Run: /home/maxwelhelp/main/bin/python3 ambiguity_experiment.py
  3"""
  4import json, math, random
  5from pathlib import Path
  6import numpy as np
  7import torch
  8from torch import nn
  9
 10SEED = 611
 11np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 12try:
 13    torch.set_num_threads(4)
 14except Exception:
 15    pass
 16
 17
 18def audit(probs, tau=0.30, delta=0.05):
 19    probs = np.asarray(probs)
 20    admissible = probs >= tau
 21    amb = admissible.sum(axis=1) >= 2
 22    dhat = float(amb.mean()); n = len(amb)
 23    eps = math.sqrt(math.log(1/delta)/(2*n))
 24    floor_lcb = max(0., dhat-eps)/2
 25    return {"Dhat": dhat, "epsilon": eps, "floor_LCB": floor_lcb,
 26            "ambiguous": amb, "admissible": admissible}
 27
 28
 29def verify_minimax_floor():
 30    # Finite population: 40% have two admissible labels, others have one.
 31    rng = np.random.default_rng(SEED)
 32    n = 2000
 33    amb = rng.random(n) < .40
 34    # All deterministic predictors are represented by their predictions;
 35    # adversary chooses the opposite admissible target pointwise.
 36    # For a randomized q=P(pred=1), worst target risk is max(q,1-q).
 37    q = rng.random(n)
 38    admissible_risk = np.where(amb, np.maximum(q, 1-q),
 39                              np.where(q >= .5, 0., 1.))
 40    # On unique points q is set to the unique admissible label (0 here).
 41    q_opt = np.where(amb, .5, 0.)
 42    opt_risk = np.where(amb, .5, 0.).mean()
 43    empirical_D = amb.mean()
 44    return {"D": float(empirical_D), "minimax_risk_fair": float(opt_risk),
 45            "D_over_2": float(empirical_D/2),
 46            "random_predictor_worst_risk": float(admissible_risk.mean())}
 47
 48
 49class MLP(nn.Module):
 50    def __init__(self):
 51        super().__init__()
 52        self.net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))
 53    def forward(self, x): return self.net(x)
 54
 55
 56def make_data(n, rng):
 57    x = rng.uniform(-1, 1, size=(n, 2)).astype(np.float32)
 58    # Ground-truth boundary is nonlinear but learnable by the small MLP.
 59    score = x[:, 0] + .55*np.sin(3*x[:, 1])
 60    y = (score > 0).astype(np.int64)
 61    # Supervisor is confident away from boundary and ambiguous near it.
 62    amb = np.abs(score) < .38
 63    p1 = np.where(amb, .50 + .03*np.sign(score),
 64                  np.where(score > 0, .92, .08))
 65    # Independent generated hard label, as with stochastic LLM decoding.
 66    yl = (rng.random(n) < p1).astype(np.int64)
 67    probs = np.column_stack([1-p1, p1]).astype(np.float32)
 68    return x, y, yl, probs, amb
 69
 70
 71def train(x, yl, probs, keep, soft=False, epochs=80):
 72    torch.manual_seed(SEED + (3 if soft else 0) + int(keep.sum()))
 73    model = MLP(); opt = torch.optim.Adam(model.parameters(), lr=.012)
 74    xt = torch.from_numpy(x); yt = torch.from_numpy(yl)
 75    pt = torch.from_numpy(probs)
 76    idx = np.flatnonzero(keep)
 77    # Full-batch makes equal-step comparison deterministic; equal retained data.
 78    for _ in range(epochs):
 79        opt.zero_grad(); logits = model(xt[idx])
 80        if soft:
 81            loss = -(pt[idx] * torch.log_softmax(logits, 1)).sum(1).mean()
 82        else:
 83            loss = nn.functional.cross_entropy(logits, yt[idx])
 84        loss.backward(); opt.step()
 85    return model
 86
 87
 88def metrics(model, x, y):
 89    with torch.no_grad():
 90        z = model(torch.from_numpy(x)); p = torch.softmax(z, 1)[:,1].numpy()
 91    pred = (p >= .5).astype(np.int64)
 92    acc = float((pred == y).mean())
 93    # Expected calibration error, 10 fixed bins.
 94    ece = 0.
 95    conf = np.maximum(p, 1-p); corr = (pred == y).astype(float)
 96    for lo, hi in zip(np.linspace(0,1,11)[:-1], np.linspace(0,1,11)[1:]):
 97        m = (conf >= lo) & (conf < hi if hi < 1 else conf <= hi)
 98        if m.any(): ece += m.mean()*abs(corr[m].mean()-conf[m].mean())
 99    return acc, float(ece)
100
101
102def main():
103    rng = np.random.default_rng(SEED)
104    xtr, ytr, yl, probs, latent_amb = make_data(1800, rng)
105    xte, yte, _, _, _ = make_data(5000, rng)
106    au = audit(probs, tau=.30, delta=.05)
107    amb = au["ambiguous"]
108    # Gate keeps certified-unambiguous examples. Confidence control retains
109    # exactly the same number, using largest max-label probability.
110    gate_keep = ~amb
111    confidence = probs.max(1)
112    conf_keep = np.zeros(len(confidence), dtype=bool)
113    conf_keep[np.argsort(confidence)[-gate_keep.sum():]] = True
114    hard_all = np.ones(len(yl), dtype=bool)
115    models = {
116      "hard_all": train(xtr, yl, probs, hard_all, False),
117      "ambiguity_gate": train(xtr, yl, probs, gate_keep, False),
118      "confidence_matched": train(xtr, yl, probs, conf_keep, False),
119      "soft_all": train(xtr, yl, probs, hard_all, True),
120    }
121    results = {k: {"accuracy": metrics(v,xte,yte)[0], "ECE": metrics(v,xte,yte)[1]}
122               for k,v in models.items()}
123    out = {"seed": SEED, "tau": .30, "n_train": len(xtr),
124           "audit": {k:(float(v) if np.isscalar(v) else None) for k,v in au.items() if k not in ("ambiguous","admissible")},
125           "retained_gate": int(gate_keep.sum()), "retention": float(gate_keep.mean()),
126           "minimax_verification": verify_minimax_floor(), "results": results}
127    print(json.dumps(out, indent=2, sort_keys=True))
128    Path("results.json").write_text(json.dumps(out, indent=2, sort_keys=True))
129
130if __name__ == '__main__': main()