1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED = 1358
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10M = 1.0
 11
 12# Klein-Gordon reduced to a single Fourier mode: u_t=v, v_t=-m sin(u).
 13def F(z):
 14    u, v = z[..., 0], z[..., 1]
 15    return np.stack([v, -M*np.sin(u)], axis=-1)
 16
 17def rk4(z, dt):
 18    k1=F(z); k2=F(z+.5*dt*k1); k3=F(z+.5*dt*k2); k4=F(z+dt*k3)
 19    return z + dt*(k1+2*k2+2*k3+k4)/6
 20
 21def math_checks():
 22    bg=np.array([.73,-.21]); r=np.array([.18,-.27]); bgdot=np.array([bg[1],-.11])
 23    defect=F(bg)-bgdot; ru=np.sin(bg[0]+r[0])-np.sin(bg[0])-np.cos(bg[0])*r[0]
 24    rhs=np.array([r[1],-M*r[0]+defect[1]-M*(np.cos(bg[0])-1)*r[0]-M*ru])
 25    identity_err=float(np.max(np.abs(F(bg+r)-bgdot-rhs)))
 26    eps=np.logspace(-5,-1,9)
 27    rem=np.array([abs(np.sin(bg[0]+e)-np.sin(bg[0])-np.cos(bg[0])*e) for e in eps])
 28    slope=float(np.polyfit(np.log(eps),np.log(rem),1)[0])
 29    coeff=float(np.median(rem[:5]/eps[:5]**2)); coeff_pred=abs(np.sin(bg[0]))/2
 30    # sigmoid gate g=1/(1+exp(-a(log d-b))); midpoint predicted d=exp(b)
 31    a=3.; b=math.log(.08); ds=np.logspace(-3,0,401)
 32    gates=1/(1+np.exp(-a*(np.log(ds)-b))); d_mid=float(ds[np.argmin(abs(gates-.5))])
 33    # logit(g) versus log(d) slope predicted a
 34    fit=float(np.polyfit(np.log(ds),np.log(gates/(1-gates)),1)[0])
 35    return {'identity_max_error':identity_err,'remainder_loglog_slope':slope,
 36            'remainder_quadratic_slope_predicted':2.,'remainder_coeff_est':coeff,
 37            'remainder_coeff_predicted':coeff_pred,'gate_midpoint_observed':d_mid,
 38            'gate_midpoint_predicted':math.exp(b),'gate_logit_slope_observed':fit,
 39            'gate_logit_slope_predicted':a}
 40
 41def make_data(n=1800, steps=25, dt=.04):
 42    x=np.random.uniform(-1.4,1.4,(n,2)); ys=[]
 43    for _ in range(steps):
 44        y=rk4(x,dt); ys.append((x.copy(),y.copy())); x=y
 45    return np.concatenate([a for a,b in ys]),np.concatenate([b for a,b in ys])
 46
 47class Net(nn.Module):
 48    def __init__(self, inp, out=2):
 49        super().__init__(); self.net=nn.Sequential(nn.Linear(inp,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,out))
 50    def forward(self,x): return self.net(x)
 51
 52def train(kind, X, Y, steps=500):
 53    dev='cuda' if torch.cuda.is_available() else 'cpu'
 54    try:
 55        tx=torch.tensor(X,dtype=torch.float32,device=dev); ty=torch.tensor(Y,dtype=torch.float32,device=dev)
 56        net=Net(2 if kind=='full' else 4).to(dev); opt=torch.optim.Adam(net.parameters(),lr=2e-3)
 57        for i in range(steps):
 58            idx=torch.randint(0,len(tx),(128,),device=dev); z=tx[idx]; target=ty[idx]
 59            if kind=='full': pred=net(z)
 60            else:
 61                # cheap background is a scaled oscillator drift; defect is explicitly supplied
 62                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
 63                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
 64                inp=torch.cat((bg,r,defect[:,1:2]),1) # 5 features
 65                if net.net[0].in_features != 5: net=Net(5).to(dev); opt=torch.optim.Adam(net.parameters(),lr=2e-3)
 66                pred=bg + dt_step(z,bg,r,defect,net(inp))
 67            loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 68        with torch.no_grad():
 69            z=tx[:512]; target=ty[:512]
 70            if kind=='full': pred=net(z)
 71            else:
 72                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
 73                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
 74                inp=torch.cat((bg,r,defect[:,1:2]),1); pred=bg+dt_step(z,bg,r,defect,net(inp))
 75            one=float(torch.mean((pred-target)**2).sqrt().cpu())
 76        return net,one,dev
 77    except Exception as e:
 78        # CPU retry is intentionally simple and deterministic.
 79        if dev!='cpu': torch.cuda.empty_cache(); torch.cuda.is_available=lambda:False; return train(kind,X,Y,steps)
 80        raise e
 81
 82def dt_step(z,bg,r,defect,closure,dt=.04):
 83    linear=torch.stack((r[:,1],-r[:,0]),1)
 84    corr=torch.stack((torch.zeros_like(r[:,0]),-(torch.cos(bg[:,0])-1)*r[:,0]),1)
 85    gate=torch.sigmoid(3.0*(torch.log(torch.linalg.vector_norm(defect,dim=1)+1e-6)-math.log(.08))).unsqueeze(1)
 86    return dt*(linear+defect+corr+gate*closure)
 87
 88def rollout(kind, net, init, horizon=80, dt=.04):
 89    z=torch.tensor(init,dtype=torch.float32,device=next(net.parameters()).device); out=[]
 90    with torch.no_grad():
 91        for _ in range(horizon):
 92            if kind=='full': z=net(z)
 93            else:
 94                bg=.82*z; r=z-bg; bgdot=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1)*.82
 95                fbg=torch.stack((bg[:,1],-torch.sin(bg[:,0])),1); defect=fbg-bgdot
 96                inp=torch.cat((bg,r,defect[:,1:2]),1); z=bg+dt_step(z,bg,r,defect,net(inp))
 97            out.append(z.cpu().numpy())
 98    return np.stack(out)
 99
100def main():
101    checks=math_checks(); X,Y=make_data(); split=int(.8*len(X));
102    results={}
103    for k in ('full','residual'):
104        net,one,dev=train(k,X[:split],Y[:split]); pred=rollout(k,net,X[split:split+64],60)
105        true=[]; z=X[split:split+64].copy()
106        for _ in range(60): z=rk4(z,.04); true.append(z.copy())
107        longerr=float(np.sqrt(np.mean((pred-np.array(true))**2)))
108        results[k]={'one_step_rmse':one,'rollout_rmse':longerr,'device':dev}
109    Path('results.json').write_text(json.dumps({'math_checks':checks,'comparison':results},indent=2))
110    print(json.dumps({'math_checks':checks,'comparison':results},indent=2))
111
112if __name__=='__main__': main()