KOopman Operator Learning : A KOOL model for long term stable prediction of dynamical systems

Long-term stable prediction is crucial for surrogate modeling of dynamical systems. Typical machine learning models are accurate in short ranges but either diverge to infinity or decay and lose dynamics after multiple autoregresisve steps. To address this challenge, we propose a Koopman-based deep learning model that effectively learns the underlying invariant statistics providing improved accuracy and dynamic stability for long-term predictions. We use state-of-the-art deep learning techniques like Adaptive Fourier Neural Operators(AFNOs) for function space approximations and a notably smaller non-linear approximation for the Koopman operator. Moreover, we implement a novel loss function that is used to alternatively learn the function spaces and the Koopman Operator. During inference, the dynamics of the state is governed only by the auto-regressive application of the Koopman operator. We demonstrate the efficacy of our approach using chaotic systems like the Kuramoto Shivashinky equation and 2D turbulence. Furthermore, extensive implementations on climate datasets, showcase its ability to forecast key climate variables with greater precision compared to other machine learning techniques for longer periods of time. We can obtain high stability without using any forcing variables or induced biases. Our results highlights the potential of Koopman-based deep learning models as a powerful tool for enhancing the reliability of long-term climate surrogate models, offering valuable insights for climate science.