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, sweep_baseline, evaluate, make_report, get_dataset
  9import orientation_track as ot
 10
 11SEEDS = tuple(range(8))
 12LR_GRID = [1e-3, 3e-3, 1e-2]
 13EPOCHS = 18
 14BATCH = 128
 15
 16
 17def seed_all(seed):
 18    np.random.seed(seed)
 19    torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        torch.cuda.manual_seed_all(seed)
 22
 23
 24def ds_torch(d):
 25    out = dict(d)
 26    out["xtr"] = torch.tensor(d["xtr"], dtype=torch.float32)
 27    out["ytr"] = torch.tensor(d["ytr"], dtype=torch.long)
 28    out["xte"] = torch.tensor(d["xte"], dtype=torch.float32)
 29    out["yte"] = torch.tensor(d["yte"], dtype=torch.long)
 30    return out
 31
 32
 33def train_one(seed, lr, invariant):
 34    seed_all(seed)
 35    raw = get_dataset("orientation_phase_quotient", seed, 400, 160)
 36    raw_np = dict(raw)
 37    for k in ("xtr", "ytr", "xte", "yte"):
 38        if hasattr(raw_np[k], "detach"):
 39            raw_np[k] = raw_np[k].detach().cpu().numpy()
 40    d = ot.invariant_dataset(raw_np) if invariant else raw_np
 41    d = ds_torch(d)
 42    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
 43    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 44    return float(metric)
 45
 46
 47def baseline_factory(cfg):
 48    return lambda seed: train_one(seed, float(cfg["lr"]), False)
 49
 50
 51def idea_factory(cfg):
 52    return lambda seed: train_one(seed, float(cfg["lr"]), True)
 53
 54
 55def trained_signature(seed, lr):
 56    seed_all(seed)
 57    raw = get_dataset("orientation_phase_quotient", seed, 400, 160)
 58    raw_np = dict(raw)
 59    for k in ("xtr", "ytr", "xte", "yte"):
 60        if hasattr(raw_np[k], "detach"):
 61            raw_np[k] = raw_np[k].detach().cpu().numpy()
 62    inv = ot.invariant_dataset(raw_np)
 63    dr, di = ds_torch(raw_np), ds_torch(inv)
 64    br = make_model("mlp_tiny", dr["input_shape"], dr["out_dim"])
 65    bi = make_model("mlp_tiny", di["input_shape"], di["out_dim"])
 66    br, _, _ = train_model(br, dr, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 67    bi, _, _ = train_model(bi, di, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 68    rng = np.random.default_rng(seed + 991)
 69    u = raw_np["xte"].reshape(-1, 4 ** 3, 3)
 70    q, r = np.linalg.qr(rng.normal(size=(3, 3)))
 71    q = q @ np.diag(np.where(np.diag(r) >= 0, 1.0, -1.0))
 72    if np.linalg.det(q) < 0:
 73        q[:, 0] *= -1
 74    transformed = (u @ q.T) * rng.choice([-1.0, 1.0], size=(len(u), 4 ** 3, 1))
 75    raw2 = dict(raw_np)
 76    raw2["xte"] = transformed.reshape(len(u), -1).astype(np.float32)
 77    inv2 = ot.invariant_dataset(raw2)
 78    with torch.no_grad():
 79        devr = next(br.parameters()).device
 80        devi = next(bi.parameters()).device
 81        pbr = torch.softmax(br(torch.tensor(raw["xte"], dtype=torch.float32, device=devr)), 1)
 82        pbr2 = torch.softmax(br(torch.tensor(raw2["xte"], dtype=torch.float32, device=devr)), 1)
 83        pbi = torch.softmax(bi(torch.tensor(inv["xte"], dtype=torch.float32, device=devi)), 1)
 84        pbi2 = torch.softmax(bi(torch.tensor(inv2["xte"], dtype=torch.float32, device=devi)), 1)
 85    return {
 86        "baseline_mean_output_change": float(torch.abs(pbr - pbr2).mean().cpu()),
 87        "idea_mean_output_change": float(torch.abs(pbi - pbi2).mean().cpu()),
 88        "baseline_max_output_change": float(torch.abs(pbr - pbr2).max().cpu()),
 89        "idea_max_output_change": float(torch.abs(pbi - pbi2).max().cpu()),
 90    }
 91
 92
 93def main():
 94    # Core numerical check is performed before any neural training.
 95    math_err = ot.math_check()
 96    grid = [{"lr": x} for x in LR_GRID]
 97    base = sweep_baseline(baseline_factory, grid, seeds=(0, 1, 2, 3))
 98    idea_sweep = []
 99    for cfg in grid:
100        short = evaluate(idea_factory(cfg), seeds=(0, 1, 2, 3))
101        idea_sweep.append({"cfg": cfg, "mean": short["mean"]})
102    best_idea_cfg = min(grid, key=lambda c: next(x["mean"] for x in idea_sweep if x["cfg"] == c))
103    idea_full = evaluate(idea_factory(best_idea_cfg), seeds=SEEDS)
104    sig = trained_signature(0, float(best_idea_cfg["lr"]))
105    sig["predicted_invariant_output_change"] = 0.0
106    sig["observed_feature_change"] = math_err
107    sig["confirmed"] = bool(math_err < 1e-6 and sig["idea_mean_output_change"] < sig["baseline_mean_output_change"])
108    sig["prediction"] = "P2 local-correlation inputs and their trained predictions should be unchanged by global rotation and independent apolar flips."
109    report = make_report(
110        "orientation_phase_quotient", "mlp_tiny",
111        base,
112        dict(idea_full, best_cfg=best_idea_cfg, sweep=idea_sweep),
113        extra=sig,
114    )
115    result = {
116        "bench_report": report,
117        "custom_track": {"name": "orientation_phase_quotient", "file": "orientation_track.py", "domain": "molecular_orientation_symmetry"},
118        "math_check": {"max_transformed_feature_error": math_err, "predicted": 0.0},
119        "protocol": {"epochs": EPOCHS, "batch": BATCH, "lr_union": LR_GRID, "paired_seeds": list(SEEDS)},
120    }
121    Path("bench_results.json").write_text(json.dumps(result, indent=2, sort_keys=True))
122    print(json.dumps(result, indent=2, sort_keys=True))
123
124
125if __name__ == "__main__":
126    main()