Kevin Horner, Daimon Byl & Benjamin Cottrell: Practical Gray-Box Universal Differential Equation Methods for Chaotic Climate-Relevant Dynamics
AUTHORS: Daimon Byl, Kevin Horner, Benjamin Cottrell, Vassili Kitsios, Andrew Ooi
Introduction: Gray-box modelling with Universal Differential Equations (UDEs) is a promising route for applying machine learning to chaotic dynamical systems. Rather than replacing physical models with black-box neural networks, UDEs retain known governing structure while learning missing, uncertain, or unresolved terms from data. This is attractive for climate science, where models must remain interpretable, physically constrained, and robust beyond short training windows. However, chaotic systems present practical difficulties: small trajectory errors grow rapidly, long-horizon losses become misleading, and learned residuals can compensate for physics in non-physical ways. This work investigates the modelling and training choices required to make gray-box UDEs behave reliably on climate-relevant chaotic systems.
Methodology: We developed a sequence of Julia/SciML experiments, beginning with forced mass-spring-damper and Duffing oscillator systems before moving to Lorenz-63 and the nine-variable Peña-Kalnay coupled atmosphere-ocean model. We compared black-box NODEs, partially constrained NODEs, and UDEs with nominal physics plus learned residual terms. Attention was given to practical training such as derivative pretraining, short-window multiple shooting, Adam-to-BFGS optimisation, residual regularisation, and diagnostics based on attractor statistics rather than only pointwise trajectory error.
Audience: This work is aimed at Earth system scientists, climate model developers, and scientific machine learning practitioners. The focus is not only on whether UDEs can fit trajectories, but on what is needed for them to become trustworthy components in climate-relevant workflows.
Keywords: Neural Networks, Neural ODEs, Universal Differential Equation, Chaotic Dynamical Systems, Multiple Shooting, Gray-Box Modelling, Scientific Machine Learning.