Leap-frog neural network for learning the symplectic evolution from partitioned data

We first review some basic mathematical background of Hamiltonian systems (Meyer, Hall & Offin, 2009). Let $gl(m,R)$ denote the set of all $m\times m$ non-singular matrices with entries in the real number field $R$. Let $z$ be a coordinate vector in $R^{2n}$, $I$ be a time interval in $R$, and the matrix $K:I\to gl(2n,R)$ be continuous and symmetric. Then a Hamiltonian system can be described by $2n$ ordinary differential equations:

where $H=H(z,t)$ is called the Hamiltonian of the system, and $J$ is the $2n\times 2n$ Poisson matrix:

where $I_n$ is the $n\times n$ identity matrix.

Let $t_0\in I$ be the initial time point. From the theory of differential equations, if the right function $A\left(t\right)z$ is $C^1$, or the Lipschitz condition in the variable $z$ is satisfied, given the initial condition $z_0=z\left(t_0\right)\in R^{2n}$, there exists a unique solution $z(t,t_0,z_0)$ of equation (1) for $t\in I$. We say the time evolution of $z$ is symplectic if and only if the Jacobi matrix of the transformation $z^{\prime }$ satisfies

As for the Hamilton’s equation (1), this condition is naturally satisfied.

In the current work, the latent evolution of a Hamiltonian system is considered. Let

where $q$ is called the coordinate and $p$ is called momentum, and these two variables are said to be conjugated. Then the system of equation (1) can be written in the form

Assuming each time series $\left(q\right(t),p(t\left)\right)$ is sampled at discretized time points, with a time step $h(>0)$, then we have

Now we consider the one step mapping

Since the time series $\left(q\right(t),p(t\left)\right)$ is symplectic, the corresponding mapping $\varphi$ is symplectic.

In our previous work (Li et al., 2022), the solution of the restricted 3-body problem is learned by multi-layer neural networks. We tried to train the ANN using the data of positions and velocities (equivalent to $q$ and $p$, respectively). Although the prediction could be very closed to the ground truth data, the Jacobi integral seems to vary greatly along the orbital evolution, i.e. the mapping $\varphi$ fails to be symplectic. While we notice that the values of positions and velocities have the same order of magnitude, and this may be a positive feature in machine learning.

Now, instead of learning $q$ and $p$ directly, similar to Chen & Tao (2021) we turn to learn a class of mappings $\varphi$. Then the time series $\left(q\right(t_i),p(t_i\left)\right)$ can be derived step-by-step via the symplectic mapping $\varphi$, thus the symplectic structure of the system evolution can be naturally kept. In addition, we need to control the accuracy of prediction of $q$ and $p$ in a small scale. Indeed, we can also say $\varphi$ in equation (7) is a canonical transformation, then we will adopt the classical technique “generating function” in the following.

For clarity, here we consider the multidimensional vectors [$q$,$p$], where $q,p\in R^n$. Corresponding to the 2$n$ Hamiltonian system, $q$ is the position variable, ${\text{q}}_i$ is the conjugate momentum variable. Then the standard symplectic form is defined to be

where the operator $\wedge$ is called the exterior product or wedge product. Let $Q$ and $P$ be new variables, and we now investigate substitutions of the form

The change of variables in equation (9) is volume-preserving if and only if

which actually indicates that the transformation $(q,p)\to (Q,P)$ is symplectic. By equation (10), we can easily have

This is sufficient to make sure the form $pdq+QdP$ is exact, i.e. if we can find a function $F$ that satisfies

Since the differential form of $F$ can be written as

we can obtain the change of variables from $(q,p)$ to $(Q,P)$, by

By the implicit function theorem, when the Hessian matrix $\frac{{\partial }^{2}F}{\partial \text{q}\partial \text{P}}$ is non-singular, the equations in (14) are solvable for $P$ as a function of $q$ and $p$ and for $p$ as a function of $Q$ and $P$. So equation (14) gives a straightforward way to construct the symplectic transformation, and the function $F$ is usually called generating function. For more details, the reader is referred to Chapter 6 in Meyer, Hall & Offin (2009).

In Chen & Tao (2021), the authors designed the network GFNN to learn the generating function, in order to obtain a symplectic mapping from one time point to the next. Then step by step, the system can continue to evolve with time. More precisely, for the one time step evolution, the mapping can be described as:

then denoting

the next step is

This process could go on until

Here, the subscript $i$ corresponds to the $i$-th time point $t_i$ in equation (7). We note that, when $N$ is large, i.e. for the long time evolution, the above prediction of the time series could cause severe accumulative error.

