import json, math, time import numpy as np import torch from torch import nn SEED = 1234 torch.manual_seed(SEED); np.random.seed(SEED) DTYPE = torch.float64 try: DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if DEVICE.type == 'cuda': torch.zeros(1, device=DEVICE) except Exception: DEVICE = torch.device('cpu') # Bogacki-Shampine 3(2), an embedded explicit RK pair. def rk23_step(f, t, y, h): k1 = f(t, y) k2 = f(t + h/2, y + h*k1/2) k3 = f(t + 3*h/4, y + 3*h*k2/4) y3 = y + h*(2*k1/9 + k2/3 + 4*k3/9) k4 = f(t + h, y3) y2 = y + h*(7*k1/24 + k2/4 + k3/3 + k4/8) return y3, y2 def rk4_step(f, t, y, h): k1=f(t,y); k2=f(t+h/2,y+h*k1/2); k3=f(t+h/2,y+h*k2/2); k4=f(t+h,y+h*k3) return y+h*(k1+2*k2+2*k3+k4)/6 def adaptive_record(f, y0, t0, t1, atol=1e-7, rtol=1e-5, h0=.08, max_steps=10000): # The record is deliberately detached: accept/reject is not part of the graph. with torch.no_grad(): y=y0.detach().clone(); t=float(t0); h=float(h0); hs=[]; ys=[y.clone()]; rejects=0 direction=1.0 if t1 >= t0 else -1.0 h=abs(h)*direction while direction*(t1-t) > 1e-13 and len(hs) < max_steps: h=direction*min(abs(h), abs(t1-t)) yn, ye=rk23_step(f,t,y,h) scale=atol+rtol*torch.maximum(y.abs(),yn.abs()) err=torch.sqrt(torch.mean(((yn-ye)/scale)**2)).item() if err <= 1.0 or abs(h) <= 1e-12: y=yn; t += h; hs.append(h); ys.append(y.clone()) fac=2.0 if err == 0 else min(2.0,max(.2,.9*err**(-1/3))) h=direction*min(abs(h)*fac, abs(t1-t) if direction*(t1-t)>0 else abs(h)*fac) else: rejects += 1 h=direction*abs(h)*max(.1,.9*err**(-1/3)) if len(hs)>=max_steps: raise RuntimeError('adaptive solver exceeded max_steps') return torch.tensor(hs, dtype=y0.dtype, device=y0.device), torch.stack(ys), rejects def replay(f, y0, hs, step='rk23'): y=y0 t=torch.zeros((), dtype=y.dtype, device=y.device) for h in hs: if step=='rk23': y,_=rk23_step(f,t,y,h) else: y=rk4_step(f,t,y,h) t=t+h return y def adaptive_differentiable(f, y0, t0=0., t1=1., atol=2e-4, rtol=2e-4, h0=.1, max_steps=1000): # Genuine adaptive baseline: accepted maps remain differentiable, while # error tests and step-size decisions use detached scalar values. y=y0; t=float(t0); h=float(h0); direction=1.0 if t1 >= t0 else -1.0 h=abs(h)*direction; accepted=0 while direction*(t1-t) > 1e-13 and accepted < max_steps: h=direction*min(abs(h), abs(t1-t)) yn, ye=rk23_step(f,t,y,h) scale=atol+rtol*torch.maximum(y.detach().abs(),yn.detach().abs()) err=torch.sqrt(torch.mean(((yn-ye)/scale)**2)).detach().item() if err <= 1.0 or abs(h) <= 1e-12: y=yn; t += h; accepted += 1 fac=2.0 if err == 0 else min(2.0,max(.2,.9*err**(-1/3))) h=direction*abs(h)*fac else: h=direction*abs(h)*max(.1,.9*err**(-1/3)) if accepted >= max_steps: raise RuntimeError('adaptive differentiable solver exceeded max_steps') return y class Field(nn.Module): def __init__(self): super().__init__() self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,2)) def forward(self,t,y): return y + 0.35*self.net(y) def math_check(): # Nonlinear scalar parameterized field gives a sensitive, independently checkable gradient. p=torch.tensor([0.7],dtype=DTYPE,requires_grad=True); y0=torch.tensor([[0.4]],dtype=DTYPE) def f(t,y): return p*y + 0.2*y**3 hs, states, rej=adaptive_record(f,y0,0.,1.,atol=1e-8,rtol=1e-7,h0=.2) yr=replay(f,y0,hs); loss=(yr**2).sum(); loss.backward(); g=p.grad.item() eps=1e-5 def evalp(v): pv=torch.tensor([v],dtype=DTYPE); return (replay(lambda t,y: pv*y+0.2*y**3,y0,hs)**2).sum().item() gfd=(evalp(.7+eps)-evalp(.7-eps))/(2*eps) # replay and recorded adaptive terminal state should coincide by construction state_err=(yr-states[-1]).abs().max().item() return {'accepted_steps':int(len(hs)), 'rejected_steps':int(rej), 'terminal_state_error':state_err, 'autodiff_gradient':g, 'finite_difference_gradient':gfd, 'relative_gradient_error':abs(g-gfd)/max(1e-12,abs(gfd))} def train_benchmark(): torch.manual_seed(SEED) n=64; y0=torch.randn(n,2,dtype=DTYPE,device=DEVICE) teacher=Field().to(DEVICE).double() with torch.no_grad(): target=replay(teacher,y0,torch.tensor([.05]*20,dtype=DTYPE,device=DEVICE),'rk4') models=[Field().to(DEVICE).double() for _ in range(3)] names=['fixed_rk4','adaptive_diff','recorded_replay']; opts=[torch.optim.Adam(m.parameters(),lr=.025) for m in models] # Equal initial weights, then record one shared mesh from that initial model. sd=models[0].state_dict(); models[1].load_state_dict(sd); models[2].load_state_dict(sd) f0=lambda t,y: models[2](t,y) hs,_,rej=adaptive_record(f0,y0,0.,1.,atol=2e-4,rtol=2e-4,h0=.1) fixed_h=torch.tensor([1/len(hs)]*len(hs),dtype=DTYPE,device=DEVICE) times=[]; final=[] for model,opt,name in zip(models,opts,names): if DEVICE.type=='cuda': torch.cuda.synchronize() t0=time.perf_counter() for _ in range(35): opt.zero_grad() ff=lambda t,y: model(t,y) if name=='fixed_rk4': pred=replay(ff,y0,fixed_h,'rk4') elif name=='adaptive_diff': pred=adaptive_differentiable(ff,y0) else: pred=replay(ff,y0,hs,'rk23') loss=((pred-target)**2).mean(); loss.backward(); opt.step() if DEVICE.type=='cuda': torch.cuda.synchronize() times.append((time.perf_counter()-t0)/35) with torch.no_grad(): ff=lambda t,y: model(t,y) pred=replay(ff,y0,hs,'rk23' if name!='fixed_rk4' else 'rk4') final.append(((pred-target)**2).mean().item()) # Gradient cosine between differentiable adaptive control and replay at common initial weights. m=Field().to(DEVICE).double(); hs2,_,_=adaptive_record(lambda t,y:m(t,y),y0,0.,1.,atol=2e-4,rtol=2e-4,h0=.1) m.zero_grad(); l1=((adaptive_differentiable(lambda t,y:m(t,y),y0)-target)**2).mean(); l1.backward(); g1=torch.cat([p.grad.flatten() for p in m.parameters()]) m.zero_grad(); l2=((replay(lambda t,y:m(t,y),y0,hs2,'rk23')-target)**2).mean(); l2.backward(); g2=torch.cat([p.grad.flatten() for p in m.parameters()]) cos=torch.nn.functional.cosine_similarity(g1,g2,dim=0).item() return {'device':str(DEVICE),'mesh_steps':len(hs),'mesh_rejections':rej,'seconds_per_update':dict(zip(names,times)), 'final_loss':dict(zip(names,final)), 'gradient_cosine_same_mesh':cos} if __name__=='__main__': out={'math_check':math_check(),'benchmark':train_benchmark()} print(json.dumps(out,indent=2))