Physically consistent and uncertainty-aware learning of spatiotemporal dynamics

TL;DR

PCNO introduces a physics-consistent neural operator that enforces mass and momentum conservation by projecting surrogate outputs onto Fourier-space constraint spaces, yielding physically consistent, accurate spatiotemporal forecasts. DiffPCNO extends this with a consistency-model diffusion mechanism to quantify and refine predictive uncertainty, improving long-term reliability. The two-stage framework delivers high-fidelity predictions across turbulent flows, real-world floods, and atmospheric dynamics while achieving resolution-invariance and robust uncertainty quantification. The work combines a momentum/mass–conserving projection layer with probabilistic residual correction, offering a practically impactful approach for physics-informed, uncertainty-aware forecasting in complex geophysical systems.

Abstract

Accurate long-term forecasting of spatiotemporal dynamics remains a fundamental challenge across scientific and engineering domains. Existing machine learning methods often neglect governing physical laws and fail to quantify inherent uncertainties in spatiotemporal predictions. To address these challenges, we introduce a physics-consistent neural operator (PCNO) that enforces physical constraints by projecting surrogate model outputs onto function spaces satisfying predefined laws. A physics-consistent projection layer within PCNO efficiently computes mass and momentum conservation in Fourier space. Building upon deterministic predictions, we further propose a diffusion model-enhanced PCNO (DiffPCNO), which leverages a consistency model to quantify and mitigate uncertainties, thereby improving the accuracy and reliability of forecasts. PCNO and DiffPCNO achieve high-fidelity spatiotemporal predictions while preserving physical consistency and uncertainty across diverse systems and spatial resolutions, ranging from turbulent flow modeling to real-world flood/atmospheric forecasting. Our two-stage framework provides a robust and versatile approach for accurate, physically grounded, and uncertainty-aware spatiotemporal forecasting.

Figures (12)

• Figure 1: Schematic illustration of the proposed PCNO and DiffPCNO frameworks. a, Embedding physics through PCNO. PCNO enforces physical consistency by projecting surrogate outputs onto function spaces that satisfy predefined physical laws via a physics-consistent projection layer. b, Schematic of the physics-consistent projection layer, which enforces mass and momentum conservation in spatiotemporal dynamics. c, Momentum-conserving projection with invariance via a rotation-invariant kernel in Fourier space. d, Mass-conserving projection by embedding the divergence-free condition in Fourier space. e, Embedding uncertainty through DiffPCNO corrects prediction residuals via a generative residual correction mechanism based on a consistency model, conditioned on the deterministic prediction $\mathbf{\hat{u}}_{t+1}$ from PCNO and the current state $\mathbf{u}_t$. It aims to capture the residual distribution $\mathbf{r}_{t+1}=\mathbf{y}-\mathbf{\hat{u}}_{t+1}$, where $\mathbf{y}$ denotes the ground-truth solution. It enables applications such as spatiotemporal forecasting, uncertainty quantification, and zero-shot super-resolution (downscaling). For instance, DiffPCNO can rapidly estimate uncertainty by generating multiple samples of the dynamical process. f, Consistency models are a class of diffusion-based generative models that generate high-quality samples in a single step while retaining the flexibility for multi-step sampling to balance computational cost and fidelity. For DiffPCNO, an improved consistency model based on the probability flow ordinary differential equation (ODE) is employed and trained using consistency training.
• Figure 1: CSI results of real-world flood inundation forecasting. a-d compare the CSI (higher values indicate better performance) at water depth thresholds of 0.05 m and 0.5 m for low-fidelity Pakistan flood forecasting, transferable Mozambique flood forecasting, high-fidelity Australia flood forecasting (30 m), and transferable UK flood forecasting (30 m). Comparisons of prediction performance at different flood depths reveal that, relative to FNO, PCNO with physical constraints and DiffPCNO with probabilistic learning consistently provide superior forecasts across varying water levels and time steps.
• Figure 1: Locations of the study areas.
• Figure 2: Results of the one-dimensional Kuramoto-Sivashinsky dynamics. a, Visualization results and uncertainty quantification of DiffPCNO and PCNO-Refiner on the 1D KSE with fixed viscosity $\nu=1$. The vertical axis of the visualization represents 256 spatial points, while the horizontal axis corresponds to the chaotic spatiotemporal evolution over 400 testing time steps. Uncertainty is generated via a stochastic recursive process with 50 sampled trajectories per test case, from which empirical standard deviations are computed (see Methods for detailed information). b, Rollout MSE of PCNO, DiffPCNO, PCNO-Refiner, and baseline methods for KSE with fixed viscosity over 400 testing time steps. c, Visualization results and uncertainty quantification of DiffPCNO and PCNO-Refiner on the 1D KSE with varying viscosity $\nu$ sampled uniformly between 0.5 and 1.5. d, Rollout MSE of PCNO, DiffPCNO, PCNO-Refiner, and baseline methods for KSE with varying viscosity over 400 testing time steps. e, Time steps where the average correlation drops below 0.9 and 0.8, indicating the temporal horizon of reliable predictions. f, Super-resolution performance evaluation of PCNO, DiffPCNO, PCNO-Refiner, and baseline methods, on the KSE with varying viscosity, with models trained on a $64 \times 64$ spatial grid and tested directly on downscaled grids at $1\times$, $2\times$, and $4\times$ resolutions.
• Figure 2: Results of transferable flood inundation forecasting. Spatial variability of final flood inundation extents and depths across flood measurements (SAR and surveyed outlines), traditional hydrodynamic models, and DiffPCNO. a, SAR-based flood map on 20 March 2019 in Mozambique. f, Surveyed flood outlines in UK. b and g, Traditional hydrodynamic-based flood maps for Mozambique and UK (flood depth$\ge$0.05 m). c and h, DiffPCNO-based flood maps for Mozambique and UK (flood depth$\ge$0.05 m). d and i, Traditional hydrodynamic-based flood maps for Mozambique and UK (flood depth$\ge$0.5 m). e and j, DiffPCNO-based flood maps for Mozambique and UK (flood depth$\ge$0.5 m). k, SAR-based flood depths for Mozambique on 20 March 2019. SAR-based flood depths are extracted using the DEM and SAR-based flood extents cohen2018estimating. p, Surveyed flood outlines in UK. l and q, Traditional hydrodynamic-based flood depths for Mozambique and UK (flood depth$\ge$0.05 m). m and r, DiffPCNO-based flood depths for Mozambique and UK (flood depth$\ge$0.05 m). n and s, Spatial distribution of errors between DiffPCNO-based and hydrodynamic-based flood depths for Mozambique and UK (flood depth$\ge$0.05 m). o, Spatial distribution of errors between DiffPCNO-based and SAR-based flood depths for Mozambique. t, Spatial distribution of errors between DiffPCNO-based and hydrodynamic-based flood depths for UK (flood depth$\ge$0.5 m).