Deep Koopman operator framework for causal discovery in nonlinear dynamical systems

TL;DR

The paper introduces Kausal, a deep Koopman operator-based framework for causal discovery in nonlinear dynamical systems. By learning optimal observable embeddings via neural encoders and estimating a finite-rank Koopman operator with dynamic mode decomposition, Kausal enables causal analysis through a marginal-vs-joint forecast-error comparison in RKHS. The approach demonstrates improved causal direction detection and magnitude estimation across coupled Rössler oscillators, a reaction-diffusion system, and ENSO dynamics, including real-world ocean data, outperforming several baselines that rely on prescribed dictionaries. The work advances scalable, data-driven causal inference in complex, nonlinear, and potentially nonstationary systems with climate science applications, and provides open-source code for replication and extension.

Abstract

We use a deep Koopman operator-theoretic formalism to develop a novel causal discovery algorithm, Kausal. Causal discovery aims to identify cause-effect mechanisms for better scientific understanding, explainable decision-making, and more accurate modeling. Standard statistical frameworks, such as Granger causality, lack the ability to quantify causal relationships in nonlinear dynamics due to the presence of complex feedback mechanisms, timescale mixing, and nonstationarity. This presents a challenge in studying many real-world systems, such as the Earth's climate. Meanwhile, Koopman operator methods have emerged as a promising tool for approximating nonlinear dynamics in a linear space of observables. In Kausal, we propose to leverage this powerful idea for causal analysis where optimal observables are inferred using deep learning. Causal estimates are then evaluated in a reproducing kernel Hilbert space, and defined as the distance between the marginal dynamics of the effect and the joint dynamics of the cause-effect observables. Our numerical experiments demonstrate Kausal's superior ability in discovering and characterizing causal signals compared to existing approaches of prescribed observables. Lastly, we extend our analysis to observations of El Niño-Southern Oscillation highlighting our algorithm's applicability to real-world phenomena. Our code is available at https://github.com/juannat7/kausal.

Figures (23)

• Figure 1: Schematic of the Kausal algorithm. We 1) estimate the embeddings using deep learning, and 2) approximate the Koopman operator with dynamic mode decomposition (DMD). Then, 3) we infer causal measures by computing the difference in prediction error between marginal (effect-only) and joint (effect-cause) models.
• Figure 2: Illustration of ODEs describing the Rössler Oscillator system (left), and the corresponding causal graph (right).
• Figure 3: Causal measure estimation of coupled Rössler oscillators by computing the difference between the true causal ($\Delta^{K^t}_{C,E}$) and non-causal direction ($\Delta^{K^t}_{E,C}$) across time shifts $t$, using different kernels to approximate the observables.
• Figure 4: Conditional forecasts in the (a) true and (b) non-causal direction using MLP kernels. In (a), the addition of $\Omega_C$ in the joint model improves the forecast of $\Omega_E$ relative to the marginal model that excludes it. In (b), however, both marginal and joint models make no qualitative difference as $\Omega_E \not \rightarrow^t_K \Omega_C$.
• Figure 5: Performance assessment of MLP and RFF kernels for varying dimensionality (i.e. increasing complexity).