import json, math, time from pathlib import Path import numpy as np import torch from torch import nn SEED = 17 np.random.seed(SEED); torch.manual_seed(SEED) try: device = 'cuda' if torch.cuda.is_available() else 'cpu' if device == 'cuda': torch.cuda.set_device(0) torch.zeros(1, device='cuda') except Exception: device = 'cpu' def wendland(t): t = np.asarray(t) return np.where(t < 1.0, (1.0-t)**4 * (4.0*t+1.0), 0.0) def centers_nested(x, levels): # C0 is all sites; subsequent sets are deterministic nested subsets. out = [x] for j in range(1, levels): stride = 2**j out.append(x[::stride]) return out def radius(c, multiplier=2.5): d = np.sqrt(((c[:,None,:]-c[None,:,:])**2).sum(-1) + np.eye(len(c))*1e9) return multiplier * np.median(d.min(1)) def eval_matrix(x, c, rho): return wendland(np.linalg.norm(x[:,None,:]-c[None,:,:], axis=-1)/rho) def frame_encode(x, u, centers, lamb=1e-5): q = np.zeros_like(u); blocks=[]; stats=[] for c in centers: rho = radius(c) A = eval_matrix(x,c,rho) gram = A.T@A + lamb*np.eye(len(c)) alpha = np.linalg.solve(gram, A.T@(u-q)) q = q + A@alpha stats.append({'n':len(c), 'rho':float(rho), 'density':float((A>0).mean()), 'cond':float(np.linalg.cond(gram)), 'residual':float(np.linalg.norm(u-q)/np.linalg.norm(u))}) blocks.append(alpha) return blocks, stats def flatten_blocks(blocks, sizes): z=[] for b,n in zip(blocks,sizes): z.append(np.pad(b, (0, max(sizes)-n))) return np.stack(z, axis=0) class CoeffOperator(nn.Module): def __init__(self, levels, width): super().__init__(); self.levels=levels; self.width=width self.net=nn.Sequential(nn.Linear(levels*width,96),nn.GELU(),nn.Linear(96,96),nn.GELU(),nn.Linear(96,levels*width)) def forward(self,z): return self.net(z.flatten(1)).reshape(-1,self.levels,self.width) class PointOperator(nn.Module): def __init__(self, m): super().__init__(); self.net=nn.Sequential(nn.Linear(m,96),nn.GELU(),nn.Linear(96,96),nn.GELU(),nn.Linear(96,m)) def forward(self,x): return self.net(x) def make_fields(n, mside): g=np.linspace(0,1,mside); X=np.stack(np.meshgrid(g,g,indexing='ij'),-1).reshape(-1,2) vals=[]; outs=[] for _ in range(n): a=np.random.randn(5)*.7; b=np.random.randn(5)*.7 v=(a[0]*np.sin(2*np.pi*X[:,0])+a[1]*np.cos(2*np.pi*X[:,1])+a[2]*np.sin(4*np.pi*(X[:,0]+X[:,1]))+ a[3]*np.cos(6*np.pi*X[:,0])+a[4]*np.sin(6*np.pi*X[:,1])) # nonlinear nonlocal-ish target, still small and deterministic u=np.tanh(v)+.15*np.sin(8*np.pi*X[:,0])*v vals.append(v); outs.append(u) return X,np.asarray(vals),np.asarray(outs) def main(): mside=12; M=mside*mside; levels=3 X,V,U=make_fields(180,mside) centers=centers_nested(X,levels); sizes=[len(c) for c in centers]; width=max(sizes) # Core math check on a held-out field: each sequential projection should not increase residual blocks, check=frame_encode(X,V[0],centers) residuals=[1.0]+[s['residual'] for s in check] monotonic=all(residuals[i+1] <= residuals[i]+1e-9 for i in range(len(residuals)-1)) sparse=all(s['density'] < .5 for s in check) # Encode all data, and decode at original points. Use same basis for target coefficients. Z=[]; T=[] for v,u in zip(V,U): zb,_=frame_encode(X,v,centers); tb,_=frame_encode(X,u,centers) Z.append(flatten_blocks(zb,sizes)); T.append(flatten_blocks(tb,sizes)) Z=np.asarray(Z,dtype=np.float32); T=np.asarray(T,dtype=np.float32) # split fixed tr=np.arange(0,130); te=np.arange(130,180) ztr=torch.tensor(Z[tr],device=device); ttr=torch.tensor(T[tr],device=device) zte=torch.tensor(Z[te],device=device); tte=torch.tensor(T[te],device=device) vtr=torch.tensor(V[tr],dtype=torch.float32,device=device); utr=torch.tensor(U[tr],dtype=torch.float32,device=device) vte=torch.tensor(V[te],dtype=torch.float32,device=device); ute=torch.tensor(U[te],dtype=torch.float32,device=device) torch.manual_seed(SEED) idea=CoeffOperator(levels,width).to(device); base=PointOperator(M).to(device) oi=torch.optim.Adam(idea.parameters(),lr=2e-3); ob=torch.optim.Adam(base.parameters(),lr=2e-3) t0=time.time() for step in range(700): oi.zero_grad(); pred=idea(ztr); ((pred-ttr)**2).mean().backward(); oi.step() ob.zero_grad(); pb=base(vtr); ((pb-utr)**2).mean().backward(); ob.step() train_seconds=time.time()-t0 with torch.no_grad(): # Reconstruct output from predicted coefficient blocks at X. pred=idea(zte).cpu().numpy(); pbase=base(vte).cpu().numpy() Aall=[eval_matrix(X,c,radius(c)) for c in centers] def decode(P): return sum(P[:,j,:sizes[j]]@Aall[j].T for j in range(levels)) pred_u=decode(pred); true_u=U[te] err=float(np.linalg.norm(pred_u-true_u)/np.linalg.norm(true_u)) berr=float(np.linalg.norm(pbase-U[te])/np.linalg.norm(U[te])) # irregular query stability/generalization: decode on jittered points against analytic target generation per sample rng=np.random.default_rng(SEED); Y=np.clip(X+rng.normal(0,.018,X.shape),0,1) AY=[eval_matrix(Y,c,radius(c)) for c in centers] qpred=sum(pred[:,j,:sizes[j]]@AY[j].T for j in range(levels)) # target evaluated directly at Y using same latent Fourier coefficients is unavailable; compare interpolation consistency interp=float(np.mean(np.linalg.norm(qpred-pred_u,axis=1)/ (np.linalg.norm(pred_u,axis=1)+1e-7))) result={'device':device,'math_check':{'residuals':residuals,'monotonic_nonincrease':monotonic,'all_levels_sparse':sparse,'levels':check}, 'experiment':{'train_seconds':train_seconds,'relative_l2_idea':err,'relative_l2_baseline':berr,'query_jitter_relative_change':interp, 'idea_parameters':sum(p.numel() for p in idea.parameters()),'baseline_parameters':sum(p.numel() for p in base.parameters()), 'active_basis_coefficients':sum(sizes),'dense_input_tokens':M}, 'seed':SEED} Path('results.json').write_text(json.dumps(result,indent=2)) print(json.dumps(result,indent=2)) if __name__=='__main__': main()