1import json, random, sys
  2import numpy as np
  3import torch
  4
  5sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8))
  9SWEEP_SEEDS = tuple(range(4))
 10# Union parity: every idea setting is also evaluated for baseline.
 11GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 12EPOCHS = 24
 13NTR, NTE = 400, 200
 14K, M, DELTA = 8, 24, 2e-3
 15
 16
 17def seed_all(seed):
 18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 19    try:
 20        if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 21    except Exception:
 22        pass
 23
 24
 25def directions(n, d, rng):
 26    q = rng.normal(size=(n, d))
 27    return q / np.maximum(np.linalg.norm(q, axis=1, keepdims=True), 1e-12)
 28
 29
 30def boundary_value(z):
 31    # Track solution u=x^3-3xy^2, harmonic in the unit disk.
 32    return z[..., 0]**3 - 3.0*z[..., 0]*z[..., 1]**2
 33
 34
 35def wos_paths(x, k, m, seed):
 36    """Walk on largest interior spheres in the unit disk.
 37
 38    The source is zero and the terminal value is evaluated after projecting
 39    near-boundary points to the unit circle. The authoritative supplied
 40    convention dt=r^2/(2d) is immaterial here because h1=0.
 41    """
 42    x = np.asarray(x, dtype=np.float32)
 43    b, d = x.shape
 44    rng = np.random.RandomState(int(seed))
 45    z = np.repeat(x[:, None, :], k, axis=1).copy()
 46    for _ in range(m):
 47        norm = np.linalg.norm(z, axis=2)
 48        r = np.maximum(1.0 - norm, 0.0)
 49        active = r > DELTA
 50        u = directions(b*k, d, rng).reshape(b, k, d)
 51        z = z + r[..., None] * u
 52        if not np.any(active):
 53            break
 54    norm = np.linalg.norm(z, axis=2, keepdims=True)
 55    z = z / np.maximum(norm, 1e-8)
 56    return boundary_value(z).astype(np.float32)
 57
 58
 59def wos_targets(x, k, m, seed):
 60    return wos_paths(x, k, m, seed).mean(axis=1, keepdims=True)
 61
 62
 63def math_check():
 64    # Sphere step identity: E[|z+rU|^2-|z|^2]=r^2.
 65    rng = np.random.RandomState(2049)
 66    z = np.array([.2, -.1]); r = .4
 67    u = directions(200000, 2, rng)
 68    inc = float(np.mean(np.sum((z+r*u)**2, axis=1)-np.sum(z*z)))
 69    # The paper's sampling prediction: variance of a K-average is ~1/K.
 70    x = np.array([[.23, .17]], dtype=np.float32)
 71    paths = wos_paths(x, 16384, M, 77)[0]
 72    rows = []
 73    for k in [1, 2, 4, 8, 16, 32, 64, 128]:
 74        q = paths[:len(paths)//k*k].reshape(-1, k).mean(axis=1)
 75        v = float(np.var(q, ddof=1))
 76        rows.append({"K": k, "variance": v, "K_times_variance": k*v})
 77    products = np.array([q["K_times_variance"] for q in rows])
 78    return {"sphere_second_moment_increment": inc, "predicted_increment": r*r,
 79            "relative_increment_error": abs(inc-r*r)/(r*r),
 80            "variance_scaling": rows,
 81            "K_product_cv": float(products.std()/max(products.mean(), 1e-12)),
 82            "confirmed": bool(abs(inc-r*r)/(r*r) < .02 and products.std()/max(products.mean(), 1e-12) < .12)}
 83
 84
 85def make_ds(seed, kind):
 86    d0 = get_dataset("poisson_boundary", seed=seed, n_train=NTR, n_test=NTE)
 87    ds = {k: d0[k].clone() if torch.is_tensor(d0[k]) else torch.as_tensor(d0[k], dtype=torch.float32)
 88          for k in ("xtr", "ytr", "xte", "yte")}
 89    ds.update({"track": "poisson_boundary", "task": "regression", "metric": "mse",
 90               "input_shape": (2,), "out_dim": 1})
 91    if kind == "idea":
 92        ds["ytr"] = torch.as_tensor(wos_targets(d0["xtr"].numpy(), K, M, seed+10000), dtype=torch.float32)
 93    return ds
 94
 95
 96def run(kind, lr, seed, return_model=False):
 97    seed_all(seed)
 98    ds = make_ds(seed, kind)
 99    model = make_model("mlp_tiny", ds["input_shape"], ds["out_dim"])
100    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=float(lr), batch=128,
101                                    log=lambda *a, **k: None)
102    if net is None:
103        raise RuntimeError("bench training failed")
104    if return_model:
105        return float(metric), net, ds
106    return float(metric)
107
108
109def base_factory(cfg):
110    return lambda seed: run("baseline", cfg["lr"], seed)
111
112
113def idea_factory(cfg):
114    return lambda seed: run("idea", cfg["lr"], seed)
115
116
117def mechanism_signature():
118    # Signature uses outputs from trained systems. Prediction: because WOS
119    # targets are unbiased harmonic boundary rollouts, larger-K training
120    # should not induce a systematic mean shift relative to exact-label MLP.
121    rows = []
122    for seed in SEEDS:
123        bm, bn, bd = run("baseline", 3e-3, seed, True)
124        im, inn, idd = run("idea", 3e-3, seed, True)
125        with torch.no_grad():
126            bdev = next(bn.parameters()).device
127            idev = next(inn.parameters()).device
128            bp = bn(bd["xte"].to(bdev)).cpu().numpy().reshape(-1)
129            ip = inn(idd["xte"].to(idev)).cpu().numpy().reshape(-1)
130        delta = float(np.mean(ip-bp))
131        rows.append({"seed": seed, "baseline_mse": bm, "idea_mse": im,
132                     "trained_prediction_mean_shift_idea_minus_baseline": delta})
133    shifts = np.array([r["trained_prediction_mean_shift_idea_minus_baseline"] for r in rows])
134    return {"prediction": "WOS stochastic supervision does not create a systematic output mean shift",
135            "predicted_shift": 0.0, "observed_shift_mean": float(shifts.mean()),
136            "observed_shift_abs_mean": float(np.mean(np.abs(shifts))),
137            "trained_systems": rows,
138            "confirmed": bool(abs(shifts.mean()) < 0.01)}
139
140
141def main():
142    math_result = math_check()
143    base = sweep_baseline(base_factory, GRID, seeds=SWEEP_SEEDS)
144    idea_trials = [{"cfg": cfg, "result": evaluate(idea_factory(cfg), SEEDS)} for cfg in GRID]
145    best_trial = min(idea_trials, key=lambda q: q["result"]["mean"])
146    report = make_report("poisson_boundary", "mlp_tiny", base, best_trial["result"], {
147        "idea_config": best_trial["cfg"],
148        "idea_sweep": idea_trials,
149        "math_check": math_result,
150        "mechanism_signature": mechanism_signature(),
151        "custom_track": {"name": "poisson_boundary", "file": "bench/custom_tracks/poisson_boundary.py", "domain": "pde"},
152        "structural_match": "PDE Dirichlet boundary-value regression on the unit disk; WOS is the intervention."})
153    with open("bench_report.json", "w") as f:
154        json.dump(report, f, indent=2)
155    print(json.dumps(report, indent=2))
156
157
158if __name__ == "__main__":
159    main()