OnsagerNet: Learning Stable and Interpretable Dynamics using a Generalized Onsager Principle

Here, we show that the generalized Onsager principle (2) or its equivalent form (3) is invariant under coordinate transformations and is a suitable structured form for model reduction. In the original Onsager principle [21, 22], the system is assumed to be not far away from the equilibrium, such that the system dissipation takes a quadratic form: $\|\dot{h}\|_{M}^{2}$, i.e. the matrix $M$ is assumed to be constant. In our generalization, we assume $M$ is a general function that depends $h$. Here we give some arguments why this is necessary and how it can be obtained if the underlying dynamic is given from the viewpoint of model reduction.

To explain nonlinear extension is necessary, we assume that the underlying high-dimensional dynamics which produces the observed trajectories satisfies the form of the generalized Onsager principle with constant $M$ and $W$. The dynamics described by an underlying PDE after spatial semi-discretization takes the form

where $M$ is a symmetric positive semi-definite constant matrix, $W$ is an anti-symmetric matrix.  For the dynamics to make sense, we need $M+W$ to be invertible. $\nabla_{u}V$ is a column vector of length $N$. By taking inner product of (4) with $\dot{u}$, we have an energy dissipation relation of form

Now, suppose that $u$ has a solution in a low-dimensional manifold that could be well described by hidden variables $h(t)\in\mathbb{R}^{m}$ with $m\ll N$. Denote the low-dimensional solution as $u=u(h(t))+\varepsilon$, and plug it into (4), we obtain

Multiplying $J^{T}:=(\nabla_{h}u)^{T}$ on both sides of the above equation, we have

So, we obtain the following ODE system with model reduction error $O(\varepsilon)$:

where

Note that as long as $\nabla_{h}u$ has a full column rank, $(\bar{M}+\bar{W})$ is invertible, so the ODE system makes sense. Now $\bar{M}$, $\bar{W}$ and $\bar{g}$ depend on $h$ in general if $\nabla_{h}u$ is not a constant matrix. Moreover, if the solutions considered all exist exactly in a low-dimensional manifold, i.e. $\varepsilon=0$ in the above derivation, then (5) is exact, which means the generalized Onsager principle is invariant to non-singular coordinate transformations.

Note that the presence of state ($h$) dependence on all the terms implies that we are not linearizing the system about some equilibrium state, as is usually done in linear response theory. Consequently, the functions $W,M,g$ and $V$ may be complicated functions of $h$, and we will learn them from data by parameterizing them as suitably designed neural networks. In this way, we preserve the physical intuition of non-equilibrium physics (dissipation term $M$, conservative term $W$, potential term $V$ and external fields $g$), yet exploit the flexibility of function approximation using data and learning.

In summary, one can view (2) and (3) both as an extension of the classical Onsager principle to handle systems far from equilibrium, or as a reduction of a high dimensional dynamical system defined by generalized Poisson brackets. Both of these dynamics are highly general in their respective domains of application, and serve as solid foundations on which we build our data-driven methodology.

From the mathematical point of view, modeling dynamics using (2) or (3) also has clear advantages, in that the learned system is well-behaved as it evolves through time, unlike unstructured approaches such as dynamic mode decomposition (DMD) (see e.g. [50, 51]) and direct parameterization by neural networks, which may achieve a short time trajectory-wise accuracy, but cannot assure mid-to-long time stability as the learned system evolves in time. In our case, we can show in the following result that under mild conditions, the dynamics described by (2) or (3) automatically ensures a dissipative structure, and remains stable as the system evolves in time.

Theorem 1 shows that the dynamics is physical under the assumed conditions and we will design our neural network parameterization so that these conditions are satisfied.

We implement the generalized Onsager principle based on equation (3).  $\tilde{M}(h)$, $\tilde{W}(h)$, $V(h)$ and $f(h)$ are represented by neural networks with shared hidden layers and are combined according to (3). The resulting composite network is named OnsagerNet. Accounting for (anti-)symmetry, the numbers of independent variables in $\tilde{M}(h)$, $\tilde{W}(h),V(h),f(h)$ are $(m+1)m/2,(m-1)m/2$, $1$, and $m$, respectively.

