From 468d05c2610c4668d54d81aaa8a28777e0c289b4 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 10 Jun 2026 05:52:18 -0700 Subject: [PATCH 1/2] Keep standardize piecewise/categorical on the device for cupy inputs (#3151) --- xrspatial/mcda/standardize.py | 20 ++++++---- xrspatial/tests/test_mcda.py | 70 +++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 8 deletions(-) diff --git a/xrspatial/mcda/standardize.py b/xrspatial/mcda/standardize.py index d33b1fd7c..1fd82ec9c 100644 --- a/xrspatial/mcda/standardize.py +++ b/xrspatial/mcda/standardize.py @@ -268,15 +268,18 @@ def _piecewise(data, *, breakpoints, values): import dask.array as da def _interp_block(block): - # Ensure block is numpy for np.interp (handles cupy chunks) - arr = np.asarray(block) - return np.interp(arr, bp, vl) + # Use the block's own array module so cupy chunks stay on + # the device (np.asarray on a cupy block raises TypeError). + xp_b = _get_xp(block) + return xp_b.interp(block, xp_b.asarray(bp), xp_b.asarray(vl)) 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 @@ -292,11 +295,12 @@ def _categorical(data, *, mapping): import dask.array as da 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) + # Use the block's own array module so cupy chunks stay on + # the device (np.asarray on a cupy block raises TypeError). + xp_b = _get_xp(block) + 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) diff --git a/xrspatial/tests/test_mcda.py b/xrspatial/tests/test_mcda.py index af7f84198..55832fa8d 100644 --- a/xrspatial/tests/test_mcda.py +++ b/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, 2)), 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, 2)), 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 # =========================================================================== From ed2c3c434572324be38dbd36bdca9ae4322a561b Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 10 Jun 2026 06:03:18 -0700 Subject: [PATCH 2/2] Address review: hoist lookup-table conversion out of chunk functions, ragged GPU chunks (#3151) --- xrspatial/mcda/standardize.py | 21 ++++++++++++++------- xrspatial/tests/test_mcda.py | 4 ++-- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/xrspatial/mcda/standardize.py b/xrspatial/mcda/standardize.py index 1fd82ec9c..e74e54f93 100644 --- a/xrspatial/mcda/standardize.py +++ b/xrspatial/mcda/standardize.py @@ -267,11 +267,16 @@ 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): - # Use the block's own array module so cupy chunks stay on - # the device (np.asarray on a cupy block raises TypeError). - xp_b = _get_xp(block) - return xp_b.interp(block, xp_b.asarray(bp), xp_b.asarray(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) @@ -294,10 +299,12 @@ 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): - # Use the block's own array module so cupy chunks stay on - # the device (np.asarray on a cupy block raises TypeError). - xp_b = _get_xp(block) out = xp_b.full(block.shape, np.nan, dtype=np.float64) for k, v in zip(keys, vals): out[block == k] = v diff --git a/xrspatial/tests/test_mcda.py b/xrspatial/tests/test_mcda.py index 55832fa8d..e9349dcad 100644 --- a/xrspatial/tests/test_mcda.py +++ b/xrspatial/tests/test_mcda.py @@ -433,7 +433,7 @@ def test_piecewise_dask_cupy(self, raw): ref = standardize(xr.DataArray(raw, dims=["y", "x"]), **self.piecewise_kw) agg = xr.DataArray( - da.from_array(cupy.asarray(raw), chunks=(1, 2)), dims=["y", "x"], + da.from_array(cupy.asarray(raw), chunks=(1, 3)), dims=["y", "x"], ) result = standardize(agg, **self.piecewise_kw) computed = result.data.compute() @@ -449,7 +449,7 @@ def test_categorical_dask_cupy(self, raw): ref = standardize(xr.DataArray(raw, dims=["y", "x"]), **self.categorical_kw) agg = xr.DataArray( - da.from_array(cupy.asarray(raw), chunks=(1, 2)), dims=["y", "x"], + da.from_array(cupy.asarray(raw), chunks=(1, 3)), dims=["y", "x"], ) result = standardize(agg, **self.categorical_kw) computed = result.data.compute()