GitHub

Hamiltonian Neural Network

The Hamiltonian Neural Network (HNN) [6] aims at building a Hamiltonian vector field with a neural network. We recall that a canonical Hamiltonian vector field on $\mathbb{R}^{2d}$ is one that can be written as:

\[ X_H(z) = \mathbb{J}_{2d}\nabla_zH,\]

where $\mathbb{J}_{2d}$ is the PoissonTensor. The idea behind a Hamiltonian neural network is to learn a vector field of this form, i.e. to learn:

\[ X_{\mathcal{NN}}(z) = \mathbb{J}_{2d}\nabla_z\mathcal{NN},\]

where $\mathcal{NN}:\mathbb{R}^{2d}\to\mathbb{R}$ is a neural network that approximates the Hamiltonian. There are then two different options to define a HNN loss, depending on the format in which the data are given.

HNN Loss for Vector Field Data

For the first loss, we assume that the given data describe the vector field of the HNN:

\[\mathcal{L}_\mathrm{HNN} = \sqrt{\sum_{i=1}^d\left(\Big|\Big|\frac{\partial\mathcal{NN}}{\partial{}q_i} + \dot{p}_i \Big|\Big|_2^2 + \Big|\Big| \frac{\partial\mathcal{NN}}{\partial{}p_i} - \dot{q}_i\Big|\Big|_2^2\right)}\]

HNN Loss for Phase Space Data

For the second loss, we assume that the given data describe points in phase space associated to a Hamiltonian system. For this approach we also need to specify a symplectic integrator [1] in order to train the neural network. In the following we use SymplecticEulerB to define this loss. This integrator does the following:

\[\mathrm{SymplecticEulerB}: (q^{(t)}, p^{(t)}) \mapsto (q^{(t+1)}, p^{(t+1)})\]

Note that this integrator is implicit in general.

\[\mathcal{L}_\mathrm{HNN} = \sqrt{\sum_t\sum_{i=1}^d \Big|\Big| \frac{\partial\mathcal{NN}}{\partial{}q_i}(q^{(t)}, p^{(t+1)}) + \frac{p_i^{t+1} - p_i^{(t)}}{\Delta{}t} \Big|\Big|_2^2 + \Big|\Big| \frac{\partial\mathcal{NN}}{\partial{}p_i}(q^{(t)}, p^{(t+1)}) + \frac{q_i^{t+1} - q_i^{(t)}}{\Delta{}t} \Big|\Big|_2^2}\]

Here the derivatives (i.e. vector field data) $\dot{q}_i^{(t)}$ and $\dot{p}_i^{(t)}$ are approximated with finite differences:

A loss of this kind is GeometricMachineLearning.SymplecticEulerLoss. Its :B variant evaluates the Hamiltonian vector field at $(q^{(t)}, p^{(t+1)})$, as above, and its :A variant at $(q^{(t+1)}, p^{(t)})$; it compares that field against the finite difference, with the sign convention $X_H = (\partial{}H/\partial{}p, -\partial{}H/\partial{}q)$.

Warning

The displayed formula is not what that loss computes. Under the convention above, the second term reads $-(q_i^{t+1} - q_i^{(t)})/\Delta{}t$, and the formula has a plus. The sign here predates the implementation and is left as it stands; the implementation is the one to read.

Info

Usually we use Zygote for computing derivatives in GeometricMachineLearning, but as the Zygote documentation itself points out: "Often using a different AD system over Zygote is a better solution [for computing second-order derivatives]." For this reason we compute the loss of the HNN with SymbolicNeuralNetworks and optionally also its gradient.

Library Functions

GeometricMachineLearning.hamiltonian_vector_field — Method
hamiltonian_vector_field(arch::HamiltonianArchitecture)

Compute an executable expression of the Hamiltonian vector field of a HamiltonianArchitecture.

Implementation

This first computes a symbolic expression of the vector field using symbolic_hamiltonian_vector_field.

The function is built with inplace = false: HNNLoss wraps it and is differentiated with Zygote, and the in-place kernel SymbolicNeuralNetworks.build_nn_function builds by default mutates its result, which Zygote does not support.

source
GeometricMachineLearning.StandardHamiltonianArchitecture — Type
StandardHamiltonianArchitecture <: HamiltonianArchitecture

