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:

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.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
GeometricMachineLearning._processing — Function
_processing(returned_pullback)

Strip returned_pullback from unnecessary Zygote-induces garbage.

These two helpers used to be SymbolicNeuralNetworks._get_params and SymbolicNeuralNetworks._get_contents; SymbolicNeuralNetworks 0.5 removed them, and they were never about symbolics in the first place — they clean up what Zygote returns.

Also see the docs for ZygotePullback.

source