cs.MSAug 24, 2026

ADDA: a Modular Framework for Representing, Simulating and Assimilating Dynamics with End-to-end Differentiability

Authors: Anthony Frion, Vien Minh Nguyen-Thanh, Ali Can Bekar, Pauleo R. Nimtz, Vadim Zinchenko, David S. Greenberg

Organizations: Model-Driven Machine Learning Institute of Coastal Systems - Analysis and Modeling Helmholtz-Zentrum Hereon Germany

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: Visualization of the true values, noisy observations and assimilated trajectory for the Lorenz-63 system, with strong-constraint 4D-Var, in a case where all variables are observed with a Gaussian noise until time 2.0.
Figure 2: Visualization of the true values, noisy observations and assimilated trajectory for the Lorenz-63 system, with strong-constraint 4D-Var, in a case where observations are entirely missing for the zz variable.
Figure 3: Visualization of the true values, noisy observations and assimilated trajectory for the Lorenz-63 system, with weak-constraint 4D-Var performed on a long time series with a sparse observation pattern.
Figure 4: Visualization of the true state values, the corresponding sparse and noisy observations, and the state estimations and errors obtained by strong- and weak-constraint 4D-Var on the Lorenz-96 dynamical system. The dashed vertical line marks the limit of assimilated observations for strong-constraint 4D-Var.
Figure 5: Comparison between the errors obtained by the assimilations of an Ensemble Kalman Filter (EnKF) and a similarly parameterized Ensemble Kalman Smoother (EnKS) on the same set of observations from the Lorenz-96 dynamical system.
Figure 6: Visualization of the groundtruth time series for the Lorenz-96 dynamics with irregular time sampling, the corresponding sparse and noisy observations and the assimilated states and associated errors for the strong- and weak-constraint 4D-Var and EnKF. For ease of visualization, we only show the first half of the whole observation domain.
Figure 7: Visualization of the true values, sparse and noisy observations, assimilated time series (with weak-constraint 4D-Var) and assimilation errors for the two-timescale Lorenz-96 system, using observations from both the slow and fast variables.
Figure 8: Visualization of the true values, sparse and noisy observations, assimilated time series (with weak-constraint 4D-Var) and assimilation errors for the slow variables of the two-timescale Lorenz-96 system when the influence of the fast variables is completely ignored.
Figure 9: Visualization of the true values, sparse and noisy observations, assimilated time series (with weak-constraint 4D-Var) and assimilation errors for all variables of the two-timescale Lorenz-96 system when the fast variables are not observed but still inferred and taken into account for the assimilation of the system.
Figure 10: Summary of our data assimilation results on the three-layer quasi-geostrophic equations. In the left part of the figure, we show the true and assimilated pressure fields of a central part of the middle layer, for three time steps contained in the observation domain. The corresponding difference fields are displayed in the top-right panel of the figure. On the bottom right part, we plot the global assimilation root mean squared errors as a function of time, both inside of the observation domain (left) and beyond it (right). The dashed horizontal lines correspond to the amplitude of the noise on sparse observations.
Figure 11: Data assimilation with 4D-Var on an inert tracer in a horizontally homogeneous setting with tides. The tracer concentration cc follows the vertical transport equation ( 18 ). In all plots, the x-axis is time tt in seconds and the y-axis is the vertical coordinate zz in meters. a) Diffusivity ν(z,t)\nu(z,t) as a function of vertical position zz and time tt . b) The tracer concentration c(z,t)c(z,t) simulated by GOTM ( 25 ) acts as a known input to the assimilated system. c) Sparse and noisy pseudo-observations of tracer concentrations with white noise (s.d. 0.20.2 , 10 %{10\,\%} of values observed). d) Initial guess of the tracer concentrations over space and time using the next observed value. e, f) Assimilated trajectories using strong- (e) and weak-constraint (f) 4D-Var. Strong-constraint 4D-Var was performed using L-BFGS with a learning rate of 0.5 for 5 steps. Weak-constraint 4D-Var used a diagonal model error covariance with diagonal coefficients 10−6{10^{-6}} , and leveraged L-BFGS with a learning rate of 1.0, for 100 steps. g, h) Errors of the assimilated trajectories for weak- (g) and strong-constraint (h) 4D-Var w.r.t. the true simulated time series.
Figure 12: Results of sliding-window strong-constraint 4D-Var on the Kolmogorov flow emulator task. Assimilated states display a lower RMSE and stay closer to the groundtruth trajectory throughout the simulation. Visual inspection of the predicted fields at the end of the assimilation window also demonstrates the improvement.
Figure 13: Assimilated results with strong-constraint 4D-Var with latent data assimilation on the Kuramoto-Sivashinsky system. Using a background prior in the latent space results in more accurate assimilation.
Figure 14: Assimilated results with strong-constraint 4D-Var on the 2D KS system. The RMSE of the predictions over time (left) and snapshots of the field variable at t=11.85st=11.85s and t=17.85st=17.85s (right). The predictions from the assimilated initial condition remain lower throughout the prediction window.
More in this topic

