Physics-Informed Koopman Network for time-series prediction of dynamical systems

Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 23:25:45

Abstract: Koopman operator theory is receiving increased attention due to its promise to linearize nonlinear dynamics. Neural networks that are developed to represent Koopman operators have shown great success thanks to their ability to approximate arbitrarily complex functions. However, despite their great potential, they typically require large training data-sets either from measurements of a real system or from high-fidelity simulations. In this work, we propose a novel architecture inspired by physics-informed neural networks, which leverage automatic differentiation to impose the underlying physical laws via soft penalty constraints during model training. We demonstrate that it not only reduces the need of large training data-sets, but also maintains high effectiveness in approximating Koopman eigenfunctions.

Main paper: 
 introduction: 
 Nonlinear dynamical systems give rise to a rich diversity of complex phenomena such as the ones arise in climate science Lorenz (1956) , neuroscience Brunton et al. (2016) , ecology Clark & Luis (2020) , finance Mann & Kutz (2016) , and epidemiology Proctor & Eckhoff (2015) . In 1931, Koopman introduced the operator-theoretic perspective of dynamical systems , complementing the traditional geometric perspectives Mezic (2020) . In this framework, a Koopman operator is defined which acts on observation functions (observables) in an appropriate function space. Under the action of this operator, the evolution of the observables are linear although the function space may be infinite-dimensional. As a consequence, approximating the Koopman operator and seeking its eigenfunctions become a key to linearize the nonlinear dynamics Mezic (1994) ; Mezić (2005) ; Brunton et al. (2021) .While there are various conventional methods to approximate Koopman operator spectra, e.g. DMD method Rowley et al. (2009) and its variants, or data-driven autoencoder-based methods, e.g. Takeishi et al. (2017) ; Morton et al. (2018) ; Lusch et al. (2018) ; Gin et al. (2021) , they all face various challenges. DMD-based methods, for instance, require an a priori, judicious selection of the observables with no guarantee that these observables span an invariant Koopman subspace Kutz et al. (2016) . Conversely, neural-network-based methods need to acquire a large enough data-set, which pose a challenge, e.g. in terms of computational cost or experimental complexities. In this work, we propose physics-informed Koopman networks (PIKNs) which assimilates the knowledge of the dynamical system and autoencoder-based Koopman networks. We demonstrate doing so reduce the need for large training data-sets and it enables the model to better predict beyond the training horizon. It should be noted that our contribution is the addition of physics-informed loss, which can be combined with any state-of-the-art data-driven approximation of the Koopman operator. 
koopman operator theory: 
 For an autonomous ordinary differential equation of d dt x(t) = f (x(t)), with x ∈ X ⊆ R n , we define the time-t flow map operator F t : X → X asEQUATIONIn Koopman framework, an alternative description for dynamical systems is in terms of evolution of functions of possible measurements g : X → C, which belongs to a set of functions G(X ) of significantly higher dimension than X . The family of Koopman operators K t : G(X ) → G(X ), parameterized by t are given by K t g(x) = g(F t (x))(2) It can be shown K t , which in general an infinite dimensional, is linear. Constructing finitedimensional representations of Koopman operator remains an open question, and is in the scope of this work. If f is sufficiently smooth, one can also define the infinitesimal generator L of the Koopman operator family asEQUATIONFrom the definition, we can easily seeEQUATIONThe generator L is sometimes referred to as the Lie operator. On the other hand, we also haveEQUATIONTherefore, we conclude Lg = ∇g • f (6) Equation 6 will be the key for the implementation of PIKN.It is straightforward to show that Koopman eigenfunctions φ(x) that satisfies K t φ(x) = λ t φ(x) for λ t ̸ = 0 are also eigenfunctions of the Lie operator, which we denote as ϕ i . It can be shown thatEQUATIONThis also implies span{φ k } M k=1 is an invariant subspace under the Koopman operator K t and can be viewed as the new coordinates on which the dynamics evolve linearly.Since our ultimate goal is to study nonlinear dynamical systems using linear theory, we do not need to restrict ourselves to Equation 7. Following Lusch et al. (2018) ; Gin et al. (2021) , we can generalize it as g(x) = ψ(φ 1 (x), φ 2 (x), . . . , φ M (x); ω)EQUATIONwhere ψ is an arbitrary transformation parameterized by ω.In physical sciences, data is scarce while governing equations are available in literature. In physicsinformed Koopman networks (PIKNs), we aim to leverage such knowledge of the dynamics, e.g., of Equation 6, to enforce the linearity constraint. The basic idea is to train the network by minimizing the quantity∥∇φ k (x) • f -µ k φ k (x)∥.More generally, a squared matrix L is used to approximate the Lie operator L, which in turn is related to the Koopman operator, and we minimize ∥Lϕ(x) -∇ϕ(x) • f ∥. Finding the eigenvalue and eigenfunction pairs of the Lie operator corresponds to performing eigendecomposition to the matrix L.We can seamlessly integrate information from measurement data X data := {x(t 0 ), x(t 1 ), • • • , x(t p )}, and the knowledge (collocation points) x := {x(t 0 ), x(t 1 ), • • • , x(t p )} such that the loss function isJ = 1 N N i=1 (ω 1 ∥Lϕ(x i ) -∇ϕ(x i ) • f (x i )∥ 2 + ω 2 ∥x i -ψ(z i )∥ 2 )+ 1 p p j=0 (ω 3 ∥e L∆tj ϕ(x(t 0 )) -ϕ(x(t j ))∥ 2 + ω 4 ∥x(t j ) -ψ(z(t j ))∥ 2 ) (9)For detailed description of each term, please see Appendix A.1-A.2.  
