1import json
  2import math
  3import sys
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8import torch.nn as nn
  9
 10sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 11from bench import make_model, train_model, evaluate, sweep_baseline, make_report
 12from sphere_track import get_dataset, META
 13
 14OUT = Path("bench_report.json")
 15
 16
 17def exp_sphere(x, v):
 18    r = torch.linalg.vector_norm(v, dim=-1, keepdim=True)
 19    return torch.cos(r) * x + torch.sinc(r / math.pi) * v
 20
 21
 22def tangent(x, u):
 23    return u - (u * x).sum(dim=-1, keepdim=True) * x
 24
 25
 26class ResidualSystem(nn.Module):
 27    """Same MLP backbone; only the residual integration rule differs."""
 28    def __init__(self, idea, h):
 29        super().__init__()
 30        self.backbone = make_model("mlp_tiny", (3,), 3)
 31        self.idea = bool(idea)
 32        self.h = float(h)
 33
 34    def forward(self, x):
 35        raw = self.backbone(x)
 36        if self.idea:
 37            return exp_sphere(x, self.h * tangent(x, raw))
 38        return x + self.h * raw
 39
 40
 41def dataset(seed):
 42    d = get_dataset(seed, 400, 400)
 43    out = {k: torch.as_tensor(v, dtype=torch.float32) for k, v in d.items()
 44           if k in ("xtr", "ytr", "xte", "yte")}
 45    out.update({"task": "regression", "metric": "mse", "input_shape": (3,), "out_dim": 3})
 46    return out
 47
 48
 49def run_one(idea, cfg, seed, return_model=False):
 50    torch.manual_seed(10000 + int(seed))
 51    np.random.seed(10000 + int(seed))
 52    net, metric, history = train_model(
 53        ResidualSystem(idea, cfg["h"]), dataset(seed),
 54        epochs=20, lr=cfg["lr"], batch=128, weight_decay=0.0,
 55        log=lambda *_args, **_kwargs: None)
 56    if metric is None:
 57        raise RuntimeError("benchmark training failed")
 58    return (float(metric), net) if return_model else float(metric)
 59
 60
 61def main():
 62    # Cheap numerical verification is performed before any training.
 63    torch.manual_seed(3)
 64    x = torch.randn(1000, 3)
 65    x = x / torch.linalg.vector_norm(x, dim=-1, keepdim=True)
 66    u = torch.randn_like(x)
 67    v = 0.12 * tangent(x, u)
 68    exp_norm_error = float((torch.linalg.vector_norm(exp_sphere(x, v), dim=-1) - 1).abs().max())
 69    add_norm_error = float((torch.linalg.vector_norm(x + v, dim=-1) - 1).abs().mean())
 70    if not (exp_norm_error < 1e-6 and add_norm_error > 1e-5):
 71        raise RuntimeError("core exponential-map norm sanity check failed")
 72
 73    lrs = [1e-3, 3e-3, 1e-2]
 74    hs = [0.06, 0.12, 0.24]
 75    grid = [{"lr": lr, "h": h} for lr in lrs for h in hs]
 76    # Baseline sees the complete union of all idea hyperparameters.
 77    base = sweep_baseline(
 78        lambda cfg: lambda seed: run_one(False, cfg, seed), grid)
 79    best = base["best_cfg"]
 80    idea_grid = [best]
 81    for cfg in grid:
 82        if cfg != best and len(idea_grid) < 3:
 83            idea_grid.append(cfg)
 84    idea_trials = []
 85    for cfg in idea_grid:
 86        r = evaluate(lambda seed, c=cfg: run_one(True, c, seed))
 87        idea_trials.append({"cfg": cfg, "result": r})
 88    best_idea_trial = min(idea_trials, key=lambda z: z["result"]["mean"])
 89    idea = best_idea_trial["result"]
 90    rep = make_report("sphere_one_step_dynamics", "mlp_tiny", base, idea, extra={
 91        "sanity_check": {
 92            "prediction": "Exp preserves unit norm while additive Euler drifts",
 93            "exp_max_norm_error": exp_norm_error,
 94            "additive_mean_norm_error": add_norm_error,
 95            "confirmed": bool(exp_norm_error < 1e-6 and add_norm_error > 1e-5),
 96        },
 97        "idea_trials": idea_trials,
 98    })
 99    rep["custom_track"] = {
100        "name": META["name"], "file": "sphere_track.py", "domain": META["domain"]}
101    # Signature is measured from independently trained systems at their selected configs.
102    cfgb = best
103    mb_metric, mb = run_one(False, cfgb, 0, return_model=True)
104    mi_metric, mi = run_one(True, best_idea_trial["cfg"], 0, return_model=True)
105    with torch.no_grad():
106        xx = dataset(0)["xte"]
107        xx_b = xx.to(next(mb.parameters()).device)
108        xx_i = xx.to(next(mi.parameters()).device)
109        pb = mb(xx_b); pi = mi(xx_i)
110        rep["mechanism_signature"] = {
111            "trained_model_seed": 0,
112            "baseline_mean_abs_norm_error": float((pb.norm(dim=1) - 1).abs().mean()),
113            "idea_mean_abs_norm_error": float((pi.norm(dim=1) - 1).abs().mean()),
114            "predicted_idea_lower_constraint_error": True,
115            "confirmed": bool((pi.norm(dim=1) - 1).abs().mean() < (pb.norm(dim=1) - 1).abs().mean()),
116        }
117    OUT.write_text(json.dumps(rep, indent=2))
118    print(json.dumps(rep, indent=2))
119
120
121if __name__ == "__main__":
122    main()