1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5from scipy.special import ndtr, ndtri, logsumexp
  6from scipy.stats import norm
  7from sklearn.metrics import roc_auc_score
  8
  9RNG = np.random.default_rng(1342)
 10
 11
 12def mixture_params(separation=4.0, prior_var=0.35**2):
 13    means = np.array([-separation / 2.0, separation / 2.0])
 14    weights = np.array([0.5, 0.5])
 15    return means, weights, prior_var
 16
 17
 18def posterior_components(y, noise_var, separation=4.0, prior_var=0.35**2):
 19    """p(x0|y) is a two-component Gaussian mixture under y=x0+N(0,noise_var)."""
 20    means, weights, _ = mixture_params(separation, prior_var)
 21    y = np.asarray(y)
 22    marginal_vars = prior_var + noise_var
 23    logw = np.log(weights)[None, :] + norm.logpdf(y[..., None], means, np.sqrt(marginal_vars))
 24    logw -= logsumexp(logw, axis=-1, keepdims=True)
 25    post_w = np.exp(logw)
 26    gain = prior_var / marginal_vars
 27    post_means = means + gain * (y[..., None] - means)
 28    post_var = prior_var * noise_var / marginal_vars
 29    return post_w, post_means, post_var
 30
 31
 32def posterior_cdf(x, y, noise_var, separation=4.0, prior_var=0.35**2):
 33    w, m, v = posterior_components(y, noise_var, separation, prior_var)
 34    return np.sum(w * ndtr((np.asarray(x)[..., None] - m) / np.sqrt(v)), axis=-1)
 35
 36
 37def posterior_logpdf(x, y, noise_var, separation=4.0, prior_var=0.35**2):
 38    w, m, v = posterior_components(y, noise_var, separation, prior_var)
 39    return logsumexp(np.log(w) + norm.logpdf(np.asarray(x)[..., None], m, np.sqrt(v)), axis=-1)
 40
 41
 42def transport_sample(y, noise_var, n, separation=4.0, prior_var=0.35**2, grid_n=5001):
 43    """Exact conditional-CDF transport T(y,z)=F^{-1}_{X|y}(Phi(z))."""
 44    y = np.asarray(y)
 45    z = RNG.normal(size=(n, len(y)))
 46    q = ndtr(z)
 47    lo = np.min(posterior_components(y, noise_var, separation, prior_var)[1]) - 8
 48    hi = np.max(posterior_components(y, noise_var, separation, prior_var)[1]) + 8
 49    # Vectorized inverse CDF by bisection; monotonicity is the key mechanism.
 50    l = np.full_like(q, lo, dtype=float); r = np.full_like(q, hi, dtype=float)
 51    yy = np.broadcast_to(y, q.shape[1:])
 52    for _ in range(45):
 53        mid = (l + r) / 2
 54        f = posterior_cdf(mid, yy[None, :], noise_var, separation, prior_var)
 55        l = np.where(f < q, mid, l); r = np.where(f < q, r, mid)
 56    return (l + r) / 2, z
 57
 58
 59def affine_params(y, noise_var, separation=4.0, prior_var=0.35**2):
 60    """Moment-matched Gaussian reverse kernel, standard affine reverse head."""
 61    w, m, v = posterior_components(y, noise_var, separation, prior_var)
 62    mean = np.sum(w * m, axis=-1)
 63    var = v + np.sum(w * (m - mean[..., None]) ** 2, axis=-1)
 64    return mean, var
 65
 66
 67def conditional_kl(y, noise_var, separation=4.0, prior_var=0.35**2):
 68    """KL(exact posterior || moment-matched Gaussian), quadrature by posterior samples."""
 69    x, _ = transport_sample(y, noise_var, 1800, separation, prior_var)
 70    mean, var = affine_params(y, noise_var, separation, prior_var)
 71    exact = posterior_logpdf(x, y[None, :], noise_var, separation, prior_var)
 72    gauss = norm.logpdf(x, mean[None, :], np.sqrt(var)[None, :])
 73    return np.mean(exact - gauss, axis=0)
 74
 75
 76def residual_auc(noise_var, separation=4.0, prior_var=0.35**2, n=5000):
 77    """Can a classifier predict y from reverse residual? AUC=0.5 means independence."""
 78    x, y = sample_joint(n, noise_var, separation, prior_var)
 79    # Exact transport residual: conditional probability integral transform.
 80    u = np.clip(posterior_cdf(x, y, noise_var, separation, prior_var), 1e-6, 1-1e-6)
 81    z_transport = ndtri(u)
 82    mu, var = affine_params(y, noise_var, separation, prior_var)
 83    z_affine = (x - mu) / np.sqrt(var)
 84    # Use |residual|, which captures state-dependent shape without fitting a classifier.
 85    # AUC is made orientation-invariant.
 86    def auc(z):
 87        a = roc_auc_score(y > np.median(y), np.abs(z))
 88        return max(a, 1-a)
 89    return auc(z_affine), auc(z_transport)
 90
 91
 92def sample_joint(n, noise_var, separation=4.0, prior_var=0.35**2):
 93    means, weights, _ = mixture_params(separation, prior_var)
 94    c = RNG.choice(2, size=n, p=weights)
 95    x = RNG.normal(means[c], np.sqrt(prior_var))
 96    y = x + RNG.normal(0, np.sqrt(noise_var), n)
 97    return x, y
 98
 99
