1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, evaluate, sweep_baseline, make_report
  9
 10N = 8
 11EDGE = np.triu_indices(N, 1)
 12SEEDS = tuple(range(8))
 13
 14
 15def tree_loss(w):
 16    b, n, _ = w.shape
 17    deg = w.sum(-1)
 18    lap = torch.diag_embed(deg) - w
 19    shift = torch.ones((b, n, n), device=w.device, dtype=w.dtype) / n
 20    sign, ld = torch.linalg.slogdet(lap + shift)
 21    if torch.any(sign <= 0):
 22        raise RuntimeError('non-positive shifted Laplacian')
 23    return (-ld + math.log(n)).mean()
 24
 25
 26def math_check():
 27    rng = np.random.default_rng(91)
 28    a = rng.uniform(.1, 1., (N, N)); w = (a + a.T) / 2
 29    np.fill_diagonal(w, 0)
 30    lap = np.diag(w.sum(1)) - w
 31    ev = np.linalg.eigvalsh(lap)
 32    pdet = np.prod(ev[1:])
 33    rdet = np.linalg.det(lap + np.ones((N, N)) / N)
 34    cof = np.linalg.det(lap[:-1, :-1])
 35    return {'relative_rank_one_pdet_error': float(abs(pdet-rdet)/pdet),
 36            'relative_matrix_tree_cofactor_error': float(abs(pdet/N-cof)/cof)}
 37
 38
 39class GraphNet(nn.Module):
 40    def __init__(self):
 41        super().__init__()
 42        self.body = nn.Sequential(nn.Flatten(), nn.Linear(16, 64), nn.ReLU(),
 43                                  nn.Linear(64, 64), nn.ReLU())
 44        self.edge = nn.Linear(64, 28)
 45
 46    def forward(self, x):
 47        z = self.body(x)
 48        e = F.softplus(self.edge(z))
 49        w = torch.zeros((x.shape[0], N, N), device=x.device, dtype=x.dtype)
 50        w[:, EDGE[0], EDGE[1]] = e
 51        w[:, EDGE[1], EDGE[0]] = e
 52        return w, e.mean(-1, keepdim=True)
 53
 54
 55def spectral_metrics(w):
 56    fiedler, comps, pdet = [], [], []
 57    for a in w:
 58        lap = np.diag(a.sum(1)) - a
 59        ev = np.linalg.eigvalsh(lap)
 60        fiedler.append(float(ev[1]))
 61        comps.append(int(np.sum(ev < 1e-7)))
 62        pdet.append(float(np.prod(np.maximum(ev[1:], 1e-30))))
 63    return {'fiedler_mean': float(np.mean(fiedler)),
 64            'components_mean': float(np.mean(comps)),
 65            'pseudo_det_mean': float(np.mean(pdet))}
 66
 67
 68def train_one(seed, alpha, lr, epochs=35, return_model=False):
 69    torch.manual_seed(seed); np.random.seed(seed)
 70    d = get_dataset('soft_graph_connectivity', seed, 400, 120)
 71    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 72    try:
 73        if device == 'cuda': torch.zeros(1, device=device).sum().item()
 74    except Exception:
 75        device = 'cpu'
 76    model = GraphNet().to(device)
 77    xtr, ytr = torch.as_tensor(d['xtr'], dtype=torch.float32, device=device), torch.as_tensor(d['ytr'], dtype=torch.float32, device=device)
 78    xte, yte = torch.as_tensor(d['xte'], dtype=torch.float32, device=device), torch.as_tensor(d['yte'], dtype=torch.float32, device=device)
 79    opt = torch.optim.Adam(model.parameters(), lr=lr)
 80    for _ in range(epochs):
 81        model.train()
 82        for start in range(0, len(xtr), 128):
 83            w, pred = model(xtr[start:start+128])
 84            loss = F.mse_loss(pred, ytr[start:start+128])
 85            if alpha:
 86                loss = loss + alpha * tree_loss(w)
 87            opt.zero_grad(); loss.backward(); opt.step()
 88    model.eval()
 89    with torch.no_grad():
 90        w, pred = model(xte)
 91        metric = float(F.mse_loss(pred, yte).cpu())
 92    if return_model:
 93        return metric, spectral_metrics(w.detach().cpu().numpy()), model
 94    return metric
 95
 96
 97def make_fn(cfg):
 98    return lambda seed: train_one(seed, cfg['alpha'], cfg['lr'])
 99
100
101def main():
102    # Union parity: every idea lr is also present in baseline sweep.
103    grid = [{'alpha': 0.0, 'lr': lr} for lr in (1e-3, 3e-3, 1e-2, 2e-3, 6e-3)]
104    base = sweep_baseline(make_fn, grid)
105    # The baseline sweep itself selects alpha=0; idea is evaluated at same lrs.
106    idea_lrs = sorted(set([base['best_cfg']['lr'], 2e-3, 6e-3]))
107    idea_cfgs = [{'alpha': a, 'lr': lr} for lr in idea_lrs for a in (1e-4, 1e-3, 1e-2)]
108    # Run all three on the full paired seeds, then retain best by the same 4-seed selection.
109    idea_sweep = []
110    for cfg in idea_cfgs:
111        r4 = evaluate(make_fn(cfg), (0,1,2,3))
112        idea_sweep.append({'cfg': cfg, 'mean': r4['mean']})
113    best_idea_cfg = min(idea_sweep, key=lambda z: z['mean'])['cfg']
114    idea = evaluate(make_fn(best_idea_cfg), SEEDS)
115    base['idea_side_sweep'] = idea_sweep
116    sig_base=[]; sig_idea=[]
117    for s in SEEDS:
118        _, mb, _ = train_one(s, 0.0, base['best_cfg']['lr'], return_model=True)
119        _, mi, _ = train_one(s, best_idea_cfg['alpha'], best_idea_cfg['lr'], return_model=True)
120        sig_base.append(mb); sig_idea.append(mi)
121    signature = {
122        'quantity': 'Fiedler eigenvalue of trained predicted adjacency',
123        'baseline_predicted_fiedler_mean': float(np.mean([z['fiedler_mean'] for z in sig_base])),
124        'idea_predicted_fiedler_mean': float(np.mean([z['fiedler_mean'] for z in sig_idea])),
125        'baseline_components_mean': float(np.mean([z['components_mean'] for z in sig_base])),
126        'idea_components_mean': float(np.mean([z['components_mean'] for z in sig_idea])),
127        'prediction': 'tree regularization should increase Fiedler connectivity',
128        'confirmed': bool(np.mean([z['fiedler_mean'] for z in sig_idea]) > np.mean([z['fiedler_mean'] for z in sig_base]))
129    }
130    report = make_report('soft_graph_connectivity', 'graphnet_custom', base, idea,
131                         {'mechanism_signature': signature,
132                          'custom_track': {'name': 'soft_graph_connectivity', 'file': 'graph_track.py', 'domain': 'graph-nn'},
133                          'math_check': math_check(), 'idea_cfg': best_idea_cfg})
134    Path('report.json').write_text(json.dumps(report, indent=2))
135    print(json.dumps(report, indent=2))
136
137if __name__ == '__main__':
138    main()