GAIA: A Foundation Model for Operational Atmospheric Dynamics

We present the GAIA (Geospatial Artificial Intelligence for Atmospheres) Foundation Model, a novel model that combines masked autoencoders (MAE) and self-distillation with no labels (DINO) for analyzing global atmospheric patterns in satellite imagery. By integrating these complementary self-supervised learning approaches, our model simultaneously captures both local features and global dependencies, addressing two critical challenges in satellite data analysis: reconstructing missing regions and estimating precipitation patterns. The model demonstrates superior attention distribution and temporal pattern capture compared to standard
MAE approaches, while maintaining robust performance in downstream tasks

README.md

@@ -0,0 +1,161 @@
+---
+title: GAIA
+emoji: 🛰️
+colorFrom: blue
+colorTo: green
+# TBD
+# sdk: gradio
+# sdk_version: 3.50.2
+# python_version: 3.8
+# app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+# GAIA: A Foundation Model for Operational Atmospheric Dynamics
+
+We present the GAIA (Geospatial Artificial Intelligence for Atmospheres) Foundation Model, a novel model that combines masked autoencoders (MAE) and Self-Distillation with NO labels (DINO) for analyzing global atmospheric patterns in satellite imagery. By integrating these complementary self-supervised learning approaches, our model simultaneously captures both local features and global dependencies, addressing two critical challenges in satellite data analysis: reconstructing missing regions and estimating precipitation patterns. The model demonstrates superior attention distribution and temporal pattern capture compared to standard MAE approaches, while maintaining robust performance in downstream tasks
+
+## Architecture Overview
+
+GAIA employs a transformer-based architecture specifically designed to handle spatio-temporal satellite data with the following attributes:
+
+- **Backbone**: Masked Autoencoder (MAE) [He et al. 2021](https://arxiv.org/pdf/2111.06377) combined with DINO self-distillation [Caron et al. 2021](https://arxiv.org/pdf/2104.14294) architecture
+
+- **Local-Global Learning**: MAE handles local patterns while DINO captures global atmospheric dynamics
+
+- **Self-Supervised Training**: Self-supervised pre-training on masked satellite imagery
+
+- **Training Strategy**: Our training strategy combines the MAE and DINO loss to learn local and global relationship in the geostationary satellite data
+
+- **Resolution**: Supports processing of medium-resolution geostationary satellite data (0.25 degree)
+
+- **Input Channels**: Single channel (long-wave infrared) from geostationary satellites
+
+![Pretraining architecture](docs/model_architecture.png "GAIA architecture")
+
+## Pre-trained Models
+
+The base GAIA model is pre-trained on a comprehensive dataset of geostationary satellite observations from 2001-2015.
+The pre-trained model weights are available at [this link](https://huggingface.co/bcg-usra-nasa-gaia/GAIA-v1/tree/main/checkpoints).
+
+```python
+model = GAIABase(
+    config_path
+)
+    
+model.configure_model()
+state_dict = torch.load("checkpoints/gaia-v1.pt")
+model.load_state_dict(state_dict)
+```
+
+## Downstream Tasks
+
+The pre-trained model is adapted to the following downstream tasks.
+
+### 1. Gap Filling in Operational Data
+
+GAIA has been fine-tuned to address operational gaps in satellite data:
+
+```python
+model = GAIAGapFill(
+    config_path
+)
+    
+model.configure_model()
+state_dict = torch.load("checkpoints/gaia_gapfill.pt")
+model.load_state_dict(state_dict)
+```
+
+The results showcase several key strengths of our approach:
+ 
+- **Large Gap Reconstruction**: The model successfully reconstructs substantial missing regions while preserving the temperature gradients and atmospheric patterns consistent with surrounding areas.
+- **Pattern Continuity**: The reconstructions maintain smooth transitions between filled regions and original data, avoiding artificial boundaries or discontinuities.
+- **Detail Preservation**: Fine-scale features such as cloud formations and temperature variations are accurately reproduced, suggesting the model has learned meaningful representations of atmospheric physics rather than simple interpolation.
+- **Visible Patch Reconstruction**: The fact that the model's reconstruction closely replicates the fine-grained details in the input image suggests that the encoder is generating rich latent representations that preserve details in the input.
+
+
+### 2. Precipitation Estimation
+
+GAIA was fine-tuned to learn the relationship between Geostationary Operational Environmental Satellite IR and precipitation data:
+
+```python
+model = GAIAPrecip(
+    config_path
+)
+    
+model.configure_model()
+state_dict = torch.load("checkpoints/gaia_precip.pt")
+model.load_state_dict(state_dict)
+```
+
+The model effectively captures the spatial distribution and intensity patterns of precipitation across a variety of atmospheric scenarios. The precipitation predictions obtain an SSIM of 0.881.
+
+
+## Demo and Inference
+
+### Quick Start
+
+1. **Create a Virtual Environment (Recommended)**
+   ```bash
+   conda create -n gaia_env python=3.10 -y
+   conda activate gaia_env
+   ```
+
+2. **Clone the Repository**
+   ```bash
+   git clone https://huggingface.co/bcg-usra-nasa-gaia/GAIA-v1
+   cd GAIA-v1
+   ```
+
+3. **Install Dependencies**
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+4. **Run Inference Notebooks**
+   Navigate to the `notebooks` directory to run the demo notebooks:
+   - For gap-filling:
+     ```bash
+     cd notebooks
+     jupyter notebook gap_filling_inference.ipynb
+     ```
+   - For precipitation estimation:
+     ```bash
+     jupyter notebook precipitation_inference.ipynb
+     ```
+
+## Feedback
+
+We welcome feedback and contributions! Please:
+- Open issues for bugs or feature requests
+- Submit pull requests for improvements
+- Share your use cases and results
+
+
+## Citation
+
+If you use GAIA in your research, please cite:
+
+```bibtex
+@article{gaia-fm,
+  title={GAIA: A Foundation Model for Operational Atmospheric Dynamics},
+  author={Ata Akbari Asanjan and Olivia Alexander and Tom Berg and Clara Zhang and Matt Yang and Jad Makki and Disha Shidham and Srija Chakraborty and William Bender and Stephen Peng and Arun Ravindran and Olivier Raiman and David Potere and David Bell},
+  year={2025},
+  eprint={2505.18179},
+  archivePrefix={arXiv},
+  primaryClass={cs.LG},
+  url={https://arxiv.org/abs/2505.18179}, 
+}
+```
+
+### Copyright Notice:
+Copyright © 2025 Boston Consulting Group and Universities Space Research Association. All rights reserved. Unauthorized use, reproduction, or distribution of this software is strictly prohibited unless it is licensed for use under terms in the “Apache 2.0” license.
+
+### Scope and limitations:
+
+GAIA is released strictly as a research prototype meant to showcase how this methodology can learn useful representations from GOES data. It is aimed at research and educational communities to demonstrate the potential of MAEs and DINO in gap filling and precipitation estimation. The model is not a production-ready forecasting tool. It is intended to let researchers experiment in controlled, academic settings.
+
+In offline experiments GAIA’s embeddings have proved valuable for gap filling of infrared imagery and precipitation estimation. This model shows promise and could be extended for further tasks, such as identifying atmospheric rivers or tropical cyclones. However, it is important to note that the model was trained only on infrared channels, covers 60° S–60° N between 2001 and 2015, and has not been tested outside of this data. The method also may lead to issues such as over-smoothed outputs due to the pattern continuity method. Consequently, any quantitative conclusions drawn from GAIA should be treated as diagnostic and must be cross-checked.
+
+Because of these limitations, GAIA must not be used in safety-critical or high-stakes settings such as flight planning, maritime routing, disaster response, or financial and insurance decisions. It should not be repurposed for other applications without extensive additional work. Anyone who wishes to deploy it operationally would need to accomplish task such as retraining or fine-tuning on up-to-date sensor data, conducting rigorous calibration and out-of-distribution testing, applying rigorous human oversight, and obtaining any required regulatory approvals (e.g., from the FAA or national meteorological agencies).

checkpoints/gaia-gapfill-v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e26453ec534a6826cdba8fd1ffee75e4289a8024654b448a9c850360a39d285
+size 1177551768

checkpoints/gaia-precip-v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d236893591f274093166025a47e7fd133ebdde5333f4c7526f763ff73f4d6fc8
+size 1076731588

checkpoints/gaia-v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31a3ea0f8dbb5a472aa6664af01e7feadfe3c2aa5edafbc644ed6752eb0b4fd8
+size 2100463430

configs/base_config.yaml

@@ -0,0 +1,48 @@
+# Model Configuration
+model:
+  name: "GAIABase"
+  # Image and patch configuration
+  img_size: [1, 480, 1440]    # Hardcoded value: time, height, width
+  patch_size: [1, 30, 30]     # Hardcoded value: time_patch, height_patch, width_patch
+  in_chans: 1               # Hardcoded: Number of input channels
+  
+  # Encoder configuration
+  encoder_embed_dim: 912      # Hardcoded based on original computation
+  encoder_depth: 24
+  encoder_num_heads: 16
+  
+  # Decoder configuration
+  decoder_embed_dim: 512
+  decoder_depth: 8
+  decoder_num_heads: 16
+  
+  # Common configuration
+  mlp_ratio: 4.0
+  norm_pix_loss: false
+  drop_channels_rate: 0
+  
+  # DINO specific configuration
+  adjacent_masking: false
+  norm_last_layer: true
+  dino_head_dim: 1024
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.04
+  warmup_teacher_temp_epochs: 1
+  student_temp: 0.1
+  center_momentum: 0.9
+  momentum_teacher: 0.996
+
+  # Masking configuration
+  mask_ratio: 0.75
+  
+training:
+  epochs: 100
+  val_step: 0
+  optimizer: 
+    learning_rate: 1.5e-4
+  scheduler: 
+    type: "CosineAnnealingLR"
+    warmup_steps: 500
+    total_steps: 10000
+    eta_min: 1.0e-6
+    last_step: -1

configs/gap_fill_config.yaml

@@ -0,0 +1,35 @@
+# Model Configuration
+model:
+  name: "GapFill"
+  # Image and patch configuration
+  img_size: [1, 480, 1440]   # [time, height, width]
+  patch_size: [1, 30, 30]    # [time_patch, height_patch, width_patch]
+  in_chans: 1                # Number of input channels
+  input_horizon: 1
+  epochs: 100
+  
+  # Encoder configuration
+  encoder_embed_dim: 912
+  encoder_depth: 24
+  encoder_num_heads: 16
+  
+  # Decoder configuration
+  decoder_embed_dim: 512
+  decoder_depth: 16
+  decoder_num_heads: 16
+  
+  # Common configuration
+  mlp_ratio: 4.0
+  norm_pix_loss: false
+  drop_channels_rate: 0.0
+  
+  # Masking configuration
+  adjacent_masking: false
+  mask_ratio: 0.10
+
+  # Pretrained model settings (if applicable)
+  pretrained:
+    enabled: true
+    checkpoint_path: "./checkpoints/foundation_model/last.ckpt"
+    freeze_encoder: false
+      

configs/precip_config.yaml

@@ -0,0 +1,32 @@
+# Model Configuration
+model:
+  name: "Precip"
+  # Image and patch configuration
+  img_size: [1, 480, 1440]   # [time, height, width]
+  patch_size: [1, 30, 30]    # [time_patch, height_patch, width_patch]
+  in_chans: 1                # Number of input channels
+  
+  # Encoder configuration
+  encoder_embed_dim: 912
+  encoder_depth: 24
+  encoder_num_heads: 16
+  
+  # Decoder configuration
+  decoder_embed_dim: 512
+  decoder_depth: 8
+  decoder_num_heads: 16
+  
+  # Common configuration
+  mlp_ratio: 4.0
+  norm_pix_loss: false
+  drop_channels_rate: 0.0
+  
+  # Masking configuration
+  adjacent_masking: false
+  mask_ratio: 0
+
+  # Pretrained model settings (if applicable)
+  pretrained:
+    enabled: true
+    checkpoint_path: "./checkpoints/foundation_model/last.ckpt"
+    freeze_encoder: false

data/IR/01/merg_2024100100_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a246acda6e82a38278723cfe63ee4939e8316567e862a8daf8a2c400904709f
+size 39025660

data/IR/01/merg_2024100101_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a239ddc43436e464dc92be0c15da922d846b7ced50ecb1672d6128a69fe5e14
+size 39076736

data/IR/01/merg_2024100102_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb21efb499080c8f2bb0a70637e134c2ee40d50dff1a43ab618898ae38a2e1db
+size 39248195

data/IR/01/merg_2024100103_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2edda29c68ba94f2b524659c392e1e9f058263676617fb6df622a972cabd1628
+size 39400438

data/IR/01/merg_2024100104_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9386034b0557609d436f857e13ae0c21739a02fb98021a6ba3bf070802c4b9a
+size 39515294

data/IR/01/merg_2024100105_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:231b9cade5687ca4ceb5e4027ece60a559c569804a14c526136ddd25dc6212f5
+size 39645681

data/IR/01/merg_2024100106_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00beb97a2433832bf923db982ff0fd85bb0fef6b0ab3af90f20507b5e56d36aa
+size 39812872

data/IR/01/merg_2024100107_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05c02d954a906736e23b976a727f3dc94bcb5fb4bf932866c2e0b9dff6631a88
+size 39979073

data/IR/01/merg_2024100108_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b85aff38412ca32dcf7f847044c9450834d1fc5c042568102b71f08008c4e7b7
+size 40163260

data/IR/01/merg_2024100109_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70aecb859a626900e782e6f1d6b75620154888184877fd7cd188a98d161b6ba8
+size 40280077

data/IR/01/merg_2024100110_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e4c83bda3ef3ce17bf9578b5d35b25af015cff961f004db11384292a1affd80
+size 40328790

data/IR/01/merg_2024100111_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47a94564d6fc99b0995c3f453b80872cd8c4ece0f3a2def5ebc9f18367218bb9
+size 40330532

data/IR/01/merg_2024100112_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f4a0d1df54bb3a2854e5754e827c170a04c0f8104509f6ac884a779eb99f0ff
+size 40341792

data/IR/01/merg_2024100113_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cd12c18c8eaaeef3d7397d9788b2a9f9ec2760867b9dbd9afc0855f443703a8
+size 40335086

data/IR/01/merg_2024100114_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe18b2874b64c3be5441abc275ddffce17d45bb7522b3c9b940a42db8518dbe1
+size 40317904

data/IR/01/merg_2024100115_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3983d18f7fc499e3da4f623878bdc779338025a13473997489a5de969372e1f4
+size 40319660

data/IR/01/merg_2024100116_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d852c46ca83bf10253c0d256204289486807fadcf36b4d77f9137e5d9ca6f213
+size 40418593

data/IR/01/merg_2024100117_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3bbaa0fc155407cdda13e2424366530f50f69aeeda5278305afc8ba5f3f4aa84
+size 40466406

data/IR/01/merg_2024100118_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0698f1e9947ed6f309a717177f3c0c11498a5a26206943c2f270ede30a37d21
+size 40421937

data/IR/01/merg_2024100119_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eaba6dc6745552270492290764fb74d6d95089e2140cdba5140822b68d7a045e
+size 40248300

data/IR/01/merg_2024100120_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:940a7a45349f39b0e51be1b67708e1736790254647688390cb0a9b96607f1a3b
+size 39987296

data/IR/01/merg_2024100121_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4ef5586bd1ea5bd320fabded1d35baa8804494240d99cda9700210686ed800a
+size 39754575

data/IR/01/merg_2024100122_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:729c4acc5e805cbad0a2f77b770287816da2ae066f53c903d9d7e2f22265c263
+size 39609695

data/IR/01/merg_2024100123_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba6f1f93a6d55ed747460d6780ce3d4ce2138e6a5f5488f3243cfe8e7b9dc5a3
+size 39573505

data/IR/02/merg_2024100200_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f45e9e4b2e0edaaad7857275adda71a8e1b52c352790afd67b66ab8828c71960
+size 39649743

data/IR/02/merg_2024100201_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84218d4d1267cc2e4dbcb5e9a7188e20315fa48eeee4c629c686ba4ccb685fcb
+size 39781066

data/IR/02/merg_2024100202_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1228f7145a2499a9606ca6a39feb0080816020df0bd289f7ec34da19f00fbb63
+size 39990716

data/IR/02/merg_2024100203_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85d7dce3e7ac6fac43c25d3daab5f1a73a53bd32cff7c21220dfb928d6d375d8
+size 40117553

data/IR/02/merg_2024100204_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0c67369d2b92c4094838230e3c71d0c8dbaed4ce1baddffe5c2cf02870f7409
+size 40171587

data/IR/02/merg_2024100205_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d475c348eb3a08505a2fca41a2c7d0f88f4700913474a3ce0ef3dbafe3ee52f
+size 40261205

data/IR/02/merg_2024100206_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09ea8e9cb6f109a9b11b040f525e8204b60db0d1870d5a9e459f696f3c44cef2
+size 40362058

data/IR/02/merg_2024100207_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa58f54426cdc2d242a6632444ee1a22b7d8978404c41e97270e5b4a92e34215
+size 40477668

data/IR/02/merg_2024100208_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c45e267ae37d280e29b38e8be3480f8a8e18826a68833a84006faccccc0b0a9f
+size 40608315

data/IR/02/merg_2024100209_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe1c24a4a5e9a19753680d3e4a592f73a1bfe07c4d47a040ebb13bf2281f833b
+size 40691246

data/IR/02/merg_2024100210_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9d99df009bbd112dd02c0ab08c7b4f8bea04b5fce1b4ab667c1025ab80929ea
+size 40695744

data/IR/02/merg_2024100211_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:306b1113a5a976cfe0aa7a0ca043c6cecd6750a377a0acdaeb56516555f890bb
+size 40651218

data/IR/02/merg_2024100212_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8708d194e721acfca47ad0668d041af2830dc4c5988976fd1ecba2b8387c80f7
+size 40633727

data/IR/02/merg_2024100213_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4623b15954d934b3b30dc5fe1bc0545f269238a5b316bfbb96468407b1260bb5
+size 40605444

data/IR/02/merg_2024100214_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25ee71e4124025a1dd212232d15fb1616882fa88dd0c5f16bc3fe40c6349c3a2
+size 40561373

data/IR/02/merg_2024100215_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fff475eb6af84270b31b3785c9f2e074bad0fa2689c9d6613afff1ebba9a709
+size 40628273

data/IR/02/merg_2024100216_4km-pixel.nc4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f22c033bbbbd09d9b33d093ff0ab84b2fcaad27c8d24329095083ac9ad2a0dca
+size 40718453