1import json
  2import random
  3import sys
  4from pathlib import Path
  5import numpy as np
  6import torch
  7from torch import nn
  8
  9sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 10from bench import train_model, evaluate, sweep_baseline, make_report
 11import cochain_track as ct
 12
 13SEEDS = tuple(range(8))
 14SWEEP_SEEDS = (0, 1, 2, 3)
 15LRS = [1e-3, 3e-3, 1e-2]
 16EPOCHS = 24
 17NTR, NTE = 400, 200
 18WIDTH = 16
 19M0 = np.ones(3)
 20M1 = np.ones(3)
 21
 22
 23def weighted_qr(a, masses):
 24    q, _ = np.linalg.qr(np.sqrt(masses)[:, None] * a)
 25    return q / np.sqrt(masses)[:, None]
 26
 27
 28def mass_norm(a, mi, mo):
 29    s = np.sqrt(mo)[:, None] * a * (1.0 / np.sqrt(mi))[None, :]
 30    return float(np.linalg.svd(s, compute_uv=False)[0])
 31
 32
 33def math_check():
 34    assert np.max(np.abs(ct.D1 @ ct.D0)) < 1e-12
 35    q = weighted_qr(ct.D0[:, :2], M1)
 36    p0 = np.eye(3)
 37    p1 = q @ q.T @ np.diag(M1)
 38    # A deliberately unrelated rank-2 edge subspace is the control.
 39    r = weighted_qr(np.array([[1., 0.], [0., 1.], [1., 1.]]), M1)
 40    bad = r @ r.T @ np.diag(M1)
 41    return {
 42        "incidence_d1_d0_max": float(np.max(np.abs(ct.D1 @ ct.D0))),
 43        "compatible_idempotence": mass_norm(p1 @ p1 - p1, M1, M1),
 44        "compatible_bound": mass_norm(p1, M1, M1),
 45        "compatible_commutation_d0": mass_norm(ct.D0 @ p0 - p1 @ ct.D0, M0, M1),
 46        "compatible_commutation_d1": mass_norm(ct.D1 @ p1, M1, np.ones(1)),
 47        "independent_commutation_d0": mass_norm(ct.D0 @ p0 - bad @ ct.D0, M0, M1),
 48        "independent_commutation_d1": mass_norm(ct.D1 @ bad, M1, np.ones(1)),
 49    }
 50
 51Q1 = torch.tensor(weighted_qr(ct.D0[:, :2], M1), dtype=torch.float32)
 52P1 = Q1 @ Q1.T @ torch.diag(torch.tensor(M1, dtype=torch.float32))
 53
 54
 55def seed_all(seed):
 56    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 57    if torch.cuda.is_available():
 58        try: torch.cuda.manual_seed_all(seed)
 59        except Exception: pass
 60
 61
 62def tensors(seed):
 63    d = ct.get_dataset(seed, NTR, NTE)
 64    for k in ("xtr", "ytr", "xte", "yte"):
 65        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 66    d["input_shape"] = tuple(d["xtr"].shape[1:])
 67    return d
 68
 69
 70class CochainNet(nn.Module):
 71    def __init__(self, projected):
 72        super().__init__()
 73        self.projected = projected
 74        self.node = nn.Linear(2, WIDTH)
 75        self.edge = nn.Linear(2, WIDTH)
 76        self.head = nn.Sequential(nn.Linear(6 * WIDTH, 32), nn.ReLU(), nn.Linear(32, 1))
 77
 78    def forward(self, x, return_hidden=False):
 79        hn = torch.relu(self.node(x[:, :3, :]))
 80        he = torch.relu(self.edge(x[:, 3:, :]))
 81        if self.projected:
 82            he = torch.einsum("ij,bjc->bic", P1.to(he.device), he)
 83        h = torch.cat((hn, he), dim=1).flatten(1)
 84        out = self.head(h)
 85        return (out, hn, he) if return_hidden else out
 86
 87
 88def run(kind, lr, seed, return_model=False):
 89    seed_all(seed)
 90    d = tensors(seed)
 91    net, metric, hist = train_model(CochainNet(kind == "idea"), d,
 92                                    epochs=EPOCHS, lr=float(lr), batch=128,
 93                                    log=lambda *a, **k: None)
 94    if return_model:
 95        return float(metric), net, d
 96    return float(metric)
 97
 98
 99def base_factory(cfg):
100    return lambda seed: run("baseline", cfg["lr"], seed)
101
102
103def idea_factory(cfg):
104    return lambda seed: run("idea", cfg["lr"], seed)
105
106
107def mechanism_signature():
108    # Retrain/evaluate actual benchmark systems, then measure the spatial
109    # derivative commutation defect on their learned hidden edge representations.
110    vals = {"baseline": [], "idea": []}
111    for seed in (0, 1, 2, 3):
112        for kind in vals:
113            _, net, d = run(kind, 0.003, seed, return_model=True)
114            net = net.to("cpu")
115            with torch.no_grad():
116                _, hn, he = net(d["xte"], return_hidden=True)
117                boundary = torch.einsum("ij,bjc->bic", torch.tensor(ct.D1, dtype=torch.float32), he)
118                vals[kind].append(float(torch.sqrt(torch.mean(boundary ** 2))))
119    b, i = np.mean(vals["baseline"]), np.mean(vals["idea"])
120    return {"prediction": "trained projected edge representations have lower face-boundary residual",
121            "predicted_direction": "idea < baseline",
122            "baseline_boundary_rms": float(b), "idea_boundary_rms": float(i),
123            "reduction_fraction": float((b - i) / b) if b else 0.0,
124            "confirmed": bool(i < b and np.isfinite(i) and np.isfinite(b))}
125
126
127def main():
128    checks = math_check()
129    grid = [{"lr": x} for x in LRS]
130    base = sweep_baseline(base_factory, grid, seeds=SWEEP_SEEDS)
131    trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
132    best = min(trials, key=lambda z: z["result"]["mean"])
133    # Full baseline at every idea-tested learning rate enforces search-space parity.
134    full_base = [{"cfg": c, "result": evaluate(base_factory(c), SEEDS)} for c in grid]
135    chosen_base = min(full_base, key=lambda z: z["result"]["mean"])
136    base["full"] = chosen_base["result"]
137    report = make_report("simplicial_cochain_regression", "cochain_mlp_tiny",
138                         base, best["result"], {
139                             "math_check": checks,
140                             "baseline_full_grid": full_base,
141                             "idea_sweep": trials,
142                             "idea_config": best["cfg"],
143                             "custom_track": {"name": "simplicial_cochain_regression",
144                                              "file": "cochain_track.py",
145                                              "domain": "simplicial_geometry_pde_like"},
146                             "mechanism_signature": mechanism_signature()})
147    Path("bench_report.json").write_text(json.dumps(report, indent=2))
148    print(json.dumps(report, indent=2))
149
150
151if __name__ == "__main__":
152    main()