Volume-Mass Diffusion GNN / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED=1150
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def sbm(n1=60,n2=60,p_in=.22,p_out=.025):
 10    n=n1+n2
 11    y=np.r_[np.zeros(n1,dtype=int),np.ones(n2,dtype=int)]
 12    R=np.random.rand(n,n)
 13    same=y[:,None]==y[None,:]
 14    A=((R < np.where(same,p_in,p_out))).astype(float)
 15    A=np.triu(A,1); A=A+A.T; np.fill_diagonal(A,0)
 16    # ensure connected enough; this seed graph is connected
 17    return A,y
 18
 19def matrices(A,v):
 20    d=A.sum(1); L=np.diag(d)-A
 21    inv=np.diag(1/np.sqrt(v)); S=inv@L@inv
 22    B=np.diag(1/v)@L
 23    return d,L,S,B
 24
 25def ipr(u): return float(np.sum(u**4)/np.sum(u**2)**2)
 26
 27def math_checks(A):
 28    d,L,S,B=matrices(A,np.ones(len(A)))
 29    lam,U=np.linalg.eigh(S); lmax=lam[-1]
 30    # Prediction 1: Euler spectral boundary eta_c=2/lambda_max.
 31    pred_eta=2/lmax
 32    etas=pred_eta*np.linspace(.80,1.20,161)
 33    rho=np.array([max(abs(1-e*lam)) for e in etas])
 34    
 35    # first grid point whose spectral radius exceeds one, ignoring eta=0
 36    crossing_idx=np.where(rho>1+1e-10)[0]
 37    cross=etas[crossing_idx[0]] if len(crossing_idx) else float('nan')
 38    # Prediction 2: exact modal decay/amplification after m steps.
 39    eta=.65*pred_eta; m=17; j=len(lam)-1
 40    x=U[:,j]; measured=np.linalg.norm(np.linalg.matrix_power(np.eye(len(A))-eta*S,m)@x)
 41    predicted=abs(1-eta*lam[j])**m
 42    # Prediction 3: volume heterogeneity/localization. use v=k^alpha, and independent lognormal.
 43    rows=[]
 44    for alpha in [-1.0,-.5,0,.5,1.0,1.5,2.0]:
 45        v=np.maximum(d,1e-3)**alpha
 46        _,_,Ss,_=matrices(A,v); ee,uu=np.linalg.eigh(Ss)
 47        r=d/v
 48        rows.append({'kind':'degree_power','alpha':alpha,'std_log_r':float(np.std(np.log(r))), 'ipr_top':ipr(uu[:,-1])})
 49    # Replicate independent lognormal volumes: the prediction is a positive
 50    # association between ratio heterogeneity and extremal-mode IPR.
 51    for s in [0,.25,.5,.75,1.0,1.25]:
 52        vals=[]
 53        for rep in range(8):
 54            rng=np.random.RandomState(SEED+1000+rep)
 55            v=np.exp(s*rng.randn(len(A))); _,_,Ss,_=matrices(A,v); ee,uu=np.linalg.eigh(Ss)
 56            r=d/v
 57            vals.append((np.std(np.log(r)),ipr(uu[:,-1])))
 58        rows.append({'kind':'lognormal_replicated','s':s,
 59                     'std_log_r_mean':float(np.mean([z[0] for z in vals])),
 60                     'ipr_top_mean':float(np.mean([z[1] for z in vals])),
 61                     'ipr_top_std':float(np.std([z[1] for z in vals]))})
 62    rep_rows=[z for z in rows if z['kind']=='lognormal_replicated']
 63    localization_corr=float(np.corrcoef([z['std_log_r_mean'] for z in rep_rows], [z['ipr_top_mean'] for z in rep_rows])[0,1])
 64    # Similarity/eigenvalue sanity
 65    simerr=np.max(np.abs(np.sort(np.real(np.linalg.eigvals(B)))-lam))
 66    return {'lambda_max':float(lmax),'predicted_eta_c':pred_eta,'observed_eta_crossing':float(cross),
 67            'boundary_relative_error':float(abs(cross-pred_eta)/pred_eta),
 68            'decay_eta':eta,'decay_steps':m,'decay_predicted':float(predicted),'decay_measured':float(measured),
 69            'similarity_eigenvalue_max_error':float(simerr),'localization':rows,'localization_correlation':localization_corr}
 70
 71def train_experiment(A,y,device='cpu'):
 72    try:
 73        import torch
 74        torch.manual_seed(SEED)
 75        if device=='cuda': torch.cuda.manual_seed_all(SEED)
 76        dev=torch.device(device)
 77        n=len(y); X=torch.eye(n,device=dev); Y=torch.tensor(y,device=dev)
 78        At=torch.tensor(A,dtype=torch.float32,device=dev)
 79        d=At.sum(1); L=torch.diag(d)-At
 80        # standard normalized GCN propagation, with self loops
 81        An=At+torch.eye(n,device=dev); dn=An.sum(1); Pn=An/dn[:,None]
 82        # mass diffusion: stable step and self-residual, same 2-layer width
 83        v=torch.tensor(np.maximum(d.cpu().numpy(),1e-3)**1.0,dtype=torch.float32,device=dev)
 84        B=L/v[:,None]
 85        lam=torch.linalg.eigvalsh(torch.diag(1/torch.sqrt(v))@[email protected](1/torch.sqrt(v)))[-1]
 86        eta=.8*2/lam
 87        Pm=torch.eye(n,device=dev)-eta*B
 88        def run(kind):
 89            torch.manual_seed(SEED+3)
 90            W1=torch.nn.Parameter(torch.randn(n,16,device=dev)*.08); W2=torch.nn.Parameter(torch.randn(16,2,device=dev)*.08)
 91            opt=torch.optim.Adam([W1,W2],lr=.03,weight_decay=1e-3)
 92            idx=torch.randperm(n,device=dev); tr=idx[:80]; te=idx[80:]
 93            vals=[]
 94            for step in range(250):
 95                opt.zero_grad()
 96                P=Pn if kind=='gcn' else Pm
 97                h=torch.relu(P@X@W1); z=P@h@W2
 98                loss=torch.nn.functional.cross_entropy(z[tr],Y[tr]); loss.backward(); opt.step()
 99                if step in (49,99,249): vals.append(float(loss.detach().cpu()))
100            with torch.no_grad():
101                h=torch.relu(P@X@W1); z=P@h@W2; acc=float((z[te].argmax(1)==Y[te]).float().mean().cpu())
102            return {'loss_steps_50_100_250':vals,'test_accuracy':acc}
103        return {'device':str(dev),'gcn':run('gcn'),'mass_diffusion':run('mass')}
104    except Exception as e:
105        return {'error':repr(e)}
106
107def main():
108    A,y=sbm(); checks=math_checks(A)
109    try:
110        import torch
111        device='cuda' if torch.cuda.is_available() else 'cpu'
112    except Exception: device='cpu'
113    exp=train_experiment(A,y,device)
114    out={'seed':SEED,'n':len(y),'math':checks,'experiment':exp}
115    Path('results.json').write_text(json.dumps(out,indent=2))
116    print(json.dumps(out,indent=2))
117
118if __name__=='__main__': main()