PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.23297v1cs.MScs.LG

TL;DR

Data assimilation research is hindered by incompatibilities between simulations, observations, and assimilation methods, especially across varied discretizations and differentiability requirements. ADDA provides a modular, automatically differentiable framework for these components, demonstrates its use across dynamical systems and DA methods, and supports PyTorch with JAX-based differentiable dynamics. The paper presents ADDA as a foundation for flexible applications and future benchmarking, rather than as a benchmark itself.

  • Problem

    Integrating and comparing DA methods across diverse simulations, state representations, observation types, discretizations, and computational capabilities remains difficult.

  • Method

    ADDA modularizes system states, simulations, observation operators, and DA methods with end-to-end automatic differentiation and interoperability across PyTorch and JAX.

  • Results

    ADDA demonstrates flexible DA across varied dynamical systems and examples, including Lorenz-63, tracer transport, latent assimilation, and a JAX-based Kuramoto-Sivashinsky setup.

  • Takeaways & Limitations

    ADDA provides a reusable foundation for applying variational and other DA methods to custom dynamics, observation operators, probability distributions, and differentiable simulations.

  • Takeaways & Limitations

    ADDA is not a benchmark, and fair comprehensive benchmarking would require community-wide method implementations, tuning, task selection, and metrics.

Abstract

from arXiv · show

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.

1 Introduction

Data assimilation combines simulations and observations to estimate system states, tune parameters, and identify conflicts between models and data. ADDA addresses the integration challenges by providing a modular, differentiable framework designed to support future comparisons across methods and tasks.

  • Data assimilation combines simulations and observations to obtain system-state sequences consistent with both.
  • Researchers face substantial difficulty integrating DA methods with diverse simulations, state representations, observation types, parallel computation, and gradient computation.
  • ADDA provides modular components for representing and simulating dynamical systems, implementing observation operators, and building, applying, and evaluating DA methods.
  • Its requirements include end-to-end autodiff, diverse state structures and observations, flexible dynamics, modularity, batched CPU/GPU computation, and modifiable examples.
  • ADDA is implemented in Python and PyTorch with an interface through which simulations and DA algorithms interoperate.
  • ADDA is intended as a foundation for future benchmarking, not as a benchmark or benchmarking study.

2 Data Assimilation: Methods and Formulations

Data assimilation estimates state trajectories from observations and known dynamical and probabilistic components, but nonlinear systems make the full posterior difficult to compute. The section introduces variational, sequential, and broader formulations that motivate different approximations and algorithmic choices.

  • 2.1 DA task formulation: Gaussian process and observation noise: The standard DA setup models states evolving under dynamics with process noise and observations generated through observation operators with observation noise.
  • 2.1 DA task formulation: Gaussian process and observation noise: The DA objective is to estimate the state trajectory x_0:T from observations while assuming knowledge of dynamics, noise, observation dependence, and the initial-state prior.
  • 2.2 Bayesian derivation of the posterior distribution of the state: The Bayesian posterior can be represented through factorized prior and likelihood terms under Markovian dynamics and temporally independent observations.
  • 2.2 Bayesian derivation of the posterior distribution of the state: For nonlinear dynamics, the posterior is generally intractable and non-Gaussian, motivating DA methods that exploit temporal structure through approximations.
  • 2.3 The 4D-Var cost and its minimization: 4D-Var estimates the posterior mode by minimizing a variational cost, with weak-constraint formulations penalizing initial-state, model-evolution, and observation deviations.
  • 2.3 The 4D-Var cost and its minimization: Strong-constraint 4D-Var enforces exact model evolution, whereas its effectiveness relative to weak constraint depends on model accuracy, noise knowledge, dynamical sensitivity, and observations.
  • 2.3 The 4D-Var cost and its minimization: Automatic differentiation replaces manual adjoint derivation for differentiable dynamics, making several 4D-Var variants more straightforward to use.
  • 2.4 Sequential methods: Sequential methods assimilate observations forward in time, with Kalman methods exact for linear-Gaussian systems and ensemble methods representing distributions with finite ensembles for nonlinear cases.

3 Implemented structures and methods

