Repository: numba/numba

Issue:
Support for numpy.select in nopython mode
It would be nice to support `numpy.select()` in nopython mode for piecewise-defined array computations.

``` python
def piece(x):
    condlist = np.array(x<3, x>5)
    choicelist = np.array(x, x**2)
    return np.select(condlist, choicelist)

result = piece(np.arange(10))
np.allclose(result, np.array([ 0,  1,  2,  0,  0,  0, 36, 49, 64, 81]))
```



Discussion:
Looks doable, though we would need to pass tuples instead of lists to the select function.

(Side note: `np.array(x<3, x>5)` doesn't seem to work on Numpy 1.9.  Is this a typo, or does this work on older versions of Numpy?)

My mistake: The arguments to `np.array` should be wrapped in brackets or parentheses.

I can adapt to the use of tuples versus lists for my use case.

As per https://github.com/sympy/sympy/issues/9747, apparently old versions of numpy require lists here, not tuples, so numba will need to support them. 

@stuartarchibald, I can pick it up if no one has yet (apparently not as I cannot find any PR referencing it). 
@luk-f-a sounds good, thanks v. much.

Resolving patch:
diff --git a/numba/targets/arraymath.py b/numba/targets/arraymath.py
--- a/numba/targets/arraymath.py
+++ b/numba/targets/arraymath.py
@@ -3518,6 +3518,7 @@ def impl(a, dtype=None):
 
     return impl
 
+
 @overload(np.extract)
 def np_extract(condition, arr):
 
@@ -3543,6 +3544,56 @@ def np_extract_impl(condition, arr):
 
     return np_extract_impl
 
+
+@overload(np.select)
+def np_select(condlist, choicelist, default=0):
+
+    def np_select_arr_impl(condlist, choicelist, default=0):
+        if len(condlist) != len(choicelist):
+            raise ValueError('list of cases must be same length as list of conditions')
+        out = default * np.ones(choicelist[0].shape, choicelist[0].dtype)
+        # should use reversed+zip, but reversed is not available
+        for i in range(len(condlist) - 1, -1, -1):
+            cond = condlist[i]
+            choice = choicelist[i]
+            out = np.where(cond, choice, out)
+        return out
+
+    # first we check the types of the input parameters
+    if not isinstance(condlist, (types.List, types.UniTuple)):
+        raise TypeError('condlist must be a List or a Tuple')
+    if not isinstance(choicelist, (types.List, types.UniTuple)):
+        raise TypeError('choicelist must be a List or a Tuple')
+    if not isinstance(default, (types.Number, types.Boolean)):
+        raise TypeError('default must be a scalar (number or boolean)')
+    # the types of the parameters have been checked, now we test the types
+    # of the content of the parameters
+    # implementation note: if in the future numba's np.where accepts tuples
+    # as elements of condlist, then the check below should be extended to
+    # accept tuples
+    if not isinstance(condlist[0], types.Array):
+        raise TypeError('items of condlist must be arrays')
+    if not isinstance(choicelist[0], types.Array):
+        raise TypeError('items of choicelist must be arrays')
+    # the types of the parameters and their contents have been checked,
+    # now we test the dtypes of the content of parameters
+    if isinstance(condlist[0], types.Array):
+        if not isinstance(condlist[0].dtype, types.Boolean):
+            raise TypeError('condlist arrays must contain booleans')
+    if isinstance(condlist[0], types.UniTuple):
+        if not (isinstance(condlist[0], types.UniTuple)
+                and isinstance(condlist[0][0], types.Boolean)):
+            raise TypeError('condlist tuples must only contain booleans')
+    # the input types are correct, now we perform checks on the dimensions
+    if isinstance(condlist[0], types.Array) and condlist[0].ndim != choicelist[0].ndim:
+        raise TypeError('condlist and choicelist elements must have the '
+                        'same number of dimensions')
+    if isinstance(condlist[0], types.Array) and condlist[0].ndim < 1:
+        raise TypeError('condlist arrays must be of at least dimension 1')
+
+    return np_select_arr_impl
+
+
 #----------------------------------------------------------------------------
 # Windowing functions
 #   - translated from the numpy implementations found in:
@@ -3713,3 +3764,4 @@ def np_kaiser_impl(M, beta):
         return _i0n(n, alpha, beta)
 
     return np_kaiser_impl
+


Test patch:
diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
--- a/numba/tests/test_np_functions.py
+++ b/numba/tests/test_np_functions.py
@@ -193,6 +193,10 @@ def array_repeat(a, repeats):
     return np.asarray(a).repeat(repeats)
 
 
+def np_select(condlist, choicelist, default=0):
+    return np.select(condlist, choicelist, default=default)
+
+
 def np_bartlett(M):
     return np.bartlett(M)
 
@@ -2768,6 +2772,105 @@ def test_repeat_exception(self):
                 with self.assertRaises(TypingError):
                     nbfunc(np.ones(1), rep)
 
+    def test_select(self):
+        np_pyfunc = np_select
+        np_nbfunc = njit(np_select)
+        test_cases = [
+            # Each test case below is one tuple.
+            # Each tuple is separated by a description of what's being tested
+
+            # test with arrays of length 3 instead of 2 and a different default
+            ([np.array([False, False, False]),
+              np.array([False, True, False]),
+              np.array([False, False, True])],
+             [np.array([1, 2, 3]),
+              np.array([4, 5, 6]),
+              np.array([7, 8, 9])], 15.3),
+            # test with arrays of length 1 instead of 2
+            ([np.array([True]),
+              np.array([False])], [np.array([1]), np.array([2])], 0),
+            # test with lists of length 100 of arrays of length 1
+            ([np.array([False])] * 100, [np.array([1])] * 100, 0),
+            # passing arrays with NaNs
+            ([np.isnan(np.array([1, 2, 3, np.nan, 5, 7]))] * 2,
+             [np.array([1, 2, 3, np.nan, 5, 7])] * 2, 0),
+            # passing lists with 2d arrays
+            ([np.isnan(np.array([[1, 2, 3, np.nan, 5, 7]]))] * 2,
+             [np.array([[1, 2, 3, np.nan, 5, 7]])] * 2, 0),
+            # passing arrays with complex numbers
+            ([np.isnan(np.array([1, 2, 3 + 2j, np.nan, 5, 7]))] * 2,
+             [np.array([1, 2, 3 + 2j, np.nan, 5, 7])] * 2, 0)
+        ]
+
+        for x in (np.arange(10), np.arange(10).reshape((5, 2))):
+            # test with two lists
+            test_cases.append(([x < 3, x > 5], [x, x ** 2], 0))
+            # test with two tuples
+            test_cases.append(((x < 3, x > 5), (x, x ** 2), 0))
+            # test with one list and one tuple
+            test_cases.append(([x < 3, x > 5], (x, x ** 2), 0))
+            # test with one tuple and one list
+            test_cases.append(((x < 3, x > 5), [x, x ** 2], 0))
+
+        for condlist,  choicelist, default in test_cases:
+            self.assertPreciseEqual(np_pyfunc(condlist, choicelist, default),
+                                    np_nbfunc(condlist, choicelist, default))
+
+    def test_select_exception(self):
+        np_nbfunc = njit(np_select)
+        x = np.arange(10)
+        self.disable_leak_check()
+        for condlist, choicelist, default, expected_error, expected_text in [
+            # Each test case below is one tuple.
+            # Each tuple is separated by the description of the intended error
+
+            # passing condlist of dim zero
+            ([np.array(True), np.array([False, True, False])],
+             [np.array(1), np.arange(12).reshape(4, 3)], 0,
+             TypingError, "condlist arrays must be of at least dimension 1"),
+            # condlist and choicelist with different dimensions
+            ([np.array(True), np.array(False)], [np.array([1]), np.array([2])],
+             0, TypingError, "condlist and choicelist elements must have the "
+                             "same number of dimensions"),
+            # condlist and choicelist with different dimensions
+            ([np.array([True]), np.array([False])],
+             [np.array([[1]]), np.array([[2]])], 0, TypingError,
+             "condlist and choicelist elements must have the "
+             "same number of dimensions"),
+            # passing choicelist of dim zero
+            ([np.array(True), np.array(False)], [np.array(1), np.array(2)], 0,
+             TypingError, "condlist arrays must be of at least dimension 1"),
+            # passing an array as condlist instead of a list or tuple
+            (np.isnan(np.array([1, 2, 3, np.nan, 5, 7])),
+             np.array([1, 2, 3, np.nan, 5, 7]), 0, TypingError,
+             "condlist must be a List or a Tuple"),
+            # default is a list
+            ([True], [0], [0], TypingError,
+             "default must be a scalar"),
+            # condlist with ints instead of booleans
+            ([(x < 3).astype(int), (x > 5).astype(int)], [x, x ** 2], 0,
+             TypingError, "condlist arrays must contain booleans"),
+            # condlist and choicelist of different length
+            ([x > 9, x > 8, x > 7, x > 6], [x, x**2, x], 0, ValueError,
+             "list of cases must be same length as list of conditions"),
+
+            # condlist contains tuples instead of arrays
+            # if in the future numba's np.where accepts tuples, the
+            # implementation of np.select should also accept them and
+            # the following two test cases should be normal tests
+            # instead of negative tests
+
+            # test with lists of length 100 of tuples of length 1 for condlist
+            ([(False,)] * 100, [np.array([1])] * 100, 0, TypingError,
+             'items of condlist must be arrays'),
+            # test with lists of length 100 of tuples of length 1 for choicelist
+            ([np.array([False])] * 100, [(1,)] * 100, 0, TypingError,
+             'items of choicelist must be arrays'),
+        ]:
+            with self.assertRaises(expected_error) as e:
+                np_nbfunc(condlist, choicelist, default)
+            self.assertIn(expected_text, str(e.exception))
+
     def test_windowing(self):
         def check_window(func):
             np_pyfunc = func