1import json, math, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=362
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float64)
  9device='cuda' if torch.cuda.is_available() else 'cpu'
 10try:
 11    if device=='cuda': torch.cuda.set_device(0)
 12except Exception:
 13    device='cpu'
 14
 15class MLP(nn.Module):
 16    def __init__(self):
 17        super().__init__(); self.net=nn.Sequential(nn.Linear(1,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,1))
 18    def forward(self,x): return self.net(x)
 19
 20def deriv(y,x,create=True):
 21    return torch.autograd.grad(y,x,torch.ones_like(y),create_graph=create,retain_graph=True)[0]
 22
 23def make_model(): return nn.ModuleList([MLP().to(device),MLP().to(device)])
 24
 25def eval_u(net,x,need2=False):
 26    x=x.clone().detach().requires_grad_(True); u=net(x); ux=deriv(u,x,True)
 27    if need2: return x,u,ux,deriv(ux,x,True)
 28    return x,u,ux
 29
 30def setup():
 31    # Gauss-like uniform quadrature; exact solution x(1-x), -u''=2.
 32    xL=torch.linspace(0.005,0.495,48,device=device).reshape(-1,1)
 33    xR=torch.linspace(0.505,0.995,48,device=device).reshape(-1,1)
 34    xf=torch.tensor([[0.5]],device=device)
 35    return xL,xR,xf
 36
 37def nitsche_loss(m, xL,xR,xf, gamma=20., ghost=1.):
 38    _,uL,dL=eval_u(m[0],xL); _,uR,dR=eval_u(m[1],xR)
 39    bulk=(0.5*(dL*dL)-2*uL).mean()*0.5 + (0.5*(dR*dR)-2*uR).mean()*0.5
 40    # exterior homogeneous Dirichlet: - n P u + gamma/(2h)u^2, h=.5
 41    xb=torch.tensor([[0.0],[1.0]],device=device)
 42    _,ul,dl=eval_u(m[0],xb[:1]); _,ur,dr=eval_u(m[1],xb[1:])
 43    # n=-1 left, +1 right
 44    ext=(dl*ul + gamma/(1.0)*ul.square()/2 - dr*ur + gamma/(1.0)*ur.square()/2).mean()
 45    # symmetric interface Nitsche for jump uL-uR, average flux; h=.5
 46    _,uiL,diL=eval_u(m[0],xf); _,uiR,diR=eval_u(m[1],xf)
 47    jump=uiL-uiR; avg=.5*(diL+diR)
 48    iface=(-avg*jump + gamma/(1.0)*jump.square()/2).mean()
 49    # p=1 ghost: gamma_A h^(1) [u']^2
 50    g=(ghost*0.5*(diL-diR).square()).mean()
 51    return bulk+ext+iface+g, {'bulk':bulk.item(),'boundary':(ul.square().mean()+ur.square().mean()).item(),'jump':jump.abs().item(),'djump':(diL-diR).abs().item()}
 52
 53def baseline_loss(m,xL,xR,xf):
 54    _,uL,dL,rL=eval_u(m[0],xL,True); _,uR,dR,rR=eval_u(m[1],xR,True)
 55    # standard PINN strong residual and pointwise penalties, including patch interface
 56    loss=(rL.add(2).square().mean()+rR.add(2).square().mean())
 57    xb=torch.tensor([[0.0],[1.0]],device=device)
 58    _,ul,_=eval_u(m[0],xb[:1]); _,ur,_=eval_u(m[1],xb[1:])
 59    _,uiL,diL=eval_u(m[0],xf); _,uiR,diR=eval_u(m[1],xf)
 60    loss=loss+100*(ul.square().mean()+ur.square().mean())+100*(uiL-uiR).square().mean()+1*(diL-diR).square().mean()
 61    return loss, {'boundary':(ul.square().mean()+ur.square().mean()).item(),'jump':(uiL-uiR).abs().item(),'djump':(diL-diR).abs().item()}
 62
 63def metrics(m):
 64    xx=torch.linspace(0,1,401,device=device).reshape(-1,1); vals=[]
 65    for a,b in [(0,.5),(.5,1)]:
 66        q=xx[(xx[:,0]>=a)&(xx[:,0]<=b)].reshape(-1,1); vals.append(m[0 if a==0 else 1](q))
 67    pred=torch.cat(vals).flatten(); q=torch.cat([xx[(xx[:,0]>=0)&(xx[:,0]<=.5)],xx[(xx[:,0]>=.5)&(xx[:,0]<=1)]])
 68    truth=q*(1-q); err=(pred-truth).abs().max().item()
 69    xL,xR,xf=setup(); _,uL,dL=eval_u(m[0],torch.tensor([[0.0]],device=device)); _,uR,dR=eval_u(m[1],torch.tensor([[1.0]],device=device)); _,a,da=eval_u(m[0],xf); _,b,db=eval_u(m[1],xf)
 70    return {'max_error':err,'boundary_rms':math.sqrt((uL.square()+uR.square()).mean().item()),'interface_value_jump':abs((a-b).item()),'interface_derivative_jump':abs((da-db).item())}
 71
 72def run(kind,steps=900):
 73    torch.manual_seed(SEED+ (0 if kind=='baseline' else 1)); m=make_model(); opt=torch.optim.Adam(m.parameters(),lr=2e-3)
 74    xL,xR,xf=setup(); curve=[]; t=time.time()
 75    for it in range(steps):
 76        opt.zero_grad(); loss,parts=(baseline_loss(m,xL,xR,xf) if kind=='baseline' else nitsche_loss(m,xL,xR,xf)); loss.backward(); opt.step()
 77        if it in [0,99,299,599,899]: curve.append(float(loss.detach().cpu()))
 78    out=metrics(m); out.update({'loss_curve':curve,'final_objective':float(loss.detach().cpu()),'seconds':time.time()-t})
 79    return out
 80
 81def math_check():
 82    # Constitutive AD identity plus a coercivity scan of the 1D symmetric
 83    # Nitsche quadratic form on u(x)=a*x+b.
 84    x=torch.tensor([[.2]],requires_grad=True)
 85    u=x*x/3; ux=deriv(u,x,True); F=1+ux
 86    W=.5*(F-1).square()
 87    P=torch.autograd.grad(W,F,torch.ones_like(W),retain_graph=True)[0]
 88    fp_err=abs((P-(F-1)).item())
 89    vals={}
 90    for gamma in [0.1,1.,2.,4.,10.,20.]:
 91        def q(v):
 92            aa,bb=v
 93            return .5*aa*aa-(aa*(aa+bb)+aa*bb)+gamma/2*(bb*bb+(aa+bb)**2)
 94        H=np.zeros((2,2)); H[0,0]=q((1,0)); H[1,1]=q((0,1))
 95        H[0,1]=H[1,0]=(q((1,1))-H[0,0]-H[1,1])/2
 96        vals[str(gamma)]=float(np.linalg.eigvalsh(H).min())
 97    jump=1.7-(-.4); ghost_energy=.8*.5*jump**2
 98    return {'F_identity_and_P_error':fp_err,
 99            'nitsche_min_eigenvalue_by_gamma':vals,
100            'ghost_energy_for_derivative_jump':ghost_energy,
101            'ghost_is_nonnegative':bool(ghost_energy>=0)}
102
103if __name__=='__main__':
104    try:
105        chk=math_check(); base=run('baseline'); idea=run('idea')
106    except Exception as e:
107        if device=='cuda':
108            device='cpu'; torch.cuda.empty_cache(); chk=math_check(); base=run('baseline'); idea=run('idea')
109        else: raise
110    print(json.dumps({'device':device,'math_check':chk,'baseline':base,'idea':idea},indent=2))