1import json, math, os, random
  2import numpy as np
  3import torch
  4from torch import nn
  5from scipy.integrate import solve_ivp
  6
  7SEED=2976
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11try:
 12    if device.type=='cuda': torch.cuda.empty_cache()
 13except Exception:
 14    device=torch.device('cpu')
 15
 16T=0.30; dt=0.03; N=int(T/dt); K=8
 17# Smooth, positive, strongly variable diffusion on a periodic 1D domain.
 18def A(x): return 0.20 + 0.15*torch.sin(x)
 19def An(x): return 0.20 + 0.15*np.sin(x)
 20
 21def ref_solution(nx=256):
 22    x=np.linspace(-np.pi,np.pi,nx,endpoint=False); dx=2*np.pi/nx
 23    # Fourier spectral second derivative, periodic reference solve.
 24    k=2*np.pi*np.fft.fftfreq(nx,d=dx)
 25    def rhs(t,v):
 26        vxx=np.fft.ifft(-(k*k)*np.fft.fft(v)).real
 27        return -0.5*An(x)*vxx
 28    sol=solve_ivp(rhs,(T,0),np.cos(2*x),rtol=2e-8,atol=2e-9,method='DOP853')
 29    return x,sol.y[:,-1]
 30
 31class MLP(nn.Module):
 32    def __init__(self):
 33        super().__init__(); self.net=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1))
 34    def forward(self,t,x): return self.net(torch.cat((t,x),1))
 35
 36def deriv_x(v,x):
 37    g=torch.autograd.grad(v.sum(),x,create_graph=True)[0]
 38    h=torch.autograd.grad(g.sum(),x,create_graph=True)[0]
 39    return g,h
 40
 41def make_batch(n, mode, model):
 42    # Time locations are grid points; network learns all time slices jointly.
 43    ni=torch.randint(0,N,(n,),device=device); t=ni.float().view(-1,1)*dt/T
 44    x=(torch.rand(n,1,device=device)*2*math.pi-math.pi)/math.pi # normalized coordinate
 45    # physical coordinate for coefficients; network coordinate is x/pi
 46    xp=(x*math.pi).detach().requires_grad_(True)
 47    tn=((ni+1).float().view(-1,1)*dt/T).detach()
 48    if mode=='raw':
 49        vn=model(tn,xp/math.pi)
 50        _,vxx=deriv_x(vn,xp)
 51        target=vn + dt*0.5*A(xp)*vxx
 52    else:
 53        # Frozen kernel is applied with local A0=A(x), using antithetic samples.
 54        z=torch.randn(n,K//2,device=device); z=torch.cat((z,-z),1)
 55        xprop=xp.detach()+torch.sqrt(A(xp.detach())*dt)*z
 56        # periodic wrap; normalized network coordinate remains in [-1,1]
 57        xwrap=(xprop+math.pi)%(2*math.pi)-math.pi
 58        tt=tn[:,None,:].expand(n,K,1).reshape(n*K,1)
 59        xx=xwrap.reshape(n*K,1).detach().requires_grad_(True)
 60        vv=model(tt,xx/math.pi)
 61        gx,hx=deriv_x(vv,xx)
 62        aq=A(xx); a0=A(xp.detach()).repeat_interleave(K).view(-1,1)
 63        q=0.5*(aq-a0)*hx
 64        # P V + dt P[(A-A0)Vxx/2]. The local source is zero at patch center,
 65        # but is retained at propagated points as required by the parametrix formula.
 66        target=(vv + dt*q).view(n,K,1).mean(1)
 67    return model(t,x), target.detach()
 68
 69def train(mode, steps=700):
 70    torch.manual_seed(SEED+ (0 if mode=='raw' else 1)); model=MLP().to(device)
 71    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 72    # Soft terminal anchoring makes the comparison a complete backward solver.
 73    hist=[]
 74    for it in range(steps):
 75        opt.zero_grad(set_to_none=True)
 76        pred,y=make_batch(96,mode,model)
 77        loss=((pred-y)**2).mean()
 78        xt=(torch.rand(96,1,device=device)*2*math.pi-math.pi)/math.pi
 79        tt=torch.ones_like(xt)
 80        term=((model(tt,xt)-torch.cos(2*xt*math.pi))**2).mean()
 81        total=loss+2.0*term
 82        total.backward(); opt.step()
 83        if it%50==0: hist.append(float(total.detach().cpu()))
 84    with torch.no_grad():
 85        xx=np.linspace(-math.pi,math.pi,256,endpoint=False)
 86        tx=torch.full((256,1),0.,device=device); xxn=torch.tensor(xx/math.pi,dtype=torch.float32,device=device).view(-1,1)
 87        pred=model(tx,xxn).cpu().numpy().ravel()
 88    return model,hist,pred
 89
 90def main():
 91    # Core math sanity: Gaussian multiplier for cos(kx), and correction exactly zero for constant A.
 92    x=torch.linspace(-math.pi,math.pi,10001); k=3.; h=.07; a=.4
 93    z=torch.randn(300000); mc=torch.cos(k*(x[5000]+math.sqrt(a*h)*z)).mean().item()
 94    exact=math.cos(k*x[5000].item())*math.exp(-.5*h*a*k*k)
 95    const_corr=float(torch.max(torch.abs((torch.full_like(x,.4)-.4)*torch.ones_like(x))).item())
 96    rx,rv=ref_solution()
 97    results={'device':str(device),'math_check_mc':mc,'math_check_exact':exact,'math_abs_error':abs(mc-exact),'constant_A_correction_max':const_corr,'reference_rms':float(np.sqrt(np.mean(rv*rv)))}
 98    allhist={}
 99    for mode in ('raw','parametrix'):
100        _,hist,pred=train(mode)
101        allhist[mode]=hist
102        results[mode+'_final_rel_l2']=float(np.linalg.norm(pred-rv)/np.linalg.norm(rv))
103        results[mode+'_loss_trace']=hist
104    results['winner_lower_error']='parametrix' if results['parametrix_final_rel_l2']<results['raw_final_rel_l2'] else 'raw'
105    with open('results.json','w') as f: json.dump(results,f,indent=2)
106    print(json.dumps(results,indent=2))
107
108if __name__=='__main__': main()