ADDA implements a shared, extensible basis for states, observation operators, dynamics, and DA algorithms. Its structures support heterogeneous representations, differentiable variational methods, flexible initialization, sliding windows, ensemble smoothers, and independent state and observation time axes.

  • ADDA implements Kalman, ensemble Kalman, and single- or sliding-window 4D-Var methods around extensible state and observation representations.
  • The State class represents state sequences with timestamps and supports named fields having different dimensions and spatial structures, including grids and unstructured meshes.
  • Observation operators evaluate p(y_0:T|x_0:T), support sampling and density operations, and may provide conditional means for ensemble Kalman methods.
  • Dynamics receive a time step, dynamic forcings or boundary conditions, and static configurable inputs, while retaining a common four-input interface.
  • Initialization heuristics include interpolation and user-provided defaults because poor initial states can impair variational optimization and ensemble posterior estimates.
  • Strong-constraint 4D-Var uses a supplied differentiable model, observations, observation operator, optional prior, and time-dependent forcing or boundary data.
  • ADDA parallelizes strong-constraint 4D-Var over time and supports non-Gaussian distributions for its variational implementation.
  • Sliding-window variational methods address long sequences and unstable or memory-intensive optimization by solving successive sub-window problems.

4 Illustrative experiments

ADDA’s illustrative experiments span diverse differentiable dynamical systems and demonstrate variational assimilation on Lorenz-63 under noisy, incomplete, and chaotic-system settings. The examples show accurate state reconstruction and forecasting, while motivating weak-constraint formulations for longer windows.

  • Illustrative experiments: ADDA demonstrates capabilities across diverse dynamical systems, observation operators, and data-assimilation techniques, including PyTorch and JAX-based implementations.The examples include linear dynamics, Lorenz systems, Kuramoto–Sivashinsky, Korteweg–De Vries, Kolmogorov flow, advection–diffusion, and quasi-geostrophic dynamics.
  • Lorenz-63 setup: Lorenz-63 provides a three-variable chaotic testbed for strong- and weak-constraint 4D-Var, with standard parameters and a Lyapunov time of roughly 1.10.The ground-truth trajectories are generated with fourth-order Runge–Kutta integration using a time step of 0.01.
  • Complete observations: Strong-constraint 4D-Var convincingly denoises complete Gaussian-noise observations and accurately extends predictions beyond the observation range.The experiment uses a 2-time-unit window, 200 time steps, five L-BFGS steps, and no background prior term.
  • Incomplete observations: Strong-constraint 4D-Var accurately reconstructs the full Lorenz-63 state, including unobserved z, from noisy observations of only x and y, and forecasts beyond the observation window.The observation operator truncates the third variable and is differentiated automatically.
  • Weak-constraint assimilation: Chaotic dynamics make long single-window strong-constraint assimilation ill-behaved, motivating weak-constraint 4D-Var, which optimizes the full state trajectory with modeled error.The weak-constraint experiment uses a diagonal Gaussian model-error distribution and 50 L-BFGS iterations.

4.2 Data assimilation on the Lorenz-96 system

The Lorenz-96 experiments compare variational and ensemble methods under sparse, noisy, and irregular observations. Weak-constraint 4D-Var remains stable over long windows, while strong-constraint 4D-Var eventually diverges and smoothing improves on filtering.

  • Setup: The Lorenz-96 experiments use a 40-variable periodic system with sparse masked observations and Gaussian noise.At each time step, 30 of 40 variables are masked and noise with variance 1 is added to the remainder.
  • Regular sampling: Strong-constraint 4D-Var initially reduces error below observation-noise amplitude but becomes unskilled after roughly 4 time units beyond its observation window.Its error grows exponentially after the observation window because of chaotic dynamics.
  • Regular sampling: Weak-constraint 4D-Var assimilates observations across several Lyapunov times with relatively low, temporally stable errors.Model errors allow the optimized trajectory to depart from the known dynamics; errors are largest near assimilation-window boundaries.
  • Ensemble methods: Ensemble smoothing produces significantly lower RMSE throughout the assimilation window than filtering under the same Lorenz-96 observation scenario.Both methods have their largest errors during the first few time steps because of imperfect initialization.
  • Irregular sampling: Irregular observation times do not significantly affect the tested methods: weak-constraint 4D-Var and EnKF reconstruct the trajectory, while strong-constraint 4D-Var eventually diverges.Weak-constraint 4D-Var again obtains the lowest assimilation errors, with edge errors for weak constraint and initial errors for EnKF.