One important fact is that $V(h)$, as an energy function, should be lower bounded. To ensure this automatically, one may define $V(h)=\frac{1}{2}U(h)^{2}+C$, where $C$ is some constant that is smaller than or equal to $V$’s lower bound. Since the constant $C$ does not affect the dynamics, we drop it in numerical implementation. The actual form we take is

where $\beta\geqslant 0$ is a positive hyper-parameter as assumed in Theorem 1 for forced systems. $U_{i}$ has a similar structure as one component of $\tilde{W}(h)$. We use $m$ terms in the form (10) to ensure that an original quadratic energy after a dimension reduction can be handled easily. To see this, suppose the potential function for the high-dimensional problem (with coordinates $u$) is defined as $V(u)=u^{T}{Au}$ with $A$ symmetric positive definite. Further, let a linear PCA $u\approx u_{0}+Jh$ be used for dimensionality reduction. Then, $V(h)\approx(u_{0}+Jh)^{T}A(u_{0}+Jh)=\|u_{0}\|^{2}+2u_{0}^{T}AJh+h^{T}J^{T}AJh=\|v_{0}+Gh\|^{2}+\text{constant}$, where $G^{T}G=J^{T}AJ,v_{0}^{T}=u_{0}^{T}AJG^{-1}$. Hence, with $(\gamma_{ij})$ representing the matrix $G$, a constant $U_{i}$ suffices to fully represent quadratic potentials. The autograd mechanism implemented in PyTorch [52] is used to calculate $\nabla V(h)$.

To ensure the positive semi-definite property of $\tilde{M}(h)$, we let $\tilde{M}(h)=L(h)L(h)^{T}+\alpha I$, where $L(h)$ is a lower triangular matrix, $I$ is the identity matrix, $\alpha\geqslant 0$. Note that the degree of freedom of $L(h)$ and $\tilde{W}(h)$ can be combined into one $m\times m$ matrix, whose upper triangular part $R(h)$ determines $\tilde{W}(h)=R(h)-R(h)^{T}$. A standard multi-layer perception neural network with residual network structure (ResNet) [53] is used to generate adaptive bases, which takes $(h_{1},\ldots,h_{m})$ as input, and outputs $\{L(h),R(h),U_{i}(h)\}$ as linear combinations of those bases. The external force ${f}_{i}(h)$ are parameterized based on a priori information, and should be of limited capacity so as not to dilute the physical structures imposed by the other terms. For forced systems considered in this paper, we typically take $f_{i}$ as affine functions of $h$. The final output of the OnsagerNet is given by

where $V(h)$ is defined by (10). Note that in an unforced system we have $\alpha=\beta=0$, since they are only introduced in forced system to ensure a stability of the learned system as required by Theorem 1. The computation procedure of OnsagerNet is described in Architecture 1. The full architecture is summarized in Fig. 1.

To learn ODE model represented by OnsagerNet based on sampled ODE trajectory data, we minimize following loss function

Here $\tau$ is the time interval of sample data. ${S}$ is the sample set. $\operatorname{RK2}$ stands for a second order Runge-Kutta method (Heun method). $n_{s}$ is the number of $\operatorname{RK2}$ steps used to forward the solution of OnsagerNet from snapshots at $t$ to $t+\tau$. Other Runge-Kutta method can be used. For simplicity, we only present results using Heun method in this paper. This Runge-Kutta embedding method has several advantages over the linear multi-step methods [5, 6], e.g. the variable time interval case and long time interval case can be easily handled by Runge-Kutta methods.

With the definition of the loss function and model architecture, we can then use standard stochastic gradient algorithms to train the network. Here we will use the Adam optimizer [54, 55] to minimize the loss function with a learning rate scheduler that halves the learning rate if the loss is not decreasing in certain numbers of iterations.

