1"""Spatial-quantile conformal bands: math check and small operator experiment."""
  2import json, math, random
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def weighted_quantile(values, weights, mass):
  8    values = np.asarray(values, float).ravel()
  9    weights = np.asarray(weights, float).ravel()
 10    order = np.argsort(values, kind="mergesort")
 11    v, w = values[order], weights[order]
 12    c = np.cumsum(w / w.sum())
 13    return float(v[np.searchsorted(c, mass, side="left")])
 14
 15
 16def spatial_score(residual, gamma, weights=None):
 17    """Smallest s whose weighted residual mass is at least 1-gamma."""
 18    residual = np.asarray(residual)
 19    if weights is None:
 20        weights = np.ones(residual.shape[-1]) / residual.shape[-1]
 21    return np.array([weighted_quantile(x, weights, 1.0 - gamma) for x in residual])
 22
 23
 24def conformal_q(scores, alpha):
 25    scores = np.sort(np.asarray(scores, float))
 26    n = len(scores)
 27    k = min(n, int(math.ceil((n + 1) * (1.0 - alpha))))
 28    return float(scores[k - 1])
 29
 30
 31def bands(center, scale, q):
 32    return center, q * scale
 33
 34
 35def math_sanity():
 36    # Directly checks weighted quantile definition, monotonicity in gamma,
 37    # and the finite-sample order-statistic indexing including k > n.
 38    x = np.array([[1., 2., 10., 20.]])
 39    w = np.array([.1, .2, .3, .4])
 40    s01 = spatial_score(x, .1, w)[0]
 41    s50 = spatial_score(x, .5, w)[0]
 42    assert s01 >= s50 and s01 == 20.0 and s50 == 10.0
 43    assert conformal_q(np.arange(1., 11.), .1) == 10.0  # ceil(11*.9)=10
 44    assert conformal_q(np.arange(1., 4.), .01) == 3.0  # clipped k
 45    # Every returned score covers at least the requested weighted spatial mass.
 46    rng = np.random.default_rng(4)
 47    r = rng.lognormal(size=(20, 37))
 48    for gamma in (.01, .05, .1, .3):
 49        ss = spatial_score(r, gamma)
 50        covered = (r <= ss[:, None]).mean(axis=1)
 51        assert np.all(covered >= 1 - gamma - 1e-12)
 52    return {"weighted_gamma_0.1": s01, "weighted_gamma_0.5": s50,
 53            "score_monotone": True, "order_statistic_checks": True}
 54
 55
 56def make_data(seed=17, n=1800, m=48, d=6):
 57    rng = np.random.default_rng(seed)
 58    x = np.linspace(0, 1, m)
 59    z = rng.normal(size=(n, d))
 60    # Smooth latent operator with heteroscedastic, localized spatial noise.
 61    truth = (z[:, 0, None] * np.sin(2*np.pi*x)[None, :] +
 62             .7*z[:, 1, None] * np.cos(4*np.pi*x)[None, :] +
 63             .4*z[:, 2, None] * np.sin(6*np.pi*x)[None, :] +
 64             .25*z[:, 3, None] * x[None, :])
 65    truth += .08 * z[:, 4, None] * np.cos(10*np.pi*x)[None, :]
 66    return z, truth, x
 67
 68
 69def predictor_fit(xtr, ytr, xall, yall, seed=19):
 70    # A deliberately small frozen operator surrogate: ridge regression on
 71    # input coefficients and fixed Fourier features. It is a neural-operator-
 72    # style field map with shared input representation, without heavy training.
 73    rng = np.random.default_rng(seed)
 74    d = xtr.shape[1]
 75    m = ytr.shape[1]
 76    phi = np.concatenate([xtr, np.tanh(xtr), np.ones((len(xtr), 1))], axis=1)
 77    # tiny random feature layer makes this nonlinear while remaining reproducible
 78    W = rng.normal(0, .55, size=(phi.shape[1], 18))
 79    H = np.tanh(phi @ W)
 80    H = np.concatenate([phi, H], axis=1)
 81    lam = 1e-2
 82    coef = np.linalg.solve(H.T @ H + lam*np.eye(H.shape[1]), H.T @ ytr)
 83    def predict(xx):
 84        p = np.concatenate([xx, np.tanh(xx), np.ones((len(xx), 1))], axis=1)
 85        hh = np.tanh(p @ W)
 86        hh = np.concatenate([p, hh], axis=1)
 87        return hh @ coef
 88    return predict
 89
 90
 91def run_experiment():
 92    z, y, grid = make_data()
 93    # Fixed disjoint train/calibration/test split.
 94    pred = predictor_fit(z[:900], y[:900], z, y)
 95    cal = slice(900, 1100)
 96    test = slice(1100, 1800)
 97    yc, yt = y[cal], y[test]
 98    pc, pt = pred(z[cal]), pred(z[test])
 99    # Scale is learned only from the training residual envelope, then frozen.
100    train_res = np.abs(y[:900] - pred(z[:900]))
101    scale = np.maximum(np.quantile(train_res, .75, axis=0), .035)
102    rc = np.abs(yc-pc) / scale[None, :]
103    rt = np.abs(yt-pt) / scale[None, :]
104    out = {"n_cal": len(yc), "n_test": len(yt), "grid": len(grid), "scale_mean": float(scale.mean())}
105    for gamma in (.01, .05, .10):
106        qscores = spatial_score(rc, gamma)
107        q = conformal_q(qscores, .10)
108        maxq = conformal_q(rc.max(axis=1), .10)
109        frac_q = (rt <= q).mean(axis=1)
110        frac_max = (rt <= maxq).mean(axis=1)
111        target = 1-gamma
112        out[f"gamma_{gamma:.2f}"] = {
113            "quantile_q": q, "max_q": maxq,
114            "quantile_mean_width": float(2*q*scale.mean()),
115            "max_mean_width": float(2*maxq*scale.mean()),
116            "quantile_mean_fraction": float(frac_q.mean()),
117            "max_mean_fraction": float(frac_max.mean()),
118            "quantile_event_rate": float((frac_q >= target).mean()),
119            "max_event_rate": float((frac_max >= target).mean()),
120            "target_event_rate": .90,
121        }
122    return out
123
124
125def main():
126    result = {"math_sanity": math_sanity(), "experiment": run_experiment()}
127    Path("results.json").write_text(json.dumps(result, indent=2))
128    print(json.dumps(result, indent=2))
129
130if __name__ == "__main__":
131    main()