1import json, math, time, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=2149
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8device='cuda' if torch.cuda.is_available() else 'cpu'
  9try:
 10    if device=='cuda': torch.cuda.set_device(0)
 11except Exception:
 12    device='cpu'
 13
 14torch.set_num_threads(8)
 15
 16# Smooth random Fourier fields and a known operator: output is a shifted,
 17# mildly nonlinear transform of the input field.
 18def fields(n, m, phase=0):
 19    x=np.linspace(0,1,m,endpoint=False)[None,:]
 20    rng=np.random.RandomState(SEED+phase)
 21    a=rng.randn(n,5)*np.array([1,.6,.4,.3,.2])[None,:]
 22    b=rng.randn(n,5)*np.array([1,.6,.4,.3,.2])[None,:]
 23    u=np.zeros((n,m));
 24    for k in range(1,6): u += a[:,k-1,None]*np.sin(2*np.pi*k*x)+b[:,k-1,None]*np.cos(2*np.pi*k*x)
 25    u=u/(np.std(u,axis=1,keepdims=True)+1e-6)*.7
 26    # operator is pointwise nonlinear plus shifted low-frequency component
 27    y=np.tanh(u)+.25*np.roll(u,2,axis=1)
 28    return x.astype(np.float32),u.astype(np.float32),y.astype(np.float32)
 29
 30def interp_periodic(x0, v, q):
 31    xx=np.mod(x0,1); order=np.argsort(xx); xx=xx[order]; vv=v[...,order]
 32    xx2=np.r_[xx,xx[0]+1]; vv2=np.concatenate([vv,vv[...,0:1]],axis=-1)
 33    qq=np.mod(q,1); j=np.searchsorted(xx2,qq,side='right')-1; j=np.clip(j,0,len(xx2)-2)
 34    t=(qq-xx2[j])/(xx2[j+1]-xx2[j]+1e-8)
 35    return vv2[...,j]*(1-t)+vv2[...,j+1]*t
 36
 37class MeasurementNet(nn.Module):
 38    def __init__(self,z=32):
 39        super().__init__(); self.enc=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,z),nn.Tanh())
 40        self.g=nn.Sequential(nn.Linear(z,64),nn.Tanh(),nn.Linear(64,z),nn.Tanh())
 41        self.dec=nn.Sequential(nn.Linear(z+1,64),nn.Tanh(),nn.Linear(64,1))
 42    def encode(self, coords, vals):
 43        h=self.enc(torch.stack((coords,vals),-1)); return h.mean(1)
 44    def forward(self, coords, vals, q):
 45        z=self.g(self.encode(coords,vals)); zz=z[:,None,:].expand(-1,q.shape[1],-1)
 46        return self.dec(torch.cat((zz,q[:,:,None]),-1)).squeeze(-1)
 47
 48class FixedNet(nn.Module):
 49    def __init__(self,m=32):
 50        super().__init__(); self.net=nn.Sequential(nn.Linear(m,96),nn.Tanh(),nn.Linear(96,128),nn.Tanh(),nn.Linear(128,64))
 51    def forward(self,v): return self.net(v)
 52
 53def train_models():
 54    x,u,y=fields(900,64,1); xt,ut,yt=fields(180,64,2)
 55    # canonical measurements are first 32 regular points, output 64 grid
 56    ci=np.linspace(0,63,32).round().astype(int); qi=np.arange(64)
 57    mc=MeasurementNet().to(device); base=FixedNet().to(device)
 58    opt=torch.optim.Adam(mc.parameters(),lr=3e-3); opb=torch.optim.Adam(base.parameters(),lr=3e-3)
 59    X=torch.tensor(x[:,ci],device=device); U=torch.tensor(u[:,ci],device=device); Q=torch.tensor(x[:,qi],device=device)
 60    Y=torch.tensor(y,device=device); UB=torch.tensor(u[:,ci],device=device)
 61    for step in range(300):
 62        ix=torch.randint(0,len(x),(96,),device=device)
 63        pred=mc(X[ix],U[ix],Q[ix]); loss=((pred-Y[ix])**2).mean()
 64        opt.zero_grad(); loss.backward(); opt.step()
 65        pb=base(UB[ix]); lb=((pb-Y[ix])**2).mean(); opb.zero_grad(); lb.backward(); opb.step()
 66    return mc.cpu().eval(),base.cpu().eval(),(x,u,y,xt,ut,yt,ci)
 67
 68def predict_mesh(model, coords, vals, q):
 69    with torch.no_grad():
 70        return model(torch.tensor(coords[None],dtype=torch.float32),torch.tensor(vals[None],dtype=torch.float32),torch.tensor(q[None],dtype=torch.float32)).numpy()[0]
 71
 72def main():
 73    t=time.time(); mc,base,data=train_models(); x,u,y,xt,ut,yt,ci=data
 74    # Mesh transfer: input sensors are alternative layouts; output is arbitrary query mesh.
 75    rng=np.random.RandomState(77); rows=[]
 76    for n in [8,16,32,64]:
 77        errs=[]
 78        ss=np.sort(rng.rand(len(xt),n),axis=1); vals=np.stack([interp_periodic(x[0],ut[i],ss[i]) for i in range(len(xt))])
 79        q=np.sort(rng.rand(len(xt),48),axis=1); truth=np.stack([interp_periodic(x[0],yt[i],q[i]) for i in range(len(xt))])
 80        with torch.no_grad():
 81            pred=mc(torch.tensor(ss,dtype=torch.float32),torch.tensor(vals,dtype=torch.float32),torch.tensor(q,dtype=torch.float32)).numpy()
 82        rows.append((n,float(np.sqrt(np.mean((pred-truth)**2)))))
 83    # baseline canonical error on its supported fixed grid, and baseline cannot accept new mesh.
 84    with torch.no_grad():
 85        bp=base(torch.tensor(ut[:,ci],dtype=torch.float32)).numpy()
 86    base_rmse=float(np.sqrt(np.mean((bp-yt)**2)))
 87    # Core prediction 1: permutation invariance of E_X.
 88    i=0; s=np.sort(rng.rand(20)); vv=interp_periodic(x[0],u[i],s)
 89    p1=predict_mesh(mc,s,vv,np.linspace(0,1,20,endpoint=False)); perm=rng.permutation(20)
 90    p1p=predict_mesh(mc,s[perm],vv[perm],np.linspace(0,1,20,endpoint=False))
 91    perm_diff=float(np.max(np.abs(p1-p1p)))
 92    # Core prediction 2: Lipschitz composition discrepancy scales linearly with reconstruction perturbation.
 93    # Hold a measurement set fixed and perturb reconstructed function by amplitude delta.
 94    s=np.linspace(0,1,32,endpoint=False); vv=interp_periodic(x[0],u[0],s); q=np.linspace(0,1,40,endpoint=False)
 95    direction=np.sin(2*np.pi*3*s); deltas=np.array([0,.01,.02,.04,.08,.16]); diffs=[]
 96    for d in deltas:
 97        a=predict_mesh(mc,s,vv,q); b=predict_mesh(mc,s,vv+d*direction,q); diffs.append(float(np.sqrt(np.mean((a-b)**2))))
 98    slope=float(np.polyfit(deltas[1:],diffs[1:],1)[0]); r2=float(np.corrcoef(deltas,diffs)[0,1])
 99    # Core prediction 3: more measurements improve reconstruction/transfer, with diminishing returns.