Data Assimilation

Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.

See all 58 papers

Explore similar work

CardsList
Mar 21, 2026

Auto-differentiable data assimilation: Co-learning of states, dynamics, and filtering algorithms

Data assimilation algorithms estimate the state of a dynamical system from partial observations, where the successful performance of these algorithms hinges on costly parameter tuning and on employing an accurate model for the dynamics. This paper introduces a framework for jointly learning the state, dynamics, and parameters of filtering algorithms in data assimilation through a process we refer to as auto-differentiable filtering. The framework leverages a theoretically motivated loss function that enables learning from partial, noisy observations via gradient-based optimization using auto-differentiation. We further demonstrate how several well-known data assimilation methods can be learned or tuned within this framework. To underscore the versatility of auto-differentiable filtering, we perform experiments on dynamical systems spanning multiple scientific domains, such as the Clohessy-Wiltshire equations from aerospace engineering, the Lorenz-96 system from atmospheric science, and the generalized Lotka-Volterra equations from systems biology. Finally, we provide guidelines for practitioners to customize our framework according to their observation model, accuracy requirements, and computational budget.

Melissa Adrian, Daniel Sanz-Alonso, Rebecca Willett

School of Mechanical Engineering, Purdue University, West Lafayette, IN 47907 · Department of Statistics, The University of Chicago, Chicago, IL 60637 · Department of Computer Science, The University of Chicago, Chicago, IL 60637 +1

Data AssimilationState Space ModelsSystem Identification
AskPDF
May 14, 2026

ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing

Data assimilation (DA) estimates the state of an evolving dynamical system from noisy, partial observations, and is widely used in scientific simulation as well as weather and climate science. In practice, filtering methods rely on frame-to-frame transition models. However, these models are fragile when observations are non-Markovian (when they form only a partial slice of a higher-dimensional latent state as in real-world weather data): they tend to accumulate errors over long horizons. At the same time, learned DA methods typically commit to a single regime, either filtering (nowcasting, real-time forecasting) or smoothing (retrospective reanalysis), which splits what should be a shared prior across application-specific pipelines. To address both issues, we introduce ForcingDAS, a unified and robust DA framework. Built on Diffusion Forcing with an independent noise level assigned to each frame, ForcingDAS learns a joint-trajectory prior instead of frame-to-frame transitions. This allows it to capture long-horizon temporal dependencies and reduce error accumulation. In addition, the same trained model spans the full filtering to smoothing spectrum at inference time. Specifically, nowcasting, fixed-lag smoothing, and batch reanalysis are selected through the inference schedule alone, without retraining. We evaluate ForcingDAS on 2D Navier-Stokes vorticity, precipitation nowcasting, and global atmospheric state estimation. Across all settings, a single model is competitive with or outperforms both learned and classical baselines that are specialized for individual regimes, with the largest gains observed on real-world weather benchmarks.

Yixuan Jia, Siyi Chen, Yida Pan +9

University of Michigan, Ann Arbor · Michigan State University · Los Alamos National Laboratory

Data Assimilation
AskPDF
Oct 2, 2026

DAWIS: Data Assimilation with Windowed Inverse Sampling via Multitask Interpolants

Flow- and diffusion-based generative models have recently emerged as flexible and highly efficient forecasting models for dynamical systems. When combined with inference-time guidance, they offer a promising route to high-dimensional non-Gaussian data assimilation (DA), the problem of combining forecasts with observations to estimate latent system states. Existing filters, however, condition on a fixed history and assimilate only the most recent observation, leaving them unable to revise past states when new observations arrive. Estimates then stay tethered to a history that later observations may contradict, and errors accumulate over the assimilation run. To this end, we introduce DAWIS, a unified DA method covering filtering, fixed-lag smoothing, and block smoothing within a single framework. DAWIS replaces the single flow time of a state-level prior with a multitask stochastic interpolant over a window of consecutive states, assigning a separate flow time to each. An assimilation cycle inverts the window to a vector of per-state turning points and regenerates it under observation guidance, with the turning points controlling how strongly each state is held fixed, revised, or generated from scratch. The same construction can also absorb the forecast into the assimilation cycle, removing the need for a separate forecasting model. Experiments on challenging nonlinear systems show that DAWIS improves on both filtering and smoothing baselines under sparse, noisy, and nonlinear observations. The code for DAWIS is available at https://github.com/Erik-Wikingsson/DAWIS

Erik Wikingsson, Martin Andrae, Tomas Landelius +1

Linköping University, Sweden · Swedish Meteorological and Hydrological Institute

Data AssimilationFlow-Based Generative ModelingStochastic Interpolants
AskPDF