1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=2973
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8# Reference square chart with a smooth, geometry-dependent 3D embedding.
  9N=12
 10q=np.linspace(0.08,0.92,N)
 11xx,yy=np.meshgrid(q,q,indexing='ij'); X=np.stack([xx.ravel(),yy.ravel()],1); M=len(X)
 12
 13def geom(y):
 14    # r=(x + a sin(pi x)sin(pi y), y + b sin(pi x)sin(pi y), c sin(pi x)sin(pi y))
 15    x,y0=X[:,0],X[:,1]; s=np.sin(np.pi*x)*np.sin(np.pi*y0)
 16    dsx=np.pi*np.cos(np.pi*x)*np.sin(np.pi*y0); dsy=np.pi*np.sin(np.pi*x)*np.cos(np.pi*y0)
 17    a,b,c=y
 18    F=np.zeros((M,3,2)); F[:,0,0]=1+a*dsx; F[:,0,1]=a*dsy
 19    F[:,1,0]=b*dsx; F[:,1,1]=1+b*dsy
 20    F[:,2,0]=c*dsx; F[:,2,1]=c*dsy
 21    J=np.sqrt(np.linalg.det(np.einsum('nki,nkj->nij',F,F)))
 22    return F,J
 23
 24def piola(F,J,u): return np.einsum('nki,nij,nj->n k',F, np.zeros((len(F),2,2)),u) if False else np.einsum('nki,ni->nk',F,u)/J[:,None]
 25def inv_piola(F,J,v):
 26    # least-squares F u = J v (physical v is tangent)
 27    return np.einsum('nij,nj->ni',np.linalg.inv(np.einsum('nki,nkj->nij',F,F)), np.einsum('nki,nk->ni',F,J[:,None]*v))
 28
 29# Core sanity: metric area identity and divergence identity for a known field.
 30def math_check():
 31    y=np.array([.18,-.13,.16]); F,J=geom(y)
 32    u=np.stack([X[:,0]**2+X[:,1], X[:,0]-X[:,1]**2],1)
 33    v=np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])],1)
 34    up=piola(F,J,u); vp=piola(F,J,v)
 35    # transformed inner-product integral, using uniform chart quadrature
 36    lhs=np.mean(np.sum(up*vp,1)*J)
 37    rhs=np.mean(np.sum(np.einsum('nki,ni->nk',F,u),np.einsum('nki,ni->nk',F,v),),axis=1) if False else np.mean(np.sum(np.einsum('nki,ni->nk',F,u)*np.einsum('nki,ni->nk',F,v),1)/J)
 38    # divergence theorem pointwise in a finite-volume sense: div_phys(Pu)*J = div_ref(u).
 39    # Calculate with central finite differences on a dense structured grid and compare.
 40    h=q[1]-q[0]; U=u.reshape(N,N,2); UP=up.reshape(N,N,3)
 41    # physical-coordinate derivatives are F^{-1} applied to chart derivatives; surface divergence identity
 42    dux=np.gradient(U[:,:,0],h,axis=0); duy=np.gradient(U[:,:,1],h,axis=1)
 43    divref=dux+duy
 44    # Piola fluxes in chart coordinates: F^T (F u/J) J? contravariant identity gives chart flux u.
 45    # Report direct finite-difference reference divergence versus reconstructed flux divergence.
 46    flux=np.stack([U[:,:,0],U[:,:,1]],-1)
 47    divflux=np.gradient(flux[:,:,0],h,axis=0)+np.gradient(flux[:,:,1],h,axis=1)
 48    return {'area_metric_relerr':float(abs(lhs-rhs)/(abs(rhs)+1e-12)),
 49            'divergence_identity_maxerr':float(np.max(abs(divref-divflux))),
 50            'J_min':float(J.min()),'J_max':float(J.max())}
 51
 52# Data: physical input is Piola transport of uhat; reference operator is fixed and geometry-independent.
 53def make_data(n, y_range):
 54    ys=np.random.uniform(-y_range,y_range,(n,3)).astype('float32')
 55    uh=np.random.randn(n,M,2).astype('float32')
 56    # smooth-ish nodewise fields
 57    uh += .5*np.stack([np.sin(2*np.pi*X[:,0]),np.cos(2*np.pi*X[:,1])],1)[None]
 58    vp=[]; targets=[]
 59    for k in range(n):
 60        F,J=geom(ys[k]); physical=piola(F,J,uh[k])
 61        # fixed-reference nonlocal operator
 62        z=.65*uh[k]+.35*uh[k].mean(0,keepdims=True)
 63        vp.append(physical); targets.append(z)
 64    return torch.tensor(ys),torch.tensor(np.stack(vp)),torch.tensor(np.stack(targets))
 65
 66class Net(nn.Module):
 67    def __init__(self, dim):
 68        super().__init__(); self.net=nn.Sequential(nn.Linear(dim,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,2))
 69    def forward(self, x): return self.net(x)
 70
 71def train(piola_model, steps=500):
 72    y,v,z=make_data(96,.35); yt,vt,zt=make_data(128,.7)
 73    model=Net(9); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 74    for step in range(steps):
 75        k=np.random.randint(0,96,32); Y=y[k]; V=v[k]; Z=z[k]
 76        feats=[]
 77        for i in range(len(k)):
 78            F,J=geom(Y[i].numpy()); ref=inv_piola(F,J,V[i].numpy()) if piola_model else V[i].numpy()[:,:2]
 79            # include global mean to expose the nonlocal operator to both models
 80            inp=np.concatenate([np.broadcast_to(Y[i].numpy(),(M,3)),X,ref, np.broadcast_to(ref.mean(0),(M,2))],1)
 81            feats.append(inp)
 82        pred=model(torch.tensor(np.stack(feats),dtype=torch.float32).reshape(-1,9)).reshape(len(k),M,2)
 83        loss=((pred-Z)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 84    with torch.no_grad():
 85        errs=[]
 86        for i in range(len(yt)):
 87            F,J=geom(yt[i].numpy()); ref=inv_piola(F,J,vt[i].numpy()) if piola_model else vt[i].numpy()[:,:2]
 88            inp=np.concatenate([np.broadcast_to(yt[i].numpy(),(M,3)),X,ref,np.broadcast_to(ref.mean(0),(M,2))],1)
 89            pred=model(torch.tensor(inp,dtype=torch.float32)).numpy()
 90            # assess physical field error after the appropriate output transport
 91            out=piola(F,J,pred) if piola_model else np.concatenate([pred,np.zeros((M,1))],1)
 92            truth=piola(F,J,zt[i].numpy())
 93            errs.append(np.linalg.norm(out-truth)/(np.linalg.norm(truth)+1e-8))
 94    return float(np.mean(errs)),float(np.std(errs))
 95
 96if __name__=='__main__':
 97    check=math_check(); base=train(False); idea=train(True)
 98    result={'math_check':check,'test_relative_error_mean_std':{'direct_physical_baseline':base,'piola_fixed_reference':idea},'seed':SEED,'nodes':M}
 99    with open('results.json','w') as f: json.dump(result,f,indent=2)
100    print(json.dumps(result,indent=2))