1import json, math, random
 2import numpy as np
 3import torch
 4from torch import nn
 5from symplectic_block import SymplecticBlock
 6
 7SEED=646
 8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 9
10def gradV(q): return q+0.4*q**3
11def H(q,p): return .5*(q*q+p*p)+.1*q**4
12
13def leap(q,p,h):
14    p=p-.5*h*gradV(q); q=q+h*p; p=p-.5*h*gradV(q); return q,p
15def euler(q,p,h): return q+h*p,p-h*gradV(q)
16def metrics(fn,q0=1.,p0=.3,h=.1,steps=10000):
17    q,p=q0,p0; e0=float(H(q,p)); mx=me=0.; div=None; samples=[]
18    for k in range(steps):
19        try: q,p=fn(q,p,h); norm=math.hypot(float(q),float(p)); en=float(H(q,p))
20        except (OverflowError,ValueError,ZeroDivisionError): div=k+1; norm=en=float('inf'); break
21        if not math.isfinite(norm) or not math.isfinite(en): div=k+1; break
22        mx=max(mx,norm); me=max(me,abs(en-e0))
23        if k in (9,99,999,4999,9999): samples.append([k+1,norm,en-e0])
24    return {'final_norm':norm,'max_norm':mx,'max_abs_energy_error':me,'diverged_at_step':div,'samples':samples}
25
26def learned_regression():
27    torch.manual_seed(SEED); n=256; x=torch.randn(n,1); q=torch.randn(n,2); p=torch.zeros_like(q)
28    target_q=q+0.1*p; target_p=p-0.1*(q+0.4*q**3)
29    sym=SymplecticBlock(2,1,substeps=2,hidden=16,step=.1)
30    opt=torch.optim.Adam(sym.parameters(),lr=3e-3)
31    for _ in range(80):
32        opt.zero_grad(); qo,po=sym(q,p,x); loss=((qo-target_q)**2+(po-target_p)**2).mean(); loss.backward(); opt.step()
33    qo,po=sym(q,p,x); symloss=float(((qo-target_q)**2+(po-target_p)**2).mean().detach())
34    base=nn.Sequential(nn.Linear(5,16),nn.Tanh(),nn.Linear(16,4)); opt=torch.optim.Adam(base.parameters(),lr=3e-3)
35    inp=torch.cat([q,p,x.expand(n,1)],1); target=torch.cat([target_q,target_p],1)
36    for _ in range(80):
37        opt.zero_grad(); loss=((base(inp)-target)**2).mean(); loss.backward(); opt.step()
38    baseloss=float(((base(inp)-target)**2).mean().detach())
39    return {'one_step_mse':{'euler_analogue':baseloss,'learned_symplectic':symloss},'parameters':{'euler':sum(p.numel() for p in base.parameters()),'symplectic':sum(p.numel() for p in sym.parameters())}}
40
41def main():
42    z=torch.tensor([.7,-.4],dtype=torch.float64,requires_grad=True)
43    def lf(v): return torch.stack(leap(v[0],v[1],.2))
44    def eu(v): return torch.stack(euler(v[0],v[1],.2))
45    dl=float(torch.linalg.det(torch.autograd.functional.jacobian(lf,z))); de=float(torch.linalg.det(torch.autograd.functional.jacobian(eu,z)))
46    qr,pr=1.,.3
47    for _ in range(100000): qr,pr=leap(qr,pr,1e-3)
48    errors={}
49    for name,fn in [('leapfrog',leap),('euler',euler)]:
50        try:
51            q,p=1.,.3
52            for _ in range(1000): q,p=fn(q,p,.1)
53            errors[name]=float(math.hypot(q-qr,p-pr))
54        except (OverflowError,ValueError): errors[name]=None
55    out={'seed':SEED,'jacobian_determinant':{'leapfrog':dl,'euler':de},'rollouts':{str(h):{'leapfrog':metrics(leap,h=h),'euler':metrics(euler,h=h)} for h in (.1,.2)},'T100_error':errors,'learned_regression':learned_regression()}
56    with open('results.json','w') as f: json.dump(out,f,indent=2)
57    print(json.dumps(out,indent=2))
58if __name__=='__main__': main()