Long-range geophysical forecasts are fundamentally limited by chaotic dynamics and numerical errors. While data assimilation can mitigate these issues, classical variational smoothers require computationally expensive tangent-linear and adjoint models. Conversely, recent efficient latent filtering methods often enforce weak trajectory-level constraints and assume fixed observation grids. To bridge this gap, we propose Latent Ensemble Variational Data Assimilation (LEVDA), an ensemble-space variational smoother that operates in the low-dimensional latent space of a pretrained differentiable neural dynamics surrogate. By performing four-dimensional ensemble-variational (4DEnVar) optimization within an ensemble subspace, LEVDA jointly assimilates states and unknown parameters without the need for adjoint code or auxiliary observation-to-latent encoders. Leveraging the fully differentiable, continuous-in-time-and-space nature of the surrogate, LEVDA naturally accommodates highly irregular sampling at arbitrary spatiotemporal locations. Across three challenging geophysical benchmarks, LEVDA matches or outperforms state-of-the-art latent filtering baselines under severe observational sparsity while providing more reliable uncertainty quantification. Simultaneously, it achieves substantially improved assimilation accuracy and computational efficiency compared to full-state 4DEnVar.
Figures & tables
Setting
Grid
Simulation steps
Obs. times (interval)
Δt
Obs. grid
Kolmogorov flow
150×150
200
40 (5)
0.04
10×10
Tsunami modeling
150×150
2000
50 (40)
≈21 s
10×10
Atmospheric modeling
512×256
30240
63 (480)
≈60 s
16×8
Table 1: Summary of experimental setups (adapted from Xiao et al. (2026) ). Obs. times represent the total number of observation times while interval in parentheses denotes the spacing (in simulation steps) between consecutive observations. Obs. grid indicates the spatial subsampling resolution.
Figure 1: Relative RMSE of the ensemble-mean trajectory over time for Kolmogorov flow (left), tsunami modeling (middle), and atmospheric modeling (right); shaded regions show the standard deviation of member-wise relative RMSE.
Figure 2: Atmospheric benchmark: final-time zonal wind u at 500 hPa. The visualization compares the ground truth (with sparse observations) to the assimilated states and errors for different methods.
Figure 3: Effect of irregular sampling on tsunami modeling: (a) moving observation locations (randomized at each observed time), (b) irregular observation times with fixed spatial locations, and (c) joint spatiotemporal irregularity. Curves depict the relative RMSE of the ensemble-mean trajectory; shaded regions indicate the standard deviation of member-wise relative RMSE.
Example
Kolmogorov Flow
Tsunami Modeling
Atmospheric Modeling
Metric
4DEnVar
Latent-EnSF
LD-EnSF
LEVDA
4DEnVar
Latent-EnSF
LD-EnSF
LEVDA
4DEnVar
Latent-EnSF
LD-EnSF
LEVDA
Encoder Training (s)
0
43,156
298
0
0
27,718
312
0
0
58,231
1429
0
Assimilation (s)
−
7.57
0.35
13.91
24,639
5.77
0.37
11.32
−
37.5
1.42
19.47
Assimilation dim.
Dx=45,000
Ds=499
Ds=10
Ds=10
Dx=67,500
Ds=400
Ds=12
Ds=12
Dx=393,216
Ds=512
Ds=52
Ds=52
Table 2: Wall-clock runtime comparison. We report the latent assimilation dimension Ds for latent methods (Latent-EnSF, LD-EnSF, LEVDA) and the full-state dimension Dx for 4DEnVar. Assimilation runtimes correspond to full-trajectory assimilation with an ensemble size of K=100 for latent methods, whereas 4DEnVar is restricted to K=40 due to the high computational cost of full-state propagation. 4DEnVar simulations use an AMD 7543 CPU (64 threads), while data-driven and latent methods utilize an Nvidia A6000 GPU.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Tsunami benchmark sensitivity analysis: LEVDA Relative RMSE as a function of (a) smoothing time horizon τ , (b) ensemble size K , and (c) spatial observation stride (sparsity).
Figure 5: Additional Tsunami benchmark ablations: (a) Impact of temporal sparsity on LEVDA Relative RMSE (using extended smoothing window τ=10 ). (b) Robustness to observation noise (5%, 10%, 20% noise-to-signal ratios) under fixed likelihood weighting. (c) Parameter estimation accuracy: Relative RMSE of the parameter u over the assimilation window.
Figure 6: Spread-Error Ratio (SER) over time for Kolmogorov flow, tsunami modeling, and atmospheric modeling (left-to-right). Values closer to 1 indicate better calibration of the ensemble spread.
Figure 7: Continuous Ranked Probability Score (CRPS) over time for Kolmogorov flow, tsunami modeling, and atmospheric modeling (left-to-right). Lower values represent higher probabilistic accuracy.
Capability
Latent-EnSF
LD-EnSF
SDA
FlowDAS
APPA
LEVDA
Smoothing
X
X
✓
✓
✓
✓
Differentiable Dynamics
X
✓
X
X
X
✓
Moving observations
X
X
✓
✓
✓
✓
Irregular-time observations
X
X
X
X
X
✓
Parameter estimation
X
✓
X
X
X
✓
Appendix
Table 3: Feature comparison across latent and score-based DA methods. Here, “Differentiable Dynamics” refers to the use of an explicit learned surrogate for state evolution (as opposed to implicit trajectory priors), while “Irregular-time Observations” denotes native support for data at arbitrary time stamps without resampling to a fixed temporal grid.
Figure 8: Visualization of surface elevation η in tsunami dynamics. The rows display: (1) ground truth; (2) LDNet predictions from a known initial condition; (3) prediction errors; and (4) sparse observations ( 10×10 subsampled from a 150×150 grid). The remaining rows show assimilation results starting from a deviated initial condition: (5) a single assimilated trajectory from the LD-EnSF ensemble ( N=20 ); (6) LD-EnSF error; (7) a single assimilated trajectory from the LEVDA ensemble ( N=20 ); and (8) LEVDA error.
Figure 9: Kolmogorov flow benchmark: vorticity field, comparing ground truth and assimilated reconstruction. The rows display: (1) ground truth; (2) perturbed trajectory with different Reynolds number; (3) error of perturbed trajectory; (4) sparse observations; ( 10×10 subsampled from a 150×150 grid (5) assimilated results of a single trajectory from the LD-EnSF ensemble ( N=20 ); (6) LD-EnSF error; (7) a single assimilated trajectory from the LEVDA ensemble ( N=20) ; (8) LEVDA error.
Data assimilation combines model forecasts with noisy, incomplete observations to estimate the evolving state of a dynamical system. Existing methods face two compounding challenges: high-dimensional nonlinear dynamics make repeated forward simulation computationally expensive, while sparse observations provide limited direct information about the full state. To address these challenges, we propose the Latent-Dynamics Ensemble Flow Filter (LD-EnFF), a sequential Bayesian filtering framework that performs both forecast propagation and filtering updates in a compact latent space. LD-EnFF combines a latent dynamics surrogate for ensemble propagation with a variational autoencoder (VAE)-based observation model that evaluates a state-dependent observation likelihood in latent space. At each assimilation step, an ensemble filtering update based on flow matching uses the forecast ensemble and this likelihood to generate posterior samples, jointly updating latent states and uncertain parameters. This design avoids repeated full-state simulation during forecasting and full-field reconstruction during likelihood evaluation. LD-EnFF substantially outperforms a broad range of data assimilation algorithms on benchmarks spanning Kolmogorov flow, tsunami propagation, and atmospheric modeling, all featuring complex dynamics and sparse, noisy observations.
Ziyu Tian, Kaichen Shen, Wenbo Hao +3
School of Mathematics, Georgia Institute of Technology · School of Computational Science and Engineering, Georgia Institute of Technology
Data assimilation (DA) uses Bayesian inference to update the state of a numerical forecast model with observed data. In this study, we propose a fundamentally different, unified approach to atmospheric data assimilation. We use latent video flow-matching to sample temporally consistent trajectories from a prior trained using ERA5 reanalysis (69 variables over an 8-day window). We also use posterior sampling to assimilate real observation sources, such as those from the NOAA Integrated Global Radiosonde Archive and the Integrated Surface Database. Because the prior generates a continuous trajectory, it naturally propagates information between observed and unobserved frames. Therefore, we can perform various DA tasks, such as filtering and smoothing, simply by changing the observed frames. Moreover, we generate full-state ensemble forecasts directly from sparse observations, achieving performance competitive with state-of-the-art observation-to-forecast models.
Dibyajyoti Chakraborty, Romit Maulik
Information Sciences and Technology The Pennsylvania State University University Park, Pennsylvania · School of Mechanical Engineering Purdue University West Lafayette, Indiana
The ensemble Kalman filter (EnKF) is widely used for data assimilation in high-dimensional systems, but its performance often deteriorates for strongly nonlinear dynamics due to the structural mismatch between the Kalman update and the underlying system behavior. In this work, we propose a latent autoencoder ensemble Kalman filter (LAE-EnKF) that addresses this limitation by reformulating the assimilation problem in a learned latent space with linear and stable dynamics. The proposed method learns a nonlinear encoder--decoder together with a stable linear latent evolution operator and a consistent latent observation mapping, yielding a closed linear state-space model in the latent coordinates. This construction restores compatibility with the Kalman filtering framework and allows both forecast and analysis steps to be carried out entirely in the latent space. Compared with existing autoencoder-based and latent assimilation approaches that rely on unconstrained nonlinear latent dynamics, the proposed formulation emphasizes structural consistency, stability, and interpretability. We provide a theoretical analysis of learning linear dynamics on low-dimensional manifolds and establish generalization error bounds for the proposed latent model. Numerical experiments on representative nonlinear and chaotic systems demonstrate that the LAE-EnKF yields more accurate and stable assimilation than the standard EnKF and related latent-space methods, while maintaining comparable computational cost and data-driven.
Xin T. Tong, Yanyan Wang, Liang Yan
Department of Mathematics, National University of Singapore, 119076, Singapore · School of Mathematics, Southeast University, Nanjing 210096, China