FNP: Fourier Neural Processes for Arbitrary-Resolution Data Assimilation

TL;DR

The paper tackles the challenge of assimilating observations at arbitrary resolutions in atmospheric data assimilation. It introduces Fourier Neural Processes (FNP), an architecture that uses a unified coordinate transformation, spatial-variable functional representations with SetConv and Neural Fourier Layers, and a dynamic alignment and merge module to fuse background and observations into a probabilistic analysis. FNP demonstrates state-of-the-art performance across multiple resolutions and shows robust out-of-domain generalization, including direct applicability to observational information reconstruction without fine-tuning. These results suggest substantial practical impact for end-to-end AI-based weather forecasting and flexible data assimilation workflows, while the authors acknowledge limitations related to synthetic data and lack of temporal modeling, with future work aimed at extending to temporal dimensions.

Abstract

Data assimilation is a vital component in modern global medium-range weather forecasting systems to obtain the best estimation of the atmospheric state by combining the short-term forecast and observations. Recently, AI-based data assimilation approaches have attracted increasing attention for their significant advantages over traditional techniques in terms of computational consumption. However, existing AI-based data assimilation methods can only handle observations with a specific resolution, lacking the compatibility and generalization ability to assimilate observations with other resolutions. Considering that complex real-world observations often have different resolutions, we propose the \textit{\textbf{Fourier Neural Processes}} (FNP) for \textit{arbitrary-resolution data assimilation} in this paper. Leveraging the efficiency of the designed modules and flexible structure of neural processes, FNP achieves state-of-the-art results in assimilating observations with varying resolutions, and also exhibits increasing advantages over the counterparts as the resolution and the amount of observations increase. Moreover, our FNP trained on a fixed resolution can directly handle the assimilation of observations with out-of-distribution resolutions and the observational information reconstruction task without additional fine-tuning, demonstrating its excellent generalization ability across data resolutions as well as across tasks.

Figures (7)

• Figure 1: Overview of the network architecture of FNP. Unified coordinate transformation ensures spatial alignment of the background and observations, and extracts the coordinates and values of the conditional points. FNP models the two components of the conditional information globally to get their respective spatial-variable functional representation through SetConv for data embedding and stacking of neural Fourier layers for deep feature extraction. The dynamic alignment and merge module integrates these functional representation based on similarity to shared features and aligns them into the target domain, resulting in a comprehensive functional representation over the target space. MLPs are finally employed to decode the functional representation and output the mean and variance of the analysis based on the coordinates of the target points.
• Figure 2: Visualization of assimilation results by different models for q700. The visualization date-time is randomly selected at 2018-04-02 06:00 UTC. The first row shows the ERA5 (ground truth), background, background error and observations with 0.25° resolution. Other rows show the assimilation results of different models.
• Figure 3: Visualization of assimilation results with 0.25° resolution for u10.
• Figure 4: Visualization of assimilation results with 0.25° resolution for t2m.
• Figure 5: Visualization of assimilation results with 0.25° resolution for z500.