🍃⚡ Aero: Robust Backend-Agnostic Spatial Slicing

environment.yml

@@ -14,6 +14,7 @@ dependencies:
   - dask
   - cubed
   - cubed-xarray
+  - zarr<3
   - netcdf4
   - pandas
   - matplotlib

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "dask",
     "cubed",
     "cubed-xarray",
+    "zarr<3",
     "netCDF4",
     "esmpy",
     "cf-xarray",

requirements.txt

@@ -5,6 +5,7 @@ scipy
 dask
 cubed
 cubed-xarray
+zarr<3
 netCDF4
 esmpy
 cf-xarray

requirements_no_esmpy.txt

@@ -6,6 +6,7 @@ scipy
 dask
 cubed
 cubed-xarray
+zarr<3
 netCDF4
 cf-xarray
 pyproj

src/xregrid/utils.py

@@ -1635,53 +1635,103 @@ def spatial_slice(
     min_y -= buffer
     max_y += buffer
 
-    # 3. Y-Slicing (Latitude or Projection Y)
-    y_dim = y_da.dims[0]
-    if obj.indexes[y_dim].is_monotonic_increasing:
-        obj = obj.sel({y_dim: slice(min_y, max_y)})
-    else:
-        obj = obj.sel({y_dim: slice(max_y, min_y)})
-
-    # 4. X-Slicing (Longitude or Projection X)
-    x_dim = x_da.dims[0]
-    if not is_geographic:
-        # Standard slice for projected coordinates
-        if obj.indexes[x_dim].is_monotonic_increasing:
-            obj = obj.sel({x_dim: slice(min_x, max_x)})
-        else:
-            obj = obj.sel({x_dim: slice(max_x, min_x)})
-        return obj
-
-    # 5. Longitude Wrapping Logic
-    # Get grid convention from eager indexes
-    lon_grid = obj.indexes[x_dim]
-    g_min = lon_grid.min()
-
-    # Normalize extent to [g_min, g_min + 360]
-    norm_min_x = (min_x - g_min) % 360 + g_min
-    norm_max_x = (max_x - g_min) % 360 + g_min
-
-    # Detect if we need a wrapped slice
-    if norm_min_x > norm_max_x:
-        # Crosses the grid boundary
-        if lon_grid.is_monotonic_increasing:
-            part1 = obj.sel({x_dim: slice(norm_min_x, g_min + 360)})
-            part2 = obj.sel({x_dim: slice(g_min, norm_max_x)})
+    # 3. Slicing Logic
+    # Determine if we can use .sel() (Rectilinear) or must use .where() (Curvilinear/Unstructured)
+    is_rectilinear = (
+        y_da.ndim == 1
+        and x_da.ndim == 1
+        and y_da.name == y_da.dims[0]
+        and x_da.name == x_da.dims[0]
+        and y_da.dims[0] != x_da.dims[0]
+        and y_da.dims[0] in obj.indexes
+        and x_da.dims[0] in obj.indexes
+    )
+
+    if is_rectilinear:
+        # Fast path for rectilinear grids using indexes (Always Lazy)
+        y_dim, x_dim = y_da.dims[0], x_da.dims[0]
+
+        # Y-Slicing
+        if obj.indexes[y_dim].is_monotonic_increasing:
+            obj = obj.sel({y_dim: slice(min_y, max_y)})
         else:
-            part1 = obj.sel({x_dim: slice(g_min + 360, norm_min_x)})
-            part2 = obj.sel({x_dim: slice(norm_max_x, g_min)})
+            obj = obj.sel({y_dim: slice(max_y, min_y)})
 
-        # Concatenate parts
-        res = xr.concat([part1, part2], dim=x_dim)
+        # X-Slicing with Longitude Wrapping
+        if not is_geographic:
+            if obj.indexes[x_dim].is_monotonic_increasing:
+                obj = obj.sel({x_dim: slice(min_x, max_x)})
+            else:
+                obj = obj.sel({x_dim: slice(max_x, min_x)})
+            res = obj
+            wrapped = False
+        else:
+            lon_grid = obj.indexes[x_dim]
+            g_min = lon_grid.min()
+            norm_min_x = (min_x - g_min) % 360 + g_min
+            norm_max_x = (max_x - g_min) % 360 + g_min
+
+            if norm_min_x > norm_max_x:
+                wrapped = True
+                if lon_grid.is_monotonic_increasing:
+                    p1 = obj.sel({x_dim: slice(norm_min_x, g_min + 360)})
+                    p2 = obj.sel({x_dim: slice(g_min, norm_max_x)})
+                else:
+                    p1 = obj.sel({x_dim: slice(g_min + 360, norm_min_x)})
+                    p2 = obj.sel({x_dim: slice(norm_max_x, g_min)})
+                res = xr.concat([p1, p2], dim=x_dim)
+            else:
+                wrapped = False
+                if lon_grid.is_monotonic_increasing:
+                    res = obj.sel({x_dim: slice(norm_min_x, norm_max_x)})
+                else:
+                    res = obj.sel({x_dim: slice(norm_max_x, norm_min_x)})
     else:
-        # Simple non-wrapped slice
-        if lon_grid.is_monotonic_increasing:
-            res = obj.sel({x_dim: slice(norm_min_x, norm_max_x)})
+        # Robust path for Curvilinear and Unstructured using masking
+        mask_y = (y_da >= min_y) & (y_da <= max_y)
+
+        if not is_geographic:
+            mask_x = (x_da >= min_x) & (x_da <= max_x)
+            wrapped = False
         else:
-            res = obj.sel({x_dim: slice(norm_max_x, norm_min_x)})
+            # Strictly Lazy Longitude Wrapping logic
+            # Works for any longitude convention (0-360 or -180 to 180)
+            range_x = (max_x - min_x) % 360
+            if range_x == 0 and max_x != min_x:
+                range_x = 360
+
+            if (max_x - min_x) >= 360:
+                mask_x = xr.DataArray(True)
+                wrapped = False
+            else:
+                mask_x = ((x_da - min_x) % 360) <= range_x
+                wrapped = (max_x - min_x) > range_x or min_x < 0 or max_x > 360
+
+        mask = mask_y & mask_x
+
+        # Aero Protocol: No Hidden Computes.
+        # Decide if we can use drop=True. We avoid it for lazy data because
+        # it forces a compute to determine the output shape.
+        is_lazy = False
+        try:
+            from dask.base import is_dask_collection
+
+            if is_dask_collection(obj) or is_dask_collection(mask):
+                is_lazy = True
+        except ImportError:
+            # Fallback for Dataset/DataArray robustness without dask installed
+            if hasattr(obj, "data_vars"):  # Dataset
+                is_lazy = any(hasattr(v.data, "dask") for v in obj.data_vars.values())
+            elif hasattr(obj, "data"):  # DataArray
+                is_lazy = hasattr(obj.data, "dask")
+
+            if not is_lazy and hasattr(mask, "data"):
+                is_lazy = hasattr(mask.data, "dask")
+
+        res = obj.where(mask, drop=not is_lazy)
 
     # Metadata update
-    msg = f"Spatially sliced to extent {extent} (wrapped={norm_min_x > norm_max_x})"
+    msg = f"Spatially sliced to extent {extent} (wrapped={wrapped})"
     update_history(res, msg)
 
     return res

tests/test_aero_spatial_slice.py

@@ -1,72 +1,104 @@
+from __future__ import annotations
 import numpy as np
-import pytest
-from xregrid.utils import spatial_slice, create_global_grid
-
-
-def test_spatial_slice_basic():
-    """Test standard slicing on a NumPy-backed Dataset."""
-    ds = create_global_grid(1.0, 1.0)
-    ds["data"] = (["lat", "lon"], np.random.rand(180, 360))
-
-    # Slice a region in the middle
-    extent = (10, 30, 10, 30)
-    sliced = spatial_slice(ds, extent)
-
-    assert sliced.lat.min() >= 10
-    assert sliced.lat.max() <= 30
-    assert sliced.lon.min() >= 10
-    assert sliced.lon.max() <= 30
-    assert not hasattr(sliced.data.data, "dask")
-
-
-def test_spatial_slice_dask():
-    """Test slicing on a Dask-backed Dataset verifying laziness."""
-    ds = create_global_grid(1.0, 1.0, chunks={"lat": 90, "lon": 90})
-    ds["data"] = (["lat", "lon"], np.random.rand(180, 360))
-    ds = ds.chunk({"lat": 90, "lon": 90})
-
-    extent = (10, 30, 10, 30)
-    sliced = spatial_slice(ds, extent)
+import xarray as xr
+from xregrid.utils import create_global_grid, create_mesh_from_coords, spatial_slice
+
+
+def test_spatial_slice_rectilinear() -> None:
+    """
+    Test spatial_slice with a rectilinear grid (Eager and Lazy).
+
+    Verifies that spatial_slice correctly subsets a rectilinear grid
+    using coordinate indexes and handles both NumPy and Dask backends.
+    """
+    # 1. Eager
+    ds = create_global_grid(res_lat=1.0, res_lon=1.0)
+    # Extent: (min_x, max_x, min_y, max_y)
+    extent = (10.5, 20.5, 30.5, 40.5)
+    ds_sliced = spatial_slice(ds, extent)
+
+    assert ds_sliced.lat.min() >= 30.5
+    assert ds_sliced.lat.max() <= 40.5
+    assert ds_sliced.lon.min() >= 10.5
+    assert ds_sliced.lon.max() <= 20.5
+    assert "history" in ds_sliced.attrs
+    assert "Spatially sliced" in ds_sliced.attrs["history"]
+
+    # 2. Lazy (Dask)
+    ds_lazy = create_global_grid(
+        res_lat=1.0, res_lon=1.0, chunks={"lat": 10, "lon": 10}
+    )
+    ds_sliced_lazy = spatial_slice(ds_lazy, extent)
+
+    # In xarray, dimension coordinates are often eager (NumPy) due to indexing.
+    # Check lat_b instead, which should remain lazy.
+    assert hasattr(ds_sliced_lazy.lat_b.data, "dask")
+    xr.testing.assert_allclose(ds_sliced, ds_sliced_lazy.compute())
+
+
+def test_spatial_slice_unstructured() -> None:
+    """
+    Test spatial_slice with an unstructured grid (Eager and Lazy).
+
+    Verifies that spatial_slice correctly subsets an unstructured grid
+    using boolean masking and maintains laziness for Dask-backed data.
+    """
+    # Create points: one in the box, one outside
+    lons = np.array([15.0, 50.0])
+    lats = np.array([35.0, 60.0])
+    ds = create_mesh_from_coords(lons, lats, crs="EPSG:4326")
+
+    extent = (10.0, 20.0, 30.0, 40.0)
+
+    # Eager path: drop=True is supported
+    ds_sliced = spatial_slice(ds, extent)
+
+    assert ds_sliced.sizes["n_pts"] == 1
+    assert ds_sliced.lon.values[0] == 15.0
+    assert ds_sliced.lat.values[0] == 35.0
+
+    # Lazy path: drop=False is used to preserve laziness
+    ds_lazy = create_mesh_from_coords(lons, lats, crs="EPSG:4326", chunks=1)
+    ds_lazy["data"] = (["n_pts"], ds_lazy.lat.data, {"units": "K"})
+    ds_sliced_lazy = spatial_slice(ds_lazy, extent)
 
     # Verify laziness
-    assert hasattr(sliced.data.data, "dask")
-
-    # Compute and verify values
-    computed = sliced.compute()
-    assert computed.lat.min() >= 10
-    assert computed.lat.max() <= 30
+    assert hasattr(ds_sliced_lazy.data.data, "dask")
 
+    # Verify results (after compute and dropna since drop=False was used)
+    ds_res = ds_sliced_lazy.compute().dropna("n_pts")
+    assert ds_res.sizes["n_pts"] == 1
+    assert ds_res.lon.values[0] == 15.0
 
-def test_spatial_slice_wrap():
-    """Test longitude wrapping (cross-meridian)."""
-    # Grid [0, 360]
-    ds = create_global_grid(1.0, 1.0)
-    ds["data"] = (["lat", "lon"], np.random.rand(180, 360))
 
-    # Requested extent [-20, 20] which crosses 0/360
-    extent = (-20, 20, -10, 10)
-    sliced = spatial_slice(ds, extent)
+def test_spatial_slice_wrapping() -> None:
+    """
+    Test spatial_slice with longitude wrapping across different grid types.
 
-    # Result should have longitudes around 340-360 and 0-20
-    assert sliced.lon.size > 0
-    assert (sliced.lon >= 340).any()
-    assert (sliced.lon <= 20).any()
-    assert "wrapped=True" in sliced.attrs["history"]
+    Verifies that spatial_slice correctly handles regions crossing the
+    meridian/dateline for both rectilinear and unstructured grids.
+    """
+    ds = create_global_grid(res_lat=1.0, res_lon=1.0)
 
+    # Slice crossing the 0/360 boundary
+    # Extent: (min_x, max_x, min_y, max_y)
+    extent = (-10.5, 10.5, -10.5, 10.5)
+    ds_sliced = spatial_slice(ds, extent)
 
-def test_spatial_slice_crs():
-    """Test CRS-aware slicing."""
-    pytest.importorskip("pyproj")
+    # Lon in ds is 0-360. -10.5 should map to 349.5
+    assert ds_sliced.lon.min() >= 0
+    assert ds_sliced.lon.max() <= 360
 
-    ds = create_global_grid(1.0, 1.0)
-    ds["data"] = (["lat", "lon"], np.random.rand(180, 360))
+    # Check that we have both the 0-10.5 and 349.5-360 parts
+    assert (ds_sliced.lon <= 10.5).any()
+    assert (ds_sliced.lon >= 349.5).any()
 
-    # Extent in Web Mercator (EPSG:3857) around (0,0)
-    # roughly 10 degrees in meters
-    extent_3857 = (-1113194, 1113194, -1113194, 1113194)
-    sliced = spatial_slice(ds, extent_3857, crs="EPSG:3857")
+    # Verify unstructured wrapping
+    lons = np.array([5.0, 355.0, 180.0])
+    lats = np.array([0.0, 0.0, 0.0])
+    ds_unstructured = create_mesh_from_coords(lons, lats, crs="EPSG:4326")
+    ds_un_sliced = spatial_slice(ds_unstructured, extent)
 
-    assert sliced.lat.min() < 0
-    assert sliced.lat.max() > 0
-    assert sliced.lon.min() < 10  # 350 in 0-360
-    assert sliced.lon.max() > 0
+    # Since it's eager, drop=True was used
+    assert ds_un_sliced.sizes["n_pts"] == 2
+    assert set(ds_un_sliced.lon.values) == {5.0, 355.0}