Keep standardize piecewise/categorical on the device for cupy inputs

xrspatial/mcda/standardize.py

@@ -267,16 +267,24 @@ def _piecewise(data, *, breakpoints, values):
     if _is_dask(data):
         import dask.array as da
 
+        # Convert the lookup tables once, to the module of the chunks
+        # (known from the dask meta), instead of once per block. cupy
+        # chunks then interpolate on the device with device-resident
+        # tables; np.asarray on a cupy block would raise TypeError.
+        xp_b = _get_xp(data._meta)
+        bp_t = xp_b.asarray(bp)
+        vl_t = xp_b.asarray(vl)
+
         def _interp_block(block):
-            # Ensure block is numpy for np.interp (handles cupy chunks)
-            arr = np.asarray(block)
-            return np.interp(arr, bp, vl)
+            return xp_b.interp(block, bp_t, vl_t)
 
         result = da.map_blocks(_interp_block, data, dtype=np.float64)
         result = da.where(da.isfinite(data), result, np.nan)
         return result
 
-    result = xp.interp(data, bp, vl)
+    # cupy.interp rejects numpy operands, so move the lookup tables to
+    # the data's array module first (a no-op copy for numpy).
+    result = xp.interp(data, xp.asarray(bp), xp.asarray(vl))
     result = xp.where(xp.isfinite(data), result, xp.nan)
     return result
 
@@ -291,12 +299,15 @@ def _categorical(data, *, mapping):
     if _is_dask(data):
         import dask.array as da
 
+        # Build each block's output with the chunks' own array module
+        # (known from the dask meta) so cupy blocks stay on the device;
+        # np.asarray on a cupy block would raise TypeError.
+        xp_b = _get_xp(data._meta)
+
         def _apply_mapping(block):
-            # Ensure block is numpy (handles cupy chunks)
-            arr = np.asarray(block)
-            out = np.full(arr.shape, np.nan, dtype=np.float64)
+            out = xp_b.full(block.shape, np.nan, dtype=np.float64)
             for k, v in zip(keys, vals):
-                out[arr == k] = v
+                out[block == k] = v
             return out
 
         result = da.map_blocks(_apply_mapping, data, dtype=np.float64)

xrspatial/tests/test_mcda.py

@@ -6,6 +6,8 @@
 import pytest
 import xarray as xr
 
+from xrspatial.tests.general_checks import cuda_and_cupy_available
+
 from xrspatial.mcda import (
     ahp_weights,
     boolean_overlay,
@@ -389,6 +391,74 @@ def test_piecewise_dask(self):
         assert float(computed.values[0, 1]) == pytest.approx(1.0)
 
 
+class TestStandardizeCupy:
+    """GPU paths for standardize (#3151).
+
+    piecewise on cupy used numpy lookup tables (cupy.interp rejects
+    them), and the piecewise/categorical dask chunk functions called
+    np.asarray on cupy blocks (implicit conversion raises TypeError).
+    """
+
+    @pytest.fixture
+    def raw(self):
+        return np.array([
+            [0.0, 25.0, 50.0, 100.0],
+            [75.0, np.nan, 10.0, 90.0],
+        ], dtype=np.float64)
+
+    piecewise_kw = dict(
+        method="piecewise", breakpoints=[0, 50, 100], values=[0.0, 1.0, 0.5],
+    )
+    categorical_kw = dict(
+        method="categorical", mapping={0: 0.1, 50: 0.5, 100: 0.9},
+    )
+
+    @cuda_and_cupy_available
+    def test_piecewise_cupy(self, raw):
+        import cupy
+        ref = standardize(xr.DataArray(raw, dims=["y", "x"]),
+                          **self.piecewise_kw)
+        agg = xr.DataArray(cupy.asarray(raw), dims=["y", "x"])
+        result = standardize(agg, **self.piecewise_kw)
+        # Result stays on the device
+        assert isinstance(result.data, cupy.ndarray)
+        np.testing.assert_allclose(
+            result.data.get(), ref.values, equal_nan=True,
+        )
+
+    @cuda_and_cupy_available
+    @pytest.mark.skipif(not HAS_DASK, reason="Requires dask")
+    def test_piecewise_dask_cupy(self, raw):
+        import cupy
+        ref = standardize(xr.DataArray(raw, dims=["y", "x"]),
+                          **self.piecewise_kw)
+        agg = xr.DataArray(
+            da.from_array(cupy.asarray(raw), chunks=(1, 3)), dims=["y", "x"],
+        )
+        result = standardize(agg, **self.piecewise_kw)
+        computed = result.data.compute()
+        assert isinstance(computed, cupy.ndarray)
+        np.testing.assert_allclose(
+            computed.get(), ref.values, equal_nan=True,
+        )
+
+    @cuda_and_cupy_available
+    @pytest.mark.skipif(not HAS_DASK, reason="Requires dask")
+    def test_categorical_dask_cupy(self, raw):
+        import cupy
+        ref = standardize(xr.DataArray(raw, dims=["y", "x"]),
+                          **self.categorical_kw)
+        agg = xr.DataArray(
+            da.from_array(cupy.asarray(raw), chunks=(1, 3)), dims=["y", "x"],
+        )
+        result = standardize(agg, **self.categorical_kw)
+        computed = result.data.compute()
+        assert isinstance(computed, cupy.ndarray)
+        np.testing.assert_allclose(
+            computed.get(), ref.values, equal_nan=True,
+        )
+
+
 # ===========================================================================
 # weights
 # ===========================================================================