{"slug": "training-a-hamiltonian-neural-network", "title": "Training a Hamiltonian Neural Network", "summary": "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.", "body_md": "## Training a Hamiltonian Neural Network: Using an NN for Simulating the Phase Space of a Harmonic Oscillator\n\n## Introduction\n\nWe 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.\n\n## Prerequisites\n\nI 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).\n\n## Objectives\n\nWe 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.\n\n## Code\n\nAll code used in this project is available on GitHub: [**ritog/harmonic**](https://github.com/ritog/harmonic).\n\n## Our Physical System\n\nWe will work with a rigid pendulum.\n\nThe related variables are:\n\n- A body of mass m on a rigid rod of length l\n- Angle q (0 is straight down)\n- Angular momentum, p\n- Gravity g\n\n## Defining the Hamiltonian of Our System\n\n1. Potential Energy(V): Height is l(1-\\cos{q}), so, V = mgl(1-\\cos{q})\n\nHow?\n\n1. Kinetic Energy (T):\n\nRotational 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}.\n\nSo, Kinetic Energy in terms of momentum, T(p) = \\dfrac{p^2}{2ml^2}\n\nRemember that, in Hamiltonian mechanics, the Hamiltonian of the system is:\n\nH(q, p) = V + T\n\nAnd the Hamiltonian equations are:\n\n1. \\dfrac{dq}{dt} = \\dfrac{\\partial H}{\\partial p}\n2. \\dfrac{dp}{dt} = -\\dfrac{\\partial H}{\\partial q}\n\nIf we calculate theses, we will get:\n\n\\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}\n\nYou can trust me, or check using a pen and paper. Remember that I said the same thing in the last article as well?\n\n## Implementing\n\nThis is how we implement this in Python. Using this function, we can find the derivatives, and using these, we can find the trajectory.\n\nHere’s a basic implementation using NumPy:\n\n``` php\ndef pendulum_dynamics(t, state: List, m, l, g) -> List:\n    \"\"\"\n    inputs:\n    t: The current time, solvers expect it.\n    state: A list or array containing [q, p].\n    m: mass\n    l: the length of the rigid pendulum\n    g: gravitational acceleration\n    output:\n    the list [dq_dt, dp_dt]\n    \"\"\"\n    dq_dt = state[1] / (m * np.power(l, 2))\n    dp_dt = -m * g * l * np.sin(state[0])\n    return [dq_dt, dp_dt]\n```\n\nOr, we can write a version using `torch`. This will make things better.\n\n``` php\ndef pendulum_dynamics_tensor(t, state: torch.Tensor, m, l, g) -> torch.Tensor:\n    dq_dt = (state[:, 1] / (m * torch.pow(torch.tensor(l), 2))).unsqueeze(1)\n    dp_dt = (-m * g * l * torch.sin(state[:, 0])).unsqueeze(1)\n    return torch.cat([dq_dt, dp_dt], dim=1)\n```\n\nWe can use the earlier function to plot the trajectory of the pendulum.\n\n``` python\nfrom pendulum_nonlinear import pendulum_dynamics\n\n# different initial states\ninit_a = [0.5, 0]  # small release\ninit_b = [3.1, 0]  # close to top\ninit_c = [0, 5.0]  # close to bottom with force\n\n# params\nm = 1.0    # mass\nl = 1.0    # length of rod\ndt = 0.05  # time-step\ng = 9.8    # gravitational acceleration\n\ndef run_simulation(init_state: List):\n    p_vals = []\n    q_vals = []\n    for i in range(1_000):\n        q, p = init_state\n        _, dp_dt = pendulum_dynamics(t=dt, state=[q, p], m=m, l=l, g=g)\n        p = p + dp_dt * dt\n        dq_dt, _ = pendulum_dynamics(t=dt, state=[q, p], m=m, l=l, g=g)\n        q = q + dq_dt * dt\n        init_state = [q, p]\n        q_vals.append(q)\n        p_vals.append(p)\n    return p_vals, q_vals\n\na_p_vals, a_q_vals = run_simulation(init_a)\nb_p_vals, b_q_vals = run_simulation(init_b)\nc_p_vals, c_q_vals = run_simulation(init_c)\n\n# Plotting\nplt.plot(a_q_vals, a_p_vals, label=\"$q_0=0, p_0=0$\")\nplt.plot(b_q_vals, b_p_vals, label=\"$q_0=3.1, p_0=0$\")\nplt.plot(c_q_vals, c_p_vals, label=\"$q_0=0, p_0=5.0$\")\nplt.xlabel(\"$q$\")\nplt.ylabel(\"$p$\")\nplt.title(\"$p v. q$ for different inital conditions\")\nplt.legend()\nplt.tight_layout()\nplt.show()\n```\n\nThis is the plot that the code generates:\n\nThe 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.\n\nOne 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.\n\nThis 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.\n\nI am not writing more about this here. Maybe in a future post I write more on this.\n\n## Hamiltonian Neural Network\n\nWith our Python function, we can generate the data from our knowledge of Physics.\n\nWith Deep Learning, we solve the opposite problem- going from data to the Physics.\n\nStandard Neural ODEs try to learn the derivatives directly from the available data.\n\nInput: (q, p), Output: (\\dfrac{dq}{dt}, \\dfrac{dp}{dt})\n\nBut, 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.\n\n\\hat{H} = NeuralNet(q, p;\\theta)\n\nHere, \\theta is the set of trainable parameters of the neural network.\n\nWe 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*.\n\n1. \\dfrac{d\\hat{q}}{dt} = \\dfrac{\\partial \\hat{H}}{\\partial p}\n2. \\dfrac{d\\hat{p}}{dt} = -\\dfrac{\\partial \\hat{H}}{\\partial q}\n\nHere is our Neural Network:\n\n``` python\nimport torch\nfrom torch import nn\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nclass HNN(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.linear_block = nn.Sequential(\n            nn.Linear(2, 200),\n            nn.Tanh(),\n            nn.Linear(200, 200),\n            nn.Tanh(),\n            nn.Linear(200, 1),\n        )\n\n    def forward(self, x):\n        H = self.linear_block(x)\n        return H\n```\n\nThere are two things to note in the code:\n\n- 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.\n- We chose the `tanh()` activation function. As the function is thoroughly differentiable at every point.\n\nWe want to feed a batch of 200 pairs of q and p to the NN.\n\n## Training the Hamiltonian Neural Network\n\nWe want to write vectorized code. We don’t want to bottleneck the model by feeding in data through naive for loops.\n\nFor that, we can write a function to get the derivatives of the model, in batches:\n\n``` python\nimport torch\nfrom HNN import HNN\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\ndef get_model_time_derivatives(model, x):\n    \"\"\"\n    Compute time derivatives [dq/dt, dp/dt] for a batch of inputs.\n    x: Tensor of shape (Batch_Size, 2)\n    \"\"\"\n    H_hat = model(x)\n\n    # sum the energy to get a scalar,\n    # the gradients separate out perfectly per row\n    grads = torch.autograd.grad(H_hat.sum(), x, create_graph=True)[0]\n\n    # grads shape: (Batch, 2) -> [dH/dq, dH/dp]\n\n    # flipping (Symplectic Swap (Hamilton's Eqs))\n    # dq/dt =  dH/dp\n    # dp/dt = -dH/dq\n\n    dH_dq = grads[:, 0].unsqueeze(1)\n    dH_dp = grads[:, 1].unsqueeze(1)\n\n    return torch.cat([dH_dp, -dH_dq], dim=1)  # - because minus\n```\n\nThere 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.\n\nI am not going deep into it for now. I hope that you know why this is the case.\n\nHere’s what the training script looks like:\n\n``` python\nimport torch\nfrom tqdm import tqdm\n\nfrom HNN import HNN\nfrom hnn_model_derivs import get_model_time_derivatives\nfrom pendulum_tensor import pendulum_dynamics_tensor\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# params\nm = 1.0    # mass\nl = 1.0    # length of rod\ndt = 0.05  # time-step\ng = 9.8    # gravitational acceleration\n\nmean = torch.tensor([-3.0, 3.0])\nstd = 0.1\ninit_states = (6 * torch.rand(1_000, 2) - 3).to(device).requires_grad_()\n\ntrue_derivatives = pendulum_dynamics_tensor(t=dt, state=init_states, m=m, l=l, g=g)\n\nhamiltonian_nn = HNN().to(device)\nloss_func = torch.nn.MSELoss()\noptimizer = torch.optim.Adam(hamiltonian_nn.parameters(), lr=1e-2)\n\nn_epochs = 1_200\n\nfor epoch in tqdm(range(n_epochs + 1)):\n    deriv_pred = get_model_time_derivatives(hamiltonian_nn, init_states)\n    loss = loss_func(deriv_pred, true_derivatives)\n\n    optimizer.zero_grad()\n    loss.backward(retain_graph=True)\n\n    optimizer.step()\n    if epoch % 100 == 0:\n        print(f\"Epoch: {epoch}\\t Loss: {loss}\")\n\ntorch.save(hamiltonian_nn.state_dict(), \"hamiltonian_nn_1.pth\")\n```\n\nNote 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)`.\n\nAfter running this script, I had a loss of 0.0009111023391596973 - which is great.\n\nNote 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.\n\nThroughout the training, and generating data, maintaining the graph of the computation is crucial.\n\n## Plotting the Trajectory as Predicted Using the Model\n\nNow, 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.\n\n``` python\n# Here, we have an already trained HNN\n# We use it to plot trajectory\nimport torch\nfrom matplotlib import pyplot as plt\n\nfrom HNN import HNN\nfrom hnn_model_derivs import get_model_time_derivatives\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# different initial states\ninit_a = torch.tensor([0.5, 0]).to(device)  # small release\ninit_b = torch.tensor([3.1, 0]).to(device)  # close to top\ninit_c = torch.tensor([0, 5.0]).to(device)  # kicked from near bottom\n\n# params\nm = 1.0  # mass\nl = 1.0  # length of rod\ndt = 0.05  # time-step\ng = 9.8  # gravitational acceleration\n\nhamiltonian_model = HNN().to(device)\nhamiltonian_model.load_state_dict(torch.load(\"hamiltonian_nn_1.pth\", weights_only=True))\n\ndef run_simulation_HNN(init_state: torch.Tensor):\n    p_vals = []\n    q_vals = []\n    for i in range(1_000):\n        curr_state_tensor = (\n            init_state.clone().detach().unsqueeze(0).requires_grad_(True)\n        )\n        derivs = get_model_time_derivatives(hamiltonian_model, curr_state_tensor)\n\n        dq_dt = derivs[0, 0]\n        dp_dt = derivs[0, 1]\n\n        q_old, p_old = init_state\n\n        p_new = p_old + dp_dt * dt\n        q_new = q_old + dq_dt * dt\n\n        init_state = torch.tensor([q_new, p_new]).to(device)\n\n        q_vals.append(q_new.item())\n        p_vals.append(p_new.item())\n    return p_vals, q_vals\n\na_p_vals, a_q_vals = run_simulation_HNN(init_a)\nb_p_vals, b_q_vals = run_simulation_HNN(init_b)\nc_p_vals, c_q_vals = run_simulation_HNN(init_c)\n\n# Plotting\nplt.plot(a_q_vals, a_p_vals, label=\"$q_0=0, p_0=0$\")\nplt.plot(b_q_vals, b_p_vals, label=\"$q_0=3.1, p_0=0$\")\nplt.plot(c_q_vals, c_p_vals, label=\"$q_0=0, p_0=5.0$\")\n\nplt.xlabel(\"$q$\")\nplt.ylabel(\"$p$\")\nplt.title(\"$p  v. q$ for different inital conditions simulated via Hamiltonian NN\")\nplt.legend()\nplt.tight_layout()\nplt.savefig(\"FIG5.png\")\n```\n\nNote 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))`.\n\nThis is the plot that we get:\n\nFor 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.\n\nThis is due to the fact that the point was an out-of-distribution data point for the model.\n\n## Conclusion\n\nWe 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.\n\n### Discuss\n\nIf 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).\n\n### Changelog\n\nThis is an Open Source blog. Feel free to inspect diffs in the GitHub [repo](https://github.com/ritog/ritog.github.io).\n\n### Cite this Article\n\n```\n@ONLINE {,\n    author = \"Ritobrata Ghosh\",\n    title  = \"Training a Hamiltonian Neural Network\",\n    month  = \"jan\",\n    year   = \"2026\",\n    url    = \"https://ritog.github.io/posts/hamiltonian_nn\"\n}\n```\n\n", "url": "https://wpnews.pro/news/training-a-hamiltonian-neural-network", "canonical_source": "https://ritoghosh.com/posts/hamiltonian_nn/", "published_at": "2026-09-20 16:24:28+00:00", "updated_at": "2026-09-20 16:53:17.275928+00:00", "lang": "en", "topics": ["machine-learning", "neural-networks", "ai-research", "developer-tools"], "entities": ["ritog/harmonic", "PyTorch", "NumPy", "GitHub"], "alternates": {"html": "https://wpnews.pro/news/training-a-hamiltonian-neural-network", "markdown": "https://wpnews.pro/news/training-a-hamiltonian-neural-network.md", "text": "https://wpnews.pro/news/training-a-hamiltonian-neural-network.txt", "jsonld": "https://wpnews.pro/news/training-a-hamiltonian-neural-network.jsonld"}}