Discovering Conservation Laws using Optimal Transport and Manifold Learning

Consider a dynamical system with states $x\in M$ that live a in $d$-dimensional phase space $M$ and evolve in time according to a system of first order ODEs

with $n$ conserved quantities $G_1\left(x\right),\dots ,G_n\left(x\right)$.

To analyze the isosurface shape space manifold $C$—i.e. the manifold of conserved quantities—using manifold learning methods, we need to place some structure on the points $X_c\in C$. Having associated each isosurface $X_c$ with a corresponding distribution defined by an ergodic physical measure ${\mu }_c$, we choose to lift the Euclidean metric on the phase space into the space of distributions using the 2-Wasserstein metric from optimal transport

where the cost function $c(x,y)=\parallel x-y{\parallel }^2$ is the squared Euclidean distance, and $\pi \in \Pi ({\mu }_c,{\mu }_{c^{\prime }})$ is a valid transport map between ${\mu }_c$ and ${\mu }_{c^{\prime }}$ [Villani2009].

For discrete samples, the 2-Wasserstein distance between two sets of sample points $\{x_1,x_2,\dots ,x_k\}$ and $\{y_1,y_2,\dots ,y_l\}$ is defined as

where the cost matrix $C_{ij}=\parallel x_i-y_j{\parallel }^2$, and the transport matrix $T$ is subject to the constraints

To efficiently compute an entropy regularized form of this optimization problem, we use the Sinkhorn algorithm [NIPS2013_af21d0c9] and estimate the Wasserstein distance as a debiased Sinkhorn divergence [pmlr-v119-janati20a].

One important subtlety in this construction is the choice of the ground metric for optimal transport. As previously mentioned, we use a Euclidean metric on the phase space, which implicitly imposes a choice of units to make the phase space dimensionless. In fact, there is no canonical choice for the ground metric, and different choices result in different Wasserstein metrics on the shape space. For example, when multiple conserved quantities are present, the relative effect of each conserved quantity on the computed Wasserstein distances will determine how prominent each conserved quantity is and how easily it is identified using manifold learning. To partially mitigate this issue and improve consistency, we normalize each component of our data to have a maximum absolute value of 1 before computing the pairwise Wasserstein distances.

Using the structure provided by the Wasserstein metric, we then use diffusion maps to generate an embedding for $C$. The diffusion map manifold learning method uses a spectral embedding algorithm applied to an affinity matrix to construct a low dimensional embedding of the data manifold [6789755, Coifman2005]. Using the pairwise Wasserstein distances $W_2({\mu }_i,{\mu }_j)$ computed from discrete samples provided by the trajectory data (Eq. 5), we first construct a kernel matrix using a Gaussian kernel

and then scale it to form an affinity matrix for our spectral embedding

where $D_i={\sum }_k{K}_{ik}$, and $\alpha$ is a hyperparameter. The spectral embedding algorithm then takes this affinity matrix and constructs a normalized graph Laplacian

where $I$ is the identity matrix. The eigenvectors $v_i$ corresponding to the smallest eigenvalues ${\lambda }_i\ge 0$ (excluding ${\lambda }_0=0$) of the Laplacian then provide an approximate low dimensional embedding of the manifold of conserved quantities $C$. In our experiments, we set $\alpha =1$ so that the Laplacian computed by the spectral embedding algorithm approximates the Laplace–Beltrami operator [6789755].

To estimate the dimensionality of $C$ and to choose which eigenvectors $v_i$ to include in our embedding, we use a heuristic score that combines a measure of relevance, given by a length scale computed from the Laplacian eigenvalues, with a previously suggested measure of “unpredictability” for minimizing redundancy [pfau2018minimally]. To construct our embedding, we only include the Laplacian eigenvectors with score above a chosen cutoff value and discard the rest as either noise or redundant embedding components. To determine the cutoff, we perform a sweep of the cutoff value looking for robust ranges and find that a cutoff of $0.6$ works well across all of our experiments, which consist of a wide variety of datasets and dynamical systems. See Appendix B for more details.

