import json, math, random import numpy as np import torch from torch import nn SEED = 17 random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) try: # CUDA is attempted, with CPU fallback below device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') except Exception: device = torch.device('cpu') # Fixed sparse directed cycle; incidence-like signed aggregation is retained. N, D = 20, 2 src = torch.arange(N, device=device) dst = (src + 1) % N # A second, non-adjacent coupling is part of the synthetic ground truth. far = (torch.arange(N, device=device) + N//2) % N def mechanism_checks(): # Exact scalar linearization z' = gamma*Lambda*z. # Prediction 1: contraction/divergence boundary gamma*Lambda=1. products = np.array([0.60, 0.80, 0.95, 1.00, 1.05, 1.20]) observed_rho, observed_ratio = [], [] for p in products: J = np.array([[p]]) observed_rho.append(float(abs(np.linalg.eigvals(J)[0]))) z = 1.0 hist = [] for _ in range(14): z *= p; hist.append(abs(z)) observed_ratio.append(float(hist[-1]/hist[-2])) # Prediction 2: perturbation at k is (gamma Lambda)^k. p = 0.80; k = 12 z = 1.0 for _ in range(k): z *= p predicted = p**k # Prediction 3: nonlocal contribution is linear in coupling alpha. alphas = np.array([0., .1, .2, .4, .8]) summary = 1.7 measured = np.abs(alphas * summary) slope = float(np.polyfit(alphas, measured, 1)[0]) return { 'boundary': {'predicted_product': 1.0, 'products': products.tolist(), 'observed_jacobian_rho': observed_rho, 'observed_step_ratio': observed_ratio, 'classification': ['contractive' if p < 1 else ('neutral' if p == 1 else 'divergent') for p in products]}, 'geometric_decay': {'product': p, 'steps': k, 'predicted': predicted, 'observed': float(z), 'relative_error': float(abs(z-predicted)/(predicted+1e-12))}, 'nonlocal_linearity': {'alphas': alphas.tolist(), 'predicted_slope': summary, 'observed_slope': slope, 'relative_slope_error': abs(slope-summary)/summary} } class LocalGNN(nn.Module): def __init__(self, hidden=32): super().__init__() self.edge = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D)) self.node = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D)) def forward(self, x): # Supports [N,D] and [B,N,D]. e = self.edge(torch.cat([x[...,src,:], x[...,dst,:]], -1)) q = torch.zeros_like(x).index_add(-2, src, e).index_add(-2, dst, -e) return self.node(torch.cat([x,q], -1)) class RecursiveFeedbackGNN(nn.Module): def __init__(self, hidden=32): super().__init__(); self.hidden=hidden self.ctx = nn.GRUCell(2*D + D, hidden) self.u = nn.Sequential(nn.Linear(2*D+hidden+D, hidden), nn.Tanh(), nn.Linear(hidden, 1)) self.edge = nn.Sequential(nn.Linear(2*D+1, hidden), nn.Tanh(), nn.Linear(hidden, D)) self.node = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D)) def forward(self, x, z): # Supports [N,D] and [B,N,D], with one context per edge/node. r=x.mean(-2, keepdim=True).expand_as(x) inp=torch.cat([x[...,src,:],x[...,dst,:],r],-1) shape=inp.shape z=self.ctx(inp.reshape(-1,shape[-1]), z.reshape(-1,z.shape[-1])).reshape(*shape[:-1],self.hidden) u=self.u(torch.cat([x[...,src,:],x[...,dst,:],z,r[...,src,:]],-1)) e=self.edge(torch.cat([x[...,src,:],x[...,dst,:],u],-1)) q=torch.zeros_like(x).index_add(-2,src,e).index_add(-2,dst,-e) return self.node(torch.cat([x,q],-1)), z def truth(x, alpha=.65): # Local ring dynamics plus a distant, same-channel interaction. local=torch.roll(x,-1,0)+torch.roll(x,1,0)-2*x nonlocal_term=torch.roll(x,N//2,0)-x return x + .18*local + alpha*.18*nonlocal_term def make_data(n=320): X=[]; Y=[] for _ in range(n): x=torch.randn(N,D,device=device) X.append(x); Y.append(truth(x)) return torch.stack(X),torch.stack(Y) def train(model, X, Y, feedback=False, steps=220): opt=torch.optim.Adam(model.parameters(),lr=3e-3) model.train() for step in range(steps): i=torch.randint(0,len(X),(min(64,len(X)),),device=device) xb=X[i] if feedback: z=torch.zeros(xb.shape[0],N,model.hidden,device=device) pred,_=model(xb,z) else: pred=model(xb) loss=((pred-Y[i])**2).mean() opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),2.0); opt.step() return model def evaluate(model,X,Y,feedback=False): model.eval(); x=X[:100].clone(); z=torch.zeros(x.shape[0],N,model.hidden,device=device) if feedback else None with torch.no_grad(): one=None for k in range(8): if feedback: x,z=model(x,z) else: x=model(x) if k==0: one=((x-Y[:100])**2).mean().sqrt().item() long=((x-Y[:100])**2).mean().sqrt().item() return one,long def jacobian_rho(model, feedback=False): x=torch.randn(N,D,device=device,requires_grad=True) if feedback: z=torch.zeros(N,model.hidden,device=device) def fn(v): return model(v,z)[0].reshape(-1) else: def fn(v): return model(v).reshape(-1) J=torch.autograd.functional.jacobian(fn,x).detach().cpu().numpy().reshape(N*D,N*D) return float(np.max(np.abs(np.linalg.eigvals(J)))) def main(): checks=mechanism_checks() X,Y=make_data(180) local=LocalGNN().to(device); rec=RecursiveFeedbackGNN().to(device) train(local,X,Y,False); train(rec,X,Y,True) lm= evaluate(local,X,Y,False); rm=evaluate(rec,X,Y,True) result={'device':str(device),'seed':SEED,'mechanism_checks':checks, 'mini_experiment':{'local_sparse':{'one_step_rmse':lm[0],'8_step_rmse':lm[1],'rho':jacobian_rho(local)}, 'recursive_nonlocal_feedback':{'one_step_rmse':rm[0],'8_step_rmse':rm[1],'rho':jacobian_rho(rec,True)}, 'parameter_counts':{'local':sum(p.numel() for p in local.parameters()),'recursive':sum(p.numel() for p in rec.parameters())}}} with open('results.json','w') as f: json.dump(result,f,indent=2) print(json.dumps(result,indent=2)) if __name__=='__main__': main()