1import json
  2import sys
  3from pathlib import Path
  4import numpy as np
  5import torch
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import make_model, train_model, evaluate, sweep_baseline, make_report
  9import cayley_graph_track as track
 10
 11SEEDS = tuple(range(8))
 12NTRAIN, NTEST = 400, 100
 13EPOCHS, BATCH = 18, 128
 14# Shared union: both systems are evaluated at every learning rate.
 15GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 16
 17
 18def ds_torch(seed, kind):
 19    raw = track.get_dataset(seed, NTRAIN, NTEST)
 20    d = track.encode(raw, kind)
 21    return {**d,
 22            "xtr": torch.tensor(d["xtr"], dtype=torch.float32),
 23            "ytr": torch.tensor(d["ytr"], dtype=torch.float32),
 24            "xte": torch.tensor(d["xte"], dtype=torch.float32),
 25            "yte": torch.tensor(d["yte"], dtype=torch.float32)}
 26
 27
 28def train_one(seed, kind, cfg, return_model=False):
 29    torch.manual_seed(10000 + int(seed))
 30    np.random.seed(10000 + int(seed))
 31    d = ds_torch(seed, kind)
 32    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
 33    model, metric, history = train_model(net, d, epochs=EPOCHS,
 34                                         lr=float(cfg["lr"]), batch=BATCH,
 35                                         log=lambda *_: None)
 36    if return_model:
 37        return model, metric, d
 38    return float(metric)
 39
 40
 41def baseline_factory(cfg):
 42    return lambda seed: train_one(seed, "baseline", cfg)
 43
 44
 45def idea_factory(cfg):
 46    return lambda seed: train_one(seed, "idea", cfg)
 47
 48
 49def signature(cfg):
 50    """Measure the trained models, not an analytic proxy.
 51    A simultaneous cyclic shift of both endpoints must preserve g_uv."""
 52    raw = track.get_dataset(0, NTRAIN, NTEST)
 53    n = raw["n_nodes"]
 54    shift = 7
 55    shifted = dict(raw)
 56    shifted["xte"] = (np.asarray(raw["xte"]) + shift) % n
 57    vals = {}
 58    for kind in ("baseline", "idea"):
 59        model, _, d = train_one(0, kind, cfg, return_model=True)
 60        if model is None:
 61            vals[kind] = float("nan")
 62            continue
 63        model.eval()
 64        a = track.encode(raw, kind)
 65        b = track.encode(shifted, kind)
 66        device = next(model.parameters()).device
 67        with torch.no_grad():
 68            pa = model(torch.tensor(a["xte"], dtype=torch.float32, device=device)).cpu()
 69            pb = model(torch.tensor(b["xte"], dtype=torch.float32, device=device)).cpu()
 70        vals[kind] = float(torch.abs(pa - pb).mean())
 71    confirmed = bool(np.isfinite(vals["idea"]) and vals["idea"] < 0.25 * max(vals["baseline"], 1e-12))
 72    return {"prediction": "relative group differences are invariant to simultaneous vertex translation",
 73            "baseline_mean_output_change": vals["baseline"],
 74            "idea_mean_output_change": vals["idea"],
 75            "predicted_idea_change": 0.0, "confirmed": confirmed}
 76
 77
 78def main():
 79    math = track.math_check()
 80    assert math["path_independence"] and math["cycle_sum_mod_n"] == 0
 81    base = sweep_baseline(baseline_factory, GRID, seeds=SEEDS)
 82    idea_trials = []
 83    for cfg in GRID:
 84        r = evaluate(idea_factory(cfg), seeds=SEEDS)
 85        idea_trials.append({"cfg": cfg, "mean": r["mean"], "per_seed": r["per_seed"]})
 86    best_cfg = min(idea_trials, key=lambda x: x["mean"])["cfg"]
 87    idea_full = evaluate(idea_factory(best_cfg), seeds=SEEDS)
 88    idea_full["best_cfg"] = best_cfg
 89    idea_full["sweep"] = [{"cfg": x["cfg"], "mean": x["mean"]} for x in idea_trials]
 90    rep = make_report("cayley_cycle_distance", "mlp_tiny", base, idea_full,
 91                      {"mechanism_signature": signature(best_cfg),
 92                       "custom_track": {"name": track.META["name"], "file": "cayley_graph_track.py", "domain": track.META["domain"]},
 93                       "math_check": math,
 94                       "protocol_notes": {"epochs": EPOCHS, "batch": BATCH, "paired_seeds": list(SEEDS), "grid_union": GRID,
 95                                          "structure_justification": "The custom task is a cyclic Cayley graph distance prediction problem, directly containing generator increments, cycle constraints, and relative group displacements."}})
 96    Path("bench_report.json").write_text(json.dumps(rep, indent=2, sort_keys=True))
 97    print(json.dumps(rep, indent=2, sort_keys=True))
 98
 99
100if __name__ == "__main__":
101    main()