Discovering Symbolic Laws Directly from Trajectories with Hamiltonian
  Graph Neural Networks
The time evolution of physical systems is described by differential
equations, which depend on abstract quantities like energy and force.
Traditionally, these quantities are derived as functionals based on observables
such as positions and velocities. Discovering these governing symbolic laws is
the key to comprehending the interactions in nature. Here, we present a
Hamiltonian graph neural network (HGNN), a physics-enforced GNN that learns the
dynamics of systems directly from their trajectory. We demonstrate the
performance of HGNN on n-springs, n-pendulums, gravitational systems, and
binary Lennard Jones systems; HGNN learns the dynamics in excellent agreement
with the ground truth from small amounts of data. We also evaluate the ability
of HGNN to generalize to larger system sizes, and to hybrid spring-pendulum
system that is a combination of two original systems (spring and pendulum) on
which the models are trained independently. Finally, employing symbolic
regression on the learned HGNN, we infer the underlying equations relating the
energy functionals, even for complex systems such as the binary Lennard-Jones
liquid. Our framework facilitates the interpretable discovery of interaction
laws directly from physical system trajectories. Furthermore, this approach can
be extended to other systems with topology-dependent dynamics, such as cells,
polydisperse gels, or deformable bodies.
Suresh Bishnoi, Ravinder Bhattoo,  Jayadeva, Sayan Ranu, N M Anoop Krishnan
# Discovering Symbolic Laws Directly from Trajectories with Hamiltonian Graph Neural Networks

###### Abstract

The time evolution of physical systems is described by differential equations, which depend on abstract quantities like energy and force. Traditionally, these quantities are derived as functionals based on observables such as positions and velocities. Discovering these governing symbolic laws is the key to comprehending the interactions in nature. Here, we present a Hamiltonian graph neural network (Hgnn), a physics-enforced Gnn that learns the dynamics of systems directly from their trajectory. We demonstrate the performance of Hgnn on \(n-\)springs, \(n-\)pendulums, gravitational systems, and binary Lennard Jones systems; Hgnn learns the dynamics in excellent agreement with the ground truth from small amounts of data. We also evaluate the ability of Hgnn to generalize to larger system sizes, and to hybrid spring-pendulum system that is a combination of two original systems (spring and pendulum) on which the models are trained independently. Finally, employing symbolic regression on the learned Hgnn, we infer the underlying equations relating the energy functionals, even for complex systems such as the binary Lennard-Jones liquid. Our framework facilitates the interpretable discovery of interaction laws directly from physical system trajectories. Furthermore, this approach can be extended to other systems with topology-dependent dynamics, such as cells, polydisperse gels, or deformable bodies.

Any system in the universe is always in a continuous state of motion. This motion, also known as the dynamics, is observed and noted in terms of the trajectory, which comprises the system's configuration (that is, positions and velocities) as a function of time. Any understanding humans have developed about the universe is through analyzing the dynamics of different systems. Traditionally, the dynamics governing a physical system are expressed as governing differential equations derived from fundamental laws such as energy or momentum conservation, which, when integrated,
provide the system's time evolution. However, these equations require the knowledge of functionals that relate abstract quantities such as energy, force, or stress with the configuration [1]. Thus, discovering these governing equations directly from the trajectory remains the key to understanding and comprehending the phenomena occurring in nature. Alternatively, several symbolic regression (SR) approaches have been used to discover free-form laws directly from observations [2, 3, 4]. However, the function space to explore in such cases is prohibitively large, and appropriate assumptions and constraints regarding the equations need to be provided to obtain a meaningful and straightforward equation [5, 6, 7].

Learning the dynamics of physical systems directly from their trajectory is a problem of interest in wide areas such as robotics, mechanics, biological systems such as proteins, and atomistic dynamics [8, 9, 10, 11, 12]. Recently, machine learning (ML) tools have been widely used to learn the dynamics of systems directly from the trajectory of systems [13, 14, 15, 16, 17, 18, 19]. Specifically, there have been three broad approaches to this extent, namely, data-driven, physics-informed, and physics-enforced approaches. Data-driven approaches try to develop models that learn the dynamics directly from ground-truth trajectories [13, 10, 12]. Physics-informed approaches rely on an additional term in the loss function, which is the governing differential equation: data loss and physics loss [9]. In contrast, physics-enforced approaches directly infuse the inductive biases in terms of the ordinary differential equations directly in the formulation as a hard constraint. These approaches are known as Hamiltonian (Hnn) [20, 21, 22, 14], and Lagrangian neural networks (Lnn) [15, 16, 17], and Graph Neural ODEs [23, 18, 24]. Adding the inductive bias in a physics-enforced fashion instead of a soft constraint in the loss function can significantly enhance the learning efficiency while also leading to realistic trajectories in terms of conservation laws [14, 22, 25]. Additionally, combining these formulations with graph neural networks (Gnns) [26, 27, 28, 25] can lead to superior properties such as zero-shot generalizability to unseen system sizes and hybrid systems unseen during the training, more efficient learning, and inference. However, although efficient in learning the dynamics, these approaches remain black-box in nature with poor interpretability of the learned function, which questions the robustness and correctness of the learned models [29].

Here, we present a framework combining Hamiltonian graph neural networks (Hgnn) and symbolic regression (SR), which enables the discovery of symbolic laws governing the energy functionals directly from the trajectory of systems. Specifically, we propose a Hgnn architecture that decouples kinetic and potential energies and, thereby, efficiently learns the Hamiltonian of a system directly from the trajectory. We evaluate our architecture on several complex systems such as \(n\)-pendulum, \(n\)-spring, \(n\)-particle gravitational, and binary LJ systems. Further, the modular nature of Hgnn enables the interpretability of the learned functions, which, when combined with SR, enables the discovery of the governing laws in a symbolic form, even for complex interactions such as binary LJ systems.

