1import json, math, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1234
  7torch.manual_seed(SEED); np.random.seed(SEED)
  8DTYPE = torch.float64
  9try:
 10    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if DEVICE.type == 'cuda': torch.zeros(1, device=DEVICE)
 12except Exception:
 13    DEVICE = torch.device('cpu')
 14
 15# Bogacki-Shampine 3(2), an embedded explicit RK pair.
 16def rk23_step(f, t, y, h):
 17    k1 = f(t, y)
 18    k2 = f(t + h/2, y + h*k1/2)
 19    k3 = f(t + 3*h/4, y + 3*h*k2/4)
 20    y3 = y + h*(2*k1/9 + k2/3 + 4*k3/9)
 21    k4 = f(t + h, y3)
 22    y2 = y + h*(7*k1/24 + k2/4 + k3/3 + k4/8)
 23    return y3, y2
 24
 25def rk4_step(f, t, y, h):
 26    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)
 27    return y+h*(k1+2*k2+2*k3+k4)/6
 28
 29def adaptive_record(f, y0, t0, t1, atol=1e-7, rtol=1e-5, h0=.08, max_steps=10000):
 30    # The record is deliberately detached: accept/reject is not part of the graph.
 31    with torch.no_grad():
 32        y=y0.detach().clone(); t=float(t0); h=float(h0); hs=[]; ys=[y.clone()]; rejects=0
 33        direction=1.0 if t1 >= t0 else -1.0
 34        h=abs(h)*direction
 35        while direction*(t1-t) > 1e-13 and len(hs) < max_steps:
 36            h=direction*min(abs(h), abs(t1-t))
 37            yn, ye=rk23_step(f,t,y,h)
 38            scale=atol+rtol*torch.maximum(y.abs(),yn.abs())
 39            err=torch.sqrt(torch.mean(((yn-ye)/scale)**2)).item()
 40            if err <= 1.0 or abs(h) <= 1e-12:
 41                y=yn; t += h; hs.append(h); ys.append(y.clone())
 42                fac=2.0 if err == 0 else min(2.0,max(.2,.9*err**(-1/3)))
 43                h=direction*min(abs(h)*fac, abs(t1-t) if direction*(t1-t)>0 else abs(h)*fac)
 44            else:
 45                rejects += 1
 46                h=direction*abs(h)*max(.1,.9*err**(-1/3))
 47        if len(hs)>=max_steps: raise RuntimeError('adaptive solver exceeded max_steps')
 48    return torch.tensor(hs, dtype=y0.dtype, device=y0.device), torch.stack(ys), rejects
 49
 50def replay(f, y0, hs, step='rk23'):
 51    y=y0
 52    t=torch.zeros((), dtype=y.dtype, device=y.device)
 53    for h in hs:
 54        if step=='rk23': y,_=rk23_step(f,t,y,h)
 55        else: y=rk4_step(f,t,y,h)
 56        t=t+h
 57    return y
 58
 59def adaptive_differentiable(f, y0, t0=0., t1=1., atol=2e-4, rtol=2e-4, h0=.1, max_steps=1000):
 60    # Genuine adaptive baseline: accepted maps remain differentiable, while
 61    # error tests and step-size decisions use detached scalar values.
 62    y=y0; t=float(t0); h=float(h0); direction=1.0 if t1 >= t0 else -1.0
 63    h=abs(h)*direction; accepted=0
 64    while direction*(t1-t) > 1e-13 and accepted < max_steps:
 65        h=direction*min(abs(h), abs(t1-t))
 66        yn, ye=rk23_step(f,t,y,h)
 67        scale=atol+rtol*torch.maximum(y.detach().abs(),yn.detach().abs())
 68        err=torch.sqrt(torch.mean(((yn-ye)/scale)**2)).detach().item()
 69        if err <= 1.0 or abs(h) <= 1e-12:
 70            y=yn; t += h; accepted += 1
 71            fac=2.0 if err == 0 else min(2.0,max(.2,.9*err**(-1/3)))
 72            h=direction*abs(h)*fac
 73        else:
 74            h=direction*abs(h)*max(.1,.9*err**(-1/3))
 75    if accepted >= max_steps: raise RuntimeError('adaptive differentiable solver exceeded max_steps')
 76    return y
 77
 78class Field(nn.Module):
 79    def __init__(self):
 80        super().__init__()
 81        self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,2))
 82    def forward(self,t,y): return y + 0.35*self.net(y)
 83
 84def math_check():
 85    # Nonlinear scalar parameterized field gives a sensitive, independently checkable gradient.
 86    p=torch.tensor([0.7],dtype=DTYPE,requires_grad=True); y0=torch.tensor([[0.4]],dtype=DTYPE)
 87    def f(t,y): return p*y + 0.2*y**3
 88    hs, states, rej=adaptive_record(f,y0,0.,1.,atol=1e-8,rtol=1e-7,h0=.2)
 89    yr=replay(f,y0,hs); loss=(yr**2).sum(); loss.backward(); g=p.grad.item()
 90    eps=1e-5
 91    def evalp(v):
 92        pv=torch.tensor([v],dtype=DTYPE); return (replay(lambda t,y: pv*y+0.2*y**3,y0,hs)**2).sum().item()
 93    gfd=(evalp(.7+eps)-evalp(.7-eps))/(2*eps)
 94    # replay and recorded adaptive terminal state should coincide by construction
 95    state_err=(yr-states[-1]).abs().max().item()
 96    return {'accepted_steps':int(len(hs)), 'rejected_steps':int(rej), 'terminal_state_error':state_err,
 97            'autodiff_gradient':g, 'finite_difference_gradient':gfd,
 98            'relative_gradient_error':abs(g-gfd)/max(1e-12,abs(gfd))}
 99
