# Towards physically consistent data-driven weather forecasting: Integrating data assimilation with equivariance-preserving deep spatial transformers

arXiv: https://arxiv.org/abs/2103.09360 · HTML: https://mlsanity.com/paper/2103.09360/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 76 references that are in the map.
- First posted on arXiv 16 Mar 2021.

## Builds on (most cited first, 23 of 23)

- [WeatherBench: A benchmark dataset for data-driven weather forecasting](https://mlsanity.com/paper/2002.00469/) (2020)
- [Gauge Equivariant Convolutional Networks and the Icosahedral CNN](https://mlsanity.com/paper/1902.04615/) (2019)
- [Learning SO(3) Equivariant Representations with Spherical CNNs](https://mlsanity.com/paper/1711.06721/) (2017)
- [Data-Driven Forecasting of High-Dimensional Chaotic Systems with Long Short-Term Memory Networks](https://mlsanity.com/paper/1802.07486/) (2018)
- [On the Universality of Invariant Networks](https://mlsanity.com/paper/1901.09342/) (2019)
- [Improving data-driven global weather prediction using deep convolutional neural networks on a cubed sphere](https://mlsanity.com/paper/2003.11927/) (2020)
- [Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models](https://mlsanity.com/paper/2102.05107/) (2021)
- [Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: a case study with the Lorenz 96 model](https://mlsanity.com/paper/2001.01520/) (2020)
- [A deep-learning-based surrogate model for data assimilation in dynamic subsurface flow problems](https://mlsanity.com/paper/1908.05823/) (2019)
- [Deep Learning for Post-Processing Ensemble Weather Forecasts](https://mlsanity.com/paper/2005.08748/) (2020)
- [Combining data assimilation and machine learning to infer unresolved scale parametrisation](https://mlsanity.com/paper/2009.04318/) (2020)
- [Ensemble Kalman Inversion: A Derivative-Free Technique For Machine Learning Tasks](https://mlsanity.com/paper/1808.03620/) (2018)
- [Analog forecasting of extreme-causing weather patterns using deep learning](https://mlsanity.com/paper/1907.11617/) (2019)
- [Achieving Conservation of Energy in Neural Network Emulators for Climate Modeling](https://mlsanity.com/paper/1906.06622/) (2019)
- [Physics-Guided Architecture (PGA) of Neural Networks for Quantifying Uncertainty in Lake Temperature Modeling](https://mlsanity.com/paper/1911.02682/) (2019)
- [Climbing down Charney's ladder: Machine Learning and the post-Dennard era of computational climate science](https://mlsanity.com/paper/2005.11862/) (2020)
- [Using Data Assimilation to Train a Hybrid Forecast System that Combines Machine-Learning and Knowledge-Based Components](https://mlsanity.com/paper/2102.07819/) (2021)
- [Data-driven super-parameterization using deep learning: Experimentation with multi-scale Lorenz 96 systems and transfer-learning](https://mlsanity.com/paper/2002.11167/) (2020)
- [A test case for application of convolutional neural networks to spatio-temporal climate data: Re-identifying clustered weather patterns](https://mlsanity.com/paper/1811.04817/) (2018)
- [Machine Learning, Deepest Learning: Statistical Data Assimilation Problems](https://mlsanity.com/paper/1707.01415/) (2017)
- [Reservoir Computing as a Tool for Climate Predictability Studies](https://mlsanity.com/paper/2103.06206/) (2021)
- [Analog ensemble data assimilation and a method for constructing analogs with variational autoencoders](https://mlsanity.com/paper/2006.01101/) (2020)
- [Designing Accurate Emulators for Scientific Processes using Calibration-Driven Deep Models](https://mlsanity.com/paper/2005.02328/) (2020)

## Built on by (most cited first, 15 of 15)

- [Learning Dissipative Dynamics in Chaotic Systems](https://mlsanity.com/paper/2106.06898/) (2021)
- [Continuum Attention for Neural Operators](https://mlsanity.com/paper/2406.06486/) (2024)
- [Foundation Models for Weather and Climate Data Understanding: A Comprehensive Survey](https://mlsanity.com/paper/2312.03014/) (2023)
- [W-MAE: Pre-trained weather model with masked autoencoder for multi-variable weather forecasting](https://mlsanity.com/paper/2304.08754/) (2023)
- [Using Sequences of Life-events to Predict Human Lives](https://mlsanity.com/paper/2306.03009/) (2023)
- [Can AI weather models predict out-of-distribution gray swan tropical cyclones?](https://mlsanity.com/paper/2410.14932/) (2024)
- [Learning Spatiotemporal Chaos Using Next-Generation Reservoir Computing](https://mlsanity.com/paper/2203.13294/) (2022)
- [Statistical Deep Learning for Spatial and Spatio-Temporal Data](https://mlsanity.com/paper/2206.02218/) (2022)
- [Data Assimilation with Machine Learning Surrogate Models: A Case Study with FourCastNet](https://mlsanity.com/paper/2405.13180/) (2024)
- [Long-term stability and generalization of observationally-constrained stochastic data-driven models for geophysical turbulence](https://mlsanity.com/paper/2205.04601/) (2022)
- [A Multi-Scale Deep Learning Framework for Projecting Weather Extremes](https://mlsanity.com/paper/2210.12137/) (2022)
- [A Novel Transformer Network with Shifted Window Cross-Attention for Spatiotemporal Weather Forecasting](https://mlsanity.com/paper/2208.01252/) (2022)
- [MODS: Multi-source Observations Conditional Diffusion Model for Meteorological State Downscaling](https://mlsanity.com/paper/2506.14798/) (2025)
- [Long-time accuracy of ensemble Kalman filters for chaotic and machine-learned dynamical systems](https://mlsanity.com/paper/2412.14318/) (2024)
- [Data Assimilation using ERA5, ASOS, and the U-STN model for Weather Forecasting over the UK](https://mlsanity.com/paper/2401.07604/) (2024)

## Abstract (arXiv)

There is growing interest in data-driven weather prediction (DDWP), for example using convolutional neural networks such as U-NETs that are trained on data from models or reanalysis. Here, we propose 3 components to integrate with commonly used DDWP models in order to improve their physical consistency and forecast accuracy. These components are 1) a deep spatial transformer added to the latent space of the U-NETs to preserve a property called equivariance, which is related to correctly capturing rotations and scalings of features in spatio-temporal data, 2) a data-assimilation (DA) algorithm to ingest noisy observations and improve the initial conditions for next forecasts, and 3) a multi-time-step algorithm, which combines forecasts from DDWP models with different time steps through DA, improving the accuracy of forecasts at short intervals. To show the benefit/feasibility of each component, we use geopotential height at 500~hPa (Z500) from ERA5 reanalysis and examine the short-term forecast accuracy of specific setups of the DDWP framework. Results show that the equivariance-preserving networks (U-STNs) clearly outperform the U-NETs, for example improving the forecast skill by 45%. Using a sigma-point ensemble Kalman (SPEnKF) algorithm for DA and U-STN as the forward model, we show that stable, accurate DA cycles are achieved even with high observation noise. The DDWP+DA framework substantially benefits from large (O(1000)) ensembles that are inexpensively generated with the data-driven forward model in each DA cycle. The multi-time-step DDWP+DA framework also shows promises, e.g., it reduces the average error by factors of 2-3.