Consider a SHO with Hamiltonian

given in terms of position $q$ and momentum $p$. The SHO energy isosurfaces form concentric ellipses in a 2D phase space. Choosing units such that $m=1$ and $\omega =1$, we obtain concentric circles with uniformly distributed samples (assuming a uniform sampling in time). The 2-Wasserstein distance between a pair of uniformly distributed circular isosurfaces is simply given by the difference in radii $|r_1-r_2|$. This is because, due to the rotational symmetry of the two distributions, the optimal transport plan for an isotropic cost function is to simply move each point on isosurface 1 radially outward (or inward) to the point on isosurface 2 with the same angle $\theta$.

This result does not meaningfully change with a different choice of units, which is equivalent to rescaling the phase space coordinates $q,p$. If we rescale $q,p$ by factors $k_q,k_p$, our cost function simply becomes

where we label points on the isosurfaces by their angle $\theta$ on the original circular isosurfaces. The SHO optimal transport plan $\Pi$ takes $\theta$ on isosurface 1 to the point with the same angle $\theta$ on isosurface 2, and $\Pi$ for the SHO is invariant to coordinate rescaling (Appendix F). Therefore, the total transport cost is

so the 2-Wasserstein distance is

i.e. the same result modulo a constant factor. While this is not a general result, we find that our approach is often fairly robust to such changes, including the extreme case of scaling some phase space coordinates all the way down to zero resulting in a partially observed phase space (Appendix C).

Once we have pairwise distances in the isosurface shape space, we can use diffusion maps to study the resulting manifold of isosurface shapes. With sufficient samples, the operator constructed by the diffusion map should converge to the Laplace–Beltrami operator on the manifold. For the SHO, the isosurface shape space is isomorphic to $R^{+}$ with each circular isosurface mapped to its radius. If we sample trajectories with radii $r\in (0,\sqrt{2E_0})$ for some maximum energy $E_0$, then the manifold is a real line segment, and the resulting Laplacian operator (with open boundary conditions) has eigenvalues ${\lambda }_n={\pi }^2{n}^2/2E_0$ and corresponding eigenvectors $v_n\left(r\right)=\cos \left(\sqrt{{\lambda }_n}r\right)$. Therefore, the first eigenvector or embedding component

successfully encodes the conserved energy and is, in fact, a monotonic function of the energy.

We first numerically test our analytic result for the SHO and obtain good agreement (Fig. 2) using both the default scaling $k_q=k_p=1$ (Figs. 2a–d) as well as the position only scaling $k_q=1$, $k_p=0$ (Figs. 2e–h), which effectively reduces the dimension of the phase space. A linear fit of the first embedding component from the diffusion map with the analytically predicted component (Eq. 14) achieves a correlation coefficient of $R^2=0.9995$ for the default scaling and $R^2=0.9961$ for the position only scaling. We also verify that the heuristic score (Appendix B) accurately determines that there is only one relevant embedding component (Figs. 2c, 2g), which corresponds to the conserved energy.

To demonstrate our method on a simple nonlinear dynamical system, we analyze a simple pendulum that has a 2D phase space consisting of the angle $\theta$ and angular momentum $\omega$ (Fig. 3a). The equations of motion are

This system has a single scalar conserved quantity

corresponding to the total energy of the pendulum, so the trajectories form 1D orbits in phase space (Fig. 3b).

Our method is able to correctly determine that there is only a single conserved quantity (Fig. 3c) corresponding to the energy of the pendulum (Fig. 3d). The single extracted embedding component is monotonically related to the energy with Spearman’s rank correlation coefficient $\rho =0.9997$. We are also able to achieve similar results ($\rho =0.9978$) with a high level of Gaussian noise (standard deviation $\sigma =0.5$) added to the raw trajectory data (Figs. 3e–h), showing that our approach is quite robust to measurement noise.

To test our method on a system with multiple conserved quantities, we simulate the gravitational system of a planet orbiting a star with much greater mass (Fig. 4a). We fix the orbits to all lie in a 2D plane, giving us an effectively 4D phase space. The resulting equations of motion are

