Add cupy.piecewise

cupy/__init__.py

@@ -7,11 +7,9 @@
 from cupy import _environment
 from cupy import _version
 
-
 if sys.platform.startswith('win32') and (3, 8) <= sys.version_info:  # NOQA
     _environment._setup_win32_dll_directory()  # NOQA
 
-
 try:
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', category=ImportWarning,
@@ -39,7 +37,6 @@
 
     raise ImportError(msg) from e
 
-
 from cupy import cuda
 # Do not make `cupy.cupyx` available because it is confusing.
 import cupyx as _cupyx
@@ -51,11 +48,11 @@ def is_available():
 
 __version__ = _version.__version__
 
-
 from cupy import binary  # NOQA
 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
@@ -69,12 +66,10 @@ def is_available():
 from cupy import util  # NOQA
 from cupy import lib  # NOQA
 
-
 # import class and function
 from cupy.core import ndarray  # NOQA
 from cupy.core import ufunc  # NOQA
 
-
 # =============================================================================
 # Constants (borrowed from NumPy)
 # =============================================================================
@@ -94,7 +89,6 @@ def is_available():
 from numpy import PINF  # NOQA
 from numpy import PZERO  # NOQA
 
-
 # =============================================================================
 # Data types (borrowed from NumPy)
 #
@@ -643,13 +637,16 @@ def isscalar(element):
 from cupy.misc import may_share_memory  # NOQA
 from cupy.misc import shares_memory  # NOQA
 
+# -----------------------------------------------------------------------------
+# Functional routines
+# -----------------------------------------------------------------------------
+from cupy.functional.piecewise import piecewise  # NOQA
 
 # -----------------------------------------------------------------------------
 # Padding
 # -----------------------------------------------------------------------------
 pad = padding.pad.pad
 
-
 # -----------------------------------------------------------------------------
 # Sorting, searching, and counting
 # -----------------------------------------------------------------------------
@@ -784,7 +781,6 @@ def get_array_module(*args):
 
 disable_experimental_feature_warning = False
 
-
 # set default allocator
 _default_memory_pool = cuda.MemoryPool()
 _default_pinned_memory_pool = cuda.PinnedMemoryPool()

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,76 @@
+import cupy
+
+from cupy import core
+from cupy.core.core import ndarray
+from cupy.core._reduction import ReductionKernel
+
+_piecewise_krnl = ReductionKernel(
+    'S x1, T x2',
+    'U y',
+    'x1 ? x2 : NULL',
+    'b == NULL? a : b',
+    'y = a',
+    'NULL',
+    'piecewise'
+)
+
+
+def piecewise(x, condlist, funclist):
+    """
+    Evaluate a piecewise-defined function.
+
+    Args:
+    :param x: input domain
+    :param condlist: conditions list ( boolean arrays or boolean scalars)
+                     Each boolean array/ scalar corresponds to a function
+                     in funclist. Length of functionlist is equal to that
+                     of condlist. If one extra function is given, it is used
+                     as otherwise condition
+    :param funclist: 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 are not supported
+
+    .. seealso:: :func:`numpy.piecewise`
+    """
+
+    if any(callable(item) for item in funclist):
+        raise NotImplementedError(
+            'Callable functions are not supported currently')
+    if cupy.isscalar(x):
+        x = cupy.asarray(x)
+    scalar = 0
+    if cupy.isscalar(condlist):
+        scalar = 1
+    if cupy.isscalar(condlist) or (
+            (not isinstance(condlist[0], (list, ndarray))) and x.ndim != 0):
+        condlist = [condlist]
+    condlist = cupy.array(condlist, dtype=bool)
+    condlen = len(condlist)
+    funclen = len(funclist)
+    if condlen + 1 == funclen:  # o.w
+        condelse = ~ cupy.any(condlist, axis=0, keepdims=True)
+        condlist = core.concatenate_method([condlist, condelse], 0)
+        condlen += 1
+    elif condlen != funclen:
+        raise ValueError('with {} condition(s), either {} or {} functions'
+                         ' are expected'.format(condlen, condlen, condlen + 1))
+
+    funclist = cupy.asarray(funclist)
+    if scalar:
+        funclist = funclist[condlist]
+        condlist = cupy.ones(shape=(1, x.size), dtype=bool)
+
+    y = cupy.zeros(x.shape, x.dtype)
+    if not x.ndim:
+        _piecewise_krnl(condlist, funclist, y)
+    else:
+        for i in range(x.size):
+            _piecewise_krnl(condlist[:, i], funclist, y[i])
+    return y

tests/cupy_tests/functional_tests/test_piecewise.py

@@ -0,0 +1,73 @@
+import unittest
+
+import pytest
+
+import cupy
+from cupy import testing
+
+
+class TestPiecewise(unittest.TestCase):
+
+    @testing.for_all_dtypes(no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_linespace(self, xp, dtype):
+        x = xp.linspace(-2.5, 2.5, 6, dtype=dtype)
+        condlist = [x < 0, x >= 0]
+        funclist = [-1, 1]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes(no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_scalar_value(self, xp, dtype):
+        x = dtype(2)
+        condlist = [x < 0, x >= 0]
+        funclist = [-10, 10]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes(no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_scalar_condition(self, xp, dtype):
+        x = xp.linspace(-2.5, 2.5, 4, dtype=dtype)
+        condlist = True
+        funclist = [-10, 10]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes(no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_otherwise_condition1(self, xp, dtype):
+        x = xp.linspace(-2, 4, 4, dtype=dtype)
+        condlist = [x < 0, x >= 0]
+        funclist = [-1, 0, 2]
+        return xp.piecewise(x, condlist, funclist)
+
+    @testing.for_all_dtypes(no_bool=True, no_complex=True)
+    @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(no_bool=True, no_complex=True)
+    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(no_bool=True, no_complex=True)
+    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(no_bool=True, no_complex=True)
+    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)