1import os, time, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED=397
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10    if device.type=='cuda': torch.zeros(1,device=device)
 11except Exception:
 12    device=torch.device('cpu')
 13
 14# Structured quad mesh. Each cell owns its four corner copies; shared_id is the
 15# only connectivity object needed by the unassembled implementation.
 16def mesh(n):
 17    cells=[]
 18    for y in range(n):
 19        for x in range(n):
 20            cells.append([y*(n+1)+x, y*(n+1)+x+1,
 21                          (y+1)*(n+1)+x+1, (y+1)*(n+1)+x])
 22    ids=np.asarray(cells,dtype=np.int64)
 23    k=np.bincount(ids.ravel(), minlength=(n+1)**2).astype(np.float32)
 24    return ids,k
 25
 26# Exact projection check, including idempotence, symmetry, and fixation.
 27def projection_check():
 28    ids,k=mesh(3); m=ids.size; N=k.size
 29    G=np.zeros((m,N),dtype=np.float64)
 30    G[np.arange(m),ids.ravel()]=1
 31    S=G@np.diag(1/k)@G.T
 32    u=np.random.randn(N,3); v=G@u
 33    z=np.random.randn(m,2); manual=G@((G.T@z)/k[:,None])
 34    return {'projection_error':float(np.max(np.abs(S@v-v))),
 35            'idempotence_error':float(np.max(np.abs(S@S-S))),
 36            'symmetry_error':float(np.max(np.abs(S-S.T))),
 37            'manual_vs_matrix':float(np.max(np.abs(S@z-manual)))}
 38
 39class LocalNet(nn.Module):
 40    def __init__(self, width=32, layers=8):
 41        super().__init__(); self.blocks=nn.ModuleList()
 42        for i in range(layers):
 43            self.blocks.append(nn.Sequential(nn.Linear(1 if i==0 else width,width),nn.Tanh()))
 44        self.out=nn.Linear(width,1)
 45    def forward(self,x):
 46        for b in self.blocks: x=b(x)
 47        return self.out(x)
 48
 49class CellNet(nn.Module):
 50    def __init__(self, ids, k, sync_every, width=32, layers=8):
 51        super().__init__(); self.ids=torch.tensor(ids,device=device); self.k=torch.tensor(k,device=device)
 52        self.sync_every=sync_every; self.blocks=nn.ModuleList()
 53        self.local_nodes=ids.shape[1]; self.width=width
 54        for i in range(layers):
 55            iw=1 if i==0 else width
 56            self.blocks.append(nn.Sequential(nn.Linear(self.local_nodes*iw,self.local_nodes*width),nn.Tanh()))
 57        self.out=nn.Linear(width,1)
 58    def sync(self,h):
 59        # h: batch, cell, local, channel; segmented reduction via index_add.
 60        B,C,J,W=h.shape; flat=self.ids.reshape(-1); z=h.reshape(B,-1,W)
 61        sums=torch.zeros(B,self.k.numel(),W,device=h.device,dtype=h.dtype)
 62        sums.index_add_(1,flat,z)
 63        return (sums[:,flat,:]/self.k[flat][None,:,None]).reshape(B,C,J,W)
 64    def forward(self,x):
 65        # input is assembled nodal scalar, duplicated into local copies
 66        h=x[:,self.ids].unsqueeze(-1)
 67        for i,b in enumerate(self.blocks,1):
 68            B,C,J,W=h.shape
 69            h=b(h.reshape(B,C,J*W)).reshape(B,C,J,self.width)
 70            if self.sync_every and i%self.sync_every==0: h=self.sync(h)
 71        h=self.sync(h)
 72        # average local outputs to assembled vector
 73        y=h[...,0]
 74        B=y.shape[0]; out=torch.zeros(B,self.k.numel(),device=y.device)
 75        out.index_add_(1,self.ids.reshape(-1),y.reshape(B,-1))
 76        return out/self.k[None,:]
 77
 78class AssembledNet(nn.Module):
 79    def __init__(self, adj, width=32, layers=8):
 80        super().__init__(); self.register_buffer('A',adj); self.blocks=nn.ModuleList()
 81        for i in range(layers): self.blocks.append(nn.Sequential(nn.Linear(1 if i==0 else width,width),nn.Tanh()))
 82        self.out=nn.Linear(width,1)
 83    def forward(self,x):
 84        h=x.unsqueeze(-1)
 85        for b in self.blocks:
 86            # standard assembled graph-GNN layer: neighbor aggregate plus self
 87            h=b(h + torch.matmul(self.A,h))
 88        return self.out(h).squeeze(-1)
 89
 90def data(n,Nsample=512):
 91    # Diffusion-like target: a fixed local operator, represented as a smooth
 92    # coefficient field mapped to its neighbor average.
 93    yy,xx=np.mgrid[0:n+1,0:n+1]; p=(n+1)**2
 94    A=np.zeros((p,p),np.float32)
 95    for q in range(p):
 96        y,x=divmod(q,n+1)
 97        ns=[]
 98        for dy,dx in ((-1,0),(1,0),(0,-1),(0,1)):
 99            if 0<=y+dy<=n and 0<=x+dx<=n: ns.append((y+dy)*(n+1)+x+dx)
100        A[q,ns]=1/len(ns)
101    X=np.random.randn(Nsample,p).astype('float32')
102    Y=X@A.T
103    return torch.tensor(X),torch.tensor(Y),torch.tensor(A)
104
105def train(model,X,Y,steps=160):
106    model.to(device); X=X.to(device); Y=Y.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
107    t=time.perf_counter();
108    for s in range(steps):
109        ix=torch.randint(0,len(X),(64,),device=device); pred=model(X[ix]); loss=((pred-Y[ix])**2).mean()
110        opt.zero_grad(); loss.backward(); opt.step()
111    if device.type=='cuda': torch.cuda.synchronize()
112    elapsed=time.perf_counter()-t
113    with torch.no_grad(): val=((model(X[-128:])-Y[-128:])**2).mean().item()
114    return val,steps*64/elapsed,elapsed
115
116def main():
117    check=projection_check(); n=8; ids,k=mesh(n); X,Y,A=data(n)
118    # Normalize adjacency to avoid an unfairly unstable assembled baseline.
119    models=[('assembled',AssembledNet(A*.5,layers=8)),
120            ('cell_sync_1',CellNet(ids,k,1,layers=8)),
121            ('cell_sync_4',CellNet(ids,k,4,layers=8)),
122            ('cell_sync_8',CellNet(ids,k,8,layers=8))]
123    results={}
124    for name,m in models: results[name]=train(m,X,Y)
125    # Disagreement is measured immediately before a scheduled synchronization.
126    h=torch.randn(16,len(ids),4,8,device=device); cm=CellNet(ids,k,4)
127    pre=cm.sync(h); disagreement_before=(h-pre).pow(2).mean().sqrt().item()
128    results['projection_check']=check; results['pre_sync_rms_disagreement']=disagreement_before
129    results['device']=str(device); results['cells']=len(ids); results['nodes']=len(k)
130    print(json.dumps(results,indent=2))
131    with open('results.json','w') as f: json.dump(results,f,indent=2)
132if __name__=='__main__': main()