100    # Compare fixed Fourier interpolation reconstruction error against true input field.
101    recon=[]
102    for n in [4,8,16,32,64]:
103        ee=[]
104        for i in range(50):
105            s=np.sort(rng.rand(n)); vv=interp_periodic(x[0],u[i],s); q=x[0]
106            ee.append(np.mean((interp_periodic(s,vv,q)-u[i])**2))
107        recon.append((n,float(np.sqrt(np.mean(ee)))))
108    out={'device':device,'seconds':round(time.time()-t,2),'baseline_fixed32_rmse':base_rmse,
109         'idea_mesh_transfer_rmse_by_input_sensors':rows,'prediction_permutation_max_difference':perm_diff,
110         'prediction_lipschitz_sweep':{'deltas':deltas.tolist(),'output_discrepancies':diffs,'linear_slope':slope,'correlation':r2},
111         'prediction_sensor_density_reconstruction_rmse':recon,
112         'note':'Measurement model uses coordinate/value mean pooling and coordinate-conditioned output decoder; baseline is fixed-grid MLP.'}
113    print(json.dumps(out,indent=2))
114if __name__=='__main__':
115    try: main()
116    except Exception as e:
117        if device=='cuda':
118            print('CUDA failed, rerun CPU:',repr(e)); device='cpu'; torch.cuda.empty_cache(); main()
119        else: raise