## Hamiltonian mechanics

Here, we briefly introduce the mathematical formulation of Hamiltonian mechanics that govern the dynamics of physical systems. Consider a system of \(n\) particles that are interacting with their positions at time \(t\) represented by the Cartesian coordinates as \(\mathbf{x}(t)=(\mathbf{x}_{1}(t),\mathbf{x}_{2}(t),...\mathbf{x}_{n}(t))\). The Hamiltonian \(H\) of the system is defined as \(H(\mathbf{p}_{\mathbf{x}},\mathbf{x})=T(\dot{\mathbf{x}})+V(\mathbf{x})\), where \(T(\dot{\mathbf{x}})\) represents the total kinetic energy and \(V(\mathbf{x})\) represents the potential energy of the system. The Hamiltonian equations of motion for this system in Cartesian coordinates are given by [30, 31, 32]

\[\dot{\mathbf{x}}=\nabla_{\mathbf{p}_{\mathbf{x}}}H,\qquad\dot{\mathbf{p}}_{ \mathbf{x}}=-\nabla_{\mathbf{x}}H \tag{1}\]

where \(\mathbf{p}_{\mathbf{x}}=\nabla_{\dot{x}}H=\mathbf{M}\dot{\mathbf{x}}\) represents the momentum of the system in Cartesian coordinates and \(\mathbf{M}\) represents the mass matrix. Assuming \(Z=[\mathbf{x};\mathbf{p}_{\mathbf{x}}]\) and \(J=[0,I;-I,0]\), the acceleration of a particle can be obtained from the Hamiltonian equations as

\[\dot{Z}=J(\nabla_{Z}H) \tag{2}\]

since \(\nabla_{Z}H+J\dot{Z}=0\) and \(J^{-1}=-J\). Sometimes systems may be subjected to constraints that depend on positions (holonomic) or velocities (Pfaffian). For example, in the case of a pendulum, the length between the bobs remains constant, or in multi-fingered grasping, the velocity of two fingers should be such that the combined geometry is able to hold the object. In such cases, the constrain equation is represented as \(\Phi(\mathbf{x})\dot{\mathbf{x}}=0\), where \(\Phi(\mathbf{x})\in\mathbb{R}^{k\times D}\) correspond to the \(k\) velocity constraints in a \(D\)-dimensional system. For instance, in the case of a pendulum, the constraint equation for two bobs located at \((0,0)\) and \((x_{1},x_{2})\) may be written as \(x_{1}\dot{x}_{1}+x_{2}\dot{x}_{2}=0\), which is the gradient of \(x_{1}^{2}+x_{2}^{2}=0\). Following this, the Hamiltonian equations of motion can be modified to feature the constraints explicitly as [16, 32]

\[\nabla_{Z}H+J\dot{Z}+(D_{Z}\Psi)^{T}\lambda=0 \tag{3}\]

where \(\Psi(Z)=(\Phi;\dot{\Phi})\), \(D_{Z}\Psi\) is the Jacobian of \(\Psi\) with respect to \(Z\), and \((D_{Z}\Psi)^{T}\lambda\) represents the effect of constraints on \(\dot{\mathbf{x}}\) and \(\dot{\mathbf{p}}_{\mathbf{x}}\)[16, 32]. Thus, \((D_{Z}\Psi)\dot{Z}=0\). Substituting for \(\dot{Z}\) from Eq. 3 and solving for \(\lambda\) yields [17, 25, 18, 30]

\[\lambda=-[(D_{Z}\Psi)J(D_{Z}\Psi)^{T}]^{-1}[(D_{Z}\Psi)J(\nabla_{Z}H)] \tag{4}\]
Substituting \(\lambda\) in the Eq. 3 and solving for \(\dot{Z}\) yields

\[\dot{Z}=J[\nabla_{Z}H-(D_{Z}\Psi)^{T}[(D_{Z}\Psi)J(D_{Z}\Psi)^{T}]^{-1}(D_{Z}\Psi )J\nabla_{Z}H] \tag{5}\]

Note that in the absence of constraint, Eq. 5 reduces to Eq. 2. In Hamiltonian mechanics, Eq.5 is used to obtain the acceleration of the particles, which, when integrated, provides the updated configuration of the system. Thus, the only unknown in the previous equation is the \(H\), which is represented as a function of \(\mathbf{p_{x}}\) and \(\mathbf{x}\).

### Hamiltonian graph neural network

Now, we introduce our ML framework proposed to learn the Hamiltonian of a system directly from the trajectory, that is, only using the time evolution of the observable quantities \((\mathbf{x},\mathbf{p_{x}})\). To this extent, we develop the Hamiltonian graph neural network (Hgnn) that parametrizes the actual \(H\) as a Gnn to obtain the learned \(\hat{H}\). Henceforth, all the terms with a hat, for example, \(\hat{\mathbf{x}}\) represent the approximate function obtained from Hgnn. Further, the \(\hat{H}\) obtained from Hgnn is substituted in the Eq.(5) to obtain the acceleration and velocity of the particles. These values are integrated using a symplectic integrator to compute the updated position.

