From Laplace to Neural Networks (Part 2)

We continue the discussion in Part 1, but now using neural networks. Can neural networks really predict a nonlinear system like double pendulum?

Well, we know from universal approximation theorem of neural networks that neural networks can approximate any smooth function. So given time-series data of the pendulum’s positions and velocities, which is governed by differential equations, we may train neural networks and forecast future dynamics.

This post is particularly motivated by a recent paper, which exploits the concepts of Hamiltonian mechanics.

Hamiltonian Mechanics

Hamiltonian of a system is expressed by generalized coordinates $q$ and momenta $p$:

The crux of Hamiltonian mechanics is Hamilton’s equations of motions, which are given by¹:

These relations have a special implication about the total energy of the system. Consider the total differential of $H$:

With Hamilton’s equations, the term in the parenthesis vanishes, and it follows that

Thus, if \(\mathcal{H}\) does not have an explicit dependency on $t$, then the Hamiltonian is a conserved quantity.

Let’s take a few examples.

Single Pendulum

In the case of single pendulum, the Hamiltonian is given as

where we defined the generalized coordinates \(q:=\theta\) and generalized momenta \(p:=mL^2\dot\theta\) ².

Hamilton’s equations of motions are:

Double Pendulum

Similarly, the Hamiltonian can be expressed as

Hamilton’s equations of motions are:

where:

Hamiltonian Neural Networks (HNN)

The key idea of HNN is learning a Hamiltonian from data. The neural networks model consumes generalized coordinates ($q$) and momenta ($p$) and produces a single “energy-like” scalar. Its training is based on the loss function that is a L2 discrepancy of Hamilton’s equations:

This approach is different from the baseline benchmark such as multilayer perceptrons, which learns the time-derivatives ($\dot q$ and $\dot p$) directly.

A striking thing about HNN is that even without hand-crafting the Hamiltonian in the model, the neural networks learn the Hamiltonian directly from data. Note that the architecture depicted above does not include anything about Hamiltonian - all we designed in HNN is to output a certain scalar ($\mathcal{H}$), and use the above loss function based on Hamilton’s equations. Nevertheless, the total-energy is preserved well in the system that is predicted by HNN models.

What if we feed physics knowledge more explicitly to the training proces? Yes, it is likely to learn even better. An additional penalty can be imposed as below.

where $H_0$ is a constant, the total-energy of the system measured from the initial condition.

Implementation of HNN

The codes are available in github.com/uriyeobi/hamiltonian_neural_networks. The implementation is inspired by that of the original paper where the authors use pytorch. Here I simplified and redesigned it in tensorflow.keras. I also added the energy-conservation parts and the task for double pendulum, which are not included in the paper.

As seen from the code below, the output dimension of HNN is 1 (for the scalar $H$), and the loss from Hamilton’s equations ($L_{data}$) is calculated by using the gradients (from forward method) and permutation matrix (M). For activation, tanh is used:

class HNN(tf.keras.Model):
    """Hamiltonian Neural Networks."""

    def __init__(self, input_dim, hidden_dims, **kwargs):
        super().__init__(**kwargs)
        self.feature_extractor = tf.keras.Sequential(
            [tf.keras.Input(shape=(input_dim,))]
            + [
                tf.keras.layers.Dense(hidden_dim, activation="tanh")
                for hidden_dim in hidden_dims
            ]
        )

        self.last_layer = tf.keras.layers.Dense(1)
        self.input_dim = input_dim

    @cached_property
    def M(self) -> npt.NDArray[Any]:
        """Permutation matrix (assuming canonical coordinates)."""
        M = np.eye(self.input_dim)
        M = np.concatenate(
            (M[self.input_dim // 2 :], -M[: self.input_dim // 2]), axis=0
        )
        return tf.constant(M, dtype="double")

    def call(self, x):
        """Call."""
        features = self.feature_extractor(x)
        outputs = self.last_layer(features)
        return outputs

    def forward(self, x):
        """Forward."""
        with tf.GradientTape() as tape:
            features = self.feature_extractor(x)
            outputs = self.last_layer(features)
        return (tape.gradient(outputs, x)) @ self.M

As for the loss function implemtation, we pre-calculate the initial Hamiltonian (ham0) and compare it with the new estimated Hamiltonian, throughout the training process. Then the loss for the energy conservation, $L_{physics}$, is added to the loss for Hamilton’s equations, $L_{data}$:

def get_loss(model, x, y, ham0, mlg, hamiltonian_method, penalty_lamb):
    """Get loss."""
    predictions = model.forward(tf.Variable(tf.stack(x)))
    ham_new = get_hamiltonian(
        x=x, mlg=mlg, hamiltonian_method=hamiltonian_method, predictions=predictions
    )

    physics_embedded_penalty = tf.reduce_mean(tf.square(ham0 - ham_new))

    return (
        tf.reduce_mean(tf.square(predictions - tf.Variable(tf.stack(y))))
        + penalty_lamb * physics_embedded_penalty
    )

Performance of HNN

The goal is to predict the time-series of generalized coordinates and momenta for the out-of-sample test data³.

Here is sample data:

• Single pendulum: $q, p, \dot q, \dot p$
• Double pendulum: $q_1, q_2, p_1, p_2, \dot q_1, \dot q_2, \dot p_1, \dot p_2$

The hyperparameters⁴ are set differently for single and double pendulums.

The out-of-sample performances are measured for the following three models:

• base: Base MLP
• hnn0: HNN with $\lambda=0$ ($L_{data}$)
• hnn1: HNN with $\lambda=0.1$ ($L_{data}+L_{physics}$)

Single Pendulum

Here are the results in the test dataset: time-series of $q, p$ and their errors. HNNs are much better than the base MLP.

Here are the errors in total energy (Note that y-axis is in log-scale). The energy-conservation is slightly better in hnn1, compared to hnn1.

For the single pendulum, both hnn0 and hnn1 predict the time-series pretty well.

Double Pendulum

Here is the predicted and actual time-series of $q1,q2,p1,p2$. As expected, it is not as perfect as the single pendulum.

We observe that hnn1 definitely outperforms the others (hnn0 is not necessarily better than base). Recall that these plots are out-of-sample forecasts starting from time step = 1600, so the predictions become less accurate as the time step increases⁵.

Lastly, the impact of embedding physics on the performance (i.e., adding the energy conservation penalty $L_{physics}$) is demonstrated in the following energy error plot - it helps conserve the total energy of the forecasted system!

Summary

In this and previous posts, we took the case of pendulum systems and explored the possibility of neural networks approaches to predict the non-linear dynamics in the out-of-sample region. In particular, we focused on the idea of Hamiltonian Neural Networks (HNN). Hamilton’s equations of motions represent very generic relations in physical systems, so the neural networks discover the physics from data. Consequently, the HNN works quite well in predicting the complex dynamics of the double pendulum. We also confirmed that the embedded physics⁶ in the model boosts up the performance.

Although HNN is a powerful tool, there is a wrinkle in it - the Hamiltonian formalism requires that the coordinates of the system be “canonical”, but many systems do not necessarily satisfy that condition. This motivated researchers to devise alternatives such as Lagrangian Neural Networks, which enables the model to learn with arbitrary coordinates.

Codes: https://github.com/uriyeobi/hamiltonian_neural_networks

Notes