The previous section describes the situation when no dimensionality reduction is sought or required, and an Onsager dynamics is learned directly on the original trajectory coordinates. On the other hand, if the data are generated from a numerical discretization of some PDE or a large ODE system, and we want to learn a small reduced order model, then a dimensionality reduction procedure is needed. One can either use linear principal component analysis (PCA) or nonlinear embedding, e.g. the autoencoder, to find a set of good latent coordinates from the high dimensional data. When PCA is used, we perform PCA and OnsagerNet training separately. When an autoencoder is used, we can train it either separately or together with the OnsagerNet. The end-to-end training loss function is taken as

where $\varphi$, $\psi$ stands for the encoder function and decoder function of the autoencoder, respectively. In the last term, $\left|\|u(t+\tau)-u(t)\|^{2}-\|\varphi\circ u(t+\tau)-\varphi\circ u(t)\|^{2}\right|$ is an estimate of the isometric loss (i.e. deviation from $\varphi$ being an isometry) of the encoder function based on trajectory data, with $\alpha_{isom}$ being a constant smaller than the average isometric loss of the PCA dimension reduction. Here, $(\cdot)_{+}$ stands for positive part and $\beta_{isom}$ is a penalty constant. $\beta_{ae}$ is a parameter to balance the autoencoder accuracy and OnsagerNet fitting accuracy.

The choice of autoencoder architecture follows from our observation that PCA performed respectably on a number of examples we studied. Thus, we build the autoencoder by extending the basic PCA. The resulting architecture, which we call PCA-ResNet, is a stacked architecture with each layer consisting of a fully connected autoencoder block with a PCA-type short-cut connection:

where $n_{k+1}<n_{k}$ and $\text{PCA}_{n_{k},n_{k+1}}$ is a PCA transform from $n_{k}$ dimension to $n_{k+1}$ dimension. $\sigma$ is a smooth activation function. The parameters $W_{2}^{k}$, $W_{1}^{k}$ $b^{k}$ in the encoder are initialized close to zero, such that the encoder becomes a small perturbation of the PCA. On can regard such an autoencoder as a nonlinear extension of PCA. The decoder is designed similarly. We note that there was a similar but different autoencoder proposed to find slow variables [56].

Here, we use the OnsagerNet to learn a deterministic damped Langevin equation

The friction coefficient $\gamma$ may be a constant or a parameter that depends on $v$. For the potential $U(x)$, we consider two cases:

• •

  Hookean spring

  $$U(x)=\frac{\kappa}{2}x^{2}.$$

  (16)

• •

  Pendulum model

  $$U(x)=\frac{4\kappa}{\pi^{2}}\left(1-\cos\left(\frac{\pi x}{2}\right)\right).$$

  (17)

Note that no dimensionality reduction is required here and the coordinate $h=(x,v)$ entails the full phase space. The goal of this toy example is to quantify the accuracy and stability of the OnsagerNet, as well as to highlight the interpretability of the learned dynamics.

We normalize the parameter $\gamma,\kappa$ by $m$, i.e. we take $m=1$. To generate data, we simulate the systems using a third order strong stability preserving Runge-Kutta method [57] for a fixed period of time with initial conditions $\{x_{0},v_{0}\}$ sampled from $\Omega_{S}=[-1,1]^{2}$. Then, we use OnsagerNet (11) to learn an ODE system by fitting the simulated data.

In particular, we will demonstrate that the energy $V$ learned has physical meaning. Note that the energy function in the generalized Onsager principle need not be unique. For example, for the heat equation $\dot{u}=\Delta u$, both $\frac{1}{2}\|u\|^{2}$ and $\frac{1}{2}\|\nabla u\|^{2}$ can serve as an energy functional governing the dynamics, with diffusion operators ($M$ matrix) being $-\Delta$ and the identity respectively. The linear Hookean model is similar. Let $V(x,v)=\frac{1}{2}\kappa x^{2}+\frac{1}{2}v^{2}$, then

