1import json, math
  2import numpy as np
  3from scipy.special import ndtr
  4
  5SEED = 1729
  6rng = np.random.default_rng(SEED)
  7
  8def f(x):
  9    return 0.85*x + 0.20*np.tanh(x)
 10
 11def jac(x):
 12    return 0.85 + 0.20/(np.cosh(x)**2)
 13
 14class Comp:
 15    __slots__ = ('w','m','v')
 16    def __init__(self, w, m, v): self.w, self.m, self.v = float(w), float(m), float(v)
 17
 18def normalize(cs):
 19    s = sum(c.w for c in cs)
 20    for c in cs: c.w /= s
 21    return cs
 22
 23def gaussian_w2(a, b):
 24    return math.sqrt((a.m-b.m)**2 + (math.sqrt(max(a.v,0))-math.sqrt(max(b.v,0)))**2)
 25
 26def merge(cs, eps):
 27    cs = list(cs)
 28    while len(cs) > 1:
 29        cand = [(gaussian_w2(cs[i],cs[j]),i,j) for i in range(len(cs)) for j in range(i+1,len(cs))]
 30        d,i,j = min(cand)
 31        if d > eps: break
 32        a,b = cs[i],cs[j]; w = a.w+b.w; m = (a.w*a.m+b.w*b.m)/w
 33        v = (a.w*(a.v+(a.m-m)**2)+b.w*(b.v+(b.m-m)**2))/w
 34        cs = [c for k,c in enumerate(cs) if k not in (i,j)] + [Comp(w,m,v)]
 35    return normalize(cs)
 36
 37def propagate(cs, d, q, eps=None):
 38    out=[]
 39    for c in cs:
 40        J=jac(c.m)
 41        out += [Comp(.5*c.w, f(c.m)-d, J*J*c.v+q), Comp(.5*c.w, f(c.m)+d, J*J*c.v+q)]
 42    return merge(out,eps) if eps is not None else normalize(out)
 43
 44def moments(cs):
 45    m=sum(c.w*c.m for c in cs)
 46    v=sum(c.w*(c.v+(c.m-m)**2) for c in cs)
 47    return m,v
 48
 49def chance(cs, threshold):
 50    return sum(c.w*ndtr((threshold-c.m)/math.sqrt(c.v)) for c in cs)
 51
 52def sample(cs, n, rg):
 53    ix=rg.choice(len(cs), n, p=[c.w for c in cs])
 54    return np.array([rg.normal(cs[k].m, math.sqrt(cs[k].v)) for k in ix])
 55
 56def empirical_w2(x,y):
 57    # In one dimension the sorted empirical coupling is optimal W2.
 58    n=min(len(x),len(y)); xx=np.sort(x[:n]); yy=np.sort(y[:n])
 59    return math.sqrt(np.mean((xx-yy)**2))
 60
 61def gaussian_component(cs):
 62    m,v=moments(cs)
 63    return [Comp(1,m,v)]
 64
 65def main():
 66    q=.0025; d=.9; initial=[Comp(1,0,.01)]
 67    # Prediction 1: after one step, mode separation is 2d, independent of covariance.
 68    sep=[]
 69    for dd in [.1,.3,.6,.9,1.2]:
 70        z=propagate(initial,dd,q)
 71        sep.append({'d':dd,'observed':z[1].m-z[0].m,'predicted':2*dd})
 72    # Prediction 2: two equal Gaussian modes merge precisely at W2 tolerance >= separation.
 73    threshold_sweep=[]
 74    raw=propagate(initial,d,q)
 75    rawdist=gaussian_w2(raw[0],raw[1])
 76    for eps in [.2,.9,1.5,1.79,1.8,1.81,2.0]:
 77        threshold_sweep.append({'eps':eps,'components':len(merge(raw,eps)),'pair_w2':rawdist,'predicted_merge':eps >= rawdist})
 78    # Prediction 3: merging at tolerance above pair W2 changes the represented law,
 79    # while tolerances below it preserve both modes.
 80    merge_error = []
 81    rg = np.random.default_rng(SEED + 77)
 82    reference = sample(raw, 30000, rg)
 83    for eps in [.9, 1.79, 1.81, 2.0]:
 84        mc = merge(raw, eps)
 85        approx = sample(mc, 30000, np.random.default_rng(SEED + int(100*eps)))
 86        merge_error.append({'eps':eps,'components':len(mc),'sample_w2_to_unmerged':empirical_w2(reference,approx),
 87                            'predicted_merged':eps >= rawdist})
 88    # Prediction 4: before merging, chance probability is mixture CDF and differs from moment Gaussian.
 89    # A symmetric threshold at zero is exactly 0.5; use a positive threshold for a quantitative sweep.
 90    chance_rows=[]
 91    for dd in [.1,.3,.6,.9,1.2]:
 92        z=propagate(initial,dd,q); mm=gaussian_component(z)
 93        p_mix=chance(z,.5); p_gauss=chance(mm,.5)
 94        chance_rows.append({'d':dd,'mixture':p_mix,'moment_gaussian':p_gauss,'absolute_gap':abs(p_mix-p_gauss)})
 95    # Small nonlinear rollout comparison against moment matching, with exact mixture as reference.
 96    horizons=[1,2,4,6,8]
 97    rollout=[]
 98    for H in horizons:
 99        exact=initial
100        mix=initial
101        for _ in range(H):
102            exact=propagate(exact,d,q)
103            mix=propagate(mix,d,q,eps=1e9) # one component after every step
104        rg=np.random.default_rng(SEED+H)
105        a=sample(exact,20000,rg); b=sample(mix,20000,rg)
106        rollout.append({'horizon':H,'components_exact':len(exact),'components_gaussian':len(mix),
107                        'sample_w2_error_gaussian_vs_mixture':empirical_w2(a,b),
108                        'exact_mean':moments(exact)[0],'gaussian_mean':moments(mix)[0]})
109    out={'seed':SEED,'parameters':{'d':d,'q':q},'prediction_mode_separation':sep,
110         'prediction_merge_transition':threshold_sweep,'prediction_chance_calibration':chance_rows,'merge_error_sweep':merge_error,
111         'rollout_comparison':rollout}
112    with open('results.json','w') as f: json.dump(out,f,indent=2)
113    print(json.dumps(out,indent=2))
114
115if __name__=='__main__': main()