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

Melissa Adrian; Daniel Sanz-Alonso; Rebecca Willett

stat.MLMar 21, 2026

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

Authors: Melissa Adrian, Daniel Sanz-Alonso, Rebecca Willett

Organizations: 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 · NSF-Simons National Institute for Theory and Mathematics in Biology, Chicago, IL 60611

Abstract

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.

Figures & tables

Alg.	$\mathcal{M}_{\theta}(x_{t-1}^{n})$	$\mathcal{K}_{\phi}(x_{t-1}^{1:N})$	$\mathcal{Y}(\cdot;R_{t})$	Ens. Size
3DVar- $C$	$\mathcal{F}_{\theta_{1}}(x_{t-1})$	$C_{\phi}H_{t}^{\top}(H_{t}C_{\phi}H_{t}^{\top}+R_{t})^{-1}$	$y_{t}$	$1$
		$C_{\phi}\in\mathbb{R}^{d_{x}\times d_{x}}$
EnKF	$\mathcal{F}_{\theta_{1}}(x_{t-1}^{n})+\tilde{Q}_{\theta_{2}}\xi_{t}^{n}$	$\hat{C}_{t}^{\theta,\phi}H_{t}^{\top}(H_{t}\hat{C}_{t}^{\theta,\phi}H_{t}^{\top}+R_{t})^{-1}$	$y_{t}+\gamma_{t}^{n}$	$N_{\text{EnKF}}$
	$\xi_{t}^{n}\sim\mathcal{N}(0,I_{d_{x}})$	( $\hat{C}_{t}^{\theta,\phi}$ defined in ( 8 ))	$\gamma_{t}^{n}\sim\mathcal{N}(0,R_{t})$
Ens3DVar	$\mathcal{F}_{\theta_{1}}(x_{t-1}^{n})+\tilde{Q}_{\theta_{2}}\xi_{t}^{n}$	$\tilde{C}_{t}^{\theta,\phi}H_{t}^{\top}(H_{t}\tilde{C}_{t}^{\theta,\phi}H_{t}^{\top}+R_{t})^{-1}$	$y_{t}+\gamma_{t}^{n}$	$N_{\text{Ens3DVar}}$
	$\xi_{t}^{n}\sim\mathcal{N}(0,I_{d_{x}})$	( $\tilde{C}_{t}^{\theta,\phi}$ defined in ( 9 ))	$\gamma_{t}^{n}\sim\mathcal{N}(0,R_{t})$

Table 1: Table of examples of named special cases of the general filtering algorithm defined in ( 4 ). We note that this list is not exhaustive.

	Forecast cost	Analysis cost
3DVar- $C$	$c_{\mathcal{F}}$	$\mathcal{O}(d_{y}^{3}+d_{x}^{2}d_{y}+d_{y}^{2}d_{x})$
EnKF	$Nc_{\mathcal{F}}$	$\mathcal{O}(d_{y}^{3}+d_{y}^{2}N+d_{y}N^{2}+d_{x}N^{2})$
Ens3DVar	$Nc_{\mathcal{F}}$	$\mathcal{O}(d_{y}^{3}+d_{x}^{2}d_{y}+d_{y}^{2}d_{x})$

Table 2: Table of computational costs for the forecast and analysis in one filtering step for straight-forward implementations of the algorithms 3DVar-

C

, EnKF, and Ens3DVar, where we assume that

N\ll d_{y}

for Ens3DVar.

Panel of Figure 1 — Figure 1: Clohessy-Wiltshire parameter estimation (Section 4.1 ). Parameter estimation error of the learned forecast parameter $\hat{\theta}_{1}$ compared to the ground truth parameter $\theta_{1}^{*}=0.0013$ across choices of $\sigma^{2}_{0}$ and variance on the diagonals of $R$ . The shaded regions correspond to the 0.25 and 0.75 quantiles of errors across 10 independent simulations. We additionally plot the median error for the initialized parameters $\{\theta_{1,j}^{(0)}\}_{j=1}^{J}$ for $j=1,\dots,J=10$ in the dashed black line (c) representing the errors in the absence of any learning, and the gray region corresponds to the 0.25 and 0.75 quantiles of test forecast RMSEs across these 10 initialized models.

