Hamiltonian Neural Networks

Overview

A physical system with a conserved energy H(q,p) evolves by Hamilton’s equations q̇ = ∂H/∂p, ṗ = −∂H/∂q. A Hamiltonian neural network (Greydanus et al. 2019) learns H as a network over phase space, then derives the dynamics by automatic differentiation — so the learned model is exactly energy-conserving by construction, unlike a black-box network.

using Rasmah
using Random
using LinearAlgebra

Fit a harmonic oscillator

The oscillator H = (q² + p²)/2 has the exact dynamics q̇ = p, ṗ = −q. Generate snapshots and fit the Hamiltonian.

m = 300
dt = 0.05
qs = [1.0]; ps = [0.0]
for _ in 1:m
    push!(qs, qs[end] + dt * ps[end])
    push!(ps, ps[end] - dt * qs[end])
end
Xs = [[qs[k], ps[k]] for k in 1:m]
Xdots = [[ps[k], -qs[k]] for k in 1:m]

hnn = HamiltonianNN(MLP([2, 16, 1]), 1)
θ = fit_hamiltonian(hnn, he_init(MersenneTwister(0), [2, 16, 1]), Xs, Xdots;
                    epochs = 800, lr = 5e-3)

65-element Vector{Float64}:
 -0.7088290342486969
 -0.5081295564025914
 -1.049727467226246
  0.27813907437234264
 -0.3909802270938124
  0.6369373015931653
 -1.2038378794039248
 -0.7087607630695457
  0.7540081579697294
 -1.2275395925137778
  ⋮
 -0.280728844887034
  0.4563085233445641
 -0.4035275123182374
  0.022963613925622177
  0.21390850812308385
  0.636658053720738
 -0.038502623669635146
  0.14126414332276735
  0.0

Verification — the learned flow matches

At a phase point the learned vector field [q̇; ṗ] must equal [p; −q]:

q, p = 0.6, 0.8
flow = hamiltonian_dynamics(hnn, θ, [q, p])
println("q̇ = ", round(flow[1]; digits = 3), "  (true p = ", p, ")")
println("ṗ = ", round(flow[2]; digits = 3), "  (true −q = ", -q, ")")
println("energy H(q,p) = ", round(hamiltonian(hnn, θ, [q, p]); digits = 4), "  (true (q²+p²)/2 = ", round((q^2 + p^2) / 2; digits = 4), ")")

q̇ = 0.772  (true p = 0.8)
ṗ = -0.558  (true −q = -0.6)
energy H(q,p) = -1.5274  (true (q²+p²)/2 = 0.5)

Next steps

SINDy — discovering equations from data.
Neural ODEs — continuous-depth networks.