1import torch
 2from torch import nn
 3
 4
 5def _grad(y, value):
 6    return torch.autograd.grad(value.sum(), y, create_graph=True, retain_graph=True)[0]
 7
 8
 9class SymplecticBlock(nn.Module):
10    """Learned separable Hamiltonian block using velocity-Verlet/leapfrog."""
11    def __init__(self, width, input_dim, substeps=2, hidden=32, step=0.1):
12        super().__init__()
13        self.width, self.substeps, self.step = width, substeps, step
14        self.potential = nn.Sequential(nn.Linear(width + input_dim, hidden), nn.Tanh(),
15                                       nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
16        self.log_mass = nn.Parameter(torch.zeros(width))
17
18    def hamiltonian(self, q, p, x):
19        u = self.potential(torch.cat([q, x], dim=-1)).squeeze(-1)
20        mass = torch.nn.functional.softplus(self.log_mass) + 1e-3
21        kinetic = 0.5 * (p * p / mass).sum(-1)
22        return u + kinetic
23
24    def forward(self, q, p, x):
25        # x stays fixed over all internal substeps, as in the proposed block.
26        for _ in range(self.substeps):
27            q = q.requires_grad_(True); p = p.requires_grad_(True)
28            h = self.hamiltonian(q, p, x)
29            gq, gp = _grad(q, h), _grad(p, h)
30            p = p - 0.5 * self.step * gq
31            q = q + self.step * gp
32            q = q.requires_grad_(True); p = p.requires_grad_(True)
33            h = self.hamiltonian(q, p, x)
34            gq = _grad(q, h)
35            p = p - 0.5 * self.step * gq
36        return q, p
37
38
39def leapfrog(q, p, step, n=1):
40    """Analytic separable H=(q^2+p^2)/2 update for cheap verification."""
41    for _ in range(n):
42        p = p - 0.5 * step * q
43        q = q + step * p
44        p = p - 0.5 * step * q
45    return q, p
46
47
48def euler(q, p, step, n=1):
49    for _ in range(n):
50        q, p = q + step * p, p - step * q
51    return q, p