1import json
  2import numpy as np
  3from scipy.optimize import minimize
  4
  5SEED = 1071
  6K = 3
  7EPS = 1e-9
  8
  9
 10def softmax(z):
 11    z = z - z.max(axis=-1, keepdims=True)
 12    e = np.exp(z)
 13    return e / e.sum(axis=-1, keepdims=True)
 14
 15
 16def pushforward(response_prob, phi, n_states):
 17    out = np.zeros(n_states)
 18    for r, j in enumerate(phi):
 19        out[j] += response_prob[r]
 20    return out
 21
 22
 23def tv(p, q):
 24    return 0.5 * np.abs(p - q).sum(axis=-1)
 25
 26
 27def make_data(n, rng):
 28    # A range of exact Bayesian posteriors, represented by logits.
 29    z = rng.normal(0, 1.25, size=(n, K))
 30    q = softmax(z)
 31    y = np.array([rng.choice(K, p=row) for row in q])
 32    return q, y
 33
 34
 35def raw_from_q(q, gamma):
 36    return q ** gamma / (q ** gamma).sum(axis=1, keepdims=True)
 37
 38
 39def fit_calibrator(p, y, ridge=1e-3):
 40    # log pi = a_j log(p_j) + b_j, fitted by held-out-style train labels.
 41    lp = np.log(np.clip(p, EPS, 1))
 42    def unpack(x):
 43        return x[:K], x[K:]
 44    def objective(x):
 45        a, b = unpack(x)
 46        logits = lp * a[None, :] + b[None, :]
 47        logz = np.logaddexp.reduce(logits, axis=1)
 48        nll = -(logits[np.arange(len(y)), y] - logz).mean()
 49        return nll + ridge * (a @ a + b @ b)
 50    res = minimize(objective, np.r_[np.ones(K), np.zeros(K)], method='BFGS',
 51                   options={'maxiter': 500, 'gtol': 1e-8})
 52    a, b = unpack(res.x)
 53    return a, b, bool(res.success)
 54
 55
 56def apply_calibrator(p, a, b):
 57    return softmax(np.log(np.clip(p, EPS, 1)) * a[None, :] + b[None, :])
 58
 59
 60def metrics(pred, q, y):
 61    n = len(y)
 62    nll = -np.log(np.clip(pred[np.arange(n), y], EPS, 1)).mean()
 63    brier = ((pred - np.eye(K)[y]) ** 2).sum(axis=1).mean()
 64    tv_mean = tv(pred, q).mean()
 65    conf = pred.max(axis=1)
 66    acc = (pred.argmax(axis=1) == y).astype(float)
 67    bins = np.linspace(0, 1, 11)
 68    ece = 0.
 69    for lo, hi in zip(bins[:-1], bins[1:]):
 70        ix = (conf >= lo) & ((conf < hi) if hi < 1 else (conf <= hi))
 71        if ix.any():
 72            ece += ix.mean() * abs(acc[ix].mean() - conf[ix].mean())
 73    # smallest top-probability set with nominal 90% mass, empirical coverage
 74    order = np.argsort(-pred, axis=1)
 75    mass = np.take_along_axis(pred, order, axis=1).cumsum(1)
 76    size = (mass < .9).sum(1) + 1
 77    covered = np.array([y[i] in order[i, :size[i]] for i in range(n)]).mean()
 78    return {'nll': float(nll), 'brier': float(brier), 'tv_to_exact_q': float(tv_mean),
 79            'ece': float(ece), 'set90_coverage': float(covered),
 80            'mean_set_size': float(size.mean())}
 81
 82
 83def run_gamma(gamma, rng):
 84    qtr, ytr = make_data(5000, rng)
 85    qte, yte = make_data(5000, rng)
 86    ptr, pte = raw_from_q(qtr, gamma), raw_from_q(qte, gamma)
 87    a, b, ok = fit_calibrator(ptr, ytr)
 88    cal = apply_calibrator(pte, a, b)
 89    return {
 90        'gamma': gamma,
 91        'predicted_a': 1.0 / gamma,
 92        'fitted_a_mean': float(a.mean()),
 93        'gamma_times_fitted_a': float(gamma * a.mean()),
 94        'raw': metrics(pte, qte, yte),
 95        'calibrated': metrics(cal, qte, yte),
 96        'optimizer_success': ok,
 97    }
 98
 99
100def main():
101    rng = np.random.default_rng(SEED)
102    # Core pushforward sanity check: response probability is conserved by phi.
103    # Each state has several synonymous responses; split mass arbitrarily.
104    phi = np.array([0, 0, 0, 1, 1, 1, 2, 2, 3])  # state 3 = abstain/malformed
105    response = rng.dirichlet(np.ones(len(phi)))
106    pushed = pushforward(response, phi, 4)
107    direct = np.zeros(4)
108    for j in range(4): direct[j] = response[phi == j].sum()
109    pushforward_error = float(np.max(np.abs(pushed - direct)))
110    # Invariance test: redistribute each state's mass among additional synonyms.
111    state_mass = pushed.copy()
112    split = np.concatenate([rng.dirichlet(np.ones(5)) * state_mass[j] for j in range(4)])
113    phi_split = np.repeat(np.arange(4), 5)
114    invariance_error = float(np.max(np.abs(pushforward(split, phi_split, 4) - state_mass)))
115
116    rows = [run_gamma(g, np.random.default_rng(SEED + int(g * 100)))
117            for g in [0.5, 0.75, 1.0, 1.5, 2.0]]
118    result = {
119        'seed': SEED,
120        'predictions': {
121            'P1_pushforward_exact': 'aggregation error is 0 up to floating point, and synonym splitting leaves state mass unchanged',
122            'P2_power_distortion': 'raw log-odds scale with gamma; fitted calibration slope a should satisfy gamma*a ~= 1',
123            'P3_calibration_effect': 'calibrated TV-to-exact-posterior should be near 0 for every gamma, while raw TV is 0 at gamma=1 and grows away from 1'
124        },
125        'sanity': {'pushforward_max_abs_error': pushforward_error,
126                   'synonym_split_max_abs_error': invariance_error,
127                   'response_mass': float(response.sum()),
128                   'pushed_mass': float(pushed.sum())},
129        'sweep': rows,
130        'comparison_gamma_2': next(r for r in rows if r['gamma'] == 2.0)
131    }
132    with open('results.json', 'w') as f:
133        json.dump(result, f, indent=2)
134    print(json.dumps(result, indent=2))
135
136if __name__ == '__main__':
137    main()