simple nonlinear system with discrete spectrum: 
 First, we consider a simple nonlinear system with a discrete eigenvalue spectrum:EQUATIONFor λ < µ < 0, this example is serves as a benchmark for Koopman related algorithms because it provides explicitly defined three-dimensional Koopman invariant subspace with µ, 2µ, λ] andϕ(x) = [x 1 , x 2 1 , x 2 -bx 2 1 ]as Koopman eigenvalues and corresponsing eigenfunction, respectively and b = λ λ-2µ . As shown in Fig 1 , PIKN is able to find three-dimensional Koopman invariant subspace. Note that these networks are trained in a purely physics-informed manner, indicating there is no need for simulation data. Refer tp Appendix A.3.1 for more details on experiments.We then study the nonlinear pendulum which exhibits a continuous eigenvalue spectrum with increasing energy:EQUATIONSince this example is broadly tested among most of alternative architectures, we use it to benchmark our method.For more setup details, refer to Appendix A.3.2.All physics-informed counterparts perform better than the original architectures, with PiFBDAEcont to be the best of all.We finally consider the one-dimensional heat equation, as an exmplar for PDEs:EQUATIONwith periodic boundary conditions. Using Fourier transform, it can be shown that discrete-time eigenvalues are µ k = -k 2 , k = 0, ±1, ±2, . . . . To approximately represent the trajectories u and the collocations u t , we discretize the spatial domain with n = 64 equally spaced grids. Therefore, we expect our network to at least mimic a discrete Fourier transform and its inverse transform, identifying the right eigenvalues after training. More specifically, we compare the results obtained from the networks that is purely physics-informed, purely data-driven and a mixture of both which we call it hybrid model. We use sums of harmonic functions with random coefficients as collocations  