This system has one scalar and two vector conserved quantities

which, in our 4D phase space, reduces to three scalar conserved quantities: the total energy $E$ (or equivalently, the semi-major axis $a=-1/2E$), the angular momentum $L=\left|L\right|$, and the orbital orientation angle $\varphi$, which is the angle of the LRL vector $A$ relative to the $x$-axis. As a result, the trajectories also form 1D orbits in the phase space (Fig. 4b).

Our approach accurately identifies the three conserved quantities (Fig. 4c), and the extracted embedding corresponds most directly to the geometric features of the orbits (Figs. 4d–f). The first two components embed the semi-major axis vector $a=(a\cos \varphi ,a\sin \varphi )$ with magnitude given by the semi-major axis $a=-1/2E$, which is related to the energy $E$, and orientation given by the orientation angle $\varphi$ of the elliptical orbit (Figs. 4d, 4e). The third relevant component (component 6) embeds the angular momentum $L$ (Fig. 4f). See Appendix B.1 for details on choosing a cutoff to identify the relevant components. A linear fit of the identified relevant embedding components with $a\cos \varphi$ ($a\sin \varphi$) has $R^2=0.987$ ($R^2=0.986$) and rank correlation $\rho =0.994$ ($\rho =0.992$). A similar linear fit with $L$ has $R^2=0.927$ and $\rho =0.970$.

This example demonstrates that, for a system with multiple conserved quantities, the ground metric for optimal transport controls the relative scale of each conserved quantity in the extracted embedding. In this case, the geometry of the shape space $C$ is dominated by changes in the semi-major axis $a$ and orientation angle $\varphi$, whereas changes in the angular momentum $L$, which controls the eccentricity of the orbit, play a more minor role and thus appear in a later embedding component with a lower score (Fig. 4c).

To test our approach on a non-integrable system with higher dimensional isosurfaces, we study the classic double pendulum system (Fig. 5a) with unit masses and unit length pendulum arms. This system has a 4D phase space, consisting of the angles ${\theta }_1,{\theta }_2$ and the angular velocities ${\omega }_1,{\omega }_2$ of the two pendulums (Fig. 5b), and only has a single scalar conserved quantity

corresponding to the total energy. However, the double pendulum system has both chaotic and non-chaotic phases. In particular, at high energies, the system is chaotic and only conserves the total energy, while at low energies, the system behaves more like two coupled harmonic oscillators with two independent (approximately) conserved energies

corresponding to the two modes of the coupled oscillator system. Therefore, we expect to see two distinct phases in our extracted embedding: one with a single conserved quantity $E$ at high energy and another with two approximately conserved quantities $E_{\pm }$ at low energy, which approximately sum to $E\approx E_{+}+E_{-}$.

At first glance, it appears as though our method has only identified a single relevant component corresponding to the conserved total energy $E$ (Figs. 5c, 5e) with rank correlation $\rho =0.996$. However, if we restrict ourselves to low energy trajectories with first embedding component $v_1<-1$, we find that there is a region of the shape space which is two-dimensional, corresponding to the two independently conserved energies $E_{\pm }$ of the low energy non-chaotic phase where the double pendulum behaves like a coupled oscillator system with two distinct modes. For the low energy trajectories, a linear fit of the now two relevant components with $E_{+}$ ($E_{-}$) has rank correlation $\rho =0.836$ ($\rho =0.937$). If we restrict ourselves to even lower energy trajectories with $v_1<-2$, a similar linear fit for $E_{+}$ ($E_{-}$) has rank correlation $\rho =0.934$ ($\rho =0.989$).

This analysis of the double pendulum shows that our method can still provide significant insight into complex dynamical systems with multiple phases involving varying numbers of conserved quantities. This manifests itself as manifolds of different dimensions in shape space that are stitched together at phase transitions, presenting a significant challenge for most manifold learning methods. In this example, this difficulty is reflected in the performance of the heuristic score, which is designed to identify relevant embedding components for a single manifold.

