1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7
  8SEEDS = [0, 1, 2]
  9N = 12
 10
 11def laplacian(w):
 12    return torch.diag(w.sum(-1)) - w
 13
 14def tree_loss(w):
 15    n = w.shape[-1]
 16    L = laplacian(w)
 17    sign, ld = torch.linalg.slogdet(L + torch.ones_like(L) / n)
 18    if sign.item() <= 0:
 19        raise RuntimeError('non-positive shifted Laplacian')
 20    return -ld + math.log(n)
 21
 22def spectral_metrics(w):
 23    a = np.asarray(w)
 24    L = np.diag(a.sum(1)) - a
 25    ev = np.linalg.eigvalsh(L)
 26    return {
 27        'pseudo_det': float(np.prod(ev[1:])),
 28        'fiedler': float(ev[1]),
 29        'components': int(np.sum(ev < 1e-7)),
 30        'eigenvalues': ev.tolist(),
 31    }
 32
 33def math_check():
 34    rng = np.random.default_rng(4)
 35    x = rng.uniform(.1, 1.2, (N, N))
 36    w = (x + x.T) / 2
 37    np.fill_diagonal(w, 0)
 38    L = np.diag(w.sum(1)) - w
 39    ev = np.linalg.eigvalsh(L)
 40    pdet = float(np.prod(ev[1:]))
 41    rankone_det = float(np.linalg.det(L + np.ones((N, N)) / N))
 42    cofactor = float(np.linalg.det(L[:-1, :-1]))
 43    eps = 1e-4
 44    derivative = (np.linalg.det(L + eps*np.eye(N)) - np.linalg.det(L - eps*np.eye(N))) / (2*eps)
 45    wt = torch.tensor(w, dtype=torch.double, requires_grad=True)
 46    tl = tree_loss(wt)
 47    tl.backward()
 48    return {
 49        'pseudo_det_vs_rank_one_det_relerr': abs(pdet-rankone_det)/pdet,
 50        'matrix_tree_pdet_over_N_vs_cofactor_relerr': abs(pdet/N-cofactor)/cofactor,
 51        'derivative_formula_relerr': abs(derivative-pdet)/pdet,
 52        'tree_loss_vs_stated_formula_abs': abs(float(tl)-(-math.log(pdet)+math.log(N))),
 53        'autograd_gradient_norm': float(wt.grad.norm()),
 54    }
 55
 56class EdgeLearner(nn.Module):
 57    def __init__(self, n):
 58        super().__init__()
 59        self.logits = nn.Parameter(torch.randn(n, n) * .15)
 60    def weights(self):
 61        s = (self.logits + self.logits.T) / 2
 62        w = F.softplus(s)
 63        return w * (1 - torch.eye(w.shape[0], device=w.device))
 64
 65def make_target(seed):
 66    rng = np.random.default_rng(seed)
 67    # Two-community target: the learning objective must preserve task-relevant structure.
 68    y = np.zeros((N, N), dtype=np.float32)
 69    groups = np.arange(N) < N//2
 70    for i in range(N):
 71        for j in range(i):
 72            base = 1.0 if groups[i] == groups[j] else .08
 73            y[i,j] = y[j,i] = base + rng.normal(0, .025)
 74    np.fill_diagonal(y, 0)
 75    return torch.tensor(y)
 76
 77def run_one(seed, alpha, device):
 78    torch.manual_seed(seed); np.random.seed(seed)
 79    target = make_target(seed).to(device)
 80    model = EdgeLearner(N).to(device)
 81    opt = torch.optim.Adam(model.parameters(), lr=.08)
 82    losses=[]; t0=time.perf_counter()
 83    for step in range(250):
 84        w = model.weights()
 85        fit = ((w-target)**2).mean()
 86        reg = tree_loss(w)
 87        loss = fit + alpha*reg
 88        opt.zero_grad(); loss.backward(); opt.step()
 89        losses.append(float(loss.detach().cpu()))
 90    elapsed=time.perf_counter()-t0
 91    w=model.weights().detach().cpu().numpy()
 92    met=spectral_metrics(w)
 93    met.update({'fit_mse': float(((torch.tensor(w)-target.cpu())**2).mean()),
 94                'final_loss': losses[-1], 'loss_std_last50': float(np.std(losses[-50:])),
 95                'seconds': elapsed})
 96    return met
 97
 98def main():
 99    device='cuda' if torch.cuda.is_available() else 'cpu'
100    try:
101        # Probe CUDA and fall back on any runtime/device error.
102        if device == 'cuda': torch.zeros(1, device='cuda').sum().item()
103    except Exception:
104        device='cpu'
105    result={'device':device, 'math_check':math_check(), 'runs':{}}
106    for alpha in [0.0, 1e-3, 1e-2]:
107        key='baseline' if alpha==0 else 'tree_alpha_'+str(alpha)
108        vals=[]
109        for seed in SEEDS:
110            try: vals.append(run_one(seed, alpha, device))
111            except Exception:
112                if device=='cuda':
113                    device='cpu'; result['device']='cpu'; vals.append(run_one(seed, alpha, device))
114                else: raise
115        result['runs'][key]=vals
116    for key, vals in result['runs'].items():
117        result['summary_'+key]={k:float(np.mean([v[k] for v in vals])) for k in ['fit_mse','fiedler','components','pseudo_det','loss_std_last50','seconds']}
118    Path('results.json').write_text(json.dumps(result, indent=2))
119    print(json.dumps(result, indent=2))
120
121if __name__=='__main__': main()