Training a Hamiltonian Neural Network

A tutorial published on GitHub repository ritog/harmonic demonstrates training a Hamiltonian Neural Network to simulate the phase space of a harmonic oscillator — a rigid pendulum defined by mass m, rod length l, angle q, angular momentum p, and gravity g — by fitting the network to the system's derivatives rather than to ground-truth outputs. The Hamiltonian H(q, p) = V + T combines potential energy V = mgl(1-cos q) with kinetic energy T = p²/(2ml²), yielding the equations dq/dt = p/(ml²) and dp/dt = -mgl sin q, which the PyTorch implementation uses as training targets. The author states the approach requires differential calculus and basic Python but no prior Hamiltonian dynamics knowledge.

Training a Hamiltonian Neural Network: Using an NN for Simulating the Phase Space of a Harmonic Oscillator Introduction We can use Neural Networks not only for predicting the class of an image, completing sentences, or classifying sentiment of a paragraph of text. We can also use Neural Networks for solving scientific problems. In this article, you will learn about using a trained Neural Network to simulate the phase space of a Harmonic Oscillator a weighted rigid pendulum . Here, there’s a twist. We will train the Neural Network not by calculating the loss between the outputs and the ground truth data, but we will train using the derivatives . We will leverage our Physics knowledge to train the Neural Network. Prerequisites I expect the reader to be knowledgeable about basics of training a Neural Network from scratch- using a library like PyTorch, Jax, etc. I will use PyTorch in this project. I expect no knowledge of Hamiltonian Dynamics or college-level Physics, but knowing the basics of Physics up to High School level. Being well-versed in Differential Calculus is required. Experience with Python is helpful. I also expect the reader to have read an earlier post on simulating a spring with Hamiltonian Mechanics and plotting the phase space as well as the trajectory: Modelling a Spring System in Hamiltonian Mechanics: Using Euler’s Method for Trajectory Plotting ../../posts/implicit euler/index.html . Objectives We will model our system using Hamiltonian Dynamics. And using the Python function, we will plot the phase space of the system. And then using the the derivatives, we will train a Neural Network using those derivatives. Then we will use the trained Neural Network to again simulate the system and plot the phase space of the system. And we will see how the NN is doing. If you don’t know about Hamiltonian Dynamics or phase spaces, then that’s okay. I will cover what we need. Code All code used in this project is available on GitHub: ritog/harmonic https://github.com/ritog/harmonic . Our Physical System We will work with a rigid pendulum. The related variables are: - A body of mass m on a rigid rod of length l - Angle q 0 is straight down - Angular momentum, p - Gravity g Defining the Hamiltonian of Our System 1. Potential Energy V : Height is l 1-\cos{q} , so, V = mgl 1-\cos{q} How? 1. Kinetic Energy T : Rotational kinetic energy is \dfrac12 I {\omega}^2. For a point mass, moment of inertia, I=ml^2, and angular velocity, \omega = \dfrac{p}{ml^2}. So, Kinetic Energy in terms of momentum, T p = \dfrac{p^2}{2ml^2} Remember that, in Hamiltonian mechanics, the Hamiltonian of the system is: H q, p = V + T And the Hamiltonian equations are: 1. \dfrac{dq}{dt} = \dfrac{\partial H}{\partial p} 2. \dfrac{dp}{dt} = -\dfrac{\partial H}{\partial q} If we calculate theses, we will get: \begin{aligned} \frac{dq}{dt} &= \frac{p}{ml^2} & \text{ Angular Velocity } \\ \frac{dp}{dt} &= -mgl \sin q & \text{ Torque due to Gravity } \end{aligned} You can trust me, or check using a pen and paper. Remember that I said the same thing in the last article as well? Implementing This is how we implement this in Python. Using this function, we can find the derivatives, and using these, we can find the trajectory. Here’s a basic implementation using NumPy: php def pendulum dynamics t, state: List, m, l, g - List: """ inputs: t: The current time, solvers expect it. state: A list or array containing q, p . m: mass l: the length of the rigid pendulum g: gravitational acceleration output: the list dq dt, dp dt """ dq dt = state 1 / m np.power l, 2 dp dt = -m g l np.sin state 0 return dq dt, dp dt Or, we can write a version using torch . This will make things better. php def pendulum dynamics tensor t, state: torch.Tensor, m, l, g - torch.Tensor: dq dt = state :, 1 / m torch.pow torch.tensor l , 2 .unsqueeze 1 dp dt = -m g l torch.sin state :, 0 .unsqueeze 1 return torch.cat dq dt, dp dt , dim=1 We can use the earlier function to plot the trajectory of the pendulum. python from pendulum nonlinear import pendulum dynamics different initial states init a = 0.5, 0 small release init b = 3.1, 0 close to top init c = 0, 5.0 close to bottom with force params m = 1.0 mass l = 1.0 length of rod dt = 0.05 time-step g = 9.8 gravitational acceleration def run simulation init state: List : p vals = q vals = for i in range 1 000 : q, p = init state , dp dt = pendulum dynamics t=dt, state= q, p , m=m, l=l, g=g p = p + dp dt dt dq dt, = pendulum dynamics t=dt, state= q, p , m=m, l=l, g=g q = q + dq dt dt init state = q, p q vals.append q p vals.append p return p vals, q vals a p vals, a q vals = run simulation init a b p vals, b q vals = run simulation init b c p vals, c q vals = run simulation init c Plotting plt.plot a q vals, a p vals, label="$q 0=0, p 0=0$" plt.plot b q vals, b p vals, label="$q 0=3.1, p 0=0$" plt.plot c q vals, c p vals, label="$q 0=0, p 0=5.0$" plt.xlabel "$q$" plt.ylabel "$p$" plt.title "$p v. q$ for different inital conditions" plt.legend plt.tight layout plt.show This is the plot that the code generates: The Blue and green loops represent the pendulum swinging back and forth. It doesn’t have enough energy to go over the top, so it stays trapped in a closed loop. The Orange loop is right on the edge If you had just a tiny bit more energy, the pendulum would stop swinging back and start spinning 360° continuously. One of the reasons that Hamiltonian Neural Networks are better than vanilla Neural ODEs comes from Hamiltonian Mechanics. There’s a special property called Liouville’s Theorem, which says that the total area under the curve for the total set of initial points will be preserved for the total set of end points. That is, if you think of the initial points as a blob of points, then the blob might stretch, skew, twist, or distort, but the total area will remain constant. This is how Neural Networks handle data in higher dimension, and maps training data to target. I recommend that you watch Alfredo Canziani’s video from NYU CDS: 02 – Neural nets: rotation and squashing https://youtu.be/0TdAmZUMj2k?si=dcQG-2jwaYiVwyHx&t=1127 to get great a visual grasp of this. These kind of transformations are modelled and studied extensively in Linear Algebra. This incompressible flow of points, and fluid-like behaviour as opposed to gas-like, as gases compress and expand are great for Neural Networks, and NNs trained using Hamiltonian mechanics are much more robust and well-behaved than simple NODEs. I am not writing more about this here. Maybe in a future post I write more on this. Hamiltonian Neural Network With our Python function, we can generate the data from our knowledge of Physics. With Deep Learning, we solve the opposite problem- going from data to the Physics. Standard Neural ODEs try to learn the derivatives directly from the available data. Input: q, p , Output: \dfrac{dq}{dt}, \dfrac{dp}{dt} But, Hamiltonian Neural Networks are smarter. Instead of training the NN to predict the data, we force the NN to predict the Hamiltonian of the system. \hat{H} = NeuralNet q, p;\theta Here, \theta is the set of trainable parameters of the neural network. We ask the Neural Network to output the Hamiltonian of the system. And, as the Neural Network is just a chain of differentiable math operations, if we calculate the gradients of the output $ , with respect to the variables q and p, then, what we have are predicted time derivatives of the Hamiltonian . 1. \dfrac{d\hat{q}}{dt} = \dfrac{\partial \hat{H}}{\partial p} 2. \dfrac{d\hat{p}}{dt} = -\dfrac{\partial \hat{H}}{\partial q} Here is our Neural Network: python import torch from torch import nn device = "cuda" if torch.cuda.is available else "cpu" class HNN nn.Module : def init self : super . init self.linear block = nn.Sequential nn.Linear 2, 200 , nn.Tanh , nn.Linear 200, 200 , nn.Tanh , nn.Linear 200, 1 , def forward self, x : H = self.linear block x return H There are two things to note in the code: - There is no activation function at the end of the final Fully Connected FC layer. Because we want the Neural Network to output a value that is perceived as the total energy of the system, and it’s a real-numbered value, not limited to the range of the \tanh{} function. - We chose the tanh activation function. As the function is thoroughly differentiable at every point. We want to feed a batch of 200 pairs of q and p to the NN. Training the Hamiltonian Neural Network We want to write vectorized code. We don’t want to bottleneck the model by feeding in data through naive for loops. For that, we can write a function to get the derivatives of the model, in batches: python import torch from HNN import HNN device = "cuda" if torch.cuda.is available else "cpu" def get model time derivatives model, x : """ Compute time derivatives dq/dt, dp/dt for a batch of inputs. x: Tensor of shape Batch Size, 2 """ H hat = model x sum the energy to get a scalar, the gradients separate out perfectly per row grads = torch.autograd.grad H hat.sum , x, create graph=True 0 grads shape: Batch, 2 - dH/dq, dH/dp flipping Symplectic Swap Hamilton's Eqs dq/dt = dH/dp dp/dt = -dH/dq dH dq = grads :, 0 .unsqueeze 1 dH dp = grads :, 1 .unsqueeze 1 return torch.cat dH dp, -dH dq , dim=1 - because minus There is a nice trick with the summing up of the predicted Hamiltonians. PyTorch can only find gradients of scalars. And here we have a tensor of predicted Hamiltonians. We can just sum them up, and then find the gradient with respect to the whole batch of the inputs. And everything gets neatly stored in rows. Gradient of a sum is equal to sum of gradients. I am not going deep into it for now. I hope that you know why this is the case. Here’s what the training script looks like: python import torch from tqdm import tqdm from HNN import HNN from hnn model derivs import get model time derivatives from pendulum tensor import pendulum dynamics tensor device = "cuda" if torch.cuda.is available else "cpu" params m = 1.0 mass l = 1.0 length of rod dt = 0.05 time-step g = 9.8 gravitational acceleration mean = torch.tensor -3.0, 3.0 std = 0.1 init states = 6 torch.rand 1 000, 2 - 3 .to device .requires grad true derivatives = pendulum dynamics tensor t=dt, state=init states, m=m, l=l, g=g hamiltonian nn = HNN .to device loss func = torch.nn.MSELoss optimizer = torch.optim.Adam hamiltonian nn.parameters , lr=1e-2 n epochs = 1 200 for epoch in tqdm range n epochs + 1 : deriv pred = get model time derivatives hamiltonian nn, init states loss = loss func deriv pred, true derivatives optimizer.zero grad loss.backward retain graph=True optimizer.step if epoch % 100 == 0: print f"Epoch: {epoch}\t Loss: {loss}" torch.save hamiltonian nn.state dict , "hamiltonian nn 1.pth" Note the calculation of loss- loss = loss func deriv pred, true derivatives . We calculate the loss between predicted and true derivatives. You are usually accustomed to see the loss being calculated between y pred and y . But the parameters of the model get updated through backpropagation, as the optimizer receives them: optimizer = torch.optim.Adam hamiltonian nn.parameters , lr=1e-2 . After running this script, I had a loss of 0.0009111023391596973 - which is great. Note that I train the model derivatives to be close to the true derivatives - unlike normal NNs - where we train the model to output values close to the ground truth. Throughout the training, and generating data, maintaining the graph of the computation is crucial. Plotting the Trajectory as Predicted Using the Model Now, I will plot the trajectory as solved from the derivatives predicted by the Neural Network. Here, we are using Semi-Implicit Euler’s Method to find points in the phase space. To learn more about this, read the previously mentioned post. python Here, we have an already trained HNN We use it to plot trajectory import torch from matplotlib import pyplot as plt from HNN import HNN from hnn model derivs import get model time derivatives device = "cuda" if torch.cuda.is available else "cpu" different initial states init a = torch.tensor 0.5, 0 .to device small release init b = torch.tensor 3.1, 0 .to device close to top init c = torch.tensor 0, 5.0 .to device kicked from near bottom params m = 1.0 mass l = 1.0 length of rod dt = 0.05 time-step g = 9.8 gravitational acceleration hamiltonian model = HNN .to device hamiltonian model.load state dict torch.load "hamiltonian nn 1.pth", weights only=True def run simulation HNN init state: torch.Tensor : p vals = q vals = for i in range 1 000 : curr state tensor = init state.clone .detach .unsqueeze 0 .requires grad True derivs = get model time derivatives hamiltonian model, curr state tensor dq dt = derivs 0, 0 dp dt = derivs 0, 1 q old, p old = init state p new = p old + dp dt dt q new = q old + dq dt dt init state = torch.tensor q new, p new .to device q vals.append q new.item p vals.append p new.item return p vals, q vals a p vals, a q vals = run simulation HNN init a b p vals, b q vals = run simulation HNN init b c p vals, c q vals = run simulation HNN init c Plotting plt.plot a q vals, a p vals, label="$q 0=0, p 0=0$" plt.plot b q vals, b p vals, label="$q 0=3.1, p 0=0$" plt.plot c q vals, c p vals, label="$q 0=0, p 0=5.0$" plt.xlabel "$q$" plt.ylabel "$p$" plt.title "$p v. q$ for different inital conditions simulated via Hamiltonian NN" plt.legend plt.tight layout plt.savefig "FIG5.png" Note that, we are not using a Python function for this plot, but load the trained weights from a .pth file: hamiltonian model.load state dict torch.load "hamiltonian nn 1.pth", weights only=True . This is the plot that we get: For the inner loops blue and green , The network did a decent job learning the “swinging” motion It captured the concentric nature of the phase space near the center. The outer “loop” orange : This trajectory starts at q=3.1. In our clean phase space figure, this was a closed loop the “eye” . In our HNN simulation, it drifts off significantly. This is due to the fact that the point was an out-of-distribution data point for the model. Conclusion We have trained a Neural Network’s parameters so that the model learns the Physics from the data, by making its gradients be close to the real derivatives. We saw that we can leverage our Physics knowledge in training of Neural Networks, and leverage Physical properties like Liouville’s Theorem to train well-behaved NNs with predictable behaviour. Discuss If you have read this post, and found it interesting or edifying, please let me know. I would like that very much. If you have any criticism, suggestion, or want to tell me anything, just add a comment or let me know privately. Discuss this post on the Fediverse https://mathstodon.xyz/@rg/115844261208692910 , Hacker News https://news.ycombinator.com/item?id=46508707 , or Twitter/X https://x.com/AllesistKode/status/2008268111871713479 . Changelog This is an Open Source blog. Feel free to inspect diffs in the GitHub repo https://github.com/ritog/ritog.github.io . Cite this Article @ONLINE {, author = "Ritobrata Ghosh", title = "Training a Hamiltonian Neural Network", month = "jan", year = "2026", url = "https://ritog.github.io/posts/hamiltonian nn" }