1import json, os, sys
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import train_model, evaluate, sweep_baseline, make_report
  8from bench.protocol import DEFAULT_SEEDS
  9from graph_track import get_dataset, META
 10
 11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 12
 13
 14class GeometricMessageNet(nn.Module):
 15    """Matched residual GNN; only the edge weighting changes."""
 16    def __init__(self, weighted=True, alpha=.5, hidden=32, depth=4):
 17        super().__init__()
 18        self.weighted, self.alpha = weighted, alpha
 19        self.layers = nn.ModuleList([nn.Linear(5, hidden)] +
 20                                     [nn.Linear(hidden, hidden) for _ in range(depth - 1)])
 21        self.head = nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1))
 22
 23    def propagation(self, x):
 24        xy = x[:, :, 2:4]
 25        d2 = ((xy[:, :, None] - xy[:, None, :]) ** 2).sum(-1)
 26        ids = d2.topk(6, dim=-1, largest=False).indices[:, :, 1:6]
 27        knn = torch.zeros_like(d2)
 28        knn.scatter_(2, ids, 1.0)
 29        knn = torch.maximum(knn, knn.transpose(1, 2))
 30        w = knn * (d2.clamp_min(1e-6) if self.weighted else 1.0)
 31        return w / w.sum(-1, keepdim=True).clamp_min(1e-8)
 32
 33    def forward(self, x):
 34        p = self.propagation(x)
 35        h = x
 36        for layer in self.layers:
 37            h = torch.relu(layer((1.0 - self.alpha) * h + self.alpha * torch.bmm(p, h)))
 38        root = x[:, :, 4:5]
 39        return self.head((h * root).sum(1) / root.sum(1).clamp_min(1.0))
 40
 41
 42def tensor_ds(seed):
 43    raw = get_dataset(seed, n_train=400, n_test=160)
 44    return {**raw, "xtr": torch.as_tensor(raw["xtr"], dtype=torch.float32),
 45            "ytr": torch.as_tensor(raw["ytr"], dtype=torch.float32).reshape(-1, 1),
 46            "xte": torch.as_tensor(raw["xte"], dtype=torch.float32),
 47            "yte": torch.as_tensor(raw["yte"], dtype=torch.float32).reshape(-1, 1)}
 48
 49
 50def train_one(seed, weighted, lr, alpha=.5, epochs=18, return_model=False):
 51    np.random.seed(seed); torch.manual_seed(seed)
 52    ds = tensor_ds(seed)
 53    net = GeometricMessageNet(weighted=weighted, alpha=alpha)
 54    trained, metric, hist = train_model(net, ds, epochs=epochs, lr=lr,
 55                                         batch=128, weight_decay=0.0, log=lambda *_: None)
 56    if trained is None:
 57        return (float("nan"), None, ds) if return_model else float("nan")
 58    return (float(metric), trained, ds) if return_model else float(metric)
 59
 60
 61def make_train_fn(cfg, weighted):
 62    return lambda seed: train_one(seed, weighted, cfg["lr"], cfg["alpha"])
 63
 64
 65def influence_signature(seed=0):
 66    result = {}
 67    for name, weighted in (("baseline", False), ("idea", True)):
 68        metric, net, ds = train_one(seed, weighted, .003, return_model=True)
 69        if net is None:
 70            result[name] = {"test_mse": float("nan"), "error": "training failed"}
 71            continue
 72        dev = next(net.parameters()).device
 73        x = ds["xte"][:1].to(dev).clone().requires_grad_(True)
 74        net.zero_grad(); net(x).sum().backward()
 75        observed = x.grad[0, :, 0].abs().detach().cpu().numpy(); observed /= observed.sum() + 1e-12
 76        with torch.no_grad():
 77            p = net.propagation(x.detach())[0].detach().cpu().numpy()
 78            root = int(np.argmax(ds["xte"][0, :, 4].numpy()))
 79            predicted = np.linalg.matrix_power(p, 4)[root]
 80            predicted /= predicted.sum() + 1e-12
 81        result[name] = {"test_mse": metric,
 82                        "predicted_l2": float(np.linalg.norm(predicted)),
 83                        "observed_gradient_l2": float(np.linalg.norm(observed)),
 84                        "predicted_observed_corr": float(np.corrcoef(observed, predicted)[0, 1])}
 85    b, i = result["baseline"], result["idea"]
 86    result["prediction"] = "distance weighting preserves more long-range influence"
 87    result["confirmed"] = bool(i["observed_gradient_l2"] > b["observed_gradient_l2"] + .01 and
 88                                i["predicted_observed_corr"] > .5)
 89    return result
 90
 91
 92def main():
 93    grid = [{"lr": lr, "alpha": .5} for lr in (.001, .003, .01)]
 94    base = sweep_baseline(lambda cfg: make_train_fn(cfg, False), grid)
 95    trials = [{"cfg": cfg, "result": evaluate(make_train_fn(cfg, True), DEFAULT_SEEDS)} for cfg in grid]
 96    best = min(trials, key=lambda z: z["result"]["mean"])
 97    rep = make_report("geometric_graph_diffusion", "GeometricMessageNet", base, best["result"],
 98                      {"track_structure": "geometric graph node-field diffusion regression",
 99                       "trained_model_measurements": influence_signature(),
100                       "idea_sweep": trials,
101                       "confirmed": False})
102    rep["custom_track"] = {"name": META["name"], "file": "graph_track.py", "domain": META["domain"]}
103    rep["protocol_notes"] = {"paired_seeds": list(DEFAULT_SEEDS),
104                             "baseline_and_idea_same_architecture": True,
105                             "only_intervention": "edge weight d^2 versus unit edge weight",
106                             "device": DEVICE}
107    os.makedirs("artifacts", exist_ok=True)
108    with open("artifacts/bench_report.json", "w") as f: json.dump(rep, f, indent=2)
109    print(json.dumps(rep, indent=2))
110
111
112if __name__ == "__main__": main()