Add Initial WeatherMesh-2 Implementation (#138)

* Push updates

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

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

* Add test for irregular grid

* Add working encoder

Lots of inspiration and code from this: https://github.com/Brayden-Zhang/WeatherMesh with some changes to work with natten and some other changes.

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

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

* Add configurable decoder and WeatherMesh2 model

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

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

* Add dacite configs

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

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

* Fix keeping the correct depth (vertical pressure level)

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

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

* Add new updates and change to pixi install

---------

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

Author: jacobbieker

.all-contributorsrc

@@ -12,11 +12,9 @@ jobs:
       fail-fast: true
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.10", "3.11", "3.12"]
-        torch-version: [2.3.0, 2.4.0]
+        python-version: ["3.11", "3.12"]
+        torch-version: [2.4.0]
         include:
-          - torch-version: 2.3.0
-            torchvision-version: 0.18.0
           - torch-version: 2.4.0
             torchvision-version: 0.19.0
     steps:

.bumpversion.cfg

@@ -43,9 +43,9 @@ def __init__(self, filepath: str, features: list, start_year: int = 2016, end_ye
         """
 
         super().__init__()
-        assert start_year <= end_year, (
-            f"start_year ({start_year}) cannot be greater than end_year ({end_year})."
-        )
+        assert (
+            start_year <= end_year
+        ), f"start_year ({start_year}) cannot be greater than end_year ({end_year})."
         assert start_year >= 2016 and start_year <= 2022, "Time data range from 2016 to 2022"
         assert end_year >= 2016 and end_year <= 2022, "Time data range from 2016 to 2022"
         self.data = xr.open_zarr(filepath)

.github/workflows/release.yaml

@@ -105,9 +105,9 @@ def __init__(
                 )
             )
             if self.res:
-                assert image_size is not None and scale_factor is not None, (
-                    "If res=True, you must provide h, w and scale_factor"
-                )
+                assert (
+                    image_size is not None and scale_factor is not None
+                ), "If res=True, you must provide h, w and scale_factor"
                 h, w = pair(image_size)
                 s_h, s_w = pair(scale_factor)
                 self.res_layers.append(

.github/workflows/workflows.yaml

@@ -309,7 +309,9 @@ def _init_khop_mesh_graph(self):
             edge_index,
             values=torch.ones_like(edge_index[0], dtype=torch.float32),
             size=(self._num_mesh_nodes, self._num_mesh_nodes),
-        ).to(self.khop_device)  # cpu is more memory-efficient, why?
+        ).to(
+            self.khop_device
+        )  # cpu is more memory-efficient, why?
 
         adj_k = adj.coalesce()
         for _ in range(self.num_hops - 1):

.gitignore

@@ -28,7 +28,7 @@ def __init__(
             device: checks for device whether it supports gpu or not
             normalize: option for normalize
         """
-        # TODO Rescale by nominal static air density at each pressure level
+        # TODO Rescale by nominal static air density at each pressure level, could be 1/pressure level or something similar
         super().__init__()
         self.feature_variance = torch.tensor(feature_variance)
         assert not torch.isnan(self.feature_variance).any()
@@ -55,15 +55,19 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor):
         """
         self.feature_variance = self.feature_variance.to(pred.device)
         self.weights = self.weights.to(pred.device)
+        print(pred.shape)
+        print(target.shape)
+        print(self.weights.shape)
 
         out = (pred - target) ** 2
-
+        print(out.shape)
         if self.normalize:
             out = out / self.feature_variance
 
         assert not torch.isnan(out).any()
         # Mean of the physical variables
         out = out.mean(-1)
+        print(out.shape)
         # Weight by the latitude, as that changes, so does the size of the pixel
         out = out * self.weights.expand_as(out)
         assert not torch.isnan(out).any()

.pre-commit-config.yaml

@@ -0,0 +1,104 @@
+"""
+Implementation based off the technical report and this repo: https://github.com/Brayden-Zhang/WeatherMesh
+"""
+
+from dataclasses import dataclass
+
+import dacite
+import einops
+import torch
+import torch.nn as nn
+from natten import NeighborhoodAttention3D
+
+from graph_weather.models.weathermesh.layers import ConvUpBlock
+
+
+@dataclass
+class WeatherMeshDecoderConfig:
+    latent_dim: int
+    output_channels_2d: int
+    output_channels_3d: int
+    n_conv_blocks: int
+    hidden_dim: int
+    kernel_size: tuple
+    num_heads: int
+    num_transformer_layers: int
+
+    @staticmethod
+    def from_json(json: dict) -> "WeatherMeshDecoder":
+        return dacite.from_dict(data_class=WeatherMeshDecoderConfig, data=json)
+
+    def to_json(self) -> dict:
+        return dacite.asdict(self)
+
+
+class WeatherMeshDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_dim,
+        output_channels_2d,
+        output_channels_3d,
+        n_conv_blocks=3,
+        hidden_dim=256,
+        kernel_size: tuple = (5, 7, 7),
+        num_heads: int = 8,
+        num_transformer_layers: int = 3,
+    ):
+        super().__init__()
+
+        # Transformer layers for initial decoding
+        self.transformer_layers = nn.ModuleList(
+            [
+                NeighborhoodAttention3D(
+                    dim=latent_dim, num_heads=num_heads, kernel_size=kernel_size
+                )
+                for _ in range(num_transformer_layers)
+            ]
+        )
+
+        # Split into pressure levels and surface paths
+        self.split = nn.Conv3d(latent_dim, hidden_dim * (2**n_conv_blocks), kernel_size=1)
+
+        # Pressure levels (3D) path
+        self.pressure_path = nn.ModuleList(
+            [
+                ConvUpBlock(
+                    hidden_dim * (2 ** (i + 1)),
+                    hidden_dim * (2**i) if i > 0 else output_channels_3d,
+                    is_3d=True,
+                )
+                for i in reversed(range(n_conv_blocks))
+            ]
+        )
+
+        # Surface (2D) path
+        self.surface_path = nn.ModuleList(
+            [
+                ConvUpBlock(
+                    hidden_dim * (2 ** (i + 1)),
+                    hidden_dim * (2**i) if i > 0 else output_channels_2d,
+                )
+                for i in reversed(range(n_conv_blocks))
+            ]
+        )
+
+    def forward(self, latent: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Needs to be (B,D,H,W,C) with Batch, Depth (vertical levels), Height, Width, Channels
+        # Apply transformer layers
+        for transformer in self.transformer_layers:
+            latent = transformer(latent)
+
+        latent = einops.rearrange(latent, "B D H W C -> B C D H W")
+        # Split features
+        features = self.split(latent)
+        pressure_features = features[:, :, :-1]
+        surface_features = features[:, :, -1:]
+        # Decode pressure levels
+        for block in self.pressure_path:
+            pressure_features = block(pressure_features)
+        # Decode surface features
+        surface_features = surface_features.squeeze(2)
+        for block in self.surface_path:
+            surface_features = block(surface_features)
+
+        return surface_features, pressure_features

Dockerfile

@@ -0,0 +1,109 @@
+"""
+Implementation based off the technical report and this repo: https://github.com/Brayden-Zhang/WeatherMesh
+"""
+
+from dataclasses import dataclass
+
+import dacite
+import einops
+import torch
+import torch.nn as nn
+from natten import NeighborhoodAttention3D
+
+from graph_weather.models.weathermesh.layers import ConvDownBlock
+
+
+@dataclass
+class WeatherMeshEncoderConfig:
+    input_channels_2d: int
+    input_channels_3d: int
+    latent_dim: int
+    n_pressure_levels: int
+    num_conv_blocks: int
+    hidden_dim: int
+    kernel_size: tuple
+    num_heads: int
+    num_transformer_layers: int
+
+    @staticmethod
+    def from_json(json: dict) -> "WeatherMeshEncoder":
+        return dacite.from_dict(data_class=WeatherMeshEncoderConfig, data=json)
+
+    def to_json(self) -> dict:
+        return dacite.asdict(self)
+
+
+class WeatherMeshEncoder(nn.Module):
+    def __init__(
+        self,
+        input_channels_2d: int,
+        input_channels_3d: int,
+        latent_dim: int,
+        n_pressure_levels: int,
+        num_conv_blocks: int = 3,
+        hidden_dim: int = 256,
+        kernel_size: tuple = (5, 7, 7),
+        num_heads: int = 8,
+        num_transformer_layers: int = 3,
+    ):
+        super().__init__()
+
+        # Surface (2D) path
+        self.surface_path = nn.ModuleList(
+            [
+                ConvDownBlock(
+                    input_channels_2d if i == 0 else hidden_dim * (2**i),
+                    hidden_dim * (2 ** (i + 1)),
+                )
+                for i in range(num_conv_blocks)
+            ]
+        )
+
+        # Pressure levels (3D) path
+        self.pressure_path = nn.ModuleList(
+            [
+                ConvDownBlock(
+                    input_channels_3d if i == 0 else hidden_dim * (2**i),
+                    hidden_dim * (2 ** (i + 1)),
+                    stride=(1, 2, 2),  # Want to keep depth the same size
+                    is_3d=True,
+                )
+                for i in range(num_conv_blocks)
+            ]
+        )
+
+        # Transformer layers for final encoding
+        self.transformer_layers = nn.ModuleList(
+            [
+                NeighborhoodAttention3D(
+                    dim=latent_dim, kernel_size=kernel_size, num_heads=num_heads
+                )
+                for _ in range(num_transformer_layers)
+            ]
+        )
+
+        # Final projection to latent space
+        self.to_latent = nn.Conv3d(hidden_dim * (2**num_conv_blocks), latent_dim, kernel_size=1)
+
+    def forward(self, surface: torch.Tensor, pressure: torch.Tensor) -> torch.Tensor:
+        # Process surface data
+        for block in self.surface_path:
+            surface = block(surface)
+
+        # Process pressure level data
+        for block in self.pressure_path:
+            pressure = block(pressure)
+        # Combine features
+        features = torch.cat(
+            [pressure, surface.unsqueeze(2)], dim=2
+        )  # B C D H W currently, want it to be B D H W C
+
+        # Transform to latent space
+        latent = self.to_latent(features)
+
+        # Reshape to get the shapes
+        latent = einops.rearrange(latent, "B C D H W -> B D H W C")
+        # Apply transformer layers
+        for transformer in self.transformer_layers:
+            latent = transformer(latent)
+        return latent