The eigenvalue of the matrix $A=\left(\begin{array}[]{cc}0&1\\ -1&-\gamma\end{array}\right)$ is $\lambda_{1,2}=-\frac{\gamma}{2}\pm\frac{1}{2}\sqrt{\gamma^{2}-4}$.  When $\gamma\geqslant 2$, we always obtain real negative eigenvalues, and the system is over-damped. From equation (18), we may define another energy $\tilde{V}(x,v)=\frac{1}{2}x^{2}+\frac{1}{2}v^{2}$ with dissipative matrix and conservative matrix $-\frac{1}{2}(A+A^{T})$, $\frac{1}{2}(A-A^{T})$ respectively. For this system, $\hat{V}(x,v):=\tilde{V}(F(x,v))$ with any non-singular linear transform $F$ could serve as energy, and the corresponding dynamics is

Hence, we use this transformation to align the learned energy before making comparison with the exact energy function.

Let us now present the numerical results. We test two cases: 1) Hookean model with $k=4$, $\gamma=3$; 2) Pendulum model with $k=4$ and $\gamma(v)=3|v|^{2}$. To generate sample data, we simulate the ODE systems to obtain 100 trajectories with uniform random initial conditions $(x,v)\in[-1,1]^{2}$. For each trajectory, 100 pairs of snapshots at $(i\,T/100,i\,T/100+0.001)$, $i=0,\ldots,99$ are used as sample data. Here $T=5$ is the chosen time period of sampled trajectories. Snapshots from the first 80 trajectories are taken as the training set, while the remaining snapshots are taken as the test set.

We test three different methods for learning dynamics: OnsagerNet, a symplectic dissipative ODE net (SymODEN [17]), and a simple multi-layer perception ODE network (MLP-ODEN) with ResNet structure. To make the numbers of trainable parameters in three ODE nets comparable, we choose $l=1$ and $n_{l}=12$ for OnsagerNet, $n_{l}=17$ for SymODEN, and MLP-ODEN with 2 hidden layers and each layer has $9$ hidden units, such that the total numbers of tunable parameters in OnsagerNet, SymODEN and MLP-ODEN are 120, 137, and 124, correspondingly.

To test the robustness of those networks paired with different activation functions, five $C^{1}$ activation functions are tested, including ReQU, ReQUr, softplus, sigmoid and $\tanh$. Here, ReQU, defined as $(x):=x^{2}$ if $x\geqslant 0$, otherwise 0, is the rectified quadratic unit studied in [58]. Since $\operatorname{ReQU}$ is not uniformly Lipschitz, we introduce ReQUr as a regularization of ReQU, defined as $\operatorname{ReQUr}(x):=\operatorname{ReQU}(x)-\operatorname{ReQU}(x-0.5)$.

The networks are trained using a new version of Adam [55] with mini-batches of size 200 and initial learning rate $0.0256$. The learning rate is halved if the loss is not decreasing in 25 epochs. The default number of iterations is set to 600 epochs.

For the Hookean case, the mean squared error (MSE) loss on testing set can be reduced to about $10^{-5}\sim 10^{-8}$ for all the three ODE nets, depending on different random seeds and activation functions used. For the nonlinear pendulum case, the MSE loss on testing set can be reduced to about $10^{-4}\sim 10^{-5}$ for OnsagerNet and $10^{-3}\sim 10^{-5}$ for MLP-ODEN, but only $10^{-2}$ for SymODEN, see Fig 2. The reason for the low accuracy of SymODEN is that in the SymODEN, the diffusion matrix is assumed to be a function of general coordinate $x$ only [17], but here in the damped pendulum problem, the diffusion term depends on $v$. From the test results presented in Figure 2(a), we see that the results of OnsagerNet are not sensitive to the nonlinear activation functions used. Moreover, 2(b) shows that OnsagerNet has much better long time prediction accuracy. Since the nonlinearities in many practical dynamical systems are of polynomial type, we will mainly use ReQU and ReQUr as activation functions for other numerical experiments in this paper.

