From dc7984f4a37c0a04081df04cdbe22f2b9b17c100 Mon Sep 17 00:00:00 2001 From: munzzyy Date: Sun, 5 Jul 2026 16:58:51 -0500 Subject: [PATCH] fix: guard division by zero in _create_grid_mapping for single-axis grids GraphWeatherForecaster._create_grid_mapping divided by (max(unique_lats) - min(unique_lats)) and the same for longitudes, with no guard for the degenerate case where all lat_lons share one latitude or one longitude. That range hits zero and raises ZeroDivisionError, so any 1D grid crashes right in __init__. Fall back to row/col 0 when the axis has no spread to normalize against. Normal 2D grids are unaffected. --- graph_weather/models/forecast.py | 18 ++++++++---------- tests/test_model.py | 25 +++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/graph_weather/models/forecast.py b/graph_weather/models/forecast.py index 15d38b76..fa238cb5 100755 --- a/graph_weather/models/forecast.py +++ b/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): diff --git a/tests/test_model.py b/tests/test_model.py index f2cb6be0..03608597 100755 --- a/tests/test_model.py +++ b/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):