1import json
  2from pathlib import Path
  3import numpy as np
  4
  5
  6def softmax(x):
  7    x = x - np.max(x)
  8    p = np.exp(x)
  9    return p / p.sum()
 10
 11
 12def weighted_moments(z, w):
 13    mu = np.sum(w[:, None] * z, axis=0)
 14    c = np.sum(w[:, None] * (z - mu) ** 2, axis=0)
 15    return mu, c
 16
 17
 18def curvature_estimate(z, w, rewards, eps=1e-10):
 19    mu, c = weighted_moments(z, w)
 20    q = (z - mu) ** 2 - c
 21    rbar = np.sum(w * rewards)
 22    # This is the weighted least-squares coefficient in Cov(q,r)/E[q^2].
 23    r2 = np.sum(w[:, None] * (rewards[:, None] - rbar) * q, axis=0)
 24    den = np.sum(w[:, None] * q * q, axis=0) + eps
 25    return r2 / den, mu, c, rbar
 26
 27
 28def verify_moments(seed=7, n=2000000):
 29    rng = np.random.default_rng(seed)
 30    mu, s, kappa, curvature, slope = 0.7, 0.8, 1.3, -1.7, 0.4
 31    x = rng.normal(mu, s, n)
 32    r = slope * (x - mu) + 0.5 * curvature * (x - mu) ** 2
 33    cov_xr = np.mean((x - np.mean(x)) * (r - np.mean(r)))
 34    q = (x - np.mean(x)) ** 2 - np.var(x)
 35    cov_qr = np.mean((q - np.mean(q)) * (r - np.mean(r)))
 36    mean_rhs = kappa * s**2 * slope
 37    var_rhs_selection = kappa * curvature * s**4
 38    # Independent diffusion check: Gaussian perturbation variance increases by 2D dt.
 39    D, dt = 0.23, 0.01
 40    x2 = x + np.sqrt(2 * D * dt) * rng.normal(size=n)
 41    diffusion_increment = np.var(x2) - np.var(x)
 42    return {
 43        "empirical_mean_derivative_rhs": float(kappa * cov_xr),
 44        "predicted_mean_derivative": float(mean_rhs),
 45        "empirical_variance_selection_rhs": float(kappa * cov_qr),
 46        "predicted_variance_selection": float(var_rhs_selection),
 47        "empirical_diffusion_variance_increment": float(diffusion_increment),
 48        "predicted_diffusion_variance_increment": float(2 * D * dt),
 49        "relative_mean_error": float(abs(kappa * cov_xr - mean_rhs) / (abs(mean_rhs) + 1e-12)),
 50        "relative_variance_error": float(abs(kappa * cov_qr - var_rhs_selection) / (abs(var_rhs_selection) + 1e-12)),
 51    }
 52
 53
 54def run_population(mode, seed=123, steps=500, n=16, d=8):
 55    rng = np.random.default_rng(seed)
 56    target = np.linspace(-1.0, 1.0, d)
 57    z = rng.normal(0.0, 1.0, (n, d))
 58    # Fixed mutation is deliberately a reasonable, nonzero control.
 59    fixed_D = 0.018
 60    dmin, dmax, gamma, target_c = 0.00015, 0.045, 0.20, 0.16
 61    kappa, dt = 2.0, 1.0
 62    records = []
 63    for t in range(steps):
 64        true_reward = -np.mean((z - target) ** 2, axis=1)
 65        noise = rng.normal(0, 0.10, n)
 66        rewards = true_reward + noise
 67        uniform = np.full(n, 1.0 / n)
 68        r2, mu, c, rbar = curvature_estimate(z, uniform, rewards)
 69        reward_var = np.var(rewards)
 70        # Adaptive diffusion follows the proposal: noisy rewards increase D;
 71        # only sufficiently reliable negative curvature permits reduction.
 72        Dcoord = np.clip(dmin + gamma * reward_var / (1.0 + np.abs(np.clip(r2, -8, 8))), dmin, dmax)
 73        ess = 1.0 / np.sum(softmax(kappa * (rewards - rbar)) ** 2)
 74        if mode == "fixed":
 75            Dcoord = np.full(d, fixed_D)
 76        elif mode == "moment":
 77            if np.all(r2 < 0) and ess > n / 2:
 78                Dcoord = np.maximum(Dcoord, dmin)
 79            # Explicit +2D safeguard for a coordinate below target diversity.
 80            Dcoord = np.maximum(Dcoord, np.maximum(0.0, (target_c - c) / (2 * dt)))
 81            Dcoord = np.clip(Dcoord, dmin, dmax)
 82        weights = softmax(kappa * (rewards - rbar))
 83        # Resample around the weighted selected center, as in a population adapter optimizer.
 84        parent_idx = rng.choice(n, size=n, p=weights)
 85        z = z[parent_idx] + np.sqrt(2.0 * Dcoord * dt)[None, :] * rng.normal(size=(n, d))
 86        if t % 10 == 0 or t == steps - 1:
 87            records.append({
 88                "step": t, "best_true_reward": float(np.max(true_reward)),
 89                "mean_true_reward": float(np.mean(true_reward)),
 90                "distance": float(np.mean((np.mean(z, axis=0) - target) ** 2)),
 91                "cov_trace": float(np.mean(np.var(z, axis=0))), "ess": float(ess),
 92                "reward_var": float(reward_var), "mutation_D": float(np.mean(Dcoord)),
 93            })
 94    return records
 95
 96
 97def main():
 98    out = Path("results")
 99    out.mkdir(exist_ok=True)
100    verification = verify_moments()
101    all_runs = {m: run_population(m) for m in ("fixed", "moment")}
102    summary = {"verification": verification, "runs": all_runs}
103    for mode, rows in all_runs.items():
104        with (out / (mode + ".jsonl")).open("w") as f:
105            for row in rows:
106                f.write(json.dumps(row) + "\n")
107    with (out / "summary.json").open("w") as f:
108        json.dump(summary, f, indent=2)
109    for mode, rows in all_runs.items():
110        a, b = rows[0], rows[-1]
111        print(mode, "final_best=%.4f final_mean=%.4f final_cov=%.4f final_D=%.4f" %
112              (b["best_true_reward"], b["mean_true_reward"], b["cov_trace"], b["mutation_D"]))
113    print("verification", json.dumps(verification))
114
115if __name__ == "__main__":
116    main()