In Fig 3 we plot the contours of learned energy functions using OnsagerNet and compare with the exact ground truth total energy function $U(x)+v^{2}/2$. We observe a clear correspondence up to rotation and scaling for the linear Hookean model case and a scaling of the nonlinear pendulum case. In all cases, the minimum $(x=0,v=0)$ is accurately captured. Note that we used a linear transform to align the learned free energy. After the alignment, the relative $L^{2}$-norm error between the learned and physical energy for the two tested cases are $6.3\times 10^{-3}$ and $8.6\times 10^{-2}$ respectively. To verify that the OnsagerNet approach is able to learn non-quadratic potentials, we also test an example with exact double-well type potential $U(x)=(x^{2}-0.5)^{2}$ and $\gamma=3$. The relative $L^{2}$-norm error between the learned and exact energy is $7.5\times 10^{-2}$, where a simple $\min$-$\max$ rescaling is used before calculating the numerical error for this example.

In Fig. 4, we plot the trajectories for exact dynamics and learned dynamics with initial values starting from four boundaries of the sample region. We see that the learned dynamics are quantitatively accurate, and the qualitative features of the phase space are also preserved over long times.

The previous oscillator models have simple Hamiltonian structures, so it is plausible that with some engineering one can obtain a special structured model that works well for such systems. In this section, we consider a target nonlinear dynamical system that may not have such a simple structure. Yet, we will show that owing to the flexibility of the generalized Onsager principle, OnsagerNet still performs well. Concretely, we consider the well-known Lorenz system [39]:

where $b>0$ is a geometric parameter, $\sigma$ is the Prandtl number, $r$ is rescaled Rayleigh number.

The Lorenz system (19)-(21) is a simple autonomous ODE system that produces chaotic solutions, and its bifurcation diagram is well-studied [59, 60]. To test the performance of OnsagerNet, we fix $b=8/3$, $\sigma=10$, and vary $r$ as commonly studied in the literature. For the $b=8/3$, $\sigma=10$ case, the first (pitchfork) bifurcation of the Lorenz system happens at $r=1$, followed by a homoclinic explosion at $r\approx 13.92$, and then a bifurcation to the Lorenz attractor at $r\approx 24.06$. Soon after, the Lorenz attractor becomes the only attractor at $r\approx 24.74$ (see e.g. [61]). To show that OnsagerNet is able to learn systems with different kinds of attractors and chaotic behavior, we present results for $r=16$ and $r=28$.

The procedure of generating sample data and training is similar to the previously discussed case of learning Langevin equations, except that here we set $\alpha=0.1$, $\beta=0.1$ and a linear representation for ${f}(h)$ in OnsagerNet (with $l=1$ and $n_{l}=20$). This is because the Lorenz system is a forced system. The result for the case $r=16$ is presented in Fig 5. We see that both the trajectories and the stable fixed points and unstable limit cycles can be learned accurately. The results for the case $r=28$ is presented in Fig 6. The largest Lyapunov indices of numerical trajectories (run for sufficient long times) are estimated using a method proposed by [62]. They are all positive, which suggests that the learned ODE system indeed has chaotic solutions.

Finally, we compare OnsagerNet with MLP-ODEN for learning the Lorenz system. The SymODEN method [17] cannot be applied since the Lorenz system is non-Hamiltonian. The OnsagerNets used here has one shared hidden layer with 20 hidden nodes, and the total number of trainable parameters is 356. The MLP-ODEN nets have 2 hidden layers with 16 hidden nodes in each layer, with a total 387 tunable parameters. In Fig. 7, we show the accuracy on the test data set for OnsagerNet and MLP-ODEN using $\operatorname{ReQU}$ and $\operatorname{ReQUr}$ as activation functions, from which we see OnsagerNet performs much better. To ensure that these results hold across different model configurations, we also tested learning the Lorenz system with larger OnsagerNets and MLP-ODENs with 3 hidden layers with 100 hidden nodes each. Some typical training and testing loss curves are given in Figure 8, from which we see OnsagerNet perform better than MLP-ODEN, and activation function ReQUr peforms better than tanh as activation function.