4.3 Weak-constraint 4D-Var on the two-timescale Lorenz-96 system

The two-timescale Lorenz-96 experiments test weak-constraint 4D-Var on heterogeneous slow and fast state variables. Accounting for fast-variable influence improves slow-variable reconstruction even when those variables are unobserved, although missing fast observations still reduce accuracy.

  • System and representation: The two-timescale Lorenz-96 system combines n slow variables with nJ fast variables of different spatial structures.ADDA represents the one-dimensional slow field and two-dimensional fast field in a single State object.
  • Assimilation scenarios: Weak-constraint 4D-Var is tested with observations of both variable sets, with slow-only observations while fast variables are either ignored or jointly inferred.The experiments vary whether fast variables are observed, omitted from the model, or included as latent state components.
  • Ignoring fast dynamics: Ignoring fast-variable influence causes a significant loss of slow-variable accuracy compared with jointly assimilating both scales.The misspecified one-scale model neglects the coupling term from the fast variables.
  • Joint inference: Jointly inferring unobserved fast variables while retaining their influence substantially improves slow-variable assimilation accuracy.The reconstruction is qualitatively accurate for both observed slow variables and inferred fast variables.
  • Joint inference: Absence of fast-variable observations remains detrimental to reconstruction quality for both fast and slow variables.The jointly inferred case improves over ignoring fast dynamics but remains worse than the case with observations at both scales.

4.4 A larger-scale example with multi-layer quasi-geostrophic equations

A three-layer quasi-geostrophic example demonstrates ADDA on a larger-scale, differentiable geophysical system using sliding-window strong-constraint 4D-Var. Assimilation errors decrease during the observation period, and forecasts remain below noise amplitude far beyond that period.

  • Experimental setup: The quasi-geostrophic experiment uses three pressure and potential-vorticity layers with synthetic masked, noisy observations over approximately 21 days.Only 0.5% of variables are observed, with observation-noise standard deviation 0.05.
  • Assimilation setup: Sliding-window strong-constraint 4D-Var processes the 3000-step observation domain in sub-windows of 500 time steps.The computation uses one A100 GPU with 80 GB of memory to avoid splitting the trajectory across multiple GPUs.
  • Assimilation results: Assimilated-state differences are 1-2 orders of magnitude smaller than the states, while RMSE decreases steadily during assimilation.RMSE discontinuities every 500 time steps correspond to transitions between sub-windows.
  • Forecast results: Forecast RMSE takes more than 20000 time steps beyond the observation domain to exceed the observation-noise amplitude.The forecast error increases gradually with forecast horizon.

4.5 Vertical diffusion and advection of an inert tracer

ADDA applies strong- and weak-constraint 4D-Var to inert-tracer transport with time-varying turbulent diffusivity in a tidal, horizontally homogeneous water column. Both methods reconstruct tracer concentrations precisely, although early-time errors remain higher.

  • Setup: The experiment assimilates sparse, noisy tracer observations in a one-dimensional tidal water column using strong- and weak-constraint 4D-Var.The tracer has constant sinking velocity and time-varying turbulent diffusivity computed with GOTM.
  • Setup: Pre-computed turbulent diffusivity serves as a forcing input for the tracer dynamics, reducing computational expense during assimilation.
  • Results: Both strong- and weak-constraint 4D-Var reconstruct the tracer concentration time series with high precision.
  • Results: Both methods have relatively high reconstruction error at the beginning of the time series, but not at the end.Diffusion disperses small-scale initial-state inaccuracies, weakening them over time.
  • Implementation: Strong-constraint 4D-Var uses 5 L-BFGS steps, whereas weak-constraint 4D-Var uses 100 steps with model-error covariance coefficients of 10−6.

4.6 Assimilation with a neural emulator of Kolmogorov flow

