Make geodesic slope memory guard backend-aware for dask

xrspatial/geodesic.py

@@ -82,6 +82,35 @@ def _check_geodesic_memory(rows, cols, func_name):
         )
 
 
+def _check_geodesic_memory_backend_aware(agg, func_name):
+    """Backend-aware front end for :func:`_check_geodesic_memory`.
+
+    Eager NumPy / CuPy backends build a ``(3, H, W)`` float64 stack for the
+    whole raster up front, so the guard sizes against the full ``(rows, cols)``.
+
+    Dask backends run the same stack chunk by chunk via ``map_overlap``, so
+    peak working memory is bounded by the largest chunk plus its 1-cell halo,
+    not the full raster. Sizing the guard against the full raster would reject
+    large-but-chunked rasters that would otherwise stream through fine, so for
+    Dask we check the largest chunk instead.
+    """
+    data = getattr(agg, "data", agg)
+
+    chunks = getattr(data, "chunks", None)
+    if chunks is None:
+        # Eager (NumPy / CuPy): size against the full raster.
+        rows, cols = agg.shape[-2], agg.shape[-1]
+        _check_geodesic_memory(rows, cols, func_name=func_name)
+        return
+
+    # Dask: size against the largest spatial chunk, plus the 1-cell halo that
+    # ``map_overlap`` adds on each side (depth (0, 1, 1) -> +2 per spatial dim).
+    row_chunks, col_chunks = chunks[-2], chunks[-1]
+    rows = (max(row_chunks) if row_chunks else 0) + 2
+    cols = (max(col_chunks) if col_chunks else 0) + 2
+    _check_geodesic_memory(rows, cols, func_name=func_name)
+
+
 # =====================================================================
 # CPU (Numba ngjit) primitives
 # =====================================================================

xrspatial/slope.py

@@ -23,7 +23,7 @@ class cupy(object):
 
 # local modules
 from xrspatial.dataset_support import supports_dataset
-from xrspatial.geodesic import (INV_2R, WGS84_A2, WGS84_B2, _check_geodesic_memory,
+from xrspatial.geodesic import (INV_2R, WGS84_A2, WGS84_B2, _check_geodesic_memory_backend_aware,
                                 _cpu_geodesic_slope, _run_gpu_geodesic_slope)
 from xrspatial.utils import (Z_UNITS, ArrayTypeFunctionMapping, _boundary_to_dask,
                              _extract_latlon_coords, _pad_array, _validate_boundary,
@@ -393,8 +393,7 @@ def slope(agg: xr.DataArray,
             )
         z_factor = Z_UNITS[z_unit]
 
-        rows, cols = agg.shape[-2], agg.shape[-1]
-        _check_geodesic_memory(rows, cols, func_name='slope')
+        _check_geodesic_memory_backend_aware(agg, func_name='slope')
 
         lat_2d, lon_2d = _extract_latlon_coords(agg)
 

xrspatial/tests/test_geodesic_slope.py

@@ -273,6 +273,47 @@ def test_planar_method_skips_guard(self, monkeypatch):
         assert result.shape == (8, 8)
 
 
+@dask_array_available
+class TestGeodesicSlopeMemoryGuardDask:
+    """The dask geodesic backend streams the raster chunk by chunk, so the
+    memory guard must size against the largest chunk, not the full raster.
+    A raster that would be rejected eagerly should be allowed once it is
+    chunked small enough to fit."""
+
+    def test_chunked_raster_allowed_when_eager_would_reject(self, monkeypatch):
+        # 1 MB available. The 200x200 raster needs ~2.2 MB eagerly (56 B/cell)
+        # and trips the eager guard, but 20x20 chunks need only ~25 KB each.
+        monkeypatch.setattr(
+            'xrspatial.geodesic._available_memory_bytes', lambda: 1024 * 1024
+        )
+        elev = _flat_surface(H=200, W=200)
+        # eager numpy of the same size is rejected (sanity check the budget).
+        r_np = _make_geo_raster(elev, 40.0, 41.0, 10.0, 11.0)
+        with pytest.raises(MemoryError, match="slope"):
+            slope(r_np, method='geodesic')
+        # same raster, chunked small — guard must let it through.
+        r_da = _make_geo_raster(
+            elev, 40.0, 41.0, 10.0, 11.0,
+            backend='dask+numpy', chunks=(20, 20),
+        )
+        result = slope(r_da, method='geodesic')
+        assert result.compute().shape == (200, 200)
+
+    def test_single_huge_chunk_still_rejected(self, monkeypatch):
+        # A dask array whose only chunk spans the whole raster has no memory
+        # advantage over eager, so the guard must still reject it.
+        monkeypatch.setattr(
+            'xrspatial.geodesic._available_memory_bytes', lambda: 1024 * 1024
+        )
+        elev = _flat_surface(H=200, W=200)
+        r_da = _make_geo_raster(
+            elev, 40.0, 41.0, 10.0, 11.0,
+            backend='dask+numpy', chunks=(200, 200),
+        )
+        with pytest.raises(MemoryError, match="slope"):
+            slope(r_da, method='geodesic')
+
+
 # ---------------------------------------------------------------------------
 # Tests — cross-backend consistency
 # ---------------------------------------------------------------------------