Merge pull request #22 from CU-ESIIL/codex/fix-spatial-subsetting-bug-in-gridmet-xjaa0t

Fix gridMET bbox slicing

Author: ttuff

src/cubedynamics/streaming/gridmet.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import io
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Sequence
 
+import numpy as np
 import requests
 import xarray as xr
 from xarray.backends.plugins import list_engines
@@ -12,6 +13,53 @@
 _AVAILABLE_ENGINES = list_engines()
 
 
+def _axis_slice(coord: Sequence[float], bound_a: float, bound_b: float) -> slice:
+    """
+    Return a slice that spans ``[bound_a, bound_b]`` regardless of axis order.
+
+    gridMET tiles have a resolution of roughly 1/24th of a degree. When an AOI
+    bounding box is smaller than that resolution, its numeric bounds can fall
+    entirely between adjacent coordinate centers. To avoid empty selections we
+    expand the slice bounds by half the native grid spacing before subsetting.
+    """
+
+    values = np.asarray(coord, dtype=float)
+    if values.size == 0:
+        val = float(min(bound_a, bound_b))
+        return slice(val, val)
+
+    lo = float(min(bound_a, bound_b))
+    hi = float(max(bound_a, bound_b))
+
+    if values.size > 1:
+        diffs = np.diff(values)
+        diffs = diffs[np.nonzero(diffs)]
+        if diffs.size:
+            spacing = float(np.min(np.abs(diffs)))
+            span = hi - lo
+            if spacing > 0 and span < spacing:
+                padding = (spacing - span) / 2.0
+                lo -= padding
+                hi += padding
+
+    descending = values[0] > values[-1]
+    if descending:
+        return slice(hi, lo)
+    return slice(lo, hi)
+
+
+def _lat_slice(lat_coord: Sequence[float], south: float, north: float) -> slice:
+    """Return a latitude slice that works for ascending or descending axes."""
+
+    return _axis_slice(lat_coord, south, north)
+
+
+def _lon_slice(lon_coord: Sequence[float], west: float, east: float) -> slice:
+    """Return a longitude slice that works for ascending or descending axes."""
+
+    return _axis_slice(lon_coord, west, east)
+
+
 def _select_stream_engine() -> Optional[str]:
     """Pick the best available xarray engine for streaming gridMET files."""
 
@@ -162,10 +210,26 @@ def stream_gridmet_to_cube(
 
     # 3) Spatial subset using the AOI bbox
     bbox = _bbox_from_geojson(aoi_geojson)
-    da = ds[variable].sel(
-        lat=slice(bbox["south"], bbox["north"]),
-        lon=slice(bbox["west"], bbox["east"]),
-    )
+    lat_coord = ds.coords.get("lat")
+    if lat_coord is None:
+        raise KeyError("gridMET dataset is missing the 'lat' coordinate")
+    lon_coord = ds.coords.get("lon")
+    if lon_coord is None:
+        raise KeyError("gridMET dataset is missing the 'lon' coordinate")
+
+    lat_slice = _lat_slice(lat_coord, bbox["south"], bbox["north"])
+    lon_slice = _lon_slice(lon_coord, bbox["west"], bbox["east"])
+    da = ds[variable].sel(lat=lat_slice, lon=lon_slice)
+
+    empty_dims = [dim for dim in ("lat", "lon") if da.sizes.get(dim, 0) == 0]
+    if empty_dims:
+        raise ValueError(
+            "gridMET subset is empty along "
+            + ", ".join(f"'{dim}'" for dim in empty_dims)
+            + ". "
+            f"Requested south={bbox['south']}, north={bbox['north']}, "
+            f"west={bbox['west']}, east={bbox['east']}"
+        )
 
     # 4) Optional resampling in time (e.g., to monthly)
     if freq != "D":

src/cubedynamics/tests/test_gridmet_streaming.py

@@ -1,6 +1,13 @@
 """Tests for the gridMET streaming helper."""
 
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
 import cubedynamics as cd
+from cubedynamics.streaming import gridmet as gridmet_mod
 
 
 def test_stream_gridmet_to_cube_default_source_smoke():
@@ -32,3 +39,92 @@ def test_stream_gridmet_to_cube_default_source_smoke():
 
     assert "time" in cube.dims
     assert cube.name == "pr"
+
+
+def test_stream_gridmet_to_cube_descending_lat(monkeypatch):
+    """Descending latitude axes should still produce non-empty subsets."""
+
+    times = pd.date_range("2000-01-01", periods=3, freq="D")
+    lat = xr.DataArray([50.0, 49.5, 49.0], dims="lat")
+    lon = xr.DataArray([-120.0, -119.5, -119.0], dims="lon")
+
+    def _fake_year(variable: str, year: int, chunks=None) -> xr.Dataset:  # pragma: no cover - test helper
+        data = xr.DataArray(
+            np.arange(times.size * lat.size * lon.size, dtype="float32").reshape(times.size, lat.size, lon.size),
+            coords={"time": times, "lat": lat, "lon": lon},
+            dims=("time", "lat", "lon"),
+            name=variable,
+        )
+        return xr.Dataset({variable: data})
+
+    monkeypatch.setattr(gridmet_mod, "_open_gridmet_year", _fake_year)
+
+    aoi = {
+        "type": "Feature",
+        "geometry": {
+            "type": "Polygon",
+            "coordinates": [[
+                [-120.0, 49.1],
+                [-120.0, 49.6],
+                [-119.4, 49.6],
+                [-119.4, 49.1],
+                [-120.0, 49.1],
+            ]],
+        },
+    }
+
+    cube = cd.stream_gridmet_to_cube(
+        aoi_geojson=aoi,
+        variable="pr",
+        start="2000-01-01",
+        end="2000-01-03",
+    )
+
+    assert cube.sizes["lat"] == 1
+    assert cube.sizes["lon"] == 2
+    assert cube.lat.item() == 49.5
+
+
+def test_stream_gridmet_to_cube_bbox_padding(monkeypatch):
+    """Tiny AOI bounding boxes should still capture at least one grid cell."""
+
+    times = pd.date_range("2000-01-01", periods=2, freq="D")
+    lat = xr.DataArray([40.020833, 39.979167], dims="lat")
+    lon = xr.DataArray([-105.375, -105.333333, -105.291666], dims="lon")
+
+    def _fake_year(variable: str, year: int, chunks=None) -> xr.Dataset:  # pragma: no cover - test helper
+        data = xr.DataArray(
+            np.arange(times.size * lat.size * lon.size, dtype="float32").reshape(times.size, lat.size, lon.size),
+            coords={"time": times, "lat": lat, "lon": lon},
+            dims=("time", "lat", "lon"),
+            name=variable,
+        )
+        return xr.Dataset({variable: data})
+
+    monkeypatch.setattr(gridmet_mod, "_open_gridmet_year", _fake_year)
+
+    aoi = {
+        "type": "Feature",
+        "geometry": {
+            "type": "Polygon",
+            "coordinates": [[
+                [-105.35, 40.00],
+                [-105.35, 40.01],
+                [-105.34, 40.01],
+                [-105.34, 40.00],
+                [-105.35, 40.00],
+            ]],
+        },
+    }
+
+    cube = cd.stream_gridmet_to_cube(
+        aoi_geojson=aoi,
+        variable="pr",
+        start="2000-01-01",
+        end="2000-01-02",
+    )
+
+    assert cube.sizes["lat"] == 1
+    assert cube.sizes["lon"] == 1
+    assert np.isclose(cube.lat.item(), float(lat.isel(lat=0)))
+    assert np.isclose(cube.lon.item(), float(lon.isel(lon=1)))