First, we describe the architecture of Hgnn (see Fig. 1(a)). The physical system is modeled as an undirected graph \(\mathcal{G}=(\mathcal{V},\mathcal{E})\) with nodes as particles and edges as connections between them. For instance, in an \(n\)-ball-spring system, the balls are represented as nodes and springs as edges. The raw node features are \(t\) (type of particle) as one-hot encoding, \(\mathbf{x}\), and \(\mathbf{x}\), and the raw edge feature is the distance, \(d=||\mathbf{x}_{j}-\mathbf{x}_{i}||\), between two particles \(i\) and \(j\). A notable difference in the Hgnn architecture from previous graph architectures is the presence of global and local features--local features participate in message passing and contribute to quantities that depend on topology. In contrast, global features do not

Figure 1: **Hamiltonian graph architecture and systems studied.** (a) Hamiltonian graph neural network (Hgnn) architecture, (b) Visualization of the systems studied, namely, 3-pendulum, 5-spring, \(75\)-particles binary Lennard Jones system, \(4\)-particle gravitational system, and a hybrid spring-pendulum system. Note that the hybrid spring-pendulum system is used only to evaluate the generalizability of Hgnn.
take part in message passing. Here, we employ the position \(\mathbf{x}\), velocity \(\dot{\mathbf{x}}\) as global features for a node, while \(d\) and \(t\) are used as local features.

For the Gnn, we employ an \(L\)-layer message passing Gnn, which takes an embedding of the node and edge features created by multi-layer perceptrons (MLPs) as input. Detailed hyper-parameters are provided in the Supplementary Material. The local features participate in message passing to create an updated node and edge embeddings. The final representations of the nodes and edges, \(\mathbf{z}_{i}\) and \(\mathbf{z}_{ij}\), respectively, are passed through MLPs to obtain the Hamiltonian of the system. The Hamiltonian of the system is predicted as the sum of kinetic energy \(T\) and potential energy \(V\) in the Hgnn. Specifically, the potential energy is predicted as \(V_{i}=\sum_{i}\texttt{MLP}_{v}(\mathbf{z}_{i})+\sum_{ij}\texttt{MLP}_{v}( \mathbf{z}_{ij}))\), where \(\texttt{MLP}_{v}\) and \(\texttt{MLP}_{e}\) represent the contribution from the node (particles themselves) and edges (interactions) toward the potential energy of the system, respectively. Kinetic energy is predicted as \(T=\sum_{i}\texttt{MLP}_{T}\left(\mathbf{h}_{i}^{0}\right)\), where \(\mathbf{h}_{i}^{0}\) is the embedding of particle \(i\).

To train the Hgnn, we use only the time evolution of positions and momenta. This approach does not assume any knowledge of the functional form or knowledge of the Hamiltonian. The training approach, purely based on direct observables, can be used for any system (for example, trajectories from experiments) where the true Hamiltonian is unavailable. Thus, the loss function of Hgnn is computed by using the predicted and actual positions at the timestep \(t+1\) in a trajectory based on positions and velocities at \(t\), which is then back-propagated to train the MLPs. Specifically, we use _mean squared error (MSE)_ on the true and predicted \(Z\), which is the concatenation of positions and velocities.

\[\mathcal{L}=\frac{1}{n}\left(\sum_{i=1}^{n}\left(Z_{i}^{t+1}-\hat{Z}_{i}^{t+1} \right)^{2}\right) \tag{6}\]

### Case studies

**Systems studied.** Now, we evaluate the ability of Hgnn to learn the dynamics directly from the trajectory. To evaluate Hgnn, we selected four different types of systems, _viz_, \(5\)-pendulums with explicit internal constraints and subjected to an external gravitational field, \(5\)-springs with harmonic inter-particle harmonic interactions, \(75\)-particle binary LJ system with two types of a particle interacting based on the Kob-Andersen LJ potential [33], and \(4\)-particle gravitational system with purely repulsive gravitational potential. Finally, in order to test the generalizability of Hgnn to completely unseen system which is combination two systems on which it is trained, a hybrid system containing spring and pendulum is also considered. In this system, while the dynamics of pendulum is governed by the external gravitational field, the dynamics of the spring system depends on the internal forces generated in the system due to the expansion and compression of the spring. Thus, the systems selected here covers a broad range of cases, that is, dynamics (i) with internal constraints (pendulum), (ii) under the influence of an external field (gravitational), (iii) harmonic interactions (springs), (iv) complex breakable interactions (LJ potential), and (v) hybrid system with and without internal constraints.

The training of Hgnn is carried out for each system separately. A training dataset of \(100\) trajectories, each having \(100\) steps, were used for each system. For spring and pendulum, a 5-particle system is considered with random initial conditions. In the pendulum system, the initial conditions are considered in such a fashion that the constraints are respected. In the spring system, each ball is connected only to two other balls forming a loop structure. For gravitational system, a 4-particle system is considered where two particles are rotating in the clockwise direction, and two remaining particles are rotating in the anti-clockwise direction about their center of mass. For LJ system, a binary Kob-Andersen system with 75 particles are considered. The initial structure is generated by randomly placing the particles in a box with periodic boundary conditions. Further, the systems are simulated in a microcanonical ensemble (NVE) with temperatures corresponding to the liquid state to obtain equilibrium structures. Only once the system is equilibrated, the training data is collected for this system. Hgnn models were trained on this dataset with a \(75:25\) split for training and validation. Further, to test the long-term stability and energy and momentum conservation error, the trained model was evaluated on a forward simulation for \(10^{5}\) timesteps on 100 random initial configurations. See Methods for detailed equations for the interactions, datasets, and training parameters.

