Capacity-Preserving Transient Message Passing / run_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, os
  2import numpy as np
  3from scipy.sparse import coo_matrix, csr_matrix
  4from scipy.sparse.linalg import spsolve
  5
  6SEED = 1387
  7rng = np.random.default_rng(SEED)
  8
  9
 10def grid_graph(dim, R):
 11    """Finite integer lattice ball (L_inf ball), with Dirichlet boundary outside radius R."""
 12    pts = [tuple(x) for x in np.ndindex(*([2*R+1]*dim))]
 13    pts = [tuple(a-R for a in p) for p in pts]
 14    idx = {p:i for i,p in enumerate(pts)}
 15    edges=[]
 16    for p in pts:
 17        for d in range(dim):
 18            q=list(p); q[d]+=1; q=tuple(q)
 19            if q in idx: edges.append((idx[p],idx[q],1.0))
 20    return pts, idx, edges
 21
 22
 23def capacity_grid(dim, R, n=3, distance_weight=False):
 24    pts, idx, edges = grid_graph(dim,R)
 25    origin=idx[tuple([0]*dim)]
 26    # Dirichlet problem: f(origin)=1, all nodes on outer L_inf shell are 0.
 27    fixed={origin:1.0}
 28    for p,i in idx.items():
 29        if max(abs(x) for x in p)==R: fixed[i]=0.0
 30    free=[i for i in range(len(pts)) if i not in fixed]
 31    pos={i:j for j,i in enumerate(free)}
 32    rows=[]; cols=[]; vals=[]; b=np.zeros(len(free))
 33    adj=[[] for _ in pts]
 34    for i,j,_ in edges:
 35        # lattice edge geometric distance is one, hence w=d^(n-2)=1
 36        w=1.0
 37        adj[i].append((j,w)); adj[j].append((i,w))
 38    for i in free:
 39        r=pos[i]; deg=sum(w for _,w in adj[i]); rows.append(r); cols.append(r); vals.append(deg)
 40        for j,w in adj[i]:
 41            if j in fixed: b[r]+=w*fixed[j]
 42            else: rows.append(r); cols.append(pos[j]); vals.append(-w)
 43    f=np.zeros(len(pts)); f[list(fixed)]=list(fixed.values())
 44    if free: f[free]=spsolve(coo_matrix((vals,(rows,cols)),shape=(len(free),len(free))).tocsr(),b)
 45    # half ordered energy = sum undirected edge w*(du)^2
 46    energy=0.0
 47    for i,j,w in edges: energy += w*(f[i]-f[j])**2
 48    return float(energy)
 49
 50
 51def path_transition(N, self_loop=0.0):
 52    W=np.zeros((N,N))
 53    for i in range(N-1): W[i,i+1]=W[i+1,i]=1.0
 54    if self_loop: W += self_loop*np.eye(N)
 55    return W/W.sum(axis=1,keepdims=True)
 56
 57
 58def influence_scaling():
 59    # Large path, start away from boundaries; fit the known 1D local CLT prediction ||p_t||_2 ~ t^-1/4.
 60    N=4001; P=path_transition(N, self_loop=1.0); x=np.zeros(N); x[N//2]=1
 61    times=np.array([4,8,16,32,64,128,256,512,1024,2048])
 62    vals=[]; t=0
 63    for target in times:
 64        for _ in range(target-t): x=P@x
 65        t=target; vals.append(np.linalg.norm(x))
 66    slope=np.polyfit(np.log(times),np.log(vals),1)[0]
 67    return {'times':times.tolist(),'l2':np.round(vals,8).tolist(),'fitted_exponent':float(slope),'predicted_exponent':-0.25}
 68
 69
 70def diffusion_checks():
 71    # Distance-weighted P on a tiny irregular geometric graph. Verify stochasticity and stationary conservation.
 72    coords=np.array([[0.,0.],[1.,0.],[0.,2.],[2.,1.],[3.,0.]])
 73    E=[(0,1),(0,2),(1,3),(2,3),(1,4),(3,4)]
 74    n=4; W=np.zeros((len(coords),len(coords)))
 75    for i,j in E:
 76        d=np.linalg.norm(coords[i]-coords[j]); w=d**(n-2); W[i,j]=W[j,i]=w
 77    P=W/W.sum(1,keepdims=True)
 78    x=rng.normal(size=(len(coords),3)); y=P@x
 79    return {'max_row_sum_error':float(np.max(abs(P.sum(1)-1))), 'max_nonnegative_error':float(max(0,-P.min())), 'feature_mass_before':float(x.sum()), 'feature_mass_after_degree_weighted':float((W.sum(1)[:,None]*y).sum()), 'feature_mass_expected':float((W.sum(1)[:,None]*x).sum())}
 80
 81
 82def classification_demo():
 83    # Same graph/parameters: noisy coordinate labels, compare unweighted GCN and d^(n-2) diffusion.
 84    rs=np.random.default_rng(SEED); N=180; xy=rs.uniform(-1,1,(N,2)); labels=(xy[:,0]+0.25*xy[:,1]>0).astype(np.float32)
 85    # kNN graph, n=4 makes longer edges receive larger weights.
 86    D=((xy[:,None,:]-xy[None,:,:])**2).sum(2); W0=np.zeros((N,N))
 87    k=8
 88    for i in range(N):
 89        for j in np.argsort(D[i])[1:k+1]: W0[i,j]=W0[j,i]=1
 90    Wd=W0*np.maximum(D,1e-12) # n-2=2, d^2 = D
 91    def norm(W): return W/W.sum(1,keepdims=True)
 92    Ps=[norm(W0),norm(Wd)]
 93    X=np.stack([xy[:,0]+rs.normal(0,.7,N),xy[:,1]+rs.normal(0,.7,N)],1)
 94    out=[]
 95    for name,P in zip(['GCN_unweighted','capacity_diffusion_weighted'],Ps):
 96        h=X.copy(); acc=[]; var=[]
 97        # fixed linear readout, labels are the task; report best depth, not train a model.
 98        for depth in range(1,21):
 99            h=.5*h+.5*(P@h)
100            score=h[:,0]; pred=(score>0).astype(np.float32)
101            acc.append(float((pred==labels).mean())); var.append(float(np.mean(np.sum((h-h.mean(0))**2,1))))
102        out.append({'model':name,'best_accuracy':max(acc),'best_depth':int(np.argmax(acc)+1),'accuracy_depth20':acc[-1],'variance_depth1':var[0],'variance_depth20':var[-1]})
103    return out
104
105
106def main():
107    caps={}
108    for dim in [1,2,3]:
109        Rs=([4,8,16,32,64,128] if dim==1 else ([4,8,16,24,32,48,64] if dim==2 else [2,3,4,5,6,8,10]))
110        vals=[capacity_grid(dim,R) for R in Rs]
111        caps[str(dim)]={'R':Rs,'capacity':vals}
112    # Predictions: 1D Cap~R^-1; 2D Cap~1/log R. Fit slopes on sufficiently large radii.
113    r1=np.array(caps['1']['R'],float); c1=np.array(caps['1']['capacity']); slope1=np.polyfit(np.log(r1),np.log(c1),1)[0]
114    r2=np.array(caps['2']['R'],float); c2=np.array(caps['2']['capacity']); invlog=np.polyfit(1/np.log(r2),c2,1)
115    # In 3D, compare late/early capacity; predicted nonzero limiting capacity.
116    c3=np.array(caps['3']['capacity'])
117    result={'seed':SEED,'prediction_checks':{
118      'path_capacity_power_law':{'predicted_exponent':-1.0,'observed_exponent':float(slope1),'relative_error':float(abs(slope1+1))},
119      'grid2_capacity_log_law':{'prediction':'Cap approximately affine in 1/log(R)','R2':float(np.corrcoef(1/np.log(r2),c2)[0,1]**2),'capacity_ratio_R64_over_R4':float(c2[-1]/c2[0]),'log_ratio_prediction':float(np.log(r2[0])/np.log(r2[-1]))},
120      'grid3_capacity_nonvanishing':{'early':float(c3[0]),'late':float(c3[-1]),'late_over_early':float(c3[-1]/c3[0])},
121      'influence_local_clt':influence_scaling(),
122      'row_stochasticity':diffusion_checks()},'capacities':caps,'classification':classification_demo()}
123    with open('results.json','w') as f: json.dump(result,f,indent=2)
124    print(json.dumps(result,indent=2))
125
126if __name__=='__main__': main()