Figure 2: Clohessy-Wiltshire filtering (Section 4.1 ). Visualization of observations, true trajectories, and estimated positions for $x_{1}$ and $x_{3}$ , and velocities $\frac{dx_{1}}{dt}$ and $\frac{dx_{3}}{dt}$ across each of the four filtering methods using the learned forecast parameters $\hat{\theta}$ and learned filtering parameters $\hat{\phi}$ , evaluated on test data. In this example, the observation noise is fixed at $R=0.1I_{d_{y}\times d_{y}}$ , and the variance of the initialization of $\theta_{1}^{(0)}$ is fixed at $\sigma_{0}^{2}=0.03$ .

Figure 3: Clohessy-Wiltshire data log-likelihood (Section 4.1 ). Plots of the forecast log-likelihood evaluated at observations from ( 10 ) for assimilating test observations into filters parameterized by the learned forecast parameters $\hat{\theta}$ and $\hat{\phi}$ learned from AD-3DVar- $C$ , AD-EnKF ( $N=25$ ), and AD-Ens3DVar ( $N=25$ ). The algorithm used for filtering corresponds to the algorithm used in the learning task (i.e., EnKF is used as the filter with the learned parameters from AD-EnKF). The log-likelihood from the optimal solution to the filtering task, the Kalman filter (detailed in Algorithm 3 ) using the true forecast parameters $\hat{\theta}$ , is shown in each plot as a comparison.

Panel of Figure 4 — Figure 4: Lorenz-96 learned forecast performance (Section 4.2 ). Test performance of the learned Lorenz-96 model in ( 22 ) learned from various methods: AD-3DVar- $C$ , AD-3DVar- $K$ , AD-EnKF with a covariance tapering radius of 5 and $N=25$ , and AD-Ens3DVar with a covariance tapering radius of 5 and $N=25$ . The shaded regions correspond to the 0.2 and 0.8 quantiles of test forecast performances across 10 independent simulations. In the plot varying $d_{y}/d_{x}$ , we additionally plot the median forecast RMSE for the initialized models $\mathcal{F}_{\beta_{j}}$ for $j=1,\dots,J=10$ in the dashed black line (c), and the gray region corresponds to the 0.2 and 0.8 quantiles of test forecast RMSEs across these 10 initialized models.

Figure 5: Lorenz-96 filtering (Section 4.2 ). Visualization of filtering on test observation data using the learned $\hat{\theta}$ and $\hat{\phi}$ values across the methods AD-3DVar- $C$ , AD-3DVar- $K$ , AD-EnKF with $N=25$ ensemble members and a tapering radius of 5, and AD-Ens3DVar with $N=25$ ensemble members and a tapering radius of 5 compared to the noisy observations and ground truth Lorenz-96 simulations. In this plot, the settings of the filtering problem are $R=I_{d_{y}\times d_{x}}$ , $d_{y}/d_{x}=0.6$ , $\sigma^{2}_{0}=1$ , and $d_{x}=200$ . For ease of visualization, only the first 50 out of 200 state components are shown in the plot. Filtering is done with the base-filtering algorithm used during training (i.e., the results presented for AD-EnKF use the learned $\hat{\theta}$ and $\hat{\phi}$ in an EnKF algorithm).

Figure 6: Example interaction matrix $A$ for 16 interacting species in the generalized Lotka-Volterra equations. The values $\pm a_{1:10}$ denote the shared parameter values within each block.