**Learning the dynamics.** Now, we evaluate the performance of the trained Hgnn models. To evaluate the long-term stability of the dynamics learned by Hgnn, we analyze the trajectory predicted by Hgnn for 100 random initial configurations. Specifically, we compare the predicted and actual phase space, trajectory, kinetic energy, potential energy, and forces on all the particles of the system during the trajectory. Note that the systems studied in this case are chaotic; hence, the exact trajectory followed by Hgnn will diverge with time. However, the phase space and the errors in energy and forces can be effective metrics to analyze whether the trajectory generated by Hgnn is statistically equivalent to that of the original system, that is, sampling the same regions of the energy landscape. Further, in contrast to purely data-driven [8] or physics-informed methods, the physics-enforced architecture of Hgnn strictly follows all the characteristics of the Hamiltonian equations of motion, such as the conservation laws of energy and momentum (see Supplementary Materials). This is due to the fact that the graph architecture only predicts the Hamiltonian of the
Figure 2: **Evaluation of Hgnn on the pendulum, spring, binary LJ, and gravitational systems.** (a) Predicted and (b) actual phase space (that is, \(x_{1}\)-position vs. \(x_{2}\)-velocity), predicted with respect to actual (c) kinetic energy, (d) potential energy, and (e) forces in 1 (blue square), and 2 (red triangle) directions of the 5-pendulum system. (f) Predicted and (g) actual phase space (that is, 1-position, \(x_{1}\) vs \(2\)-velocity, \(\dot{x}_{2}\)), predicted with respect to actual (h) kinetic energy, (i) potential energy, and (j) forces in 1 (blue square) and 2 (red triangle) directions of the \(5\)-spring system. (k) Predicted and (l) actual positions (that is, \(x_{1}\) and \(x_{2}\) positions), predicted with respect to actual (m) kinetic energy, (n) pair-wise potential energy, \(V_{\rm ij}\) for the (0-0), (0-1), and (1-1) interactions, and (o) forces in 1 (blue square), 2 (red triangle), and 3 (green circle) directions of the \(75\)-particle LJ system. (p) Predicted and (q) actual positions (that is, \(x_{1}-\) and \(x_{2}-\)positions), predicted with respect to actual (r) kinetic energy, (s) potential energy, and (t) forces in 1 (blue square), and 2 (red triangle) directions of the gravitational system.
system, which is then substituted in the Hamiltonian equations of motion to obtain the updated configuration. Due to this feature, the trajectory predicted by the Hgnn is more realistic and meaningful in terms of the system's underlying physics.

Fig. 2 shows the performance of Hgnn for the pendulum (Figs. 2(a)-(e), first row), spring (Figs. 2(f)-(j), second row), binary LJ (Figs. 2(k)-(o), third row), and gravitational systems (Figs. 2(p)-(t), fourth row). For pendulum and spring systems, we observe that the phase space represented by the positions in 1-direction (\(x_{1}\)) and velocities in the orthogonal direction (\(z_{2}\)) predicted by Hgnn (Figs. 2(a) and (f)) exhibit an excellent match with the ground truth trajectory. It is interesting to note that Hgnn trained only on a trajectory of a single step (\(t\) to \(t+1\)) is able to learn the dynamics accurately and simulate a long-term stable trajectory of \(10^{5}\) timesteps that exactly matches the simulated trajectory. Similarly, for the binary LJ and gravitational systems, we observe that the predicted (Figs. 2(k) and (p)) and actual (Figs. 2(j) and (q)) positions in the trajectory of random unseen initial configurations explored by the systems exhibit an excellent match. Further, we observe that the predicted kinetic (Figs. 2(c), (h), (m), (r)) and potential (Figs. 2(d), (i), (n), and (s)) energies and forces (Figs. 2(e), (j), (o), and (t)) exhibit an excellent match with the ground truth values with a mean squared error almost close to zero. Additional evaluation of the Hgnn architecture is performed by comparing it with two baselines, namely, Hnn (which is a physics-enforced MLP) and Hgn, which does not decouple potential and kinetic energies (see Supplementary Materials) and on additional metrics such as energy and momentum error. We observe that Hgnn significantly outperforms Hgn and Hnn in terms of rollout and energy error (see Supplementary Materials). These results confirm that the Hgnn architecture can learn the systems' dynamics directly from the trajectory and hence can be used for systems where the Hamiltonian is unknown or inaccessible (such as experimental or coarse-grained systems).

**Zero-shot generalizability.** Now, we evaluate the generalizability of the Hgnn to unseen systems, for instance, systems larger than those on which Hgnn is trained or a completely new system that is a combination of two systems on which it is independently trained. While traditional neural networks based on approaches are restricted to the system sizes on which it is trained, Hgnn is inductive to larger (and smaller) systems than those on which they are trained. This is due to the modular nature of the Hgnn, thanks to the graph-based approach, where the learning occurs at the node and edge level. Fig. 3 shows the generalizability of Hgnn to larger system sizes than those on which it is trained. Specifically, we evaluate Hgnn on \(10-\)endulum (Fig. 3(a)-(e)), \(50-\)spring (Fig. 3(f)-(j)), and \(600-\)particle binary LJ systems (Fig. 3(k)-(o)). We observe that Hgnn is able to generalize to larger system sizes accurately without any additional training or fine-tuning, exhibiting excellent match with the ground truth trajectory in terms of positions, energies, and forces. Additional results on \(50\)-pendulum systems and \(500\)-spring systems are included in the Supplementary Material.

