mixle.models.hamiltonian module

HamiltonianNet – a learned dynamical system that conserves its own energy, by construction.

An ordinary neural net fit to a system’s (q, p) -> (dq/dt, dp/dt) derivatives has no reason to conserve anything: it is free to predict a vector field whose flow spirals in or out, even if the true physical system is energy-conserving. A Hamiltonian network (Greydanus, Dzamba & Yosinski, 2019) instead learns a single SCALAR function H(q, p) (an MLP) and reads the dynamics off its symplectic gradient:

dq/dt =  dH/dp
dp/dt = -dH/dq

Along any trajectory generated by exactly this vector field, dH/dt = dH/dq . dq/dt + dH/dp . dp/dt = dH/dq . dH/dp - dH/dp . dH/dq = 0 identically – H is conserved along the flow for ANY choice of the network’s weights, trained or not. This is the same “hard constraint, not a soft penalty” idea as PINNRegression (which only penalizes a residual) and make_monotonic_mlp()/ build_convex_energy_net() (which enforce a shape by reparameterizing weights): here the architecture itself makes conservation a property of the dynamics, not of the density.

leapfrog_rollout integrates the learned vector field with a leapfrog scheme. That method is symplectic only for a separable Hamiltonian, H(q, p) = V(q) + T(p). HamiltonianNet therefore uses a structurally separable network by default. A general Hamiltonian can still be represented with separable=False for derivative learning, but leapfrog_rollout rejects it instead of silently applying an unsupported integrator.

Requires torch.

class HamiltonianNet(dim, hidden=(64, 64), *, separable=True)[source]

Bases: object

Learns a scalar H(q, p) and exposes the symplectic-gradient dynamics it implies.

dim is the dimension of q (and of p, always equal); hidden sizes the network computing H. By default the network is masked into independent V(q) and T(p) branches whose scalar outputs are summed, making the supported leapfrog premise structural rather than advisory. Set separable=False only when learning a general Hamiltonian that will be integrated by another method. module (the underlying torch.nn.Module) is exposed directly for training with an ordinary optimizer loop against derivative-matching data.

Parameters:

  • dim (int)

  • hidden (Sequence[int])

  • separable (bool)

hamiltonian(q, p)[source]

H(q, p), shape (...,) – squeezes the module’s scalar output dimension.

Parameters:
Return type:

Any

time_derivative(q, p)[source]

(dq/dt, dp/dt) = (dH/dp, -dH/dq) via autograd – the symplectic gradient of hamiltonian.

q/p must be leaf tensors (requires_grad_(True)) or already part of an active graph; the returned derivatives carry gradients back through self.module’s parameters, so this composes into a training loop that fits derivative-matching data end to end.

Parameters:
Return type:

tuple[Any, Any]

leapfrog_rollout(net, q0, p0, dt, n_steps)[source]

Leapfrog integration of net’s learned (or untrained) Hamiltonian flow.

Returns (qs, ps), each shaped (n_steps + 1, *q0.shape) – the trajectory including the initial state. This scheme is available only when net.is_separable is true. A symplectic integrator is the right numerical counterpart to a conservative continuous-time system: unlike a generic (e.g. Euler) integrator, its energy error stays bounded and oscillates rather than drifting monotonically over long rollouts.

Parameters:
Return type:

tuple[Any, Any]