1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 17
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    # CUDA is attempted, with CPU fallback below
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device = torch.device('cpu')
 13
 14# Fixed sparse directed cycle; incidence-like signed aggregation is retained.
 15N, D = 20, 2
 16src = torch.arange(N, device=device)
 17dst = (src + 1) % N
 18# A second, non-adjacent coupling is part of the synthetic ground truth.
 19far = (torch.arange(N, device=device) + N//2) % N
 20
 21def mechanism_checks():
 22    # Exact scalar linearization z' = gamma*Lambda*z.
 23    # Prediction 1: contraction/divergence boundary gamma*Lambda=1.
 24    products = np.array([0.60, 0.80, 0.95, 1.00, 1.05, 1.20])
 25    observed_rho, observed_ratio = [], []
 26    for p in products:
 27        J = np.array([[p]])
 28        observed_rho.append(float(abs(np.linalg.eigvals(J)[0])))
 29        z = 1.0
 30        hist = []
 31        for _ in range(14):
 32            z *= p; hist.append(abs(z))
 33        observed_ratio.append(float(hist[-1]/hist[-2]))
 34    # Prediction 2: perturbation at k is (gamma Lambda)^k.
 35    p = 0.80; k = 12
 36    z = 1.0
 37    for _ in range(k): z *= p
 38    predicted = p**k
 39    # Prediction 3: nonlocal contribution is linear in coupling alpha.
 40    alphas = np.array([0., .1, .2, .4, .8])
 41    summary = 1.7
 42    measured = np.abs(alphas * summary)
 43    slope = float(np.polyfit(alphas, measured, 1)[0])
 44    return {
 45        'boundary': {'predicted_product': 1.0, 'products': products.tolist(),
 46                     'observed_jacobian_rho': observed_rho, 'observed_step_ratio': observed_ratio,
 47                     'classification': ['contractive' if p < 1 else ('neutral' if p == 1 else 'divergent') for p in products]},
 48        'geometric_decay': {'product': p, 'steps': k, 'predicted': predicted, 'observed': float(z),
 49                            'relative_error': float(abs(z-predicted)/(predicted+1e-12))},
 50        'nonlocal_linearity': {'alphas': alphas.tolist(), 'predicted_slope': summary,
 51                               'observed_slope': slope, 'relative_slope_error': abs(slope-summary)/summary}
 52    }
 53
 54class LocalGNN(nn.Module):
 55    def __init__(self, hidden=32):
 56        super().__init__()
 57        self.edge = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D))
 58        self.node = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D))
 59    def forward(self, x):
 60        # Supports [N,D] and [B,N,D].
 61        e = self.edge(torch.cat([x[...,src,:], x[...,dst,:]], -1))
 62        q = torch.zeros_like(x).index_add(-2, src, e).index_add(-2, dst, -e)
 63        return self.node(torch.cat([x,q], -1))
 64
 65class RecursiveFeedbackGNN(nn.Module):
 66    def __init__(self, hidden=32):
 67        super().__init__(); self.hidden=hidden
 68        self.ctx = nn.GRUCell(2*D + D, hidden)
 69        self.u = nn.Sequential(nn.Linear(2*D+hidden+D, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 70        self.edge = nn.Sequential(nn.Linear(2*D+1, hidden), nn.Tanh(), nn.Linear(hidden, D))
 71        self.node = nn.Sequential(nn.Linear(2*D, hidden), nn.Tanh(), nn.Linear(hidden, D))
 72    def forward(self, x, z):
 73        # Supports [N,D] and [B,N,D], with one context per edge/node.
 74        r=x.mean(-2, keepdim=True).expand_as(x)
 75        inp=torch.cat([x[...,src,:],x[...,dst,:],r],-1)
 76        shape=inp.shape
 77        z=self.ctx(inp.reshape(-1,shape[-1]), z.reshape(-1,z.shape[-1])).reshape(*shape[:-1],self.hidden)
 78        u=self.u(torch.cat([x[...,src,:],x[...,dst,:],z,r[...,src,:]],-1))
 79        e=self.edge(torch.cat([x[...,src,:],x[...,dst,:],u],-1))
 80        q=torch.zeros_like(x).index_add(-2,src,e).index_add(-2,dst,-e)
 81        return self.node(torch.cat([x,q],-1)), z
 82
 83def truth(x, alpha=.65):
 84    # Local ring dynamics plus a distant, same-channel interaction.
 85    local=torch.roll(x,-1,0)+torch.roll(x,1,0)-2*x
 86    nonlocal_term=torch.roll(x,N//2,0)-x
 87    return x + .18*local + alpha*.18*nonlocal_term
 88
 89def make_data(n=320):
 90    X=[]; Y=[]
 91    for _ in range(n):
 92        x=torch.randn(N,D,device=device)
 93        X.append(x); Y.append(truth(x))
 94    return torch.stack(X),torch.stack(Y)
 95
 96def train(model, X, Y, feedback=False, steps=220):
 97    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 98    model.train()
 99    for step in range(steps):
100        i=torch.randint(0,len(X),(min(64,len(X)),),device=device)
101        xb=X[i]
102        if feedback:
103            z=torch.zeros(xb.shape[0],N,model.hidden,device=device)
104            pred,_=model(xb,z)
105        else: pred=model(xb)
106        loss=((pred-Y[i])**2).mean()
107        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),2.0); opt.step()
108    return model
109
110def evaluate(model,X,Y,feedback=False):
111    model.eval(); x=X[:100].clone(); z=torch.zeros(x.shape[0],N,model.hidden,device=device) if feedback else None
112    with torch.no_grad():
113        one=None
114        for k in range(8):
115            if feedback: x,z=model(x,z)
116            else: x=model(x)
117            if k==0: one=((x-Y[:100])**2).mean().sqrt().item()
118        long=((x-Y[:100])**2).mean().sqrt().item()
119    return one,long
120
121def jacobian_rho(model, feedback=False):
122    x=torch.randn(N,D,device=device,requires_grad=True)
123    if feedback:
124        z=torch.zeros(N,model.hidden,device=device)
125        def fn(v): return model(v,z)[0].reshape(-1)
126    else:
127        def fn(v): return model(v).reshape(-1)
128    J=torch.autograd.functional.jacobian(fn,x).detach().cpu().numpy().reshape(N*D,N*D)
129    return float(np.max(np.abs(np.linalg.eigvals(J))))
130
131def main():
132    checks=mechanism_checks()
133    X,Y=make_data(180)
134    local=LocalGNN().to(device); rec=RecursiveFeedbackGNN().to(device)
135    train(local,X,Y,False); train(rec,X,Y,True)
136    lm= evaluate(local,X,Y,False); rm=evaluate(rec,X,Y,True)
137    result={'device':str(device),'seed':SEED,'mechanism_checks':checks,
138            'mini_experiment':{'local_sparse':{'one_step_rmse':lm[0],'8_step_rmse':lm[1],'rho':jacobian_rho(local)},
139                               'recursive_nonlocal_feedback':{'one_step_rmse':rm[0],'8_step_rmse':rm[1],'rho':jacobian_rho(rec,True)},
140                               'parameter_counts':{'local':sum(p.numel() for p in local.parameters()),'recursive':sum(p.numel() for p in rec.parameters())}}}
141    with open('results.json','w') as f: json.dump(result,f,indent=2)
142    print(json.dumps(result,indent=2))
143
144if __name__=='__main__': main()