ADDA uses a Fourier neural operator emulator as the forward model for strong-constraint 4D-Var on Kolmogorov flow. Assimilation closely follows the groundtruth velocity field despite sparse observations and an initialization whose errors grow over time.

  • Method: A Fourier neural operator is trained to emulate Kolmogorov flow and used as the forward model in strong-constraint 4D-Var.
  • Method: The emulator is trained on data temporally coarsened by approximately 600 and spatially coarsened by 64 along both spatial axes, without forcing inputs.
  • Assimilation setup: Assimilation observes 5 % of state variables for 64 time steps with Gaussian observation noise σ = 1 and uses non-overlapping 8-step sub-windows.
  • Results: Assimilated states closely follow the groundtruth velocity field, while the initialization develops steadily increasing errors and loses correlation with the groundtruth by the observation-domain end.

4.7 Latent data assimilation on the Kuramoto-Sivashinsky system

ADDA incorporates latent data assimilation for the chaotic Kuramoto-Sivashinsky system using a learned latent representation and strong-constraint 4D-Var. The experiment compares latent assimilation with and without a background prior and with non-assimilated initializations.

  • Motivation: Latent data assimilation addresses the computational difficulty of constructing background-error covariance matrices in high-dimensional data spaces.
  • Method: The implementation applies strong-constraint 4D-Var in the latent space of a β-VAE for the one-dimensional Kuramoto-Sivashinsky system.
  • Method: The latent background prior uses an identity covariance matrix and a zero mean, motivated by the VAE’s regularization toward a normal latent distribution.
  • Assimilation setup: Assimilation uses a 100-step window, observes 5 % of state variables with Gaussian noise of standard deviation 1.0, and optimizes with L-BFGS for 20 steps.
  • Evaluation: The comparison includes latent assimilation with and without a background prior, plus encode-decode and naive initializations without assimilation.The decoder supplies an implicit soft prior even without an explicit prior term.

4.8 Strong-constraint 4D-Var on the 2D Kuramoto-Sivashinsky system with a JAX-Torch bridge

ADDA bridges JAX-based differentiable solvers with its PyTorch implementation for strong-constraint 4D-Var on chaotic two-dimensional Kuramoto-Sivashinsky dynamics. Assimilated states maintain lower prediction error and remain close to ground truth.

  • JAX-Torch bridge: ADDA uses a custom bridge between JAX simulations and PyTorch-based assimilation, passing gradients with jax.vjp and arrays with zero-copy DLPack conversion.
  • JAX-Torch bridge: The bridge introduces computational overhead relative to native PyTorch but avoids re-implementing differentiable scientific software.
  • System: The two-dimensional Kuramoto-Sivashinsky system exhibits chaotic dynamics at the selected domain size.
  • Assimilation setup: The experiment observes 25 % of state variables for 16 time steps with Gaussian noise σ = 1 and optimizes strong-constraint 4D-Var for 5 steps.
  • Results: Assimilated trajectories have lower prediction error throughout the prediction window, and predictions from the assimilated initial condition remain close to ground truth.

5 Discussion

ADDA is presented as enabling a shift toward automatic differentiation in data assimilation, while leaving important method and dynamical-system coverage gaps for future work.

  • 5 Discussion: Automatic differentiation can reduce the historical development difficulty of variational methods by avoiding manually derived adjoint models.The discussion also identifies GPU computation as a further advantage of autodiff frameworks such as PyTorch.
  • 5 Discussion: The authors frame the work as contributing to a broader shift in data assimilation practice toward automatic differentiation.This motivation is linked to the need for a flexible software basis and efficient GPU computation.
  • 5 Discussion: ADDA does not yet implement several popular classical and machine-learning-based data assimilation methods.Examples include ensemble Kalman filter variants, weak-constraint 4D-Var, analog data assimilation, 4DVarNet and score-based data assimilation.
  • 5 Discussion: Expanding ADDA to additional methods could support an extensive benchmark of data assimilation techniques.Joint state and parameter estimation is also identified as a natural extension enabled by automatic differentiation.

6 Conclusion

ADDA is introduced as a modular package for data assimilation across varied system and observation settings. Its examples demonstrate versatility, while its modular design supports custom extensions.

  • 6 Conclusion: ADDA handles heterogeneous state fields, forced dynamics and observations sampled irregularly in time.These examples are presented as evidence of the package's versatility across data assimilation contexts.
  • 6 Conclusion: ADDA's modular design supports adaptation to custom dynamics, observation operators and probability distributions.The framework is described as a software package for performing data assimilation in a large variety of contexts.
Loading 2608.23297v1…