fixed cupy failures in balanced_allocation and corridor (#985)

Author: brendancol

xrspatial/balanced_allocation.py

@@ -44,6 +44,13 @@ def _to_numpy(arr):
     return np.asarray(arr)
 
 
+def _as_numpy(arr):
+    """Convert a computed array (numpy or cupy) to numpy."""
+    if hasattr(arr, 'get'):
+        return arr.get()
+    return np.asarray(arr)
+
+
 def _extract_sources(raster, target_values):
     """Return sorted array of unique source IDs from the raster."""
     data = _to_numpy(raster.data)
@@ -116,7 +123,7 @@ def _allocate_from_costs(cost_stack, source_ids, fill_value=np.nan):
         # Replace NaN with inf for argmin
         stacked_clean = da.where(da.isnan(stacked), np.inf, stacked)
         best_idx = da.argmin(stacked_clean, axis=0).compute()
-        best_idx = np.asarray(best_idx)
+        best_idx = _as_numpy(best_idx)
     elif hasattr(first, 'get'):  # cupy
         import cupy as cp
         stacked = cp.stack(cost_stack, axis=0)
@@ -133,7 +140,7 @@ def _allocate_from_costs(cost_stack, source_ids, fill_value=np.nan):
     # Mark cells that are unreachable from all sources
     if da is not None and isinstance(first, da.Array):
         all_nan = da.all(da.isnan(da.stack(cost_stack, axis=0)), axis=0)
-        all_nan = np.asarray(all_nan.compute())
+        all_nan = _as_numpy(all_nan.compute())
     elif hasattr(first, 'get'):
         import cupy as cp
         all_nan = cp.asnumpy(
@@ -164,7 +171,7 @@ def _allocate_biased(cost_stack, biases, source_ids, fill_value=np.nan):
             layers.append(layer)
         stacked = da.stack(layers, axis=0)
         best_idx = da.argmin(stacked, axis=0).compute()
-        best_idx = np.asarray(best_idx)
+        best_idx = _as_numpy(best_idx)
     elif hasattr(first, 'get'):
         import cupy as cp
         layers = []
@@ -188,7 +195,7 @@ def _allocate_biased(cost_stack, biases, source_ids, fill_value=np.nan):
     # Mark unreachable cells
     if da is not None and isinstance(first, da.Array):
         all_nan = da.all(da.isnan(da.stack(cost_stack, axis=0)), axis=0)
-        all_nan = np.asarray(all_nan.compute())
+        all_nan = _as_numpy(all_nan.compute())
     elif hasattr(first, 'get'):
         import cupy as cp
         all_nan = cp.asnumpy(

xrspatial/corridor.py

@@ -28,6 +28,16 @@
 from xrspatial.utils import _validate_raster
 
 
+def _scalar_to_float(da_scalar):
+    """Extract a Python float from a scalar DataArray (numpy/cupy/dask)."""
+    data = da_scalar.data
+    if hasattr(data, 'compute'):
+        data = data.compute()
+    if hasattr(data, 'get'):
+        data = data.get()
+    return float(data)
+
+
 def _compute_corridor(
     cd_a: xr.DataArray,
     cd_b: xr.DataArray,
@@ -36,7 +46,7 @@ def _compute_corridor(
 ) -> xr.DataArray:
     """Sum two cost-distance surfaces, normalize, and optionally threshold."""
     corridor = cd_a + cd_b
-    corridor_min = float(corridor.min())
+    corridor_min = _scalar_to_float(corridor.min())
 
     if not np.isfinite(corridor_min):
         # Sources are mutually unreachable -- return all-NaN.
@@ -49,7 +59,20 @@ def _compute_corridor(
             cutoff = threshold * corridor_min
         else:
             cutoff = threshold
-        normalized = normalized.where(normalized <= cutoff)
+        data = normalized.data
+        try:
+            import dask.array as _da
+            if isinstance(data, _da.Array):
+                data = _da.where(data <= cutoff, data, np.nan)
+            else:
+                raise ImportError
+        except ImportError:
+            if hasattr(data, 'get'):  # cupy
+                import cupy as cp
+                data = cp.where(data <= cutoff, data, cp.nan)
+            else:
+                data = np.where(data <= cutoff, data, np.nan)
+        normalized = normalized.copy(data=data)
 
     return normalized