Bug fixes and updates

Author: jacobbieker

graph_weather/data/IFSAnalysis_dataloader.py

@@ -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)

graph_weather/data/__init__.py

@@ -1,5 +1,5 @@
 """Dataloaders and data processing utilities"""
 
-from .anemoi_dataloder import AnemoiDataset
+from .anemoi_dataloader import AnemoiDataset
 from .nnja_ai import SensorDataset, collate_fn
 from .weather_station_reader import WeatherStationReader

graph_weather/data/anemoi_dataloder.py graph_weather/data/anemoi_dataloader.pygraph_weather/data/anemoi_dataloder.py renamed to graph_weather/data/anemoi_dataloader.py

@@ -330,9 +330,9 @@ def _convert_to_synopticpy(self, observations: xr.Dataset) -> Optional[Dict]:
                     values = station_data[var_name].values
 
                     # Add to observations
-                    synoptic_data["STATION"][str(station)]["OBSERVATIONS"][
-                        var_name
-                    ] = values.tolist()
+                    synoptic_data["STATION"][str(station)]["OBSERVATIONS"][var_name] = (
+                        values.tolist()
+                    )
 
         return synoptic_data
 

graph_weather/data/weather_station_reader.py

@@ -43,7 +43,8 @@ def forward(self, points: torch.Tensor, features: torch.Tensor) -> torch.Tensor:
 
         # Normalize coordinates to [-1, 1] range
         normalized_points = torch.stack(
-            [points[..., 0] / 180.0, points[..., 1] / 90.0], dim=-1  # longitude  # latitude
+            [points[..., 0] / 180.0, points[..., 1] / 90.0],
+            dim=-1,  # longitude  # latitude
         )
 
         # Separately encode coordinates and features

graph_weather/models/aurora/model.py

@@ -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(

graph_weather/models/fengwu_ghr/layers.py

@@ -309,9 +309,7 @@ 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):

graph_weather/models/gencast/graph/graph_builder.py

@@ -50,7 +50,7 @@ def __init__(
         self.transformer_layers = nn.ModuleList(
             [
                 NeighborhoodAttention3D(
-                    dim=latent_dim, num_heads=num_heads, kernel_size=kernel_size
+                    embed_dim=latent_dim, num_heads=num_heads, kernel_size=kernel_size
                 )
                 for _ in range(num_transformer_layers)
             ]

graph_weather/models/weathermesh/decoder.py

@@ -76,7 +76,7 @@ def __init__(
         self.transformer_layers = nn.ModuleList(
             [
                 NeighborhoodAttention3D(
-                    dim=latent_dim, kernel_size=kernel_size, num_heads=num_heads
+                    embed_dim=latent_dim, kernel_size=kernel_size, num_heads=num_heads
                 )
                 for _ in range(num_transformer_layers)
             ]

graph_weather/models/weathermesh/encoder.py

@@ -31,7 +31,7 @@ def __init__(self, latent_dim, n_layers=10, kernel=(5, 7, 7), num_heads=8):
         self.layers = nn.ModuleList(
             [
                 NeighborhoodAttention3D(
-                    dim=latent_dim,
+                    embed_dim=latent_dim,
                     num_heads=num_heads,
                     kernel_size=kernel,
                 )