discussion and conclusion: 
 In this work, we presented an effective deep learning framework for identifying Koopman eigenvalue and eigenfunction pairs for nonlinear dynamics. In order to validate our method, we carefully went through three examples. In first example, we found the Koopman eigendecomposition only using the knowledge of the system (no data). In second example, we demonstrated adding physics informed loss can improve various autoencoder-based methods. In third example, we show the strength of a hybrid method for a PDE with known solution for the Koopman operator. 
 
 Appendix: In physical sciences, data is scarce while governing equations are available in literature. In physicsinformed Koopman networks (PIKNs), we aim to leverage such knowledge of the dynamics, e.g., of Equation 6, to enforce the linearity constraint. The basic idea is to train the network by minimizing the quantity∥∇φ k (x) • f -µ k φ k (x)∥, ∀k = 1, 2, 3, • • • , M .More generally, a squared matrix L is used to approximate the Lie operator L, which in turn is related to the Koopman operator, and we minimize ∥Lϕ(x) -∇ϕ(x) • f ∥. Finding the eigenvalue and eigenfunction pairs of the Lie operator corresponds to performing eigendecomposition to the matrix L.We first consider an ordinary differential equation. We start from sampling a set of collocation pointsX := {x 1 , x 2 , • • • , x N }.This set of collocation points does not need to come from any trajectories of the dynamics but they can be sampled randomly, avoiding a bulk of simulations or measurement data collections.The objective of the network is to identify a few key coordinates z = ϕ(x) spannned by a set of Koopman eigenfunctions φ k (x) : R n → R, k = 1, 2, • • • , M along with a dynamical system ż = Lz. Objective function isEQUATIONif the decoder is linear (where C represents the reconstruction coefficients), or for generic decoderEQUATIONHere, ω 1 and ω 2 are the weights for each loss term. The first term encourages linear dynamics within the latent space and the second term makes it a valid auto-encoder. One can further diagonalize L such that the diagonal elements approximate the Koopman eigenvalues and the corresponding outputs of the encoder approximate the Koopman eigenfunctions, respectively. But this constraint is not necessary as it is equivalent to performing eigendecomposition of a general-structured L after training. Once trained, it can be used for state predictions beyond the time horizon used in training.For application to partial differential equations of the formEQUATIONsimilar to ODEs, the goal is to seek coordinates v = ϕ(u) that has linear evolution v = Lv and can be used to reconstruct the original measurements û = ψ(v). The main difference is, however, the input and output of the network are functions of spatio-temporal variables u(x, t) instead of the temporal variables x(t) as in the ODE cases. Therefore, we need to sample a set of "collocation points" in an appropriate function space where u t can be cheaply evaluated. We provide examples of such suitable function families in the following chapters.Like PINNs, we can seamlessly integrate information from measurement data. Suppose we have snapshots of measurements X data := {x(t 0 ), x(t 1 ), (ω 3 ∥e L∆tj ϕ(x(t 0 )) -ϕ(x(t j ))∥ 2 + ω 4 ∥x(t j ) -ψ(z(t j ))∥ 2 ) (for ODE)EQUATIONwe can penalize the network predictions that do not match the real measurements, where ∆t j = t j -t 0 , ∀j = 1, 2, • • • , p. Again, the first term is the linearity loss and the second term is the reconstruction loss. It should be noted that Eq. 16 is consistent with previous literature Takeishi et al. (2017) ; Morton et al. (2018) ; Lusch et al. (2018) on using autoencoders to find approximation of Koopman eigenfunctions and can be seen as a baseline on how physics-informed loss improve the performance of PIKN, although we use several other data-driven models in section 3.2.For all networks in the experiments, L is set to a zero matrix initially and all other parameters are randomly initialized using the default initializer of Pytorch. The exponential linear unit (ELU) is used as the nonlinear activation function, unless stated otherwise, as we expect the transformations in PIKN to be relatively smooth.In this example, we have trained two different types of PIKN: a PIKN with a linear decoder and a PIKN with a nonlinear decoder. In both experiments, 1000 collocation points were uniformly sampled from [-1, 1] × [-1, 1] for training, an Adam optimizer with a learning rate of 1e -4 has been applied. The total number of training epochs is set to 50000 and the weights in the loss function are set to ω 1 = ω 2 = 1. The encoder part of both architectures are the same: a 2-layer fully-connected neural network with hidden layer containing 50 neurons. The major differences between the two architectures are as follows:• The linear decoder is simply a linear layer without a bias term whereas the nonlinear decoder is symmetric to the encoder: a 2-layer fully-connected neural network with hidden width 50. • For the PIKN with a linear decoder, we have a three-dimensional latent space whereas for the PIKN with a nonlinear decoder, it is two-dimensional.In this example, we implement three different types of architectures from existing literature and benchmark against their physics-informed counterparts (our methods). The first type of architecture assumes a discrete eigenvalue spectrum of Koopman operator (DAE-disc) and has different versions of implementations Takeishi et al. (2017) ; Morton et al. (2018) . The main difference is on how we obtain the matrix L of the linear dynamics. We follow the latter and treat L as a trainable parameter in Pytorch. We also experiment with different latent dimensions (2, 4, 8, 16) to study the effects. The second type of architecture is proposed by Lusch et al. (2018) which uses an auxiliary network to parameterize the continuous spectrum (DAE-cont), the matrix L varies according to the inputs which effectively captures the frequency shifts. The third type further improves upon it by considering the inverse dynamics (FBDAE-cont). The consistency between forward and backward dynamics is enforced by an extra loss term Azencot et al. (2020) . In our implementations, decoder is linear, encoder is a 3-layer fully connected neural network with hidden width 20 and the auxiliary network is a 4-layer fully connected network with hidden width 10.For training, 1000 collocation points and 20 initial conditions are uniformly sampled from [-1, 1] × [-1, 1]. The initial conditions are further used to generate sequence of snapshots with a temporalAnother way to get around this is to learn the pseudo-inverse of L instead. Then the loss term reads ∥ϕ(x i ) -L † ∇ϕ(x i ) • f (x i )∥. In that case ω1 and ω2 can be both set to 1 because the two loss terms are approximately on the same scale. 


Decision: Dear Authors,

Thank you for submitting the draft.

Both reviewers agree that the presented work presents some interesting strengths. However, both reviewers do also raise some major points of concern, especially regarding the clarity of the presentation and some of the claims made. It is expected that authors will be addressing comments by the reviewers in the final draft.

regards

AC

Reviews: review: title: title: Review of "Physics-Informed Koopman Network for time-series prediction of dynamical systems"
main_review: main_review: # Summary

The paper proposes a physics-inspired soft constraint for training Koopman networks.
The constraint in question follows from the definition of the Koopman operator, and is given by $\mathcal{L}g = \nabla g \cdot \mathbf f$, where $\mathcal{L}$ is the Lie operator, $g$ is the observable function, and $\mathbf f$ is the ground truth ODE given by $\frac{d}{dt}\mathbf x(t) = \mathbf{f}(\mathbf{x}(t))$. 
In practice, the Lie operator is approximated by a square matrix $L$, which is learned during training and whose eigenfunctions and eigenvalues can then be estimated.
In experiments on a number of prototypical systems, the proposed method is able to recover the correct eigenfunctions and eigenvalues of the Lie operator.

# Strengths
1. (**significance**) The work has the potential to be impactful, due to apparently widespread interest in obtaining linear representations of nonlinear dynamics.
1. (**significance**) The proposed soft constraint is, to the best of my knowledge, a novel contribution.

# Weaknesses
1. (**significance**, **clarity**) It is not entirely clear what the paper is trying to achieve with the proposed method.
Unlike the data-driven methods referenced in the introduction, such as those based on dynamic mode decomposition (DMD) and variational autoencoders (VAEs), the soft constraint proposed here requires complete *a priori* knowledge of the ground truth dynamics $\mathbf f$.
In other words, the proposed method incorporates the ground truth dynamics into the training procedure.
It is then reasonable to ask: if we already know $\mathbf f$, why would we want to train a Koopman network to learn (a linear representation of) $\mathbf f$?
Despite the title of the paper, it seems that the proposed method is not very interesting in a purely time-series prediction paradigm, since we necessarily already know the ground truth dynamics.
We are therefore left to speculate about other settings in which the proposed method might be useful.
For example, and this is not discussed much in the paper, one might suppose that obtaining a linear representation of some known, nonlinear $\mathbf f$ is a worthwhile aim in and of itself.
According to Brunton and Kutz [1], “Expressing nonlinear dynamics in a linear framework is appealing because of the wealth of optimal estimation and control techniques available for linear systems.”
However, if this were the aim of the method, then this needs to be stated clearly and the present work put into the context of related methods for obtaining the Koopman modes from $\mathbf f$, including purely analytic approaches such as Laurent series expansions [1].
Indeed, while all of the systems studied in this paper have analytically tractable eigenfunctions and eigenvalues, there may exist systems where they are intractable [2], offering a potentially useful application of the proposed method.
1. (**quality**) The introduction claims that "We demonstrate doing so reduce the need for large training data-sets and it enables the model to better predict beyond the training horizon," but neither of these claims is backed up clearly by the results.
The first claim would require a comparison to related methods while using varying amounts of training data.
The second claim would require a comparison to related methods using some kind of forecasting metric such as relative error (but I am still not convinced that this would be an interesting application of this method, as discussed in the previous point).
1. (**quality**, **clarity**) Further to the previous point, Section 3.2 claims that "All physics-informed counterparts perform better than the original architectures", yet the results in Table 1 suggest no significant difference between the proposed method and the baselines (although the exact meaning of the error quantity, i.e. the number after $\pm$, is not stated anywhere).
2. (**clarity**, **completeness**) There is no discussion in the paper about how to choose the dimension of the Koopman invariant subspace when it is not known *a priori*.
3. (**clarity**) Unfortunately, there are issues throughout the text with spelling, punctuation, formatting, and citations. 
As a result, presentation is a serious issue for the paper in its current form, which is hard to follow in places.
To give some concrete examples:
    - The citations at the beginning of the introduction are not formatted correctly.
    - Symbols and subscripts in the methods section are not well defined or consistent. 
    For example, in Section 2.2, $p$ is used to mean both the number of measurements and the number of collocation points (we have to go to the Appendix to find out that $N$ is in fact the number of collocation points).
    In the same section, $\mu_k$ appears out of nowhere, which is presumably supposed to be $\lambda_k$.
    - Table 1 is not referenced anywhere in the text, although it apparently contains the results for Section 3.2. 

# Conclusion
While the paper has serious issues in its current form, I think most of them can be addressed fairly easily.
In particular, with regard to the first weakness discussed above, I would encourage the authors to think carefully about what problem is solved by the proposed method and how it relates to existing work.

# Citations
[1] Brunton SL, Kutz JN. Data-Driven Science and Engineering: Machine Learning, Dynamical Systems, and Control. Cambridge University Press; 2019.

[2] Lusch, B., Kutz, J.N. & Brunton, S.L. Deep learning for universal linear embeddings of nonlinear dynamics. Nat Commun 9, 4950 (2018).
title: title: A new loss for data-driven approximation of the Koopman operator
main_review: main_review: This paper aims to improve the approximation of Koopman operators by introducing a new loss function based on mathematical properties of the infinitesimal generator of the Koopman operator. The paper is clear, and offers a simple, generic alteration for practitioners approximating Koopman operators.

The main pitfall of this paper is the lack of comparison to existing loss functions for the same task. At least one should be included in the main paper, in order to showcase the difference between existing literature and this paper.

Overall, the paper is well-written and makes a compelling case for practitioners to adopt this loss.

Some minor nitpicks:
- Each of the terms of this loss are immediately interpretable by someone with a background in the theory of linear dynamical systems, except for the terms $z$ and $z_i$. Therefore, consider pulling the definition out of the appendix so it's possible to understand the loss without going to the appendix. Especially since this is the central contribution of the paper.
- $\omega$ is used both as parameters in $\psi$, and as weight terms in the loss function, ie. $\omega_1, \omega_2, ...$ Consider using a different notation for these.
- In the third sentence after equation (2), the authors state "If, $f$ is sufficiently smooth, ..." without a reference to the smoothness conditions.
- In the second sentence after equation (6), the authors state "It is straightforward to show that Koopman eigenfunctions $\phi(x)$ that satisfies $\mathcal{K}^t \phi(x) = \lambda^t \phi(x)$ for $\lambda^t \neq 0$ are also eigenfunctions of the Lie operator..." The wording here is not entirely clear, is this to say eigenfunctions which satisfy “for all t there exists $\lambda^t$ such that $\mathcal{K}^t\phi(x) = \lambda^t \phi^t...$?“
- In the paragraph preceding equation (8), it reads "Since our ultimate goal is to study nonlinear dynamical systems using linear theory, we do not need to restrict ourselves to Equation 7." This is again not entirely clear. Why should one not limit themselves to studying linear functions of the Koopman eigenfunctions?
- In the first sentence following equation (10), "$\mu, 2\mu, \lambda]$" is missing an opening bracket.
- In the first sentence of section 3.3, "exmplar" is misspelled.