Suppose we have the training set containing a time series in the interval $[t_0,t_{99}]$, i.e. for $N=99$, with a fixed time step of $h$. Different from Chen & Tao (2021), now we consider a finite partition of this time series, which is adopted to be 4 here. Thus we would get 4 subsets of the time series $(q_i,p_i)$:

Our new idea is to learn the generating function $F$ from the respective partitioned data, then $F$ is used to predict the time series. This process is visualized in Figs. 1 and partition. Let Partition $0$ be the initial condition, we first study the generating function $F^1$ between Partition $0$ and Partition $1$. When we have $F^1(q,P)=q\cdot P+S^1(q,P)$, and $\frac{{\partial }^2{F}^1}{\partial \text{q}\partial \text{P}}$ is non-singular, we can learn $S^1$ by the neural network with respect to the loss function:

Then, Partition $1$ can be generated by:

i.e.

and accordingly we would obtain the symplectic mappings

If the generating function $S^1$ is approximated with error $\le ϵ$ in the first derivative, by the designed ANN, the $L_2$ norm of the deviation between the predicted sequence

and the ground truth sequence

satisfies, for any $i\in [25,49]$,

and

where $C$ is some constant.

It should be noted that, if we are predicting Partition $1$ by the method we described above, we need to know the momenta $p_{25},\cdots ,p_{49}$. Since we do not have the ground truth for $p_i(i\ge 25)$, only the approximate values ${\tilde{p}}_{25},\dots ,{\tilde{p}}_{49}$ could be estimated. In general, there are two ways to obtain them. The first one is that, starting with the initial condition $\text{Partition}0:(q_0,p_0),\cdots ,(q_{24},p_{24})$, we use the Newton’s method to solve the following nonlinear equations:

By this procedure we could get ${\tilde{p}}_{25},\dots ,{\tilde{p}}_{49}$. The other way is more direct, as we can just train an additional fully connected neural network to predict these approximate momenta by using the supervised learning like we did in Li et al. (2022).

Exactly the same, the Partition $2$ and Partition $3$ can also be predicted from Partition $0$, respectively, by introducing the individual generating functions $F^2(q,P)=q\cdot P+S^2(q,P)$ and $F^3(q,P)=q\cdot P+S^3(q,P)$. Then we could obtain the remaining two groups of symplectic mappings

and

In summary, for Partition $1$, $2$ and $3$, if the respect generating functions $S^1$, $S^2$ and $S^3$ are approximated with errors $\le ϵ$, the loss between the ground truth and the prediction satisfies the following inequalities:

and

where $i\in [25,99]$.

The learning scheme designed above would be called the leap-frog neural network (LFNN). The main advantage of training the LFNN is that, we may avoid the issue of accumulative error since the symplectic mappings (15), (16) and (17) actually have no iteration, as we will show in the Section 3 by comparing our results with Chen & Tao (2021). Besides, the global error is rather small and the predictions of the time series in more general cases seem quite accurate.

The considered sequence $(q_i,p_i)$ has 100 time points in total (i.e. $i=0$-$99$), and is divided into 4 partitions. Thus, for each of the mappings (15), (16) and (17), the LFNN model is predicting a time series with a length of 25, e.g. $i=25$-$49$ for $\text{Partition}1$; in addition, since we will later consider the planar dynamical models, either $q_i$ or $p_i$ has two components.

Then for each of the partitions ($j=1,2,3$), we can define two $25\times 2$ matrices as:

where the subscript “sim” denote the ground truth data from numerical simulations. Similarly, we can also define the matrices predicted by the LFNN as:

If we have the $25\times 2$ matrices $A$ and $B$, the distance between them can be defined as:

Then, for a specific partition, we can calculate the average loss between prediction ($Q_{LFNN}^j$ and $P_{LFNN}^j$) and the ground truth ($Q_{sim}^j$ and $P_{sim}^j$). When we consider just one trajectory, the averaged losses are:

And when we consider $M$ trajectories, the average losses are:

Hereafter, without the superscript $j$, we refer to the loss calculated for all the three partitions, i.e. $i=25$-$99$.

The said losses will be monitored on both the training and validation sets, and the the hyperparameters of our LFNN would be refined to make sure the global errors can decrease to small values. As a results, the precision of prediction may reach an acceptable level. In the next section, we will evaluate the prediction errors for Partitions $1$, $2$ and $3$, respectively.

