1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5# Ordered Diffusion Message Passing MVP.
  6# The toy grid makes the local Gaussian calculation transparent; the forecasting
  7# task uses the same row-normalized edge operator.
  8SEED = 7
  9rng = np.random.default_rng(SEED)
 10
 11
 12def kernel_on_offsets(offsets, eps, beta):
 13    """One row of K/P at x_i=0, with s(x)=x and f evaluated on offsets."""
 14    logits = -offsets**2 / (4.0 * eps) + beta * offsets
 15    logits -= logits.max()  # numerically stable softmax
 16    p = np.exp(logits)
 17    p /= p.sum()
 18    return p
 19
 20
 21def core_sweeps():
 22    # Fine grid and broad enough window make truncation error negligible.
 23    dx = 0.005
 24    offsets = np.arange(-6.0, 6.0 + dx / 2, dx)
 25    drift_rows = []
 26    for eps in [0.015, 0.03, 0.06, 0.12]:
 27        for beta in [0.0, 0.25, 0.5, 1.0]:
 28            p = kernel_on_offsets(offsets, eps, beta)
 29            mean = float(p @ offsets)
 30            second = float(p @ (offsets**2))
 31            drift_rows.append({
 32                "eps": eps, "beta": beta,
 33                "observed_mean": mean,
 34                "predicted_mean": 2 * beta * eps,
 35                "relative_drift_error": abs(mean - 2*beta*eps) / max(abs(2*beta*eps), 1e-12),
 36                "observed_generator_x": mean / eps,
 37                "predicted_generator_x": 2 * beta,
 38                "observed_generator_x2": second / eps,
 39                "predicted_generator_x2": 2 + 4 * beta * beta * eps,
 40            })
 41
 42    # Prediction 1: drift is zero at beta=0 and linear in beta.
 43    linear = [r for r in drift_rows if r["eps"] == 0.06]
 44    b = np.array([r["beta"] for r in linear])
 45    m = np.array([r["observed_mean"] for r in linear])
 46    slope = float(np.polyfit(b, m, 1)[0])
 47    # Prediction 2: at fixed beta, displacement scales linearly in epsilon.
 48    beta_rows = [r for r in drift_rows if r["beta"] == 0.5]
 49    e = np.array([r["eps"] for r in beta_rows])
 50    mm = np.array([r["observed_mean"] for r in beta_rows])
 51    eps_slope = float(np.polyfit(e, mm, 1)[0])
 52    # Prediction 3: generator of x^2 tends to 2 as eps -> 0 (at beta fixed).
 53    small = [r for r in drift_rows if r["beta"] == 0.5 and r["eps"] <= 0.06]
 54    gen_errors = [abs(r["observed_generator_x2"] - r["predicted_generator_x2"]) for r in small]
 55    # Exact row-stochastic/nonexpansive check on a finite random row and repeated rows.
 56    p = kernel_on_offsets(offsets, 0.06, 1.0)
 57    v = rng.normal(size=len(p))
 58    stability = {
 59        "row_sum": float(p.sum()),
 60        "max_entry": float(p.max()),
 61        "inf_norm_input": float(np.max(np.abs(v))),
 62        "inf_norm_one_step": float(abs(p @ v)),
 63        "inf_norm_bound_holds": bool(abs(p @ v) <= np.max(np.abs(v)) + 1e-12),
 64    }
 65    summary = {
 66        "drift_beta_slope_observed": slope,
 67        "drift_beta_slope_predicted": 2 * 0.06,
 68        "drift_eps_slope_observed_beta_0.5": eps_slope,
 69        "drift_eps_slope_predicted_beta_0.5": 1.0,
 70        "x2_generator_max_abs_error_small_eps": max(gen_errors),
 71        "x2_generator_prediction_tolerance": 0.02,
 72        "rows": drift_rows,
 73        "stability": stability,
 74    }
 75    return summary
 76
 77
 78def directed_forecast():
 79    """Advect a smooth signal on a randomly sampled 1-D point cloud.
 80
 81    h(x)=sin(x), target is h(x+shift). A positive ordered tilt shifts the
 82    local average forward, while symmetric diffusion only blurs h.
 83    """
 84    n = 700
 85    x = np.sort(rng.uniform(-5.0, 5.0, size=n))
 86    h = np.sin(x)
 87    target = np.sin(x + 0.22)
 88    k = 35
 89    # kNN in 1-D, same neighbors for both methods.
 90    dist = np.abs(x[:, None] - x[None, :])
 91    nn = np.argpartition(dist, kth=k, axis=1)[:, :k]
 92    eps = float(np.median(np.sort(dist, axis=1)[:, k])**2 / 4.0)
 93    eps = max(eps, 1e-5)
 94    delta = x[nn] - x[:, None]
 95    base = -delta**2 / (4 * eps)
 96
 97    def aggregate(beta):
 98        logits = base + beta * delta
 99        logits -= logits.max(axis=1, keepdims=True)
100        p = np.exp(logits)
101        p /= p.sum(axis=1, keepdims=True)
102        return np.sum(p * h[nn], axis=1), p
103
104    pred0, p0 = aggregate(0.0)
105    # Oracle scalar ordering s(x)=x is the simplest learned-ordering target.
106    # Sweep beta and report the best value without changing neighbors or width.
107    candidates = np.linspace(0, 1.8, 19)
108    mses = []
109    entropies = []
110    for beta in candidates:
111        pred, pp = aggregate(float(beta))
112        mses.append(float(np.mean((pred - target)**2)))
113        entropies.append(float(np.mean(-np.sum(pp * np.log(pp + 1e-12), axis=1))))
114    ib = int(np.argmin(mses))
115    best_beta = float(candidates[ib])
116    pred, pb = aggregate(best_beta)
117    return {
118        "n": n, "k": k, "epsilon": eps, "shift": 0.22,
119        "symmetric_mse": float(np.mean((pred0-target)**2)),
120        "ordered_best_mse": float(mses[ib]),
121        "ordered_best_beta": best_beta,
122        "ordered_improvement_fraction": float(1 - mses[ib]/max(mses[0], 1e-12)),
123        "symmetric_entropy": float(np.mean(-np.sum(p0*np.log(p0+1e-12), axis=1))),
124        "ordered_entropy": float(entropies[ib]),
125        "all_beta_mse": {str(float(b)): float(m) for b,m in zip(candidates, mses)},
126    }
127
128
129def main():
130    out = {"seed": SEED, "core_verification": core_sweeps(), "forecast": directed_forecast()}
131    Path("results.json").write_text(json.dumps(out, indent=2))
132    print(json.dumps(out, indent=2))
133
134if __name__ == "__main__":
135    main()