We also evaluate the ability of Hgnn to simulate a hybrid spring-pendulum system (see Fig. 1(b) Hybrid system). To this extent, we model the Hamiltonian of the hybrid as the superposition of the Hamiltonian of spring and pendulum systems. Further, we model two graphs based on the spring and pendulum elements and use the Hgnn trained on the spring and pendulum systems to obtain the Hamiltonian of the system. Fig. 3(p)-(t) shows the performance of Hgnn on the hybrid system. Hgnn provides the dynamics in excellent agreement with the ground truth for the unseen hybrid system as well in terms of positions, energies, and forces. Additional results on the force predicted on each particle by Hgnn in comparison to the ground truth for a trajectory of 100 steps is shown in Supplementary Material. These results confirm that Hgnn is able to learn the dynamics of systems directly from their trajectory and simulate the long-term dynamics for new initial conditions and system sizes. This is a highly desirable feature as Hgnn can be used to learn the Hamiltonian from sparse experimental data of physical systems or _ab-initio_ simulations of atomic systems. This learned model can then be used to simulate larger system sizes to investigate phenomena with higher length scales.

## Interpretability and discovering symbolic laws

Neural networks, while exhibiting excellent capability to learn functions, are notorious for their black-box nature allowing poor or no interpretability to the learned function. In contrast, we demonstrate the interpretability of the learned Hgnn. Thanks to the modular nature of Hgnn, we analyze the functions learned by the individual MLPs that represent the node and edge level potential energies (\(\texttt{MLP}_{v}\) and \(\texttt{MLP}_{e}\), respectively) and kinetic energy (\(\texttt{MLP}_{T}\)) of the particles as a function of the learned embeddings. Fig. 4(a)-(f) show the learned functions with respect to the input features such as positions, velocities, or inter-particle distances. We observe that learned functions by Hgnn for the potential energies for (i) pendulum bob (\(mgx_{2}\); Fig. 4(a)), (ii) spring (\(0.5k(r_{ij}-1)^{2}\); Fig. 4(c)), and (iii) binary LJ systems (0-0, 0-1, 1-1; Figs. 4(d)-(f), respectively) and kinetic energy of particles (\(0.5m|\dot{\textbf{x}}_{i}|^{2}\); Fig. 4(b)) exhibits a close match with the known governing equations. This shows the interpretability of the Hgnn and the additional ability to provide insights into the nature of interactions between the particles directly from their trajectory. Thus, Hgnn can be used to discover interaction laws directly from their trajectory, even when they are not accessible or available.
Figure 3: **Generalizability to unseen systems.** (a) Predicted and (b) actual phase space (that is, \(1-\)position vs. \(2-\)velocity) and predicted with respect to actual (c) kinetic energy, (d) potential energy, and (e) forces of the \(10-\)pendulum system, using Hønn trained on \(5-\)pendulum system. (f) Predicted and (g) actual phase space (that is, \(1-\)position, \(x_{1}\) vs. \(2-\)velocity, \(\dot{x}_{2}\)) and predicted with respect to actual (h) kinetic energy, (i) potential energy, and (j) forces of the \(50-\)spring system using Hønn trained on \(5-\)spring system. (k) Predicted and (l) actual positions (that is, \(1-\) and \(2-\)positions; blue and red represent type 0 and type 1 particles), and predicted with respect to actual (m) kinetic energy, (n) pair-wise potential energy, Vij and (o) forces, of the \(600-\)particle binary LJ system, using Hønn trained on \(75\) particle binary LJ system. (p) Predicted and (q) actual positions (that is, \(1-\) and \(2-\)positions), and predicted with respect to actual (r) kinetic energy, (s) potential energy, and (t) forces of the 10-particle hybrid system, using Hønn trained on \(5\)-spring and \(5\)-pendulum system.
While the interpretability of Hgnn can provide insights into the nature of energy functionals, abstracting it further as a symbolic expression can enable discovering the underlying interaction laws and energy functions. Such functionals can then be used for simulating the system or understanding the dynamics independent of the Hgnn. Thus, beyond learning the dynamics of systems, Hgnn can be used to discover underlying energy functionals and interaction laws. To this extent, we apply SR [2, 3, 4] on the learned functions by Hgnn. Specifically, we focus on the kinetic energy function, the harmonic function of the spring, gravitational potential, and the binary LJ systems. Specifically, we employ simple operations such as addition, multiplication, and polynomials to identify the governing equations that minimize the error between the values predicted by the discovered equation and those predicted by the Hgnn. The optimal equation is identified based on a score that balances complexity and loss of the equation (see Methods for details).