Next, we consider an oscillating Turing pattern system that is both dissipative and has a much higher dimensional phase space than our previous examples. In particular, we study the Barrio–Varea–Aragón–Maini (BVAM) model [BARRIO1999483, PhysRevE.86.026201]

with $D=0.08$, $C=-1.5$, and $H=3$, following PhysRevE.86.026201 who showed that this set of parameters results in a spatial Turing pattern that also exhibits chaotic oscillating temporal dynamics, on a periodic domain with size $8$. The phase space of the BVAM system consists of two functions $u\left(x\right)$ and $v\left(x\right)$ ¹¹1In our method, each trajectory $\left[u\right(x,t_i),v(x,t_i\left)\right],i\in \{1,2,\dots ,N\}$ is treated as an unordered set of sample points in phase space, so we refer to the phase space as $\left[u\right(x),v(x\left)\right]$ in a slight abuse of notation. which we discretize on a mesh of size $50$, giving us an effective phase space dimension of $100$. Because this system is dissipative, we will focus on characterizing the long term behavior of the dynamics, i.e. the oscillating Turing pattern, which appears to have a conserved spatial phase $\eta$ for our chosen set of parameters corresponding to the spatial position of the Turing pattern.

Our method successfully identifies the spatial phase $\eta$ but embeds the angle as a circle in a 2D embedding space (Fig. 6)—a result of the periodic topology of $\eta$. While this shows that the number of relevant components determined by our heuristic score may not always match the true manifold dimensionality, such cases are often easily identified by examining the components directly (Fig. 6c) or by cross checking with an intrinsic dimensionality estimator [https://doi.org/10.48550/arxiv.2106.04018]. A linear fit of the two relevant components with $\cos \eta$ ($\sin \eta$) has $R^2=0.9991$ ($R^2=0.9997$) and $\rho =0.9993$ ($\rho =0.9992$). This example both tests our method on a high dimensional phase space and demonstrates how our approach can be applied to dissipative systems to study long term behavior.

For many spatiotemporal dynamical systems, the conservation laws are local in nature. Locality can significantly simplify the analysis of the conserved quantities and suggests a way to restrict the type of conserved quantities identified by our method. Specifically, we can adapt our approach to focus on local conserved quantities by replacing the raw states (Fig. 7a) by a distribution of local features (Fig. 7b), removing the explicit spatial label and providing a fully translation invariant representation of the state. Then, instead of using the Euclidean metric in the original phase space, we use the energy distance [https://doi.org/10.1002/wics.1375, pmlr-v89-feydy19a] between the distributions of local features as the ground metric for optimal transport.

To demonstrate this method for identifying local conserved quantities, we consider the Korteweg–De Vries (KdV) equation

The KdV equation is fully integrable [PhysRevLett.19.1095] and has infinitely many conserved quantities [doi:10.1063/1.1664701], the most robust of which are the most local conserved quantities expressible in terms of low order spatial derivatives. To focus on these robust local conserved quantities, we use finite differences (i.e. $u\left(x\right),\Delta u\left(x\right)=u(x+\Delta x)-u\left(x\right),{\Delta }^{2}u\left(x\right),\dots$) as our local features, allowing us to restrict the spatial derivative order of the identified conserved quantities. In this experiment, we only take $u\left(x\right)$ and $\Delta u\left(x\right)$, meaning that the identified local conserved quantities will only contain up to first order spatial derivatives. For the KdV equation, there are three such local conserved quantities:

which also have direct analogues in generalized KdV-type equations [1937-1632_2018_4_607].

Our method successfully identifies three relevant components (Fig. 7c) corresponding to (d–f) the three local conserved quantities (Eq. 23). Linear fits of these components to $c_1$, $c_2$, and $c_3$ have rank correlations $\rho =0.995$, $0.994$, and $0.985$, respectively. This result shows how our approach can be adapted to incorporate known structure, such as locality and translation symmetry, in applications to complex high dimensional dynamical systems.