The classical KBOs have semimajor axes roughly between 42 AU to 48 AU, and are defined to be objects outside Neptune’s MMRs. In the PCR3BP, we consider the particles having initial non-dimensional semimajor axes $a=1.44$. Since the distance between the Sun and Neptune, i.e. $\sim 30$ AU, is set to be 1, the particles’ actually semimajor axes can be calculated to be about 43.2 AU. At this location, none of the MMRs with order lower than 8th exists. Besides, the KBO particles are placed on the nearly circular orbits with eccentricities $e=0.01$, to additionally make sure that they are non-resonant.

As in Section 3.1, the ground truth orbits are generated by using the 4th order Runge-Kutta integrator; and the the numbers of the orbits in the training and validation sets are ${10}^5$ and ${10}^2$, respectively. For the 100 time points with a step of $h=0.1$, the total evolution time is 10 in the non-dimensional unit. According to Kepler’s third law, the orbital period of a particle is $\sqrt[3]{31.44/1}\times 2\pi \approx 7.1$. It suggests that the particles complete more than 1 full orbit over the evolution time of 10.

For the said particles starting with $a=1.44$ and $e=0.01$, their initial longitudes of perihelia $\varpi$ are set to be ${60}^{\circ }$, while the mean longitudes $\lambda$ are randomly selected between 0 and $2\pi$. Given such initial conditions, we generate the training and validation samples, which are out of the resonance regions.

Then we are about to re-trained our best LFNN model built on the generating function. Bearing in mind that, in order to make the symplectic transformations via the obtained generating function $S(q,P)$, we need to know the new momentum $P$, i.e. $\{{\tilde{p}}_{25},\dots ,{\tilde{p}}_{99}\}$. If we apply the Newton’s method to the 100 validation orbits, in order to get the losses of $P$ and $Q$, it is not easy to find the proper values of $P$. This difficulty comes from the fact that, when we adopt the leap-frog learning, the gradient of the generating function could be not accurate enough. As we mentioned in Section 2.3, alternatively we could train an additional supervised neural network to learn the approximate $P$. Then by $Q=q+{\partial }_{P}S(q,P)$, the new position $Q$ can be predicted.

In Fig. 3, for each of the 100 trajectories from the validation set, we plot averaged losses of the predicted positions $\tilde{q}$ (i.e. loss(Q), see blue dots) and momenta $\tilde{p}$ (i.e. loss(P), see yellow dots). We find that the maximal loss is less than 0.0009, which is associated to $\tilde{q}$. We further notice that the losses of $\tilde{p}$ are generally even more smaller, on the order of loss$\left(Q\right)\sim$0.0001-0.0005. Comparing with the case of “binary + particle” in Section 3.1, where loss$\left(Q\right)$ is an order of magnitude larger, the LFNN seems to work better on learning the evolution of non-resonant KBOs.

The population of resonant KBOs is believed to be a solid evidence for the orbital migration of Jovian planets. The study of the resonant KBOs is very important for us to understand the early shake-up of the outer Solar system, and a lot of works have been carried out even since early 1990s (Malhotra, 1993, 1995; Hahn & Malhotra, 2005; Li, Zhou & Sun, 2014; Nesvorný & Vokrouhlický, 2016; Pike & Lawler, 2017; Lawler et al., 2019). For this reason, in (Li et al., 2022) we provide a machine learning method that can fast predict the orbital evolution of the KBOs in the 2:3 MMR with Neptune. In the framework of the PCR3BR, since the trajectories of the 2:3 resonators are restricted on the invariant tori, their motions are regular and were learned rather accurately.

As we stated in Section 1, we intended to use the position $q$ and momentum $p$ as training data because the same order of magnitude of these variables could be more suitable for the machine learning. However, for $q$ and $p$, our designed ANN in Li et al. (2022) can not keep the system symplectic by conserving the Jacobi integral. Now we would test how the new LFNN performs on predicting the time series $(q_i,p_i)$.

For the samples particles, their initial conditions are exactly the same as we used in (Li et al., 2022): $a_0=1.31$ (the location of the nominal 2:3 resonance), $e_0=0.1$ (for regular motion everywhere in the phase space), ${\varpi }_0={60}^{\circ }$ (unchanged by rotating the coordinate system); and, the initial resonant angles ${\sigma }_0$ are chosen in the region of $[{180}^{\circ }-{120}^{\circ },{180}^{\circ }+{120}^{\circ }]$, where ${180}^{\circ }$ is the stable libration centre. Then the initial value of $\lambda$ can be determined from a given $\sigma$ by

