Add cupy.piecewise

cupy/__init__.py

@@ -56,6 +56,7 @@ def is_available():
 import cupy.core.fusion  # NOQA
 from cupy import creation  # NOQA
 from cupy import fft  # NOQA
+from cupy import functional  # NOQA
 from cupy import indexing  # NOQA
 from cupy import io  # NOQA
 from cupy import linalg  # NOQA
@@ -278,6 +279,11 @@ def is_available():
 from cupy.creation.matrix import tril  # NOQA
 from cupy.creation.matrix import triu  # NOQA
 
+# -----------------------------------------------------------------------------
+# Functional routines
+# -----------------------------------------------------------------------------
+from cupy.functional.piecewise import piecewise  # NOQA
+
 # -----------------------------------------------------------------------------
 # Array manipulation routines
 # -----------------------------------------------------------------------------

cupy/functional/__init__.py

@@ -0,0 +1,2 @@
+# "NOQA" to suppress flake8 warning
+from cupy.functional import piecewise  # NOQA

cupy/functional/piecewise.py

@@ -0,0 +1,68 @@
+import cupy
+
+from cupy import core
+
+_piecewise_krnl = core.ElementwiseKernel(
+    'U condlist, raw S funclist, int64 xsiz',
+    'raw I prev, raw T y',
+    '''
+        if (condlist){
+          __syncthreads();
+          if (prev[i % xsiz] < i)
+            prev[i % xsiz] = i;
+          __syncthreads();
+          if (prev[i % xsiz] == i)
+            y[i % xsiz] = funclist[i / xsiz];
+        }
+        ''',
+    'piecewise_kernel'
+)
+
+
+def piecewise(x, condlist, funclist):
+    """Evaluate a piecewise-defined function.
+
+        Args:
+            x (cupy.ndarray): input domain
+            condlist (cupy.ndarray or list):
+                conditions list ( boolean arrays or boolean scalars).
+                Each boolean array/ scalar corresponds to a function
+                in funclist. Length of funclist is equal to that
+                of condlist. If one extra function is given, it is used
+                as otherwise condition
+            funclist (cupy.ndarray or list): list of scalar functions.
+
+        Returns:
+            cupy.ndarray: the result of calling the functions in funclist
+                on portions of x defined by condlist.
+
+        .. warning::
+
+            This function currently doesn't support callable functions,
+            args and kw parameters.
+
+        .. seealso:: :func:`numpy.piecewise`
+        """
+
+    if any(callable(item) for item in funclist):
+        raise NotImplementedError(
+            'Callable functions are not supported currently')
+    if cupy.isscalar(condlist):
+        condlist = cupy.full(shape=x.shape, fill_value=condlist)
+    if not isinstance(condlist[0], (list, cupy.ndarray)) and x.ndim != 0:
+        condlist = [condlist]
+    condlist = cupy.array(condlist, dtype=bool)
+    condlen = len(condlist)
+    funclen = len(funclist)
+    if condlen == funclen:
+        y = cupy.zeros(shape=x.shape, dtype=x.dtype)
+    elif condlen + 1 == funclen:  # o.w
+        y = cupy.full(shape=x.shape, fill_value=funclist[-1], dtype=x.dtype)
+        funclist = funclist[:-1]
+    else:
+        raise ValueError('with {} condition(s), either {} or {} functions'
+                         ' are expected'.format(condlen, condlen, condlen + 1))
+    funclist = cupy.asarray(funclist)
+    prev = cupy.full(shape=x.shape, fill_value=-1, dtype=int)
+    _piecewise_krnl(condlist, funclist, x.size, prev, y)
+    return y

docs/source/reference/functional.rst

@@ -0,0 +1,10 @@
+Functional programming
+======================
+
+.. https://docs.scipy.org/doc/numpy/reference/routines.functional.html
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   cupy.piecewise

docs/source/reference/routines.rst

@@ -16,6 +16,7 @@ These functions cover a subset of
    binary
    dtype
    fft
+   functional
    indexing
    io
    linalg

tests/cupy_tests/functional_tests/test_piecewise.py

@@ -0,0 +1,81 @@
+import unittest
+
+import pytest
+
+import cupy
+from cupy import testing
+
+
+class TestPiecewise(unittest.TestCase):
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_piecewise(self, xp, dtype):
+        x = xp.linspace(2.5, 12.5, 6, dtype=dtype)
+        condlist = [x < 0, x >= 0, x < 5, x >= 1.5]
+        funclist = [-1, 1, 2, 5]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_scalar_input(self, xp, dtype):
+        x = dtype(2)
+        condlist = [x < 0, x >= 0]
+        funclist = [-10, 10]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_scalar_condition(self, xp, dtype):
+        x = testing.shaped_random(shape=(2, 3, 5), xp=xp, dtype=dtype)
+        condlist = True
+        funclist = [-10, 10]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_signed_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_otherwise_condition1(self, xp, dtype):
+        x = xp.linspace(-2, 20, 12, dtype=dtype)
+        condlist = [x > 15, x <= 5, x == 0, x == 10]
+        funclist = [-1, 0, 2, 3, -5]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_otherwise_condition2(self, xp, dtype):
+        x = cupy.array([-10, 20, 30, 40], dtype=dtype)
+        condlist = [[True, False, False, True], [True, False, False, True]]
+        funclist = [-1, 1, 2]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_piecewise_ndim(self, xp, dtype):
+        x = testing.shaped_random(shape=(2, 3, 5), xp=xp, dtype=dtype)
+        condlist = [x < 0, x > 0]
+        funclist = [-1, 1, 2]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    def test_mismatched_lengths(self, dtype):
+        x = cupy.linspace(-2, 4, 6, dtype=dtype)
+        condlist = [x < 0, x >= 0]
+        funclist = [-1, 0, 2, 4, 5]
+        with pytest.raises(ValueError):
+            cupy.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    def test_callable_funclist(self, dtype):
+        x = cupy.linspace(-2, 4, 6, dtype=dtype)
+        condlist = [x < 0, x > 0]
+        funclist = [lambda x: -x, lambda x: x]
+        with pytest.raises(NotImplementedError):
+            cupy.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes()
+    def test_mixed_funclist(self, dtype):
+        x = cupy.linspace(-2, 2, 6, dtype=dtype)
+        condlist = [x < 0, x == 0, x > 0]
+        funclist = [-10, lambda x: -x, 10, lambda x: x]
+        with pytest.raises(NotImplementedError):
+            cupy.piecewise(x, condlist, funclist)