1import json
  2import random
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 1524
  7rng = np.random.default_rng(SEED)
  8
  9
 10def cosine(a, b):
 11    den = np.linalg.norm(a) * np.linalg.norm(b)
 12    return float(np.dot(a, b) / den) if den else 1.0
 13
 14
 15def make_population(n=12000):
 16    x = rng.uniform(-3.0, 3.0, size=(n, 2))
 17    z = 2.0 * x[:, 0] + 0.7 * x[:, 1] - 0.3
 18    fail_prob = 0.02 + 0.93 / (1.0 + np.exp(-z))
 19    y = rng.binomial(1, fail_prob)
 20    c = np.clip(fail_prob, 0.01, 0.99)
 21    features = np.column_stack([np.ones(n), x[:, 0], x[:, 1], x[:, 0] ** 2,
 22                                 x[:, 0] * x[:, 1], x[:, 1] ** 2])
 23    target = np.sin(x[:, 0]) + 0.35 * x[:, 1] + 0.15 * x[:, 0] ** 2
 24    return features, target, y, c
 25
 26
 27def fit_critic(x, y, epochs=800, lr=0.08):
 28    # Lightweight auxiliary predictor C_phi(s), fit only from state/failure pairs.
 29    X = np.column_stack([np.ones(len(x)), x])
 30    phi = np.zeros(X.shape[1])
 31    for _ in range(epochs):
 32        pred = 1.0 / (1.0 + np.exp(-np.clip(X @ phi, -30, 30)))
 33        phi -= lr * (X.T @ (pred - y)) / len(y)
 34    pred = 1.0 / (1.0 + np.exp(-np.clip(X @ phi, -30, 30)))
 35    bce = float(-np.mean(y * np.log(pred + 1e-8) + (1-y) * np.log(1-pred + 1e-8)))
 36    return np.clip(pred, 0.01, 0.99), bce
 37
 38
 39def proposal(c, alpha, eps=0.02):
 40    a = (eps + c) ** alpha
 41    z = float(a.mean())
 42    q = a / (len(a) * z)
 43    w = z / a
 44    ess_frac = float(1.0 / np.dot(q, w * w))
 45    return a, q, w, ess_frac
 46
 47
 48def exact_quantities(y, c, alpha, eps=0.02):
 49    a, q, w, ess_frac = proposal(c, alpha, eps)
 50    fail_p = float(y.mean())
 51    fail_q = float(np.dot(q, y))
 52    return a, q, w, fail_q / fail_p, ess_frac
 53
 54
 55def verify(alpha_grid, y, features, target, c, batches=400, batch_size=96):
 56    per_grad = -2.0 * target[:, None] * features
 57    uniform_grad = per_grad.mean(axis=0)
 58    rows = []
 59    for alpha in alpha_grid:
 60        _, q, w, enrich_exact, ess_exact = exact_quantities(y, c, alpha)
 61        failures, weighted_cos, weighted_err, unweighted_cos = [], [], [], []
 62        for _ in range(batches):
 63            ids = rng.choice(len(y), size=batch_size, replace=True, p=q)
 64            g = per_grad[ids]
 65            gw = (w[ids, None] * g).mean(axis=0)
 66            gu = g.mean(axis=0)
 67            failures.append(float(y[ids].mean()))
 68            weighted_cos.append(cosine(gw, uniform_grad))
 69            unweighted_cos.append(cosine(gu, uniform_grad))
 70            weighted_err.append(float(np.linalg.norm(gw - uniform_grad)))
 71        # Since q and w are explicitly known on this finite population, this is
 72        # an exact check of E_q[w g] = E_p[g], independent of Monte Carlo noise.
 73        exact_weighted = np.sum((q * w)[:, None] * per_grad, axis=0)
 74        exact_identity_error = float(np.linalg.norm(exact_weighted - uniform_grad))
 75        rows.append({"alpha": alpha, "predicted_enrichment": enrich_exact,
 76                     "observed_enrichment": float(np.mean(failures) / y.mean()),
 77                     "predicted_ess_fraction": ess_exact,
 78                     "weighted_gradient_cosine": float(np.mean(weighted_cos)),
 79                     "unweighted_gradient_cosine": float(np.mean(unweighted_cos)),
 80                     "weighted_gradient_l2_error": float(np.mean(weighted_err)),
 81                     "exact_identity_l2_error": exact_identity_error})
 82    return rows
 83
 84
 85def train_replay(features, target, c, alpha, weighted, steps=450, batch_size=96, lr=0.003):
 86    _, q, w, _ = proposal(c, alpha)
 87    theta = np.zeros(features.shape[1])
 88    for _ in range(steps):
 89        ids = rng.choice(len(target), size=batch_size, replace=True, p=q)
 90        xb, tb = features[ids], target[ids]
 91        err = xb @ theta - tb
 92        grad = 2.0 * (err[:, None] * xb)
 93        if weighted:
 94            grad *= w[ids, None]
 95        theta -= lr * grad.mean(axis=0)
 96    return float(np.mean((features @ theta - target) ** 2))
 97
 98
 99def main():
100    random.seed(SEED)
101    np.random.seed(SEED)
102    features, target, y, oracle_c = make_population()
103    x = features[:, 1:3]
104    c, critic_bce = fit_critic(x, y)
105    alpha_grid = [0.0, 0.5, 1.0, 2.0, 4.0]
106    verification = verify(alpha_grid, y, features, target, c)
107    uniform = train_replay(features, target, c, 0.0, False)
108    training = []
109    for alpha in [0.5, 1.0, 2.0]:
110        training.append({"alpha": alpha, "uniform_mse": uniform,
111                         "weighted_mse": train_replay(features, target, c, alpha, True),
112                         "unweighted_mse": train_replay(features, target, c, alpha, False)})
113    result = {"seed": SEED, "n_states": len(y), "base_failure_rate": float(y.mean()),
114        "critic": {"bce": critic_bce, "mean_prediction": float(c.mean()),
115                   "mean_label": float(y.mean())},
116        "predictions": {"enrichment": "rho_fail=E_q[y]/E_p[y]",
117                         "ess": "ESS/B approaches 1/E_q[w^2]",
118                         "gradient_identity": "E_q[w g]=E_p[g]"},
119        "verification": verification, "training": training,
120        "notes": "Predictions are finite-population exact; observed values are Monte Carlo.",
121        "importance_identity_check": {"max_exact_l2_error": float(max(r["exact_identity_l2_error"] for r in verification)),
122                                     "criterion": "exact error should be numerical roundoff; finite-batch cosine degrades as ESS falls"}}
123    Path("results.json").write_text(json.dumps(result, indent=2))
124    print(json.dumps(result, indent=2))
125
126
127if __name__ == "__main__":
128    main()