1import json
  2import numpy as np
  3from scipy.optimize import nnls
  4from scipy.stats import wasserstein_distance
  5
  6SEED = 1457
  7rng = np.random.default_rng(SEED)
  8
  9def T(x):
 10    return 0.78*x + 0.28*np.sin(2.7*x)
 11
 12class RBFCodec:
 13    def __init__(self, lo=-2.0, hi=2.0, m=25, sigma=.22, ngrid=1601):
 14        self.x = np.linspace(lo, hi, ngrid)
 15        self.dx = self.x[1] - self.x[0]
 16        self.centers = np.linspace(lo, hi, m)
 17        raw = np.exp(-.5*((self.x[:, None]-self.centers[None, :])/sigma)**2)
 18        # Normalize each basis to unit mass: q_i = integral psi_i = 1 exactly up to grid precision.
 19        self.A = raw / (raw.sum(axis=0)[None, :] * self.dx)
 20        self.q = np.trapz(self.A, self.x, axis=0)
 21
 22    def fit(self, samples):
 23        edges = np.linspace(self.x[0], self.x[-1], len(self.x))
 24        h, _ = np.histogram(samples, bins=edges, density=False)
 25        dens = np.empty_like(self.x)
 26        dens[1:] = h / (len(samples) * self.dx)
 27        dens[0] = dens[1]
 28        c = nnls(self.A, dens)[0]
 29        return c / max(self.q @ c, 1e-12)
 30
 31    def density(self, c):
 32        p = np.maximum(self.A @ c, 0.)
 33        return p / max(np.trapz(p, self.x), 1e-12)
 34
 35    def sample(self, c, n):
 36        p = self.density(c)
 37        cdf = np.maximum.accumulate(np.cumsum(p) * self.dx)
 38        cdf /= cdf[-1]
 39        return np.interp(rng.random(n), np.r_[0., cdf[:-1]], self.x)
 40
 41    def moments(self, c):
 42        p = self.density(c)
 43        mu = np.trapz(self.x*p, self.x)
 44        va = np.trapz((self.x-mu)**2*p, self.x)
 45        return mu, va
 46
 47def mass_project(K, q):
 48    residual = q - q @ K
 49    return K + np.outer(q / (q @ q), residual)
 50
 51def train_operator(codec, ntrain=260, npart=1800, ridge=1e-5):
 52    C, Y = [], []
 53    for _ in range(ntrain):
 54        w = rng.dirichlet(np.ones(len(codec.centers)) * .7)
 55        C.append(w)
 56        Y.append(codec.fit(T(codec.sample(w, npart))))
 57    C, Y = np.asarray(C).T, np.asarray(Y).T
 58    Kraw = Y @ C.T @ np.linalg.inv(C @ C.T + ridge*np.eye(C.shape[0]))
 59    K = mass_project(Kraw, codec.q)
 60    return K, Kraw
 61
 62def learned_roll(codec, K, c0, horizons=(1, 5, 10, 25), ntruth=12000):
 63    truth = codec.sample(c0, ntruth)
 64    c = c0.copy(); out = []
 65    for k in range(1, max(horizons)+1):
 66        truth = T(truth); c = K @ c
 67        if k in horizons:
 68            out.append(wasserstein_distance(truth, codec.sample(c, ntruth)))
 69    return out
 70
 71def gaussian_roll(c0, codec, horizons=(1, 5, 10, 25), ntruth=12000):
 72    truth = codec.sample(c0, ntruth)
 73    mu, var = codec.moments(c0); out = []
 74    for k in range(1, max(horizons)+1):
 75        truth = T(truth)
 76        s = np.sqrt(max(var, 1e-10)); z = np.array([mu-s, mu, mu+s]); wt = np.array([.25, .5, .25])
 77        mz = np.sum(wt*T(z)); var = np.sum(wt*(T(z)-mz)**2); mu = mz
 78        if k in horizons:
 79            out.append(wasserstein_distance(truth, rng.normal(mu, np.sqrt(max(var, 1e-10)), ntruth)))
 80    return out
 81
 82def spectral_sweep():
 83    # Prediction: for M=gamma*A, transition occurs at gamma*rho(A)=1 and
 84    # asymptotic log norm slope is log(gamma*rho(A)).
 85    lam = .93; A = np.diag([lam] + [.45]*7); c = np.ones(8); rows=[]
 86    for gamma in [.70, .90, 1.00, 1.075, 1.10, 1.30]:
 87        M = gamma*A; z = c.copy(); norms = []
 88        for _ in range(100): norms.append(np.linalg.norm(z)); z = M @ z
 89        slope = np.polyfit(np.arange(50, 100), np.log(np.maximum(norms[50:], 1e-300)), 1)[0]
 90        rows.append({'gamma': gamma, 'predicted_rho': gamma*lam,
 91                     'observed_log_slope': float(slope), 'predicted_log_slope': float(np.log(gamma*lam)),
 92                     'observed_bounded': bool(norms[-1] <= norms[0])})
 93    return rows
 94
 95def separation_sweep(codec, K):
 96    # Prediction: a multimodal PF representation should retain an advantage as
 97    # mode separation grows, while a single Gaussian loses shape information.
 98    rows = []
 99    for left, right in [(-.25, .25), (-.55, .55), (-.85, .85), (-1.15, 1.15)]:
100        c0 = np.exp(-.5*((codec.centers-left)/.13)**2) + np.exp(-.5*((codec.centers-right)/.13)**2)
101        c0 /= codec.q @ c0
102        pf = learned_roll(codec, K, c0, horizons=(5,), ntruth=10000)[0]
103        gauss = gaussian_roll(c0, codec, horizons=(5,), ntruth=10000)[0]
104        rows.append({'separation': right-left, 'PF_W1': pf, 'Gaussian_W1': gauss, 'PF_minus_Gaussian': pf-gauss})
105    return rows
106
107def main():
108    codec = RBFCodec(); K, Kraw = train_operator(codec)
109    q = codec.q
110    raw_mass_error = float(np.max(np.abs(q @ Kraw - q)))
111    projected_mass_error = float(np.max(np.abs(q @ K - q)))
112    c0 = np.zeros(len(q)); c0[5] = .5; c0[19] = .5
113    horizons = (1, 5, 10, 25)
114    pf = learned_roll(codec, K, c0, horizons)
115    gauss = gaussian_roll(c0, codec, horizons)
116    eig = np.linalg.eigvals(K)
117    result = {
118      'seed': SEED, 'basis': len(q), 'basis_integral_min_max': [float(q.min()), float(q.max())],
119      'mass_error_raw': raw_mass_error, 'mass_error_after_projection': projected_mass_error,
120      'spectral_radius_fitted_K': float(max(abs(eig))), 'spectral_prediction_sweep': spectral_sweep(),
121      'separation_sweep': separation_sweep(codec, K),
122      'w1_horizons': {'horizons': list(horizons), 'PF_RBF': pf, 'Gaussian_moment': gauss},
123      'mean_w1_PF': float(np.mean(pf)), 'mean_w1_Gaussian': float(np.mean(gauss))}
124    print(json.dumps(result, indent=2))
125
126if __name__ == '__main__': main()