CGKN: A Deep Learning Framework for Modeling Complex Dynamical Systems and Efficient Data Assimilation

TL;DR

The paper introduces CGKN, a neural stochastic differential equation framework that converts nonlinear dynamical systems into a conditional Gaussian nonlinear system via a learned latent representation. By encoding unobserved states to latent variables with conditional linear dynamics, CGKN enables closed-form data assimilation formulae and integrates DA performance into training, improving efficiency over ensemble methods. The approach is demonstrated on PSBSE, Lorenz 96, and ENSO SPDEs, showing state forecasting accuracy comparable to or better than baselines and superior DA performance, with clear computational advantages. This work advances scientific machine learning by coupling Koopman-inspired latent linearity with conditional Gaussian structure, enabling robust forecasting and efficient DA in high-dimensional, non-Gaussian systems, with potential extensions to control and inverse problems.

Abstract

Deep learning is widely used to predict complex dynamical systems in many scientific and engineering areas. However, the black-box nature of these deep learning models presents significant challenges for carrying out simultaneous data assimilation (DA), which is a crucial technique for state estimation, model identification, and reconstructing missing data. Integrating ensemble-based DA methods with nonlinear deep learning models is computationally expensive and may suffer from large sampling errors. To address these challenges, we introduce a deep learning framework designed to simultaneously provide accurate forecasts and efficient DA. It is named Conditional Gaussian Koopman Network (CGKN), which transforms general nonlinear systems into nonlinear neural differential equations with conditional Gaussian structures. CGKN aims to retain essential nonlinear components while applying systematic and minimal simplifications to facilitate the development of analytic formulae for nonlinear DA. This allows for seamless integration of DA performance into the deep learning training process, eliminating the need for empirical tuning as required in ensemble methods. CGKN compensates for structural simplifications by lifting the dimension of the system, which is motivated by Koopman theory. Nevertheless, CGKN exploits special nonlinear dynamics within the lifted space. This enables the model to capture extreme events and strong non-Gaussian features in joint and marginal distributions with appropriate uncertainty quantification. We demonstrate the effectiveness of CGKN for both prediction and DA on three strongly nonlinear and non-Gaussian turbulent systems: the projected stochastic Burgers-Sivashinsky equation, the Lorenz 96 system, and the El Niño-Southern Oscillation. The results justify the robustness and computational efficiency of CGKN.

Figures (11)

• Figure 2.1: Schematic diagram of conditional Gaussian Koopman network (CGKN) for a partially observed system. The complex nonlinear dynamical system is transformed into the conditional Gaussian nonlinear system via a generalized usage of Koopman theory, which maps between the unobserved state variables and the latent state variables featured by conditional linear dynamics. The unknown dynamics of the conditional Gaussian nonlinear system are constructed by neural networks and jointly learned with the unknown nonlinear mappings. The analytic DA formulae for the conditional Gaussian nonlinear system can significantly accelerate DA process and also allow a computationally affordable DA loss to be incorporated into the CGKN training process. The trained CGKN model can be used for state forecasts and efficient DA, for which the computations are mainly performed in the conditional Gaussian state space and then mapped back to the original state space.
• Figure 2.2: The architecture of CGKN and its applications for state prediction and data assimilation. The CGKN approximates the true nonlinear complex dynamical system by a modeled system with a conditional linear structure, which is also known as a conditional Gaussian nonlinear system. The conditional linear structure is enabled by a proper choice of latent state variables $\mathbf{v}$ as illustrated in (a). The nonlinear mappings between the original state variables $\mathbf{u}_2$ and the latent state variables $\mathbf{v}$ are achieved via the encoder $\boldsymbol{\varphi}$ and decoder $\boldsymbol{\psi}$, which are jointly learned with the unknown functions $\mathbf{f}_1$, $\mathbf{g}_1$, $\mathbf{f}_2$, $\mathbf{g}_2$ of the conditional Gaussian nonlinear system. The pre-trained CGKN can be used for state prediction and efficient data assimilation as illustrated in (b) and (c).
• Figure 3.1: Simulation results of the projected stochastic Burgers–Sivashinsky equation, with (a) time series of each state variable, (b) the PDFs of the corresponding state variable, and (c) the ACFs of the corresponding state variable. It should be noted that the PDFs and ACFs are estimated from much longer simulations than the one presented in panel (a).
• Figure 3.2: DA results of the projected stochastic Burgers–Sivashinsky equation. EnKBF is used for the true model and analytic formulae are used for the other three models. The uncertainties are indicated by the grey-colored regions, which correspond to two estimated standard deviations from the posterior mean. The uncertainty area of the CGKN model and the standard KoopNet model are estimated based on the method of residual analysis.
• Figure 3.3: Simulation results and statistics of the Lorenz 96 system, with (a) the Hovmoller diagram of the spatiotemporal patterns, (b) time series of $x_2$, (c) the PDF of $x_2$, and (d) the ACF of $x_2$. Note that the PDF and ACF are estimated from a simulation longer than the one presented in panel (b). The behavior of other state variables is similar to $x_2$ since the system is statistically homogeneous.