A realization of the standard Hamiltonian neural network (HNN) [6].

Also see GeneralizedHamiltonianArchitecture.

Constructor

The constructor takes the following input arguments:

  1. dim: system dimension,
  2. width = dim: width of the hidden layer. By default this is equal to dim,
  3. nhidden = 1: the number of hidden layers,
  4. activation = AbstractNeuralNetworks.TanhActivation(): the activation function used in the HNN.
source
GeometricMachineLearning.HNNLoss — Type
HNNLoss <: NetworkLoss

The loss for a Hamiltonian neural network.

Constructor

This can be called with a NeuralNetwork, built with a HamiltonianArchitecture, as the only input arguemtn, i.e.:

HNNLoss(nn)

where nn is a NeuralNetwork, that is built with a HamiltonianArchitecture, gives the corresponding Hamiltonian loss.

Functor

loss(c, ps, input, output)
loss(ps, input, output) # equivalent to the above
source
GeometricMachineLearning.SymplecticEulerLoss — Type
SymplecticEulerLoss <: NetworkLoss

The loss that trains a Hamiltonian neural network on a trajectory, through one step of a symplectic Euler method.

HNNLoss needs $(\dot{q}, \dot{p})$ in the data. This one does not: it needs two consecutive states and the timestep between them, and it asks that one symplectic Euler step of the learned Hamiltonian carry the first state to the second.

Variant :A evaluates the Hamiltonian vector field at $(q_{n+1}, p_n)$ and variant :B at $(q_n, p_{n+1})$. Writing $X_H$ for that vector field and $\Delta{}t$ for the timestep, the residual of variant :A is

\[ X_H(q_{n+1}, p_n) - \frac{1}{\Delta{}t} \begin{pmatrix} q_{n+1} - q_n \\ p_{n+1} - p_n \end{pmatrix},\]

and the loss is its norm relative to the norm of the finite difference, so it is scale invariant in the same way HNNLoss is.

Constructor

SymplecticEulerLoss(arch, timestep)                 # variant :A
SymplecticEulerLoss(arch, timestep; variant = :B)

where arch is a HamiltonianArchitecture.

Functor

loss(model, ps, input, output)
loss(ps, input, output) # equivalent to the above

input is $(q_n, p_n)$ stacked and output is $(q_{n+1}, p_{n+1})$ stacked, so both have $2n$ rows.

source
GeometricMachineLearning._symplectic_euler_evaluation_point — Function
_symplectic_euler_evaluation_point(loss, input, output)

The staggered state at which the Hamiltonian vector field is evaluated: $(q_{n+1}, p_n)$ for variant :A, $(q_n, p_{n+1})$ for variant :B.

Taking one half of the state from input and the other from output is what makes the step implicit, and is the whole difference from an explicit Euler residual.

source
GeometricMachineLearning.symbolic_hamiltonian_vector_field — Method
symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork)

Get the symbolic expression for the vector field belonging to the HNN nn.

Implementation

This is calling SymbolicNeuralNetworks.Jacobian and then multiplies the result with a Poisson tensor.

source
SymbolicNeuralNetworks.SymbolicPullback — Method
SymbolicPullback(arch::HamiltonianArchitecture)

Make a SymbolicPullback based on a HamiltonianArchitecture.

Implementation

Internally this is calling SymbolicNeuralNetwork and HNNLoss.

SymbolicNeuralNetworks.SymbolicPullback(nn, loss) cannot be used here. It sizes the symbolic target of the loss with output_dimension(nn.model), which for a Hamiltonian network is 1 — the scalar Hamiltonian. HNNLoss does not compare against that: it compares against the Hamiltonian vector field, which has the dimension of the network's input. The construction below is the upstream one with that one dimension corrected.

source
GeometricMachineLearning.GeneralizedHamiltonianArchitecture — Type
GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture

A realization of generalized Hamiltonian neural networks (GHNNs) as introduced in [27].

Also see StandardHamiltonianArchitecture.

Constructor

The constructor takes the following input arguments:

  1. dim: system dimension,
  2. width = dim: width of the hidden layer. By default this is equal to dim,
  3. nhidden = 1: the number of hidden layers,
  4. activation = AbstractNeuralNetworks.TanhActivation(): the activation function used in the GHNN,
  5. integrator = nothing: the integrator that is used to design the GHNN.
source