100def math_check():
101    # Bayes identity: integrate the joint density relation at random points.
102    x, y = sample_joint(3000, 0.5)
103    prior = 0.5 * norm.pdf(x, -2, .35) + 0.5 * norm.pdf(x, 2, .35)
104    likelihood = norm.pdf(y, x, np.sqrt(.5))
105    marginal = .5 * norm.pdf(y, -2, np.sqrt(.35**2+.5)) + .5 * norm.pdf(y, 2, np.sqrt(.35**2+.5))
106    post = np.exp(posterior_logpdf(x, y, .5))
107    bayes_relerr = np.median(np.abs(post - likelihood*prior/marginal) / (post + 1e-12))
108    # PIT uniformity and transport conditional moment agreement.
109    u = posterior_cdf(x, y, .5)
110    ks_like = float(np.max(np.abs(np.sort(u) - (np.arange(len(u)) + .5) / len(u))))
111    return {"bayes_median_relative_error": float(bayes_relerr), "pit_sup_deviation": ks_like}
112
113
114def run():
115    out = {"math_check": math_check(), "separation_sweep": [], "noise_sweep": [], "independence": []}
116    # Prediction 1: mixture separation past posterior SD produces increasing affine KL; transport stays ~0.
117    for d in [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]:
118        ys = np.linspace(-3.5, 3.5, 15)
119        kl = conditional_kl(ys, .35**2, d)
120        out["separation_sweep"].append({"separation": d, "affine_kl_mean": float(np.mean(kl)), "transport_kl_expected": 0.0})
121    # Prediction 2: as forward noise increases, posterior becomes less multimodal and affine gap shrinks.
122    for nv in [.03, .10, .25, .5, 1.0, 2.0, 4.0]:
123        ys = np.linspace(-4, 4, 17)
124        kl = conditional_kl(ys, nv, 4.0)
125        out["noise_sweep"].append({"noise_var": nv, "affine_kl_mean": float(np.mean(kl)), "transport_kl_expected": 0.0})
126    # Prediction 3: exact PIT transport gives state-independent residual, affine residual does not.
127    for nv in [.03, .1, .5, 1.0, 2.0]:
128        a, t = residual_auc(nv, 4.0)
129        out["independence"].append({"noise_var": nv, "affine_auc": float(a), "transport_auc": float(t), "ideal_auc": .5})
130
131    # Quantitative transition prediction: equal-weight posterior components become
132    # visibly separated when their mean gap exceeds two posterior standard deviations.
133    # gap/sd = d*sqrt(nv)/(sqrt(prior_var)*sqrt(prior_var+nv)).
134    prior = .35**2
135    d = 4.0
136    # Solve d^2*nv/(prior*(prior+nv)) = 4.
137    predicted_nv = (4 * prior**2) / (d**2 - 4 * prior)
138    out["bimodality_transition"] = {"predicted_noise_var_gap_eq_2sd": predicted_nv, "sweep": []}
139    for nv in [0.01, 0.03, 0.05, predicted_nv, 0.10, 0.25, 1.0]:
140        w, m, v = posterior_components(np.array([0.0]), nv, d, prior)
141        ratio = float((m[0, 1] - m[0, 0]) / np.sqrt(v))
142        out["bimodality_transition"]["sweep"].append({"noise_var": float(nv), "posterior_gap_over_sd": ratio})
143
144    # Marginal sampling check: transport samples should match the mixture; affine
145    # samples have the same first two moments but lose mixture shape.
146    xtrue, ytest = sample_joint(6000, .5, 4.0, prior)
147    xt, _ = transport_sample(ytest[:600], .5, 10, 4.0, prior)
148    am, av = affine_params(ytest[:600], .5, 4.0, prior)
149    xa = am[None, :] + np.sqrt(av)[None, :] * RNG.normal(size=xt.shape)
150    out["sample_moments"] = {
151        "true_mean": float(np.mean(xtrue)), "transport_mean": float(np.mean(xt)), "affine_mean": float(np.mean(xa)),
152        "true_second_moment": float(np.mean(xtrue**2)), "transport_second_moment": float(np.mean(xt**2)), "affine_second_moment": float(np.mean(xa**2)),
153        "true_excess_kurtosis": float(np.mean((xtrue-np.mean(xtrue))**4)/np.var(xtrue)**2-3),
154        "transport_excess_kurtosis": float(np.mean((xt-np.mean(xt))**4)/np.var(xt)**2-3),
155        "affine_excess_kurtosis": float(np.mean((xa-np.mean(xa))**4)/np.var(xa)**2-3)}
156    Path("results.json").write_text(json.dumps(out, indent=2))
157    print(json.dumps(out, indent=2))
158
159if __name__ == "__main__":
160    run()