Aurora model implementation (#136)

* aurora model implementation

* Temp commit for switching branches

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* a new implementation of Aurora (3d+Unstructured)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* checkpointing configuration with smaller nums

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* EINOPS and processor configuration added

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* EINOPS implementation in encoder and processor for clear reshaping

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* checkpointing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: yuvraajnarula

graph_weather/models/aurora/__init__.py

@@ -0,0 +1,87 @@
+"""
+Aurora: A Foundation Model for Earth System Science
+- Combines 3D Swin Transformer encoding
+- Perceiver processing for efficient computation
+- 3D decoding for spatial-temporal predictions
+"""
+
+from .decoder import Decoder3D
+from .encoder import Swin3DEncoder
+from .model import AuroraModel, EarthSystemLoss
+from .processor import PerceiverProcessor
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AuroraModel",
+    "EarthSystemLoss",
+    "Swin3DEncoder",
+    "Decoder3D",
+    "PerceiverProcessor",
+]
+
+# Default configurations for different model sizes
+MODEL_CONFIGS = {
+    "tiny": {
+        "in_channels": 1,
+        "out_channels": 1,
+        "embed_dim": 48,
+        "latent_dim": 256,
+        "spatial_shape": (16, 16, 16),
+        "max_seq_len": 2048,
+    },
+    "base": {
+        "in_channels": 1,
+        "out_channels": 1,
+        "embed_dim": 96,
+        "latent_dim": 512,
+        "spatial_shape": (32, 32, 32),
+        "max_seq_len": 4096,
+    },
+    "large": {
+        "in_channels": 1,
+        "out_channels": 1,
+        "embed_dim": 192,
+        "latent_dim": 1024,
+        "spatial_shape": (64, 64, 64),
+        "max_seq_len": 8192,
+    },
+}
+
+
+def create_model(config="base", **kwargs):
+    """
+    Create an Aurora model with specified configuration.
+
+    Args:
+        config (str): Model size configuration ('tiny', 'base', or 'large')
+        **kwargs: Override default configuration parameters
+
+    Returns:
+        AuroraModel: Initialized model with specified configuration
+    """
+    if config not in MODEL_CONFIGS:
+        raise ValueError(
+            f"Unknown configuration: {config}. Choose from {list(MODEL_CONFIGS.keys())}"
+        )
+
+    # Start with default config and update with any provided kwargs
+    model_config = MODEL_CONFIGS[config].copy()
+    model_config.update(kwargs)
+
+    return AuroraModel(**model_config)
+
+
+def create_loss(alpha=0.5, beta=0.3, gamma=0.2):
+    """
+    Create an EarthSystemLoss instance with specified weights.
+
+    Args:
+        alpha (float): Weight for MSE loss
+        beta (float): Weight for gradient loss
+        gamma (float): Weight for physical consistency loss
+
+    Returns:
+        EarthSystemLoss: Initialized loss function
+    """
+    return EarthSystemLoss(alpha=alpha, beta=beta, gamma=gamma)

graph_weather/models/aurora/decoder.py

@@ -0,0 +1,47 @@
+"""
+3D Decoder:
+- Takes processed latent representations and reconstructs output.
+- Uses transposed convolution to upscale back to spatial-temporal format.
+"""
+
+import torch.nn as nn
+
+
+class Decoder3D(nn.Module):
+    """
+    3D Decoder:
+    - Takes processed latent representations and reconstructs the spatial-temporal output.
+    - Uses transposed convolutions to upscale latent features to the original format.
+    """
+
+    def __init__(self, output_channels=1, embed_dim=96, target_shape=(32, 32, 32)):
+        """
+        Args:
+            output_channels (int): Number of channels in the output tensor (e.g., 1 for grayscale).
+            embed_dim (int): Dimension of the latent features (matches the encoder's output).
+            target_shape (tuple): The desired shape of the reconstructed 3D tensor (D, H, W).
+        """
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.target_shape = target_shape
+        self.deconv1 = nn.ConvTranspose3d(
+            embed_dim, output_channels, kernel_size=3, padding=1, stride=1
+        )
+
+    def forward(self, x):
+        """
+        Forward pass for the decoder.
+
+        Args:
+            x (torch.Tensor): Input latent representation, shape (batch, seq_len, embed_dim).
+
+        Returns:
+            torch.Tensor: Reconstructed 3D tensor, shape (batch, output_channels, *target_shape).
+        """
+        batch_size = x.shape[0]
+        depth, height, width = self.target_shape
+        # Reshape latent features into 3D tensor
+        x = x.view(batch_size, self.embed_dim, depth, height, width)
+        # Transposed convolution to upscale to the final shape
+        x = self.deconv1(x)
+        return x

graph_weather/models/aurora/encoder.py

@@ -0,0 +1,76 @@
+"""
+Swin 3D Transformer Encoder:
+- Uses a 3D convolution for initial feature extraction.
+- Applies layer normalization and reshapes data.
+- Uses a transformer-based encoder to learn spatial-temporal features.
+"""
+
+import torch.nn as nn
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+
+class Swin3DEncoder(nn.Module):
+    def __init__(self, in_channels=1, embed_dim=96):
+        super().__init__()
+        self.conv1 = nn.Conv3d(in_channels, embed_dim, kernel_size=3, padding=1, stride=1)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.swin_transformer = nn.Transformer(
+            d_model=embed_dim,
+            nhead=8,
+            num_encoder_layers=4,
+            num_decoder_layers=4,
+            dim_feedforward=embed_dim * 4,
+        )
+        self.embed_dim = embed_dim
+
+        # Define rearrangement patterns using einops
+        self.to_transformer_format = Rearrange("b d h w c -> (d h w) b c")
+        self.from_transformer_format = Rearrange("(d h w) b c -> b d h w c", d=None, h=None, w=None)
+
+    # To use rearrange function directly instead of the Rearrange layer
+    def forward(self, x):
+        # 3D convolution with einops rearrangement
+        x = self.conv1(x)
+
+        # Rearrange for normalization using einops
+        x = rearrange(x, "b c d h w -> b d h w c")
+        x = self.norm(x)
+
+        # Store spatial dimensions for later reconstruction
+        d, h, w = x.shape[1:4]
+
+        # Transform to sequence format for transformer
+        x = rearrange(x, "b d h w c -> (d h w) b c")
+        x = self.swin_transformer.encoder(x)
+
+        # Restore original spatial structure
+        x = rearrange(x, "(d h w) b c -> b (d h w) c", d=d, h=h, w=w)
+
+        # Reshape to the expected output format (batch, seq_len, embed_dim)
+        x = rearrange(x, "b (d h w) c -> b (d h w) c", d=d, h=h, w=w)
+
+        return x
+
+    def convolution(self, x):
+        """Apply 3D convolution with clear shape transformation."""
+        return self.conv1(x)  # b c d h w -> b embed_dim d h w
+
+    def normalization_layer(self, x):
+        """Apply layer normalization with einops rearrangement."""
+        x = rearrange(x, "b c d h w -> b d h w c")
+        return self.norm(x)
+
+    def transformer_encoder(self, x, spatial_dims):
+        """
+        Apply transformer encoding with proper shape handling.
+
+        Args:
+            x (torch.Tensor): Input tensor
+            spatial_dims (tuple): Original (depth, height, width) dimensions
+        """
+        d, h, w = spatial_dims
+        x = self.to_transformer_format(x)
+        x = self.swin_transformer.encoder(x)
+        x = self.from_transformer_format(x, d=d, h=h, w=w)
+        return x