🔢 mathematics

Multimodal Spatiotemporal Atmospheric Data Assimilation with Latent Flow-matching

This paper proposes a unified, latent flow-matching framework for multimodal atmospheric data assimilation that leverages a prior trained on ERA5 reanalysis to generate temporally consistent full-state trajectories, enabling flexible filtering, smoothing, and competitive ensemble forecasting from sparse observations.

Original authors: Dibyajyoti Chakraborty, Romit Maulik

Published 2026-08-06

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Dibyajyoti Chakraborty, Romit Maulik

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine trying to predict the weather is like trying to guess the plot of a movie you've only seen a few scattered frames of. The atmosphere is a giant, chaotic system where wind, temperature, and pressure dance together across the entire globe. To forecast what happens next, scientists usually rely on massive, complex computer models that simulate physics from scratch. But these models are heavy, slow, and need a perfect starting picture to work well. The problem is, we can never see the whole picture at once; our instruments only give us a few snapshots from specific spots on the ground. This creates a puzzle: how do we fill in the missing gaps and guess the future without running a supercomputer simulation every single time? This is where "data assimilation" comes in. Think of it as a smart detective who takes a rough guess about the story (the model's background) and updates it using the few clues they actually have (the observations). For decades, this has been done with rigid math that requires endless calculations. But recently, a new kind of artificial intelligence called "generative machine learning" has entered the chat. Instead of solving equations, these AI models learn the "vibe" or the general pattern of how the atmosphere behaves, allowing them to imagine new, realistic weather scenarios that fit the clues we have.

This paper introduces a new, playful way to solve that weather puzzle using a technique called "latent video flow-matching." The researchers, Dibyajyoti Chakraborty and Romit Maulik, trained an AI to watch a "movie" of the Earth's atmosphere. They didn't just feed it one photo; they fed it 8-day-long video clips (32 frames) of 69 different weather variables, like temperature and wind, taken from a global dataset called ERA5. The AI learned to generate these videos from pure noise, essentially learning the rules of how weather moves and changes over time. Once trained, this AI acts as a "prior," a background story that knows how the atmosphere should behave.

The magic happens when they mix this AI with real-world data. Instead of retraining the AI for every new situation, they use a clever sampling trick. Imagine the AI is drawing a picture of the weather from a blurry cloud of noise. If you tell the AI, "Hey, at this specific spot, the temperature is actually 20°C," the AI doesn't just ignore the rest of the cloud; it gently nudges its drawing to match that fact while keeping the rest of the picture consistent with the movie it learned. Because the AI learned a continuous video, it can naturally fill in the gaps between the spots where we have data and even predict what happens in the future frames, all in one go.

The team tested this on three main challenges. First, they tried "super-resolution," where they gave the AI a very low-quality, blocky version of the weather map and asked it to restore the fine details. The AI successfully reconstructed sharp, detailed weather patterns that were missing from the blurry input. Second, they used real, sparse data from weather stations (like radiosondes and surface stations) to fill in the entire globe. Even though the data was patchy, the AI generated a full, consistent 3D picture of the atmosphere that matched the real world well. Finally, they tested if this could work as a forecast. By feeding the AI only the first few days of observations and letting it "watch" the rest of the video unfold, they created a six-day forecast. They found that this method was competitive with other top-tier AI forecasting models, even without using satellite data, which are usually considered essential for long-range predictions.

The paper also shows that this single AI model can do different types of "detective work" just by changing which frames are observed. If you show it the beginning and end of a video, it acts as a "smoother," filling in the middle perfectly. If you only show it the beginning, it acts as a "filter" or forecaster, guessing the future. This suggests that a single, trained AI can replace the need for separate, complex numerical models for different tasks. However, the authors are careful to note that while their results are promising and competitive, they haven't yet replaced the massive operational systems used by weather agencies for daily forecasts. They also point out that their model struggles a bit with the most chaotic, small-scale details of wind, likely because the real atmosphere is even more complex than the data they trained on. Still, this work suggests that generative AI can be a powerful, flexible tool for understanding and predicting our dynamic planet, turning a few scattered clues into a full, moving story of the sky.

Technical Summary: Multimodal Spatiotemporal Atmospheric Data Assimilation with Latent Flow-Matching

Problem Statement

Atmospheric dynamics span vast spatial and temporal scales, making forecasting difficult due to the need for accurate initial conditions and the inherent limitations of partial observations. Classical Data Assimilation (DA) addresses this by fusing numerical forecast models (priors) with observed data using Bayesian inference. However, traditional variational and ensemble methods are computationally expensive, requiring multiple iterations of complex numerical models. Furthermore, existing deep learning approaches for atmospheric reconstruction often rely on deterministic models that lack uncertainty quantification (UQ), require retraining for different observation sources, or treat temporal consistency as an external constraint rather than an internal property of the generative model.

Methodology

The authors propose a unified, generative approach to atmospheric DA that replaces the numerical forecast model with a latent video flow-matching prior. The methodology consists of three primary components:

1. Latent Representation and Autoencoder

A 3D-convolutional autoencoder compresses 8-day (32 timesteps, 6-hourly) windows of global ERA5 reanalysis data (69 variables, $128 \times 256$ grid) into a lower-dimensional latent video space ( $128 \times 8 \times 32 \times 64$ ). This compression allows for efficient training and sampling while preserving spatiotemporal structures. The decoder maps these latents back to physical space, where likelihoods are evaluated.

2. Generative Prior: TrigFlow