100def train_benchmark():
101    torch.manual_seed(SEED)
102    n=64; y0=torch.randn(n,2,dtype=DTYPE,device=DEVICE)
103    teacher=Field().to(DEVICE).double()
104    with torch.no_grad(): target=replay(teacher,y0,torch.tensor([.05]*20,dtype=DTYPE,device=DEVICE),'rk4')
105    models=[Field().to(DEVICE).double() for _ in range(3)]
106    names=['fixed_rk4','adaptive_diff','recorded_replay']; opts=[torch.optim.Adam(m.parameters(),lr=.025) for m in models]
107    # Equal initial weights, then record one shared mesh from that initial model.
108    sd=models[0].state_dict(); models[1].load_state_dict(sd); models[2].load_state_dict(sd)
109    f0=lambda t,y: models[2](t,y)
110    hs,_,rej=adaptive_record(f0,y0,0.,1.,atol=2e-4,rtol=2e-4,h0=.1)
111    fixed_h=torch.tensor([1/len(hs)]*len(hs),dtype=DTYPE,device=DEVICE)
112    times=[]; final=[]
113    for model,opt,name in zip(models,opts,names):
114        if DEVICE.type=='cuda': torch.cuda.synchronize()
115        t0=time.perf_counter()
116        for _ in range(35):
117            opt.zero_grad()
118            ff=lambda t,y: model(t,y)
119            if name=='fixed_rk4': pred=replay(ff,y0,fixed_h,'rk4')
120            elif name=='adaptive_diff': pred=adaptive_differentiable(ff,y0)
121            else: pred=replay(ff,y0,hs,'rk23')
122            loss=((pred-target)**2).mean(); loss.backward(); opt.step()
123        if DEVICE.type=='cuda': torch.cuda.synchronize()
124        times.append((time.perf_counter()-t0)/35)
125        with torch.no_grad():
126            ff=lambda t,y: model(t,y)
127            pred=replay(ff,y0,hs,'rk23' if name!='fixed_rk4' else 'rk4')
128            final.append(((pred-target)**2).mean().item())
129    # Gradient cosine between differentiable adaptive control and replay at common initial weights.
130    m=Field().to(DEVICE).double();
131    hs2,_,_=adaptive_record(lambda t,y:m(t,y),y0,0.,1.,atol=2e-4,rtol=2e-4,h0=.1)
132    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()])
133    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()])
134    cos=torch.nn.functional.cosine_similarity(g1,g2,dim=0).item()
135    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}
136
137if __name__=='__main__':
138    out={'math_check':math_check(),'benchmark':train_benchmark()}
139    print(json.dumps(out,indent=2))