Fix conversion from cupy in categorical rescale_discrete_levels (#1179)

* Fix conversion from cupy in categorical rescale_discrete_levels

* Fix typo

datashader/tests/test_transfer_functions.py

@@ -9,7 +9,7 @@
 import pytest
 from collections import OrderedDict
 import datashader.transfer_functions as tf
-from datashader.tests.test_pandas import assert_eq_xr
+from datashader.tests.test_pandas import assert_eq_ndarray, assert_eq_xr
 
 coords = OrderedDict([('x_axis', [3, 4, 5]), ('y_axis', [0, 1, 2])])
 dims = ['y_axis', 'x_axis']
@@ -47,10 +47,12 @@ def create_dask_array_np(*args, **kwargs):
     import cupy
     aggs = [build_agg(np), build_agg(cupy), build_agg_dask()]
     arrays = [np.array, cupy.array, create_dask_array_np]
+    array_modules = [np, cupy]
 except ImportError:
     cupy = None
     aggs = [build_agg(np), build_agg_dask()]
     arrays = [np.array, create_dask_array_np]
+    array_modules = [np]
 
 int_span = [11, 17]
 float_span = [11.0, 17.0]
@@ -515,6 +517,22 @@ def test_shade_rescale_discrete_levels(agg, attr, rescale):
     assert_eq_xr(img, sol)
 
 
+@pytest.mark.parametrize('array_module', array_modules)
+def test_shade_rescale_discrete_levels_categorical(array_module):
+    arr = array_module.array([[[1, 2], [0, 1]],
+                              [[0, 0], [0, 0]],
+                              [[1, 0], [3, 0]],
+                              [[1, 0], [2, 1]]], dtype='u4')
+    agg = xr.DataArray(data=arr, coords=dict(y=[0, 1, 2, 3], x=[0, 1], cat=['a', 'b']))
+    img = tf.shade(agg, how='eq_hist', rescale_discrete_levels=True)
+
+    sol = np.array([[0xff845c70, 0x6fb87e37],
+                    [0x006a4c8d, 0x006a4c8d],
+                    [0x6f1c1ae4, 0xff1c1ae4],
+                    [0x6f1c1ae4, 0xff503baa]])
+    assert_eq_ndarray(img.data, sol)
+
+
 empty_arrays = [
     np.zeros((2, 2, 2), dtype=np.uint32),
     np.full((2, 2, 2), np.nan, dtype=np.float64),
@@ -1160,7 +1178,8 @@ def test_shade_with_discrete_color_key():
         cupy = None
 
 
-def test_interpolate_alpha_discrete_levels_None():
-    data = np.array([[0.0, 1.0], [1.0, 0.0]])
+@pytest.mark.parametrize('array_module', array_modules)
+def test_interpolate_alpha_discrete_levels_None(array_module):
+    data = array_module.array([[0.0, 1.0], [1.0, 0.0]])
     # Issue #1084: this raises a ValueError.
     tf._interpolate_alpha(data, data, None, "eq_hist", 0.5, None, 0.4, True)

datashader/transfer_functions/__init__.py

@@ -435,11 +435,11 @@ def _interpolate_alpha(data, total, mask, how, alpha, span, min_alpha, rescale_d
 
     if cupy and isinstance(data, cupy.ndarray):
         from ._cuda_utils import interp, masked_clip_2d
-        array = cupy.array
+        array_module = cupy
     else:
         from ._cpu_utils import masked_clip_2d
         interp = np.interp
-        array = np.array
+        array_module = np
 
     # if span is provided, use it, otherwise produce a span based off the
     # min/max of the data
@@ -488,8 +488,13 @@ def _interpolate_alpha(data, total, mask, how, alpha, span, min_alpha, rescale_d
         if isinstance(norm_span, (list, tuple)):
             norm_span = norm_span[0]  # Ignore discrete_levels
 
+    # Issue 1178. Convert norm_span from 2-tuple to numpy/cupy array.
+    # array_module.hstack() tolerates tuple of one float and one cupy array,
+    # whereas array_module.array() does not.
+    norm_span = array_module.hstack(norm_span)
+
     # Interpolate the alpha values
-    a = interp(a_scaled, array(norm_span), array([min_alpha, alpha]),
+    a = interp(a_scaled, norm_span, array_module.array([min_alpha, alpha]),
                left=0, right=255).astype(np.uint8)
     return a