Figure 7: Generalized Lotka-Volterra parameter estimation (Section 4.3 ). Parameter estimation mean absolute errors (MAEs) of the learned generalized Lotka-Volterra model in ( 26 ) learned from the methods AD-3DVar- $C$ , AD-EnKF with $N=d_{x}$ , and AD-Ens3DVar with $N=d_{x}$ . The shaded regions correspond to the 0.2 and 0.8 quantiles of test forecast performances across 10 independent simulations. In each plot, we plot the median MAEs for the initialized models $\mathcal{F}_{\theta_{1}^{j}}$ for $j=1,\dots,J=10$ in the dashed black line, and the gray region corresponds to the 0.2 and 0.8 quantiles of test forecast RMSEs across these 10 initialized models. The top row of plots correspond to the error in the estimated interaction matrix $\hat{A}$ component of $\hat{\theta}_{1}$ , and the second row of plots correspond to the error in the estimated population change vector $\hat{r}$ component of $\hat{\theta}_{1}$ .

Figure 8: Generalized Lotka-Volterra filtering (Section 4.3 ). Visualization of filtering on test observational data using the learned $\hat{\theta}$ and $\hat{\phi}$ values across the methods AD-3DVar- $C$ , AD-EnKF with $N=500$ ensemble members, and AD-Ens3DVar with $N=500$ ensemble members where $d_{x}=500,$ $R_{t}=0.05I_{d_{y_{t}}\times d_{x}}$ , $\sigma_{0}^{2}=0.1,$ and $d_{y_{t}}/d_{x}=0.5$ . For ease of visualization, only 10 randomly selected species abundance trajectories out of the 500 total species are visualized. In each plot in the first row, the true trajectories are given by the dashed black lines, the sparse, noisy observation are plotted as a scatterplot, and the analyses are plotted as solid colored lines. The second row plots the errors in the analyses compared to the ground truth trajectories.

Algorithm	$H_{t}\equiv H$	varied $H_{t}$	dimensional $d_{x}$	Memory	Learning $\mathcal{F}_{\theta_{1}}$
		Time-	High-
AD-EnKF (Modified)	Reasonably	✓	May require	High; scales	Accurate
( Chen et al.,, 2022 )	recovers		non-trivial	with $N$
	all states		modifications
AD-Ens3DVar	Reasonably	✓	May require	High; scales	Accurate
(Ours)	recovers		modifications	with $N$

Table 3: Summary table of particular examples of auto-filtering algorithms in terms of their ability to learn accurate filtering parameters when

H_{t}\equiv H

when

d_{y}<d_{x}

, ability to filter with a time-varying

H_{t}

, relative scalability to high-dimensional systems, relative memory footprint, and relative accuracy in improving the forecast model

\mathcal{F}_{\theta_{1}}

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Experiment based on the Lorenz-96 system described in Section 4.2 . (Left) Median test forecast RMSE for learning a low rank update to the background covariance such that $C_{\phi}=C_{0}+B_{\phi}B_{\phi}^{\top}$ , where $B_{\phi}\in\mathbb{R}^{d_{x}\times p}$ and $p<d_{x}$ . In this experiment, $d_{x}=200,$ $d_{y}/d_{x}\in[0.3,0.5,0.8]$ , $\sigma_{0}^{2}=1$ , and observations at each time $t$ are randomly sampled across state dimensions (i.e., $H_{t}$ varies with time). Each training task was replicated for $J=10$ independent simulations. We additionally plot the median forecast RMSE for the initialized models $\mathcal{F}_{\beta^{(j)}}$ for $j=1,\dots,M$ in the dashed black line, and the gray region corresponds to the 0.2 and 0.8 quantiles of test forecast RMSEs across these 10 initialized models. (Right) Plot of millions of floating point operations per second (MFLOPS) across values of $p$ for processing one subsequence (both the forward and backward pass) for the setting where $d_{x}=200$ , $p\in[50,75,100,150,175,200],L=20,$ and $d_{y}/d_{x}\in[0.3,0.5,0.8]$ .

Figure 10: Heatmap of the EnKF filter RMSE values when filtering with the true $\mathcal{F}_{\theta_{1}^{*}}$ on unseen test data across different choices of tapering radius and covariance inflation for $d_{x}=40$ , $d_{y}/d_{x}=1$ , and $\sigma_{0}^{2}=1$ . The red trajectory corresponds to the convergence of the analysis parameter vector $\phi\in\mathbb{R}^{2}$ for the covariance tapering and inflation for AD-EnKF (Modified), which was initialized at tapering = 5 and inflation = 1.2. The red star corresponds to location of the converged value for the (tapering, inflation) pair.