Instead of training on raw ERA5 fields, the authors train a TrigFlow prior on the autoencoder latents. TrigFlow is a flow-matching formulation that defines a spherical interpolation path between data ( $z_0$ ) and Gaussian noise ( $\epsilon$ ):
$z_t = \cos(t) z_0 + \sin(t) \epsilon$
A 3D Diffusion Transformer (DiT3D) is trained to regress the velocity field ( $d z_t / dt$ ) along this path. Unlike diffusion models that reverse a noising process, flow-matching learns a velocity field to directly interpolate from noise to data. The prior is trained unconditionally on the ERA5 windows and remains frozen during all assimilation tasks, enabling zero-shot conditioning.

3. Guided Posterior Sampling

To assimilate observations, the method samples from the posterior distribution $p(z | y)$ by combining the unconditional prior score with a likelihood score derived from observations. The sampling process utilizes a Guided TrigFlow sampler incorporating:

Adaptive Noise-Level Weighting: A loss function modification that predicts a scalar log-variance to balance gradients across different noise levels, preventing high-noise levels from dominating training.
Measurement Guidance: The sampler injects a guidance term derived from the gradient of the likelihood (e.g., $-\nabla \frac{1}{2\sigma^2_{lik}} \|y - A(D(\hat{z}_0))\|^2$ ) into the ODE integration steps.
Annealed Likelihood Variance: Following Score-based Data Assimilation (SDA), the likelihood variance is annealed ( $\sigma^2_{lik} = \sigma^2_y + \gamma (\sigma/\sigma_d)^2$ ) to weaken guidance at high noise levels where the Tweedie estimate is blurry.
Predictor-Corrector Structure: The sampler employs a DPM-Solver++ predictor for state advancement and a Langevin corrector (active only at low noise levels) to refine the sample toward the posterior, improving accuracy without excessive computational cost.
Gradient Rescaling: Techniques like Diffusion with Spherical-Gaussian constraint (DSG) and momentum accumulation are used to stabilize and normalize the guidance gradients.

Key Contributions

Unified Latent Video Prior: The authors train an unconditional latent video flow-matching prior (TrigFlow with a DiT3D backbone) capable of generating 8-day, 69-variable global atmospheric trajectories. This serves as a background model for diverse DA tasks without retraining.
Internal Temporal Propagation: Unlike previous snapshot-based generative models that rely on external emulators for temporal consistency, this prior inherently propagates information between observed and unobserved frames within the video trajectory.
Flexible DA Regimes: By simply changing the set of observed frames (the observation mask), the same frozen prior recovers three classical DA regimes:
- Filtering: Observing leading frames and forecasting the tail.
- Smoothing: Observing middle frames to reconstruct the full window.
- Fixed-Interval Reconstruction: Observing frames at regular intervals (simulating reanalysis).
Multimodal Data Fusion: The framework successfully assimilates heterogeneous real-world data sources, including NOAA Integrated Global Radiosonde Archive (IGRA), Integrated Surface Database (ISD), and ICOADS (ships/buoys), to reconstruct the full atmospheric state, including unobserved variables.
Direct Observation-to-Forecast: The method generates full-state ensemble forecasts directly from sparse in-situ observations, competing with state-of-the-art observation-to-forecast models like GraphDOP.

Results

Super-Resolution (SR): The model recovers 2m temperature with an RMSE of 1.02 K from 4x temporal and 8x spatial subsampled observations. The ensemble spread is well-calibrated (Spread-Skill Ratio $\approx$ 0.72–0.75 for SR, 0.85 for IGRA) and captures fine-scale structures absent in coarse observations.
Physical Consistency: Posterior samples satisfy geostrophic balance and hypsometric consistency checks, performing comparably to the ERA5 reanalysis itself.
Real-World Data Fusion: When assimilating IGRA and ISD data, the model reconstructs the full 69-channel field. The ensemble spread correctly reflects uncertainty, being lower near observation stations and higher in sparse regions. The reconstruction error against held-out observations is comparable to or better than ERA5 deviations at those specific sites.
Observation-to-Forecast: In a direct observation-to-forecast experiment (assimilating IGRA/ISD/ICOADS and forecasting 6 days), the model improves upon GraphDOP for 2m temperature at all leads where dense surface observations are helpful. For other surface fields, errors initially exceed GraphDOP but converge by days 5–7 as forecast error saturates. The authors note that GraphDOP ingests satellite measurements alongside surface data, whereas this model relies solely on in-situ observations; however, the method demonstrates stability over a three-week autoregressive chain.
Case Study (Hurricane Laura): The method successfully tracks Hurricane Laura (2020) and performs a hindcast of its genesis. While the grid resolution (1.4°) cannot fully resolve the hurricane's eye, the ensemble captures the storm's trajectory and wind distribution, with confidence increasing as the observation window lengthens.

Significance and Claims

The paper claims that latent video flow-matching provides a fundamentally different, unified approach to atmospheric data assimilation. By internalizing temporal propagation within the generative prior, the method eliminates the need for separate numerical forecast models or emulators to maintain temporal consistency. The ability to switch between filtering, smoothing, and forecasting regimes simply by altering the observation mask highlights the flexibility of the learned prior.

The authors emphasize that their approach achieves competitive performance with operational and state-of-the-art data-driven models (like GraphDOP) using only in-situ observations, without requiring a learned encoder for observations or a separate forecast model. The framework provides a probabilistic, uncertainty-quantified reconstruction of the full atmospheric state from sparse, multimodal data sources. The authors note that while the results are promising, a full end-to-end comparison with operational assimilation systems is beyond the scope of this work, and future directions include extending the sequential cycling for sub-seasonal forecasts and integrating satellite observation operators.