where ${\lambda }_N$ is the mean longitude of Neptune. For a given set of the orbital elements $(a,e,\varpi ,\lambda )$, the initial conditions of the resonators are generated in terms of the position $x$ and velocity $v$. Then we numerically propagate the equations of motion, i.e., equation (21), by adopting the 8th-order Runge-Kutta integrator with a time step of $\sim 0.01$ yr and the local error tolerance of ${10}^{-20}$.

In order to learn the behaviour of the resonant KBOs, the time evolution should cover at least a full libration cycle. For the 2:3 resonators, the libration period is about 20000 yr, corresponding to the non-dimensional time of $\sim 762.5$. While for the two applications considered before, the evolution time is only 10. This is the largest difference in predicting the particle’s orbit by our LFNN. In order to fulfill this task, we here increase the time step $h$ to be 9.65. Then the 100 time points lead to a total of evolution time of 965, which corresponds to the period of 25000 yr that we adopted in (Li et al., 2022). We additionally note that, the LFNN uses Partition $0$ as the initial condition to predict Partitions $1$, $2$ and $3$, this is also consistent to our previous work that utilize the known first 1/4 period of the orbit to predict the subsequent 3/4 period data.

orbital elements. For the case of e = 0.1, three examples of the temporal evolution of Jacobi integral for particle with ${\sigma }_0={60}^{\circ }$ from the validation set. The initial values for the first 6250 yrs are provided by the numerical integration (black dots). For the following 18750 yrs, the orbital parameters are predicted by the trained LFNN (green dots) and calculated by current method (red dots).

Similar to the LFNN designed for the non-resonant KBOs, but for the time series $(q_i,p_i)=(x_i,v_i)$ with a much larger step $h$, our new machine-learning method is further applied to the 2:3 resonant KBOs. The losses on the validation set for 100 orbits with different ${\sigma }_0$ are presented in Fig. 4. The blue and yellow dots indicate $loss\left(Q\right)$ and $loss\left(P\right)$, respectively, calculated for the 75 predicted time points $({\tilde{q}}_{25},{\tilde{p}}_{25}),\dots ,({\tilde{q}}_{99},{\tilde{p}}_{99})$. We can see that: (1) the maximum loss($Q$)$\sim 0.08$ and loss($P$)$\sim 0.05$ are at both ends of the curves; (2) there are two valleys where the loss reaches the minimum of the order of $\sim 0.01$; (3) there are two peaks as the local maximum. The profile of the loss curve looks very similar to that in Fig. 3 from Li et al. (2022). While, although the magnitude of the loss shown here is also similar, we can not directly compare the performance of the LFNN with the ANN in our previous work. This is because the data sequences are consisted of positions and velocities, but not orbital elements any more.

To show the improvement of the LFNN, we pick the sample “tp1” starting with ${\sigma }_0={60}^{\circ }$, which has the largest loss of loss($Q\sim 0.08$) and loss($P\sim 0.05$) (see Fig. 4). Then we present in Fig. 5 the evolution for this sample from the validation set, and the associated orbital parameters can be simply obtained from the position and velocity variables. The black dots represent the ground truth values calculated by numerical integration. While the red dots denote the prediction made by the LFNN for the time period from $t_{25}$ to $t_{99}$, where $t_{99}$ corresponds to the real time of 25000 yr. As a comparison, we apply the fully connected ANN in (Li et al., 2022) to an additional prediction, and the outputs of this ANN are indicated by the green dots.

The upper panel of Fig. 5 shows the time evolution of the Jacobi integral $C$ for the sample tp1. It is quite clear that the values of $C$ deduced from our previous method (green dots) seem to considerably depart from the ground truth (black dots), indicating that the conservation of $C$ can hardly hold. But, with regard to our new LFNN, the approximate $C$ values (red dots) are different from the ground truth values only by less than 0.01 at any time point. Therefore, the first improvement of the LFNN on predicting the motion of the 2:3 resonators is that, the symplecticity of the orbital evolution can be well maintained.

As shown in the lower four panels of Fig. 5, the second improvement is the increased accuracy on the predictions of the orbital elements $a$, $e$, $\varpi$ and $M$ ($M=\lambda -\varpi$ is the mean anomaly). We can see that, the red dots are much closer to the black dots representing the ground truth than the green dots, especially for the slow variables $a$ and $e$. In this sense, the designed LFNN can help to significantly reduce the errors for making the long time prediction, and thus is better than the simple fully connected ANN.

In a brief summary, we address that the LFNN works well even when the step $h$ of the time series $(q_i,p_i)$ is relatively large.