AI-Driven Data Assimilation of Geophysical Fluids
Key words: stochastic dynamical system, generative artificial intelligence (genAI)
Complex, Chaotic Fluid Motions in the Universe
Whether you are looking at the Earth’s atmospheric circulations, or eddies and vortices in the Earth’s ocean, or even the massive zonal jets of Jupiter's atmosphere, these systems are all governed by a beautiful but highly intricate balance between the Coriolis force due to planetary rotation and pressure gradients of the fluids. This is the domain of geostrophic turbulence—a pervasive, quasi-two-dimensional phenomenon.
Challenges in Climate Modeling
Modeling the Earth's and planetary climate systems accurately and timely is an enormous challenge because of several critical bottlenecks:
Climate systems are inherently nonlinear, chaotic, and stochastic.
Both remote sensing and in-situ measurements suffer from severe observational sparsity and noise.
Numerical models based on discrete grids also suffer from inherent biases, uncertainties, and heavy computational demands.
A tiny perturbation in the initial state can trigger massive discrepancies in long-term predictions—a phenomenon widely referred to as the butterfly effect. Model errors can rapidly accumulate and hinder predictability, a limitation primarily driven by uncertainties arising from unresolved subgrid-scale processes, as well as the internal and external variabilities embedded deep within the climate system.
Data Assimilation (DA)
To mitigate error accumulation, observations are used to constrain numerical models through data assimilation. In this framework, climate models solve stochastic differential equations (SDEs) to track the evolution of the full model state x, which is sequentially corrected by observations y that are typically incomplete and noisy.
Specifically, the snapshot of x at time step k+1 is modeled as a transition function Φ of the previous snapshot x at step k, augmented by a stochastic force ξ. This evolution trajectory of the state x is continuously constrained by observations y, which map to the state x through a measurement operator A and are typically sparse and corrupted by observational noise η.
FlowDAS: a Gen-AI Framework for DA
Standard data-driven approaches, like traditional score-based diffusion models, have recently been applied to data assimilation. However, classical diffusion models come with a major bottleneck: they rely on a thermodynamic process that requires transforming data all the way into an unstructured Gaussian noise distribution and reversing it. Because the forward process is driven by Brownian motion, the generative reverse trajectories are highly curved, computationally expensive, and struggle to cleanly preserve tight physical consistency when guided by sparse real-world data.
We introduce FlowDAS, a generative, autoregressive diffusion model built on the mathematics of stochastic interpolants (Chen et al., 2025, NeurIPS). Instead of routing our data through an artificial Gaussian noise bottleneck, stochastic interpolants allow us to define a direct velocity-based transport mechanism between arbitrary probability distributions. By prescribing a linear interpolation profile between states, the generative trajectories behave like straight lines in the state space rather than highly curved diffusion paths. As you can see in the geometric illustration in the right panel of the figure, we construct a generative manifold at each interpolation step. To ensure that our reconstructed state x̂ remains strictly consistent with observations y, we incorporate a deterministic velocity field augmented by a gradient guidance term ▽|| y - A(x̂) ||.
Results: DA with 5% Sparse Observations
In a sparse-observation task, where only 5% of the domain is observed and the state evolution must be inferred from previous snapshots and partial observations, FlowDAS (highlighted in green) reconstructs the vorticity field with remarkable pixel-level fidelity. It accurately recovers even the finest turbulent structures, whereas competing models suffer from pronounced spectral attenuation and substantial variance distortions.
By leveraging the generative trajectories of stochastic interpolants, FlowDAS offers a computationally efficient, structurally accurate, and physically consistent framework for data assimilation. It represents a significant advance for high-resolution weather forecasting, ocean state estimation, and the study of turbulent geophysical flows.
Y. Jia, S. Chen, Y. Pan, X. Li, L. Shi, C. Jung, H. Yuan, I. Alkhouri, Y. C. Wu, S. Ravishankar, J. A. Fessler, and Q. Qu, "ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing," 2026, arXiv:2605.14285.
S. Chen, Y. Jia, Q. Qu, H. Sun, and J. A. Fessler, "FlowDAS: A Stochastic Interpolant-based Framework for Data Assimilation," in Advances in Neural Information Processing Systems (NeurIPS), vol. 38, 2025.