Capacity-Preserving Transient Message Passing / run_experiment.py
Failed on benchmark
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()