1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6
  7def seed_all(s=7):
  8    random.seed(s); np.random.seed(s); torch.manual_seed(s)
  9
 10
 11def get_device():
 12    if torch.cuda.is_available():
 13        try:
 14            torch.zeros(1, device='cuda')
 15            return torch.device('cuda')
 16        except Exception:
 17            pass
 18    return torch.device('cpu')
 19
 20
 21def leapfrog_matrix(h):
 22    # H=(q^2+p^2)/2; the displayed split leapfrog is a linear map.
 23    return np.array([[1-h*h/2, h], [-h*(1-h*h/4), 1-h*h/2]], dtype=float)
 24
 25
 26def symplectic_check():
 27    J=np.array([[0.,1.],[-1.,0.]])
 28    vals=[]
 29    for h in (0.05, 0.2, 0.8):
 30        A=leapfrog_matrix(h)
 31        vals.append({'h':h, 'det':float(np.linalg.det(A)),
 32                     'symplectic_error':float(np.linalg.norm(A.T@J@A-J)),
 33                     'energy_after_1000':float(np.linalg.norm(np.linalg.matrix_power(A,1000)@np.array([1.,0.]))**2/2)})
 34    # Euler control, which does not preserve the form.
 35    E=np.array([[1.,.2],[-.2,1.]])
 36    vals.append({'euler_h':.2, 'det':float(np.linalg.det(E)),
 37                 'symplectic_error':float(np.linalg.norm(E.T@J@E-J)),
 38                 'energy_after_1000':float(np.linalg.norm(np.linalg.matrix_power(E,1000)@np.array([1.,0.]))**2/2)})
 39    return vals
 40
 41
 42class HamNet(nn.Module):
 43    def __init__(self, width=32):
 44        super().__init__()
 45        self.net=nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,1))
 46    def forward(self,z): return self.net(z).squeeze(-1)
 47
 48
 49def ham_step(model,z,h,create_graph=True):
 50    # Two fixed-point iterations implement the implicit p half-step.
 51    q,p=z[:,0],z[:,1]
 52    ph=p
 53    for _ in range(2):
 54        zz=torch.stack((q,ph),1).requires_grad_(True)
 55        gq=torch.autograd.grad(model(zz).sum(),zz,create_graph=create_graph)[0][:,0]
 56        ph=p-.5*h*gq
 57    zz=torch.stack((q,ph),1).requires_grad_(True)
 58    gp=torch.autograd.grad(model(zz).sum(),zz,create_graph=create_graph)[0][:,1]
 59    qn=q+h*gp
 60    zz2=torch.stack((qn,ph),1).requires_grad_(True)
 61    gq2=torch.autograd.grad(model(zz2).sum(),zz2,create_graph=create_graph)[0][:,0]
 62    pn=ph-.5*h*gq2
 63    return torch.stack((qn,pn),1)
 64
 65
 66class Residual(nn.Module):
 67    def __init__(self, width=32, h=.15):
 68        super().__init__(); self.h=h
 69        self.net=nn.Sequential(nn.Linear(2,width),nn.Tanh(),nn.Linear(width,2))
 70    def forward(self,z): return z+self.h*self.net(z)
 71
 72
 73def make_data(n=256, length=31, h=.15, device='cpu'):
 74    # Exact oscillator rotations; each sequence has a random phase/amplitude.
 75    x=[]
 76    for _ in range(n):
 77        z=torch.randn(2)*.7
 78        seq=[z]
 79        for _ in range(length-1):
 80            q,p=seq[-1]; c=math.cos(h); s=math.sin(h)
 81            seq.append(torch.stack((c*q+s*p, -s*q+c*p)))
 82        x.append(torch.stack(seq))
 83    return torch.stack(x).to(device)
 84
 85
 86def train_model(model, data, steps=500, h=.15, device='cpu'):
 87    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 88    model.train()
 89    for it in range(steps):
 90        idx=torch.randint(0,data.shape[0],(64,),device=device)
 91        z=data[idx,0]; target=data[idx,1]
 92        pred=model(z)
 93        loss=((pred-target)**2).mean()
 94        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10.); opt.step()
 95    return float(loss.detach().cpu())
 96
 97
 98def rollout(model, initial, horizon=200):
 99    z=initial.clone(); out=[z]
100    # Hamiltonian evaluation needs autograd for input gradients, but no
101    # parameter graph is needed.
102    for _ in range(horizon):
103        # enable_grad is required by Hamiltonian input derivatives; detaching
104        # prevents graph growth and is harmless for rollout evaluation.
105        with torch.enable_grad():
106            z=model(z)
107        z=z.detach()
108        out.append(z)
109    return torch.stack(out,1)
110
111
112def evaluate(model, test, horizons=(10,50,200)):
113    pred=rollout(model,test[:,0],max(horizons))
114    result={}
115    for k in horizons:
116        result[str(k)]=float(((pred[:,k]-test[:,k])**2).mean().sqrt().cpu())
117    energy=(pred.pow(2).sum(-1)/2)
118    result['energy_abs_drift_200']=float((energy[:,200]-energy[:,0]).abs().mean().cpu())
119    return result
120
121
122def main():
123    seed_all(7); dev=get_device(); h=.15
124    check=symplectic_check()
125    train=make_data(256,31,h,dev); test=make_data(128,201,h,dev)
126    # Same architecture width, optimizer, steps and data for the learned transition.
127    seed_all(11); base=Residual(32,h); train_model(base,train,500,h,dev); b=evaluate(base,test)
128    seed_all(11); idea=HamNet(32); train_model(lambda z: ham_step(idea,z,h), train, 0, h, dev) if False else None
129    idea.to(dev); opt=torch.optim.Adam(idea.parameters(),lr=3e-3); idea.train()
130    for _ in range(500):
131        idx=torch.randint(0,train.shape[0],(64,),device=dev); z=train[idx,0]; target=train[idx,1]
132        pred=ham_step(idea,z,h); loss=((pred-target)**2).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(idea.parameters(),10.); opt.step()
133    idea.eval(); a=evaluate(lambda z: ham_step(idea,z,h),test)
134    print(json.dumps({'device':str(dev),'symplectic_check':check,'baseline':b,'idea':a},indent=2))
135
136if __name__=='__main__': main()