PDE foundation models are skillful AI weather emulators for the Martian atmosphere
Authors: Johannes Schmude, Sujit Roy, Liping Wang, Theodore van Kessel, Levente Klein, Marcus Freitag, Eloisa Bentivegna, Robert Manson-Sawko, Bjorn Lutjens, Manil Maskey, Campbell Watson, Rahul Ramachandran, Juan Bernabe-Moreno
Organizations: IBM Research · NASA Marshall Space Flight Center, Huntsville, Al, USA · Earth System Science Center, University of Alabama in Huntsville, Al, USA · Department of Electrical and Computer Engineering, Colorado State University, Fort Collins, CO, USA
We show that AI foundation models that are pretrained on numerical solutions to a diverse corpus of partial differential equations can be adapted and fine-tuned to obtain skillful predictive weather emulators for the Martian atmosphere. We base our work on the Poseidon PDE foundation model for two-dimensional systems. We develop a method to extend Poseidon from two to three dimensions while keeping the pretraining information. Moreover, we investigate the performance of the model in the presence of sparse initial conditions. Our results make use of four Martian years (approx.~34 GB) of training data and a median compute budget of 13 GPU hours. We find that the combination of pretraining and model extension yields a performance increase of 34.4% on a held-out year. This shows that PDEs-FMs can not only approximate solutions to (other) PDEs but also anchor models for real-world problems with complex interactions that lack a sufficient amount of training data or a suitable compute budget.
Table 1: Sigma level choices for different experiments
Figure 1: Extending scOT to three dimension
MODEL INIT.
SINGLE LEVEL
FULL VERTICAL
Random
0.078625
0.142879
PDE-FM (ours)
0.062886
N/A
Mixed (ours)
N/A
0.093776
Improvement
20.0%
34.4%
Table 2: Efficacy of PDE Foundation Models for atmospheric physics tasks. Loss in model units. The “mixed” initialization uses the PDE-FM weights on the 2D layers while randomly initializing the vertically acting layers.
Figure 2: When extending the training to 40,000 steps, the randomly initialized model not only does not close the performance gap, but rather shows clear signs of overfitting. The differences between the 20,000 and 40,000 step runs stem from the learning rate scheduler.
MODEL INIT.
80%
95%
Random
0.161899
0.179392
PDE-FM (ours)
0.098258
0.118802
Improvement
39.3%
33.8%
Table 3: Efficacy of multi PDE pre-training in presence of missing data.
Figure 3: Model performance after 20,000 training steps by level and masking ratio.
Figure 4: Impact of pretraining and dimensional extension on physical quantities.
Figure 5: Sample prediction for temperature at σ=0.87007517 .
Figure 6: Sample prediction for u component of wind at σ=0.87007517 .
Figure 7: Forecast error for temperature at σ=0.87007517 .
Figure 8: Forecast error for u component of wind at σ=0.87007517 .
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Impact of pretraining and dimensional extension on physical quantities; northward wind.
We investigate the transferability of Earth weather foundation models to planetary atmospheres by adapting the GraphCast graph neural weather forecasting model to Mars. While GraphCast achieves state-of-the-art performance for terrestrial forecasting, its applicability to non-Earth environments remains unexplored. Using the Mars Climate Database (MCD), which provides global atmospheric fields across vertical altitude levels (similar to Earth pressure levels), we evaluate zero-shot and fine-tuned GraphCast predictions of Martian temperature and wind fields. Zero-shot forecasts produce a surprisingly accurate depiction of current conditions but fail to reproduce diurnal variability and rapidly decay toward climatological mean states. To address this limitation, we fine-tune GraphCast using MCD variables and top-of-atmosphere solar radiation forcing while holding humidity constant. Fine-tuning enables rapid learning of Martian thermal variability. Within as few as 10 training epochs, the model begins to capture the diurnal cycle and forecasts up to 10 days reproduce seasonal and vertical temperature structure. Prediction quality improves with training sample size and exhibits sensitivity to seasonal initialization. These results demonstrate that Earth-trained AI weather models can be adapted to simulate Martian atmospheric dynamics, providing a pathway toward rapid planetary weather prediction to support mission operations, dust storm risk mitigation, and future human exploration.
M. L. Carroll, J. Li, S. D. Guzewich +3
1) NASA Goddard Space Flight Center · 2) ADNet Systems, Inc.
The martian atmosphere hosts dynamical phenomena ranging from planet-encircling dust storms to mesoscale orographic clouds and nocturnal low-level jets. General circulation model show capability to simulate these phenomena, but is computationally expensive at resolution needed to resolve mesoscale features. While assimilation of satellite remote sensing observation enable forecasting capabilities using such models, observation record is often sparse, short and fragmented across instrument generators. These constraints motivate the development of a data-driven foundation model for the Martian atmosphere. Foundation models live in a complex design landscape. There is an interplay between the available data, the physics of the underlying processes and corresponding developments in AI. Even though the idea of a foundation model is to address multiple use cases in a data- and compute-efficient manner, it is important to have a clear picture what applications can sensibly addressed by a single model. The purpose of this paper is to elucidate this design landscape. We discuss available data ranging from atmospheric retrievals to reanalysis datasets as well as existing physical models. Moreover, we identify a wide range of candidate downstream applications. Finally, we consider relevant recent developments in artificial intelligence (AI) that can be leveraged in this context. Here, we put a particular emphasis on AI models for atmospheric physics, data-driven approaches to data assimilation as well as methods to work in a limited data setting.
Sujit Roy, Udayshankar Nair, Yuling Wu +16
Earth System Science Center, University of Alabama in Huntsville, AL, USA · NASA Marshall Space Flight Center, Huntsville, AL, USA · Department of Electrical & Computer Engineering, Colorado State University, Fort Collins, CO, USA +7
General circulation models (GCMs) are the primary tool for simulating planetary atmospheres. They play a vital role in understanding Mars's atmosphere, as forecasting its unique weather is mission-critical for operations such as entry, descent, and landing. Mars poses unusual challenges for these models, as observations are sparse compared to Earth. In addition, a thin \co{} atmosphere alongside a radiatively active dust cycle creates a volatile atmosphere with large diurnal temperature swings and no true terrestrial analog for validation. Existing Mars GCMs, including the LMD PCM, the NASA Ames Mars GCM, and PlanetWRF, are mature and physically detailed but are implemented in legacy Fortran with finite-difference or finite-volume solvers, and they do not expose gradients for calibration or machine learning. Here we present AEGIS, a modular differentiable Mars climate model that couples Mars's unique atmospheric physics to the Dinosaur dynamical core, with interfaces for neural closures. We showcase stable ten-Mars-year simulations that reproduce the seasonal \co{} cycle while conserving the total \co{} inventory, capture realistic large-scale surface-temperature structure, and produce surface pressure that follows Mars Orbiter Laser Altimeter (MOLA) topography. Gradients through coupled trajectories agree with finite differences and support physical calibration and neural training. We compare with conventional GCMs, highlighting the framework's computational efficiency and differentiability.
Sameera S Kashyap, Victor Cruz, Angel Yepez +1
Department of Computer Science UC Santa Cruz Santa Cruz, CA 95064, USA