Panel of Figure 11 — Figure 11: Varying the ensemble size for AD-EnKF (Modified) and AD-Ens3DVar for the Clohessy-Wiltshire (Section 4.1 ), Lorenz-96 (Section 4.2 ), and generalized Lotka-Volterra (Section 4.3 ) examples. The shaded regions correspond to the 0.2 and 0.8 quantiles of forecast performance or forecast parameter estimation error across 10 independent simulations.

Figure 12: Visualization of the test filter RMSEs across examples and methods. The lines in each plot correspond to the test filter RMSE across parameter choices for either the covariance mixing parameter $\phi_{1}$ in Ens3DVar or the covariance inflation factor $\phi$ in EnKF. The top row corresponds to filtering results using EnKF, the bottom row corresponds to filtering results using Ens3DVar, and each column corresponds to results from one of our experimental examples. The star in each plot corresponds to the value chosen by the learning algorithm. The EnKF evaluations use the learned forecast parameters $\hat{\theta}$ in the filter, and the Ens3DVar evaluations use both the learned forecast parameters $\hat{\theta}$ and the learned 3DVar background covariance parameters $\phi_{2}$ in the filter.

Figure 13: Clohessy-Wiltshire filtering (Section 4.1 ). Visualization of observations, true trajectories, and estimated positions for $y$ and $z$ , and velocities $\frac{dy}{dt}$ and $\frac{dz}{dt}$ across each of the four filtering methods using the learned forecast parameters $\hat{\theta}$ and learned filtering parameters $\hat{\phi}$ , evaluated on test data. In this example, the observation noise is fixed at $R=0.1I_{d_{y}\times d_{y}}$ , and the variance of the initialization of $\theta_{1}^{(0)}$ is fixed at $\sigma_{0}^{2}=0.03$ .

Figure 14: Visualization of each component of the forecast log-likelihood evaluated at observations, given in ( 14 ) for filtering with the learned $\hat{\theta}$ and $\hat{\phi}$ from AD-3DVar- $C$ , AD-EnKF, and AD-Ens3DVar compared to the optimal Kalman filter. (Left) Component of the log-likelihood that computes $-\frac{1}{2}\log\det S_{t}$ for each time $t$ . (Right three) Residual component of the loss that computes $-\frac{1}{2}\|y_{t}-H_{t}\hat{m}_{t}^{\theta}\|_{S_{t}}^{2}$ for each time $t$ .

Data Assimilation

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

See all 58 papers

ADDA: a Modular Framework for Representing, Simulating and Assimilating Dynamics with End-to-end Differentiability
Aug 24, 2026Anthony Frion, Vien Minh Nguyen-Thanh, Ali Can Bekar +3Data Assimilation Differentiable Programming
Model-Driven Machine Learning Institute of Coastal Systems - Analysis and Modeling Helmholtz-Zentrum Hereon Germany
ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing
May 14, 2026Yixuan Jia, Siyi Chen, Yida Pan +9Data Assimilation
University of Michigan, Ann Arbor · Michigan State University · Los Alamos National Laboratory
Amortized Filtering and Smoothing with Conditional Normalizing Flows
Apr 8, 2026Tiangang Cui, Xiaodong Feng, Chenlong Pei +2Bayesian Filtering Data Assimilation
School of Mathematics and Statistics, The University of Sydney, Camperdown, NSW 2006, Australia · Institute of Advanced Study, Beijing Normal-Hong Kong Baptist University, Zhuhai 519087, China · SKLMS & Institute of Computational Mathematics and Scientific/Engineering Computing, Academy of Mathematics and Systems Science, Chinese Academy of Sciences +1

Aug 24, 2026J/K move · Enter open · S save

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+2

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

Data Assimilation Differentiable Programming

Open paper AskPDF arXiv

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.

Abstract

Data Assimilation

Explore similar work

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

ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing

Amortized Filtering and Smoothing with Conditional Normalizing Flows