1import sys, json
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6import bench
 7from bench import train_model, evaluate, sweep_baseline, make_report
 8
 9TRACK = 'geometric_graph_diffusion'
10
11def graph_operator(x):
12    # x[...,2:4] are the observed node coordinates in the registered track.
13    xy = x[..., 2:4]
14    d2 = ((xy[:, :, None, :] - xy[:, None, :, :]) ** 2).sum(-1)
15    b, n, _ = d2.shape
16    adj = torch.zeros_like(d2)
17    # Five nearest neighbors, symmetrized, matching the registered generator.
18    ind = torch.argsort(d2, dim=-1)[:, :, 1:6]
19    adj.scatter_(2, ind, 1.0)
20    adj = torch.maximum(adj, adj.transpose(1, 2))
21    w = adj * torch.clamp(d2, min=1e-5)
22    p = w / torch.clamp(w.sum(-1, keepdim=True), min=1e-8)
23    return 0.55 * torch.eye(n, device=x.device).unsqueeze(0) + 0.45 * p
24
25class GraphDiffusionNet(nn.Module):
26    # The two systems share every learnable parameter and differ only in alpha.
27    def __init__(self, alpha):
28        super().__init__()
29        self.alpha = float(alpha)
30        self.fc1 = nn.Linear(5, 32)
31        self.fc2 = nn.Linear(32, 1)
32    def forward(self, x):
33        q = graph_operator(x)
34        if self.alpha:
35            n = x.shape[1]
36            q = (1.0-self.alpha)*q + self.alpha*torch.ones_like(q)/n
37        h = torch.relu(self.fc1(x))
38        h = torch.bmm(q, h)
39        z = self.fc2(h).squeeze(-1)
40        # Root marker makes this a scalar graph readout, as required by target.
41        root = x[..., 4]
42        return (z * root).sum(1, keepdim=True)
43
44def run(cfg, seed, return_model=False):
45    torch.manual_seed(9000 + int(seed)); np.random.seed(9000 + int(seed))
46    ds = bench.get_dataset(TRACK, int(seed), n_train=400, n_test=160)
47    net = GraphDiffusionNet(cfg['alpha'])
48    net, metric, history = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'],
49                                       batch=128, weight_decay=cfg['wd'], log=lambda *_: None)
50    return (float(metric), net, ds) if return_model else float(metric)
51
52def factory(cfg):
53    return lambda seed: run(cfg, seed)
54
55def main():
56    # Shared union of all learning-rate/regularization settings on both sides.
57    common = [(lr, wd) for lr in (0.003, 0.01, 0.03) for wd in (0.0, 0.0001)]
58    baseline_grid = [{'alpha': 0.0, 'lr': lr, 'wd': wd, 'epochs': 20} for lr,wd in common]
59    baseline = sweep_baseline(factory, baseline_grid, seeds=(0,1,2,3))
60    best = baseline['best_cfg']
61    idea_grid = []
62    for alpha in (0.03, 0.1, 0.3):
63        c = dict(best); c['alpha'] = alpha; idea_grid.append(c)
64    idea_candidates = [{'cfg': c, 'result': evaluate(factory(c))} for c in idea_grid]
65    chosen = min(idea_candidates, key=lambda z: z['result']['mean'])
66    idea = chosen['result']
67
68    # Behavioural signature from trained systems, not an analytical toy graph.
69    rows=[]
70    for alpha in (0.0, 0.03, 0.1, 0.3):
71        cfg=dict(best); cfg['alpha']=alpha
72        _, net, ds = run(cfg, 0, return_model=True)
73        net.eval(); x=ds['xte'].to(next(net.parameters()).device)
74        with torch.no_grad():
75            q=graph_operator(x)
76            h=torch.relu(net.fc1(x)); hc=h-h.mean(1,keepdim=True)
77            before=hc.norm().item()
78            n=x.shape[1]; qt=(1-alpha)*q + alpha*torch.ones_like(q)/n
79            after=torch.bmm(qt,hc).norm().item()
80        rows.append({'alpha':alpha,'observed_mean_centered_norm_ratio':after/max(before,1e-12), 'predicted':1-alpha})
81    err=max(abs(r['observed_mean_centered_norm_ratio']-r['predicted']) for r in rows)
82    extra={'prediction':'uniform teleportation contracts mean-zero feature modes approximately by 1-alpha',
83           'trained_model_rows':rows,'max_abs_error':float(err),'confirmed':bool(err < 0.10),
84           'idea_sweep':idea_candidates,
85           'track_match':'registered geometric_graph_diffusion: node-field regression on geometric kNN graphs'}
86    report=make_report(TRACK,'graph_diffusion_net',baseline,idea,extra)
87    report['idea_candidates']=idea_candidates
88    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
89    print(json.dumps(report,indent=2))
90
91if __name__=='__main__': main()