1import json, math
  2from pathlib import Path
  3import numpy as np
  4from scipy.special import logsumexp
  5
  6SEED = 3134
  7rng = np.random.default_rng(SEED)
  8D, MODES, M = 16, 8, 1200
  9SIGMA = 0.55
 10LOG2PI = math.log(2.0 * math.pi)
 11
 12# Eight separated modes, with a common translation changing over annealing levels.
 13base_means = np.zeros((MODES, D))
 14for k in range(MODES):
 15    a = 3.0 if k < 4 else -3.0
 16    b = 3.0 if (k % 4) < 2 else -3.0
 17    base_means[k, :4] = [a, b, 2.2 * (1 if k % 2 else -1), 1.5 * (1 if k in (0, 3, 4, 7) else -1)]
 18# A shared target shift gives a clean, controllable nonstationarity signal.
 19shifts = []
 20for i in range(28):
 21    t = i / 27.0
 22    speed = 0.15 if i < 7 or i >= 21 else 0.95
 23    # piecewise rapid motion, plus a return, so consecutive q drift is observable
 24    x = (0.0 if i == 0 else sum(0.15 if j < 7 or j >= 21 else 0.95 for j in range(1, i+1)))
 25    shifts.append(np.array([x, -0.7*x, 0.35*x] + [0.] * (D-3)))
 26target_means = np.array([base_means + s for s in shifts])
 27
 28def log_mix(x, means):
 29    # equal-weight spherical Gaussian mixture
 30    z = -0.5 * (((x[:, None, :] - means[None, :, :]) / SIGMA) ** 2).sum(axis=2)
 31    z -= D * (math.log(SIGMA) + 0.5 * LOG2PI)
 32    return logsumexp(z, axis=1) - math.log(MODES)
 33
 34def sample_mix(means, n, rg):
 35    idx = rg.integers(MODES, size=n)
 36    return means[idx] + SIGMA * rg.normal(size=(n, D))
 37
 38def ess_from_logw(logw):
 39    a = logw - logsumexp(logw)
 40    return float(np.exp(-logsumexp(2*a)) / len(logw))
 41
 42def sym_kl_mc(a, b, n=5000):
 43    xa, xb = sample_mix(a, n, rng), sample_mix(b, n, rng)
 44    return float(0.5 * (np.mean(log_mix(xa, a)-log_mix(xa, b)) +
 45                        np.mean(log_mix(xb, b)-log_mix(xb, a))))
 46
 47def importance_stats(target, proposal, n=M):
 48    x = sample_mix(proposal, n, rng)
 49    lw = log_mix(x, target) - log_mix(x, proposal)
 50    return ess_from_logw(lw), float(abs(np.log(np.mean(np.exp(lw)))))
 51
 52def verify_math():
 53    a = base_means.copy()
 54    b = a + np.array([1.4, -0.5, 0.2] + [0.]*(D-3))
 55    # ESS is invariant to a common log-weight shift and lies in [1/M,1].
 56    lw = rng.normal(size=1000)
 57    e1 = ess_from_logw(lw)
 58    e2 = ess_from_logw(lw + 37.0)
 59    kl = sym_kl_mc(a, b, 12000)
 60    kl_rev = sym_kl_mc(b, a, 12000)
 61    # For this symmetric definition, swapping distributions should agree up to MC noise.
 62    return {"ess_in_range": bool(1/1000 - 1e-12 <= e1 <= 1.0),
 63            "ess_shift_invariance_abs_error": abs(e1-e2),
 64            "symmetric_kl_forward_reverse_abs_error": abs(kl-kl_rev),
 65            "symmetric_kl_positive": kl > 0, "symmetric_kl": kl}
 66
 67def run(adaptive):
 68    q = base_means.copy()
 69    previous_q = q.copy()
 70    K, R = 8, 1
 71    records = []
 72    # persistent FIFO of target means emulates stale flow-training data.
 73    history = []
 74    for i in range(len(target_means)):
 75        history.append(target_means[i].copy())
 76        if len(history) > K: history.pop(0)
 77        # A gradient-like flow update toward the FIFO training distribution.
 78        train_target = np.mean(history, axis=0)
 79        for _ in range(int(R)):
 80            q += 0.30 * (train_target - q)
 81        drift = sym_kl_mc(q, previous_q, 1800) if i else 0.0
 82        eta, evidence_error = importance_stats(target_means[i], q)
 83        # Small local Langevin correction is represented by a conservative mode-center
 84        # pull only when ESS is poor and drift is not high; it is counted separately.
 85        correction = 0
 86        if adaptive and eta < 0.10 and drift <= 0.45:
 87            correction = 6
 88            q += 0.35 * (target_means[i] - q)
 89            eta, evidence_error = importance_stats(target_means[i], q)
 90        records.append({"level": i, "eta": eta, "drift": drift, "K": K, "R": R,
 91                        "correction_steps": correction, "log_evidence_abs_error": evidence_error})
 92        if adaptive:
 93            if drift > 0.45:
 94                K = max(2, K // 2)
 95            elif drift < 0.08 and eta > 0.45:
 96                K = min(8, K + 1)
 97            if eta < 0.10:
 98                R = min(8, max(1, int(math.ceil(R * (1 + 2.0*(0.10-eta))))))
 99            elif eta > 0.55:
100                R = max(1, R-1)
101        previous_q = q.copy()
102    return records
103
104def summarize(rs):
105    eta = np.array([r["eta"] for r in rs])
106    return {"mean_eta": float(eta.mean()), "min_eta": float(eta.min()),
107            "levels_eta_below_0.1": int((eta < .1).sum()),
108            "levels_eta_below_0.02": int((eta < .02).sum()),
109            "mean_drift": float(np.mean([r["drift"] for r in rs])),
110            "total_R": int(sum(r["R"] for r in rs)),
111            "total_correction_steps": int(sum(r["correction_steps"] for r in rs)),
112            "mean_log_evidence_abs_error": float(np.mean([r["log_evidence_abs_error"] for r in rs])),
113            "max_log_evidence_abs_error": float(max(r["log_evidence_abs_error"] for r in rs))}
114
115def main():
116    mathcheck = verify_math()
117    fixed, adaptive = run(False), run(True)
118    out = {"seed": SEED, "math_check": mathcheck,
119           "fixed_summary": summarize(fixed), "adaptive_summary": summarize(adaptive),
120           "fixed_levels": fixed, "adaptive_levels": adaptive}
121    Path("results.json").write_text(json.dumps(out, indent=2))
122    print(json.dumps({"math_check": mathcheck, "fixed": summarize(fixed), "adaptive": summarize(adaptive)}, indent=2))
123
124if __name__ == "__main__": main()