Skip to content

balanced_allocation: TypeError on dask+cupy backend from cupy source_ids #1568

@brendancol

Description

@brendancol

Describe the bug

balanced_allocation raises TypeError: Implicit conversion to a NumPy array is not allowed when called on a dask+cupy raster.

File "xrspatial/balanced_allocation.py", line 359, in balanced_allocation
    float(np.sum(fric_weight[alloc == sid]))
                 ^^^^^^^^^^^^^^^^^^^^^^^^^
cupy/_core/core.pyx:1581: TypeError

Root cause: _extract_sources calls da.unique(data).compute(), which on a dask+cupy array returns a cupy array. np.sort(uniq[mask]) does not convert, so source_ids stays cupy. Downstream, source_ids[best_idx] in _allocate_biased produces a cupy alloc, and the masked fric_weight[alloc == sid] tries an implicit numpy conversion.

Expected behavior

balanced_allocation should return a valid allocation raster for all four backends (numpy, cupy, dask+numpy, dask+cupy), matching the test test_two_sources_uniform_friction.

Reproduction

pytest xrspatial/tests/test_balanced_allocation.py::test_two_sources_uniform_friction

The [dask+cupy] parametrization fails; the other three pass.

Fix

Force source_ids to numpy at the end of _extract_sources so all downstream indexing produces numpy arrays.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workinggpuCuPy / CUDA GPU supportproximity toolsProximity, allocation, direction, cost distance

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions