AuroraCFD

A Computational Fluid Dynamics (CFD) model based on Microsoft's Aurora architecture for weather prediction.

---

🚀 New to AuroraCFD?

Start here: GETTING_STARTED.md - Complete guide for transfer learning with Aurora and Poseidon data

This comprehensive guide walks you through:

1. Setting up the repository
2. Downloading Poseidon CFD datasets
3. Loading Aurora pretrained checkpoints
4. Running your first model
5. Training and fine-tuning

Additional documentation: See docs/ folder for technical reference materials.

---

Overview

AuroraCFD adapts Aurora's foundation model architecture for CFD applications, implementing surface variable encoding and decoding with a Swin3D Transformer backbone.

Features

• Surface Variable Processing: Encoder-decoder architecture for CFD variables (u, v, w, p, T)
• Swin3D Transformer Backbone: Efficient 3D attention mechanism with U-Net structure
• Patch-based Tokenization: Spatial patch embeddings for efficient processing
• LoRA Adaptation: Fine-tuning support via Low-Rank Adaptation
• Multi-timestep History: Supports temporal context from multiple previous timesteps

Installation

# Clone the repository
cd aurora-cfd

# Install dependencies with uv
uv sync

Quick Start

import torch
from src.aurora_cfd import AuroraCFD

# Define CFD surface variables
surf_vars = ("u", "v", "w", "p", "T")  # velocity, pressure, temperature

# Initialize model
model = AuroraCFD(
    surf_vars=surf_vars,
    patch_size=4,
    latent_levels=4,
    embed_dim=512,
)

# Create sample input (batch=2, time=2, vars=5, height=64, width=128)
x_surf = torch.randn(2, 2, 5, 64, 128)

# Forward pass
model.eval()
with torch.no_grad():
    predictions = model(x_surf, surf_vars=surf_vars)

# Access predictions for each variable
u_pred = predictions["u"]  # Shape: (2, 64, 128)
v_pred = predictions["v"]  # Shape: (2, 64, 128)

Architecture

┌─────────────────────────────────────────────────────────────┐
│                      AuroraCFD Pipeline                     │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  Input: (B, T, V, H, W)                                    │
│     ↓                                                       │
│  ┌──────────────────────────────────────────────┐         │
│  │ Surface Encoder                              │         │
│  │ • Patch embedding                            │         │
│  │ • Surface level encoding                     │         │
│  │ • MLP + LayerNorm                            │         │
│  │ • Level dimension expansion                  │         │
│  └──────────────────────────────────────────────┘         │
│     ↓ (B, C×L, D)                                          │
│  ┌──────────────────────────────────────────────┐         │
│  │ Swin3D Transformer Backbone                  │         │
│  │ • Encoder stages (down-sampling)             │         │
│  │ • Decoder stages (up-sampling)               │         │
│  │ • Skip connections                           │         │
│  │ • Time conditioning                          │         │
│  └──────────────────────────────────────────────┘         │
│     ↓ (B, C×L, 2D)                                         │
│  ┌──────────────────────────────────────────────┐         │
│  │ Surface Decoder                              │         │
│  │ • Dimension unwrapping                       │         │
│  │ • Surface level extraction                   │         │
│  │ • Per-variable reconstruction heads          │         │
│  │ • Unpatchify to spatial grid                 │         │
│  └──────────────────────────────────────────────┘         │
│     ↓                                                       │
│  Output: dict[var: (B, H, W)]                              │
│                                                             │
└─────────────────────────────────────────────────────────────┘

Project Structure

aurora-cfd/
├── src/
│   └── aurora_cfd/
│       ├── __init__.py          # Package exports
│       ├── aurora_cfd.py        # Main model class
│       ├── encoder.py           # Surface variable encoder
│       └── decoder.py           # Surface variable decoder
├── scripts/
│   ├── load_scot_data.py        # Load scOT datasets
│   ├── advanced_temporal_loading.py  # Temporal sequences & rollouts
│   └── README.md                # Data loading documentation
├── main.py                       # Example usage
├── pyproject.toml               # Project dependencies
├── README.md                    # This file
└── SURFACE_VARIABLES.md         # Detailed technical documentation

Components

Surface Encoder

• Converts surface variables to patch embeddings
• Adds learnable surface level encoding
• Applies Perceiver-style MLP processing
• Expands to match backbone input requirements

Swin3D Transformer Backbone

• Hierarchical vision transformer with 3D shifted windows
• U-Net architecture with skip connections
• Efficient attention computation
• LoRA support for fine-tuning

Surface Decoder

• Per-variable reconstruction heads
• Extracts surface level from latent representation
• Unpatchifies to original spatial resolution
• Outputs dictionary of predicted fields

Model Parameters

Default Configuration

AuroraCFD(
    surf_vars=("u", "v", "w", "p", "T"),  # CFD variables
    patch_size=4,                          # Spatial patch size
    latent_levels=4,                       # Number of latent levels
    embed_dim=512,                         # Embedding dimension
    window_size=(2, 6, 12),               # Swin window size (C, H, W)
    encoder_depths=(6, 10, 8),            # Encoder blocks per stage
    encoder_num_heads=(8, 16, 32),        # Encoder attention heads
    decoder_depths=(8, 10, 6),            # Decoder blocks per stage
    decoder_num_heads=(32, 16, 8),        # Decoder attention heads
    mlp_ratio=4.0,                        # MLP hidden/embedding ratio
    drop_path=0.0,                        # Drop path rate
    drop_rate=0.0,                        # Dropout rate
    use_lora=True,                        # Enable LoRA
    lora_steps=40,                        # LoRA steps
    lora_mode="single",                   # LoRA mode
    max_history_size=2,                   # Max history timesteps
    timestep=timedelta(hours=1),          # Prediction timestep
)

Example Output

Running python main.py:

============================================================
AuroraCFD Example: Surface Variable Processing
============================================================

1. Initializing AuroraCFD model...
   Surface variables: ('u', 'v', 'w', 'p', 'T')
   Model initialized with 378,534,480 parameters

2. Creating sample input...
   Shape: (B=2, T=2, V=5, H=64, W=128)

3. Running forward pass...
   Encoding: (B, T, V, H, W) -> (B, L, D)
   Backbone: (B, L, D) -> (B, L, D)
   Decoding: (B, L, D) -> dict[var: (B, H, W)]

4. Output predictions:
   u: shape=torch.Size([2, 64, 128]), mean=0.1963, std=0.9469
   v: shape=torch.Size([2, 64, 128]), mean=0.0099, std=0.9447
   w: shape=torch.Size([2, 64, 128]), mean=0.0620, std=0.9415
   p: shape=torch.Size([2, 64, 128]), mean=-0.0202, std=0.9625
   T: shape=torch.Size([2, 64, 128]), mean=0.0598, std=0.9469

Transfer Learning from Aurora

Loading Pretrained Aurora Weights

AuroraCFD supports loading pretrained weights from Aurora weather models. This enables transfer learning from weather prediction to CFD applications.

Quick Start with AuroraCFDSmall

The easiest way to use pretrained Aurora weights is with AuroraCFDSmall, which is pre-configured to match Aurora Small architecture:

from src.aurora_cfd import AuroraCFDSmall

# Initialize with Aurora Small-compatible architecture
model = AuroraCFDSmall(surf_vars=("u", "v", "w", "p", "T"))

# Load pretrained Aurora Small checkpoint (backbone only)
result = model.load_aurora_checkpoint("aurora-0.25-small-pretrained.ckpt")

print(f"Loaded {len(result['loaded'])} backbone parameters")
print(f"Skipped {len(result['skipped'])} encoder/decoder parameters")

# Now use for CFD predictions with pretrained backbone!
predictions = model(x_surf)

Custom Architecture

For custom architectures, use the base AuroraCFD class:

from src.aurora_cfd import AuroraCFD

# Initialize with custom parameters
model = AuroraCFD(
    surf_vars=("u", "v", "w", "p", "T"),
    embed_dim=768,  # Custom embedding dimension
    encoder_depths=(8, 12, 10),
    # ... other custom parameters
)

# Load compatible checkpoint (may have missing/unexpected keys)
result = model.load_aurora_checkpoint("custom-checkpoint.ckpt")

The load_aurora_checkpoint() method:

• ✓ Loads backbone (Swin3D Transformer) weights from Aurora
• ↷ Skips encoder/decoder weights (CFD-specific, trained from scratch)
• Reports detailed loading statistics

This is useful for transfer learning, where Aurora's atmospheric dynamics knowledge can be adapted to CFD problems.

Download Aurora Checkpoints:

• Aurora on Hugging Face
• Available models: aurora-0.25-small.ckpt, aurora-0.25-finetuned.ckpt

Example:

# Run the example script
uv run python scripts/load_aurora_checkpoint_example.py

For detailed documentation on checkpoint loading, fine-tuning strategies, and troubleshooting, see CHECKPOINT_LOADING.md.

Loading CFD Data

Using scOT (Poseidon) Datasets

Load CFD data from the scOT library and ingest into AuroraCFD:

# With synthetic data (no data files needed)
uv run python scripts/load_scot_data.py --synthetic

# With actual scOT datasets (requires data files)
uv run python scripts/load_scot_data.py --dataset fluids.incompressible.ShearLayer

# Specify custom data directory
uv run python scripts/load_scot_data.py --dataset fluids.incompressible.ShearLayer --data-dir /path/to/data

# Or set environment variable
export SCOT_DATA_DIR=/path/to/data
uv run python scripts/load_scot_data.py --dataset fluids.incompressible.ShearLayer

# Advanced: Temporal sequences and multi-step rollout
uv run python scripts/advanced_temporal_loading.py --data-dir /path/to/data

Configuring Data Locations

You can specify where scOT/Poseidon data files are stored in three ways:

1. Command-line argument: --data-dir /path/to/data
2. Environment variable: export SCOT_DATA_DIR=/path/to/data
3. Alternative env variable: export POSEIDON_DATA_DIR=/path/to/data

If not specified, scOT will use its default data locations.

Available datasets include:

• Incompressible flows: ShearLayer, Gaussians, KolmogorovFlow, etc.
• Compressible flows: KelvinHelmholtz, RayleighTaylor, Riemann, etc.
• Wave propagation: Acoustic waves, Gaussian packets
• Other physics: Poisson, Allen-Cahn, etc.

See docs/DATA_LOADING.md for detailed documentation on data loading.

Documentation

• GETTING_STARTED.md ⭐ START HERE - Complete transfer learning guide from setup to training
• docs/SURFACE_VARIABLES.md - Technical architecture details and shape transformations
• docs/DATA_LOADING.md - Data loading scripts reference and API documentation

Dependencies

• Python 3.11 (required, as numpy<2 compatibility is needed)
• PyTorch
• einops
• microsoft-aurora (from GitHub)
• poseidon/scot (from GitHub)
• Additional dependencies in pyproject.toml

Technical Details

Shape Transformations

Input:      (B, T, V, H, W)     # Batch, Time, Variables, Height, Width
Encoder:    (B, C×L, D)          # Batch, Levels×Patches, Embedding
Backbone:   (B, C×L, 2D)         # Doubled dimension from skip connections
Decoder:    {var: (B, H, W)}     # Per-variable predictions

Where:

• B: Batch size
• T: History timesteps
• V: Number of variables
• C: Latent levels
• L: Spatial patches = (H/patch_size) × (W/patch_size)
• D: Embedding dimension
• H, W: Spatial dimensions

Key Features from Aurora

1. Patch Embedding: Efficient spatial tokenization via convolution
2. Surface Level Encoding: Distinguishes surface from atmospheric variables
3. Perceiver Architecture: MLP + LayerNorm processing
4. Variable-specific Heads: Independent reconstruction per variable
5. Unpatchify: Converts patches back to spatial grid

References

• Aurora: Microsoft Aurora Weather Foundation Model
• Poseidon: ETH Zurich CFD Framework
• Swin Transformer: Liu et al., "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows"
• Perceiver: Jaegle et al., "Perceiver: General Perception with Iterative Attention"

License

This project builds upon Aurora (MIT License) and adapts it for CFD applications.

Citation

If you use AuroraCFD in your research, please cite the original Aurora paper:

@article{bodnar2024aurora,
  title={Aurora: A Foundation Model of the Atmosphere},
  author={Bodnar, Cristian and others},
  journal={arXiv preprint arXiv:2405.13063},
  year={2024}
}