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 (or true) vector field with a symplectic (leapfrog) scheme, which keeps the discretized trajectory’s energy error bounded and oscillating rather than drifting away over long rollouts – the numerical-integration counterpart of the continuous-time conservation property above.

Requires torch.

class HamiltonianNet(dim, hidden=(64, 64))[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 MLP computing H from the concatenated (q, p). 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])

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]

Symplectic 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. 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]