Table 1 shows the original equation and the equation discovered based on SR of the learned Hgnn functionals. Note for each system, the equation that exhibits the maximum score is chosen as the final equation (see Methods for details). All the equations discovered by SR with their loss, complexity, polynomials used, and other hyper-parameters are included in the Supplementary material. We observe that the recovered equations exhibit a close match for kinetic energy, harmonic spring, gravitational potential, and binary LJ. In the case of the binary LJ system, we observe that the equations reproduced for (0-0) and (1-1) interactions are very close to the original equation, while for (0-1) interaction, the equation is slightly different, although it exhibits low loss. Interestingly, we observe that for LJ (0-1) interaction, one of the equations provided by SR given by \(V_{ij}=\left(\frac{0.203}{r_{ij}^{12}}-\frac{0.773}{r_{ij}^{2}}\right)\) is closer to the original equation in its functional form. However, this predicted equation has a score of \(2.22\) with a loss of \(0.000109\). Thus, both the loss and the score of the equation are higher and lower, respectively, than the best equation obtained in Table 1. This also suggests that for more complex interactions, an increased number of data points, especially along the inflection points, might be required to improve the probability of discovering the original equation.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Functions & Original Eq. & Discovered Eq. & Loss & Score \\ \hline Kinetic energy & \(T_{i}=0.5m|\dot{\mathbf{x}}_{i}|^{2}\) & \(T_{i}=0.500m|\dot{\mathbf{x}}_{i}|^{2}\) & \(7.96\times 10^{-10}\) & \(22.7\) \\ Harmonic spring & \(V_{ij}=0.5(r_{ij}-1)^{2}\) & \(V_{ij}=0.499\left(r_{ij}-1.00\right)^{2}\) & \(1.13\times 10^{-9}\) & \(3.15\) \\ Binary LJ (0-0) & \(V_{ij}=\left(\frac{2.0}{r_{ij}^{12}}-\frac{2.0}{r_{ij}^{12}}\right)\) & \(V_{ij}=\left(\frac{1.90}{r_{ij}^{12}}-\frac{1.95}{r_{ij}^{12}}\right)\) & \(0.00159\) & \(2.62\) \\ Binary LJ (0-1) & \(V_{ij}=\left(\frac{0.275}{r_{ij}^{12}}-\frac{0.786}{r_{ij}^{12}}\right)\) & \(V_{ij}=\left(\frac{2.33}{r_{ij}^{2}}-\frac{2.91}{r_{ij}^{2}}\right)\) & \(3.47\times 10^{-5}\) & \(5.98\) \\ Binary LJ (1-1) & \(V_{ij}=\left(\frac{0.216}{r_{ij}^{12}}-\frac{0.464}{r_{ij}^{12}}\right)\) & \(V_{ij}=\left(\frac{0.215}{r_{ij}^{12}}-\frac{0.464}{r_{ij}^{12}}\right)\) & \(1.16\times 10^{-5}\) & \(5.41\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: **Discovering governing laws with symbolic regression.** Original equation and the best equation discovered by symbolic regression based on the score for different functions. The loss represents the mean squared error between the data points from Hgnn and the predicted equations.

Figure 4: **Interpreting the learned functions in Hgnn.** (a) Potential energy of pendulum system with the \(2\)-position of the bobs. (b) Kinetic energy of the particles with respect to the velocity for the pendulum bobs. (c) Potential energy with respect to the pair-wise particle distance for the spring system. (d) The pair-wise potential energy of the binary LJ system for 0-0, 0-1, and 1-1 type of particles. The results from Hgnn are shown with the markers, while the original function is shown as dotted lines.
## Outlook

Altogether, in this work, we present a framework Hgnn that allows the discovery of energy functionals directly from the trajectory of physical systems. The Hgnn could be extended to address several challenging problems where the dynamics depends on the topology such as the dynamics of polydisperse gels [34], granular materials [35], biological systems such as cells [36], or even rigid body dynamics. A topology to graph mapping can be developed in such cases which can then be used to learn the dynamics and further abstracted it out in terms of the governing interaction laws. At this juncture, it is worth mentioning some outstanding questions the present work raises. Although Hgnn presents a promising approach, it is applied to only particle-based systems with at most two-body interactions. Extending Hgnn to more complex systems, such as complex atomic structures with multi-body interactions or to deformable bodies in continuum mechanics could be addressed as future challenges. Further, the graph architecture presented in Hgnn could be enhanced by adding additional inductive biases such as equivariance [37]. Finally, extending the framework to non-Hamiltonian systems such as colloidal systems [38] exhibiting Brownian or Langevin dynamics could be pursued to widen the scope of the Hgnn framework to capture realistic systems.

## Methods

### Experimental systems

To simulate the ground truth, physics-based equations derived using Hamiltonian mechanics are employed. The equations for \(n\)-pendulum and spring systems are given in detail below.

#### \(n\)-Pendulum

For an \(n\)-pendulum system, \(n\)-point masses, representing the bobs, are connected by rigid (non-deformable) bars. These bars, thus, impose a distance constraint between two point masses as

\[||\mathbf{x}_{i}-\mathbf{x}_{i-1}||^{2}=l_{i}^{2} \tag{7}\]

where, \(l_{i}\) represents the length of the bar connecting the \((i-1)^{th}\) and \(i^{th}\) mass. This constraint can be differentiated to write in the form of a _Pfaffian_ constraint as

\[(\mathbf{x}_{i}-\mathbf{x}_{i-1})(\hat{\mathbf{x}}_{i}-\hat{\mathbf{x}}_{i-1} )=0 \tag{8}\]

Note that such constraint can be obtained for each of the \(n\) masses considered to obtain the constraint matrix.

The Hamiltonian of this system can be written as

\[H=\sum_{i=1}^{n}\sum_{j=1}^{2}\left(1/2m_{i}{x_{i,j}}^{2}-m_{i}gx_{i,2}\right) \tag{9}\]

where \(j=1,2\) represents the dimensions of the system, \(m_{i}\) represents the mass of the \(i^{th}\) particle, \(g\) represents the acceleration due to gravity in the \(2-\)direction and \(x_{i,2}\) represents the position of the \(i^{th}\) particle in the \(2-\) direction. Here, we use \(l_{i}=1.0\) m, \(m_{i}=1.0\) kg, and \(g=10.0\,m/s^{2}\).

#### \(n\)-spring system

Here, \(n\)-point masses are connected by elastic springs that deform linearly (elastically) with extension or compression. Note that similar to the pendulum setup, each mass \(m_{i}\) is connected to two masses \(m_{i-1}\) and \(m_{i+1}\) through springs so that all the masses form a closed connection. The Hamiltonian of this system is given by

\[H=\sum_{i=1}^{n}\sum_{j=1}^{2}\left(1/2m_{i}{x_{i,j}}^{2}\right)-\sum_{i=1}^{ n}1/2k(||\mathbf{x}_{i-1}-\mathbf{x}_{i}||-r_{0})^{2} \tag{10}\]

where \(r_{0}\) and \(k\) represent the undeformed length and the stiffness, respectively, of the spring, and \(j=1,2\) represents the dimensions of the system. Here, we use \(r_{0}=1.0\) m, \(m_{i}=1.0\) kg and \(k=1.0\) N/m.
### \(n\)-body gravitational system

Here, \(n\) point masses are in a gravitational field generated by the point masses themselves. The Hamiltonian of this system is given by

\[H=\sum_{i=1}^{n}\sum_{j=1}^{2}\left(1/2m_{i}x_{i,j}^{\ 2}\right)+\sum_{i=1}^{n} \sum_{k=1,j\neq i}^{n}Gm_{i}m_{j}/(||\mathbf{x}_{i}-\mathbf{x}_{j}||) \tag{11}\]

where \(G\) represents the gravitational constant, and \(j=1,2\) represents the dimension of the system. Here, we use \(G=1.0\) Nm2kg\({}^{-2}\), \(m_{i}=1.0\) kg and \(m_{j}=1.0\) kg \(\forall\ i,j\).

### Binary Lennard Jones system

Here, we consider a binary LJ system known as the Kob-Andersen mixture [33] composed of 80% particles of type 0 and 20% particles of type 1. The particles in this system interact based on a 12-6 LJ potential with the pair-wise potential energy \(V_{ij}\)given by

\[V_{ij}=\epsilon\left[\left(\frac{\sigma}{r_{ij}}\right)^{12}-\left(\frac{ \sigma}{r_{ij}}\right)^{6}\right] \tag{12}\]

where \(r_{ij}=||\mathbf{x}_{i}-\mathbf{x}_{j}||\) and \(\sigma\) and \(\epsilon\) are the LJ parameters, which takes the values as \(\epsilon_{0-0}=1.0\), \(\epsilon_{0-1}=1.5\), \(\epsilon_{1-1}=0.5\) and \(\sigma_{0-0}=1.00\), \(\sigma_{0-1}=0.80\), \(\sigma_{1-1}=0.88\), and \(r_{ij}\) represents the distance between particles \(i\) and \(j\). The pair-wise interaction energy between all the particles is summed to obtain the total energy of the system. For the LJ system, all the simulations are conducted at a temperature of 1.2 in the microcanonical (NVE) ensemble, ensuring the system is in a liquid state. The system is initialized with atoms placed in random positions avoiding overlap in a cubic box with periodic boundary conditions with box size \(3.968\) and cutoff for atom type \(0-0=2.5\), \(0-1=2.0\) and \(1-1=2.2\). Further, the system is equilibrated in the NVE ensemble until the memory of the initial configuration is lost. The equations of motion are integrated with the velocity Verlet algorithm.

### Gnn architecture

**Pre-Processing:** In the pre-processing layer, we generate a compact vector representation for particle and their interactions \(e_{ij}\) by employing Multi-Layer Perceptrons.

\[\mathbf{h}_{i}^{0} =\texttt{squareplus}(\texttt{MLP}_{em}(\texttt{one-hot}(t_{i}))) \tag{13}\] \[\mathbf{h}_{ij}^{0} =\texttt{squareplus}(\texttt{MLP}_{em}(e_{ij})) \tag{14}\]

Here, \(\texttt{squareplus}\) is an activation function. In our implementation, we use different \(\texttt{MLP}_{em}\)s for node representation corresponding to kinetic energy, potential energy, and drag. For brevity, we do not separately write the \(\texttt{MLP}_{em}\)s in Eq. 13.

### Kinetic energy and drag prediction.

Given that the graph employs Cartesian coordinates, the mass matrix can be represented as a diagonal matrix. Consequently, the kinetic energy (\(\tau_{i}\)) of a particle relies exclusively on the velocity (\(\hat{\mathbf{x}}_{i}\)) and mass (\(m_{i}\)) of said particle. In this context, the parameterized masses for each particle type are acquired through the utilization of the embedding (\(\mathbf{h}_{i}^{0}\)). As such, the predicted value of \(\tau_{i}\) for a given particle is determined by \(\tau_{i}=\texttt{squareplus}(\texttt{MLP}_{T}(\mathbf{h}_{i}^{0}\ ||\ \hat{\mathbf{x}}_{i}))\), where the symbol \(\|\) denotes the concatenation operator. In this equation, \(\texttt{MLP}_{T}\) denotes a multilayer perceptron responsible for learning the kinetic energy function, while \(\texttt{squareplus}\) represents the activation function employed. The overall kinetic energy of the system, denoted by \(T\), is calculated as the sum of individual kinetic energies: \(T=\sum_{i=1}^{n}\tau_{i}\).

**Potential energy prediction.** Typically, the potential energy of a system exhibits significant dependence on the topology of its underlying structure. In order to effectively capture this information, we utilize a multiple layers of message-passing among interacting particles (nodes). During the \(l^{th}\) layer of message passing, the node embeddings are iteratively updated according to the following expression:

\[\mathbf{h}_{i}^{l+1}=\texttt{squareplus}\left(\texttt{MLP}\left(\mathbf{h}_{i }^{l}+\sum_{j\in\mathcal{N}_{i}}\mathbf{W}_{\mathcal{V}}^{l}\cdot\left(\mathbf{ h}_{j}^{l}||\mathbf{h}_{ij}^{l}\right)\right)\right) \tag{15}\]

where, \(\mathcal{N}_{i}=\{u_{j}\in\mathcal{V}\ |\ (u_{i},u_{j})\in\mathcal{E}\}\) is the set of neighbors of particle \(u_{i}\). \(\mathbf{W}_{\mathcal{V}}^{l}\) is a layer-specific learnable weight matrix. \(\mathbf{h}_{ij}^{l}\) represents the embedding of incoming edge \(e_{ij}\) on \(u_{i}\) in the \(l^{th}\) layer, which is computed as follows.

\[\mathbf{h}_{ij}^{l+1}=\texttt{squareplus}\left(\texttt{MLP}\left(\mathbf{h}_{ ij}^{l}+\mathbf{W}_{\mathcal{E}}^{l}\cdot\left(\mathbf{h}_{i}^{l}||\mathbf{h}_{j}^{l} \right)\right)\right) \tag{16}\]
Similar to \(\mathbf{W}_{\mathcal{V}}^{l}\), \(\mathbf{W}_{\mathcal{E}}^{l}\) is a layer-specific learnable weight matrix specific to the edge set. The message passing is performed over \(L\) layers, where \(L\) is a hyper-parameter. The final node and edge representations in the \(L^{th}\) layer are denoted as \(\mathbf{z}_{i}=\mathbf{h}_{i}^{L}\) and \(\mathbf{z}_{ij}=\mathbf{h}_{ij}^{L}\) respectively.

The total potential energy of an \(n\)-body system is represented as \(V=\sum_{i}v_{i}+\sum_{ij}v_{ij}\). Here, \(v_{i}\) denotes the energy associated with the position of particle \(i\), while \(v_{ij}\) represents the energy arising from the interaction between particles \(i\) and \(j\). For instance, \(v_{i}\) corresponds to the potential energy of a bob in a double pendulum, considering its position within a gravitational field. On the other hand, \(v_{ij}\) signifies the energy associated with the expansion and contraction of a spring connecting two particles. In the proposed framework, the prediction for \(v_{i}\) is given by \(v_{i}=\texttt{squareplus}(\texttt{MLP}_{v_{i}}(\mathbf{h}_{i}^{0}\parallel \mathbf{x}_{i}))\). Similarly, the prediction for the pair-wise interaction energy \(v_{ij}\) is determined by \(v_{ij}=\texttt{squareplus}(\texttt{MLP}_{v_{ij}}(\mathbf{z}_{ij}))\).

The parameters of the model are trained end-to-end using the MSE loss discussed in Eq. 6.

#### Model architecture and training setup

For Hgnn, all the MLPs are two layers deep. A square plus activation function is used for all the MLPs. We used 10000 data points from 100 trajectories divided into 75:25 (train: validation) to train all the models. The timestep used for the forward simulation of the pendulum system is \(10^{-5}s\), for the spring and gravitational system is \(10^{-3}s\), and for the LJ system is 0.0001 LJ units. All the equations of motion are integrated with the velocity-Verlet integrator. Detailed training procedures and hyper-parameters are provided in the Supplementary material. All models were trained until the decrease in loss saturates to less than 0.001 over 100 epochs. The model performance is evaluated on a forward trajectory, a task it was not explicitly trained for, of \(10s\) in the case of the pendulum and \(20s\) in the case of spring. Note that this trajectory is 2-3 orders of magnitude larger than the training trajectories from which the data has been sampled. The dynamics of \(n\)-body system are known to be chaotic for \(n\geq 2\). Hence, all the results are averaged over trajectories generated from 100 different initial conditions.

#### Symbolic regression

SR refers to an approach to search for equations that fit the data points and fit them rather than a parametric approach where an equation is chosen apriori to fit the data. Here, we employ the PySR package to perform the SR [7]. PySR employs a tree-based approach for fitting the governing equation based on the operations and variables provided. Since the parametric space available for SR can be too large with every additional operation, it is important to carefully provide the minimum required input features and the operations while providing meaningful constraints on the search space.

In the present work, we choose the addition and multiplication operation. Further, we allow polynomial fit based on a set containing (square, cube, pow(n)) operations, where pow(n) refers to power from four to ten. The loss function to fit the SR is based on the mean squared error between the predicted equation and the data points obtained from Hgnn. Further, the equations are selected based on a score \(S\) that balances complexity \(C\) and loss \(L\). Specifically, the score is defined as \(S=\frac{dL}{dC}\), that is, the gradient of the loss with respect to complexity. For each set of hyperparameters, we select the top 10 equations based on the scores. Further, the equation having the best score among these equations is chosen as the optimal equation. All the hyperparameters associated with the SR and the corresponding equations obtained are included in the Supplementary material.

#### Simulation environment

All the simulations and training were carried out in the JAX environment [39, 40]. The graph architecture was developed using the jraph package [41]. The experiments were conducted on a machine with Apple M1 chip having 8GB RAM and running MacOS Monterey. **Software packages:** numpy-1.22.1, jax-0.3.0, jax-md-0.1.20, jaxlib-0.3.0, jraph-0.0.2.dev **Hardware:** Chip: Apple M1, Total Number of Cores: 8 (4 performance and 4 efficiency), Memory: 8 GB, System Firmware Version: 7459.101.3, OS Loader Version: 7459.101.3
