PyKoopman: A Python Package for Data-Driven Approximation of the Koopman Operator

PyKoopman is a Python package for the data-driven approximation of the Koopman operator associated with a dynamical system. The Koopman operator is a principled linear embedding of nonlinear dynamics and facilitates the prediction, estimation, and control of strongly nonlinear dynamics using linear systems theory. In particular, PyKoopman provides tools for data-driven system identification for unforced and actuated systems that build on the equation-free dynamic mode decomposition (DMD) and its variants. In this work, we provide a brief description of the mathematical underpinnings of the Koopman operator, an overview and demonstration of the features implemented in PyKoopman (with code examples), practical advice for users, and a list of potential extensions to PyKoopman. Software is available at http://github.com/dynamicslab/pykoopman


Statement of need
Engineers have long relied on linearization to bridge the gap between simplified descriptions where powerful analytical tools exist, and the intricate complexities of nonlinear dynamics where analytical solutions are elusive (Ljung, 2010;Wright et al., 1999).Local linearization, implemented via first-order Taylor series approximation, has been widely used in system identification (Ljung, 2010), optimization (Wright et al., 1999), and many other fields to make problems tractable.However, many real-world systems are fundamentally nonlinear and require solutions outside of the local neighborhood where linearization is valid.Rapid progress in machine learning and big data methods are driving advances in the data-driven modeling of such nonlinear systems in science and engineering (Steven L. Brunton & Kutz, 2022).
In the diverse landscape of data-driven modeling approaches, Koopman operator theory has received considerable attention in recent years (Steven L. Brunton et al., 2017;Budišić et al., 2012;Klus et al., 2018;Li et al., 2017;Mezić, 2013;Williams, Kevrekidis, et al., 2015).The main idea is illustrated in Fig. 1.This methodology enables the application of closed-form, convergence-guaranteed methods from linear system theory to general nonlinear dynamics.To fully leverage the potential of data-driven Koopman theory across a diverse range of scientific and engineering disciplines, it is critical to have a central toolkit to automate state-of-the-art Koopman operator algorithms.
Figure 1: Illustration of Koopman operator for a 2D nonlinear system.This system can be linearized into a 3D linear system with the nonlinear transform (x).In PyKoopman, searching such (x) is facilitated in a data-driven manner.
As a result, PyKoopman has been developed as a Python package for approximating the Koopman operator associated with natural and actuated dynamical systems from measurement data.Compared to implementation of DMD (e.g., PyDMD (Demo et al., 2018),) which can be viewed as a linear projection of the Koopman operator, PyKoopman offers a comprehensive set of nonlinear projection methods.Specifically, PyKoopman offers tools for designing the observables (i.e., functions of the system state) and inferring a finite-dimensional linear operator that governs the dynamic evolution of these observables in time.These steps can either be performed sequentially (Williams, Rowley, et al., 2015;Williams, Kevrekidis, et al., 2015) or combined, as demonstrated in more recent neural network models (Lusch et al., 2018;Mardt et al., 2018;Otto & Rowley, 2019;Takeishi et al., 2017).In addition, we also support data from multiple trajectories.Once a linear embedding is discovered from the data, the linearity of the transformed dynamical system can be leveraged for enhanced interpretability (Pan et al., 2021) or for designing near-optimal observers (Surana & Banaszuk, 2016) or controllers for the original nonlinear system (Kaiser et al., 2021;Korda & Mezić, 2020;Mauroy et al., 2020;Peitz et al., 2020;Peitz & Klus, 2019).The core component of the PyKoopman package is the Koopman model class.We used several base classes from scikit-learn (Pedregosa et al., 2011) to build the machine learning pipeline.We used pytorch (Paszke et al., 2019) and lightning (Falcon & The PyTorch Lightning team, 2019) to implement deep learning methods for Koopman operator.We also used PyDMD (Demo et al., 2018) to incorporate some existing implementation for regression after nonlinear observables are chosen.Finally, we used derivative (Kaptanoglu et al., 2022) to obtain time derivative to deal with non-uniformly sampled data.To summarize, the external package dependencies are depicted in Fig. 2.

Features
As illustrated in Fig. 3, PyKoopman is designed to lift nonlinear dynamics into a linear system with linear actuation.Specifically, our PyKoopman implementation involves two major steps: • observables: the nonlinear observables used to lift x to z, and reconstruct x from z; • regression: the regression used to find the optimal A. Additionally, we have a differentiation module that evaluates the time derivative from a trajectory and the analytics module for sparsifying arbitrary approximations of the Koopman operator.

Example
The PyKoopman GitHub repository provides several helpful Jupyter notebook tutorials.Here, we briefly demonstrate a typical workflow using the PyKoopman package to approximate Koopman operator of a 2D nonlinear system.First, consider the dynamical system Given X and Xnext matrices are two one-step away trajectories of the above nonlinear system, we can choose polynomial observables and EDMD regressor to learn the Koopman operator by feeding X and Xnext to the .fitmethod of an instance of pykoopman.Koopman class.Once the Koopman object model has been fit, we can use the model.simulatemethod to make predictions over an arbitrary time horizon.Fig. 4 displays the excellent agreement between ground truth and the EDMD prediction from the aforementioned Koopman model on randomly generated unseen test data.

Conclusion
Our goal of the PyKoopman package is to provide a central hub for education, application and research development of learning algorithms for Koopman operator.The PyKoopman package is aimed at researchers and practitioners alike, enabling anyone with access to discover linear embeddings of nonlinear systems from data.Following PySINDy (Silva et al., 2020) and Deeptime (Hoffmann et al., 2021), PyKoopman is designed to be accessible to users with basic knowledge of linear systems, adhering to scikit-learn standards, while also being modular for more advanced users.We hope that researchers and practioners will use PyKoopman as a platform for algorithms developement and applications of linear embedding.

Figure 3 :
Figure 3: Broad categorization of model types that can be identified with current PyKoopman.While the dotted parts (marked with "⋅'') can be simultaneously discovered within the framework, they are typically ignored for control purposes.

Figure 4 :
Figure 4: Example of learning Koopman operator for a 2D nonlinear system.Left: distribution of training data consists of multiple trajectories.Right: prediction on unseen test trajectories.