We now discuss the main application of OnsagerNet in this paper, namely learning reduced order models for the 2-dimensional Rayleigh-Bénard convection (RBC) problem, described by the coupled partial differential equations

where $\vec{u}=(u,u,0)$ is the velocity field, $g_{0}$ the gravitational acceleration, $\hat{y}$ is the unit vector opposite to gravity, $\theta$ the departure of the temperature from the linear temperature profile $\bar{\theta}(y)=\bar{\theta}_{H}-\Gamma y$, $\Gamma=(\bar{\theta}_{H}-\bar{\theta}_{L})/d$. $d$ is the depth of the channel between two plates. The constants $\alpha_{0},\nu$, and $\kappa$ denote, respectively, the coefficient of thermal expansion, the kinematic viscosity, and the thermal conductivity. The system is assumed to be homogeneous in $z$ direction, and periodic in $x$ direction, with period $L_{x}$. The dynamics depends on three non-dimensional parameters: the Prandtl number $Pr=\nu/\kappa$, the Rayleigh number $Ra=d^{4}\frac{g_{0}\alpha_{0}\Gamma}{\nu\kappa}$, and aspect ratio $a=2d/L_{x}$. The critical Rayleigh number to maintain the stability of zero solution is $R_{c}=\pi^{4}(1+a^{2})^{3}/a^{2}$. The terms $\alpha_{0}g_{0}\hat{y}\theta$ in (22) and $\Gamma v$ in (23) are external forcing terms. For given divergence-free velocity $\vec{u}$ with no-flux or periodic boundary conditions, the convection operator $\vec{u}\cdot\nabla$ is skew-symmetric. Thus, the RBC equations satisfy the generalized Onsager principle with potential $\frac{1}{2}\|u\|^{2}+\frac{1}{2}\|\theta\|^{2}$. Therefore, the OnsagerNet approach, which places the prior that the reduced system also satisfy such a principle, is appropriate in this case. The velocity $\vec{u}$ can be represented by a stream function $\phi(x,y)$:

By eliminating pressure, one gets the following equations for $\phi$ and $\theta$:

The solutions $\phi,\theta$ to (25) and (26) have representations in Fourier sine series as

where $\theta_{k_{1}k_{2}}=\bar{\theta}_{-k_{1}k_{2}}$ and $\phi_{k_{1}k_{2}}=\bar{\phi}_{-k_{1}k_{2}}$ since $\theta$ and $\phi$ are real. In the Lorenz system, only 3 most important modes $\phi_{11}$, $\theta_{11}$ and $\theta_{02}$ are retained, in this case, the solution have following form

The Lorenz ’63 equations (19)-(21) for the evolution of $X,Y,Z$ is obtained by a Galerkin approach [39] with time rescaling $\tau=\pi^{2}\frac{(1+a^{2})}{d^{2}}\kappa t$, the rescaled Rayleigh number $r=\operatorname{Ra}/R_{c}$. $b=4/(1+a^{2})$, and $\sigma=\nu/\kappa$ is the Prandtl number.

Since Lorenz system is aggressively truncated from the original RBC problem, it is not expected to give quantitatively accurate prediction of the dynamics of the original system when $r\gg 1$. Some extensions of Lorenz system to dimensions higher than 3 are proposed, e.g. Curry’s 14-dimensional model [63]. But numerical experiments have shown that much large numbers of spectral coefficients need be retained to get quantitatively good results in a Fourier-Galerkin approach [40]. In fact, [40] used more than 100 modes to obtain convergent results for a parameter region similar to $b=8/3$, $\sigma=10$, $r=28$ used by Lorenz. In the following, we show that by using OnsagerNet, we are able to directly learn reduced order models from RBC solution data that are quantitatively accurate and require much fewer modes than the Galerkin projection method.