Fix division-by-zero in GraphWeatherForecaster grid mapping for single-axis grids

graph_weather/models/forecast.py

@@ -178,17 +178,15 @@ def __init__(
     def _create_grid_mapping(self, unique_lats, unique_lons):
         """Create (row,col) mapping for original node order"""
         self.node_to_grid = []
+        lat_min, lat_max = min(unique_lats), max(unique_lats)
+        lon_min, lon_max = min(unique_lons), max(unique_lons)
+        lat_range = lat_max - lat_min
+        lon_range = lon_max - lon_min
         for lat, lon in self.original_lat_lons:
-            row = int(
-                (lat - min(unique_lats))
-                / (max(unique_lats) - min(unique_lats))
-                * (len(unique_lats) - 1)
-            )
-            col = int(
-                (lon - min(unique_lons))
-                / (max(unique_lons) - min(unique_lons))
-                * (len(unique_lons) - 1)
-            )
+            # A degenerate axis (single unique lat or lon) has no spread to
+            # normalize against, so every node maps to the same row/col.
+            row = int((lat - lat_min) / lat_range * (len(unique_lats) - 1)) if lat_range else 0
+            col = int((lon - lon_min) / lon_range * (len(unique_lons) - 1)) if lon_range else 0
             self.node_to_grid.append((row, col))
 
     def graph_to_grid(self, graph_tensor):

tests/test_model.py

@@ -133,6 +133,31 @@ def test_forecaster():
     assert not torch.isnan(out).any()
 
 
+def test_forecaster_single_lat():
+    # Grid with only one unique latitude used to raise ZeroDivisionError in
+    # _create_grid_mapping, since max(unique_lats) == min(unique_lats).
+    lat_lons = [(45.0, lon) for lon in range(0, 20, 5)]
+    model = GraphWeatherForecaster(lat_lons)
+    assert model.grid_shape == (1, len(lat_lons))
+    assert model.node_to_grid == [(0, i) for i in range(len(lat_lons))]
+
+    features = torch.randn((1, len(lat_lons), 78 + 24))
+    out = model(features)
+    assert not torch.isnan(out).any()
+
+
+def test_forecaster_single_lon():
+    # Same degenerate-axis bug, mirrored onto longitude.
+    lat_lons = [(lat, 10.0) for lat in range(-10, 10, 5)]
+    model = GraphWeatherForecaster(lat_lons)
+    assert model.grid_shape == (len(lat_lons), 1)
+    assert model.node_to_grid == [(i, 0) for i in range(len(lat_lons))]
+
+    features = torch.randn((1, len(lat_lons), 78 + 24))
+    out = model(features)
+    assert not torch.isnan(out).any()
+
+
 def test_assimilator_model():
     obs_lat_lons = []
     for lat in range(-90, 90, 7):