ADDA: a Modular Framework for Representing, Simulating and Assimilating Dynamics with End-to-end Differentiability
2026-08-24 • Mathematical Software
Mathematical SoftwareMachine Learning
AI summaryⓘ
The authors developed ADDA, a software framework to help combine computer simulations and real-world observations for better predictions in geosciences. ADDA makes it easier to use different data assimilation methods by supporting many types of data layouts and being compatible with automatic differentiation and parallel processing. It works with popular machine learning tools like PyTorch and JAX, making it flexible and powerful. The authors also provide example simulations and data assimilation setups to show how ADDA can be used.
Data assimilationAutomatic differentiationGeosciencesVariational methodsEnsemble methodsParallel computingPyTorchJAXDynamical systemsObservation operators
Authors
Anthony Frion, Vien Minh Nguyen-Thanh, Ali Can Bekar, Pauleo R. Nimtz, Vadim Zinchenko, David S. Greenberg
Abstract
Data assimilation (DA) is an essential tool for prediction and understanding in the geosciences. DA combines simulation programs representing scientific knowledge with observations that constrain system dynamics, resulting in analyses and forecasts that incorporate both knowledge and data. DA tasks can be addressed with a diverse toolset, including variational, ensemble and learning-based methods. In particular, many recent works have proposed using automatic differentiation tools for variational, learning-based or hybrid methods. However, comprehensive comparisons across algorithms and dynamical systems remain challenging, due to the incompatibility of simulation and assimilation codes, inflexible handling of spatial and temporal discretizations, specialization of DA methods to specific simulations, and limited support for automatic differentiation and parallel computation in simulations. To address this challenge, we introduce Automatic Differentiation for Data Assimilation (ADDA), a software framework for defining and working with system states, simulations, observation schemes and DA methods. ADDA provides a powerful and flexible set of base classes for representing dynamical systems and observation operators, with support for collocated and staggered grids, unstructured meshes, Lagrangian state variables and irregular or continuous-time observations. Parallel processing and differentiability are first-class features, with support for batch axes and automatic differentiation throughout. ADDA is implemented in PyTorch library, but supports DA for JAX-based computation of dynamics and their gradients. To demonstrate its features, we further provide differentiable, ADDA-compatible implementations of 10 dynamical systems of various dimensionalities and scales, from which we design multiple illustrative DA examples. All of our code is publicly available at https://github.com/m-dml/ADDA.