From 4712875840b5413a47bd8dcbed9bedeaefe829cd Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Sat, 21 Oct 2017 14:39:49 -0700 Subject: [PATCH 1/4] MAINT/BUG: Remove special-casing for 0d arrays, now that indexing with a single boolean is ok Also fix the test added in gh-4792, which didn't make sense, but passed anyway --- numpy/lib/function_base.py | 23 ++++++++--------------- numpy/lib/tests/test_function_base.py | 2 +- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 3a73409fcca3..573516f3e145 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1323,21 +1323,15 @@ def piecewise(x, condlist, funclist, *args, **kw): """ x = asanyarray(x) n2 = len(funclist) - if (isscalar(condlist) or not (isinstance(condlist[0], list) or - isinstance(condlist[0], ndarray))): - if not isscalar(condlist) and x.size == 1 and x.ndim == 0: - condlist = [[c] for c in condlist] - else: - condlist = [condlist] + + # undocumented: single condition is promoted to a list of one condition + if isscalar(condlist) or ( + not isinstance(condlist[0], (list, ndarray)) and x.ndim != 0): + condlist = [condlist] + condlist = array(condlist, dtype=bool) n = len(condlist) - # This is a hack to work around problems with NumPy's - # handling of 0-d arrays and boolean indexing with - # numpy.bool_ scalars - zerod = False - if x.ndim == 0: - x = x[None] - zerod = True + if n == n2 - 1: # compute the "otherwise" condition. condelse = ~np.any(condlist, axis=0, keepdims=True) condlist = np.concatenate([condlist, condelse], axis=0) @@ -1352,8 +1346,7 @@ def piecewise(x, condlist, funclist, *args, **kw): vals = x[condlist[k]] if vals.size > 0: y[condlist[k]] = item(vals, *args, **kw) - if zerod: - y = y.squeeze() + return y diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 39edc18b4a47..2a42c44e6386 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2538,7 +2538,7 @@ def test_0d(self): assert_(y == 0) x = 5 - y = piecewise(x, [[True], [False]], [1, 0]) + y = piecewise(x, [True, False], [1, 0]) assert_(y.ndim == 0) assert_(y == 1) From 4729550aa44c97b11d4c9bd39af9bc093c2ce0f8 Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Sat, 21 Oct 2017 14:55:13 -0700 Subject: [PATCH 2/4] DOC: piecewise callables take 1d arrays --- numpy/lib/function_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 573516f3e145..0f221fe056a6 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1258,8 +1258,8 @@ def piecewise(x, condlist, funclist, *args, **kw): is the default value, used wherever all conditions are false. funclist : list of callables, f(x,*args,**kw), or scalars Each function is evaluated over `x` wherever its corresponding - condition is True. It should take an array as input and give an array - or a scalar value as output. If, instead of a callable, + condition is True. It should take a 1d array as input and give an 1d + array or a scalar value as output. If, instead of a callable, a scalar is provided then a constant function (``lambda x: scalar``) is assumed. args : tuple, optional From 303941c929ee2938dc55ad633c9fc063610941b9 Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Sun, 22 Oct 2017 16:17:51 -0700 Subject: [PATCH 3/4] TST: Add test for 0d conditions in np.piecewise --- numpy/lib/tests/test_function_base.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 2a42c44e6386..e25962da9ea6 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2556,6 +2556,12 @@ def test_0d_comparison(self): y = piecewise(x, [x <= 3, (x > 3) * (x <= 5), x > 5], [1, 2, 3]) assert_array_equal(y, 2) + def test_0d_0d_condition(self): + x = np.array(3) + c = np.array(x > 3) + y = piecewise(x, [c], [1, 2]) + assert_equal(y, 2) + def test_multidimensional_extrafunc(self): x = np.array([[-2.5, -1.5, -0.5], [0.5, 1.5, 2.5]]) From c875b1391286a982c958d349d58eb720eb081069 Mon Sep 17 00:00:00 2001 From: Eric Wieser Date: Sun, 22 Oct 2017 16:45:05 -0700 Subject: [PATCH 4/4] BUG: Throw an error if too many functions are given to piecewise Especially necessary given the strange heuristics that decay the number of conditions to 1 --- numpy/lib/function_base.py | 7 ++++++- numpy/lib/tests/test_function_base.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 0f221fe056a6..498853d32107 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1254,7 +1254,7 @@ def piecewise(x, condlist, funclist, *args, **kw): The length of `condlist` must correspond to that of `funclist`. If one extra function is given, i.e. if - ``len(funclist) - len(condlist) == 1``, then that extra function + ``len(funclist) == len(condlist) + 1``, then that extra function is the default value, used wherever all conditions are false. funclist : list of callables, f(x,*args,**kw), or scalars Each function is evaluated over `x` wherever its corresponding @@ -1336,6 +1336,11 @@ def piecewise(x, condlist, funclist, *args, **kw): condelse = ~np.any(condlist, axis=0, keepdims=True) condlist = np.concatenate([condlist, condelse], axis=0) n += 1 + elif n != n2: + raise ValueError( + "with {} condition(s), either {} or {} functions are expected" + .format(n, n, n+1) + ) y = zeros(x.shape, x.dtype) for k in range(n): diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index e25962da9ea6..8381c2465b42 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2514,6 +2514,11 @@ def test_simple(self): x = piecewise([0, 0], [[False, True]], [lambda x:-1]) assert_array_equal(x, [0, -1]) + assert_raises_regex(ValueError, '1 or 2 functions are expected', + piecewise, [0, 0], [[False, True]], []) + assert_raises_regex(ValueError, '1 or 2 functions are expected', + piecewise, [0, 0], [[False, True]], [1, 2, 3]) + def test_two_conditions(self): x = piecewise([1, 2], [[True, False], [False, True]], [3, 4]) assert_array_equal(x, [3, 4]) @@ -2556,6 +2561,11 @@ def test_0d_comparison(self): y = piecewise(x, [x <= 3, (x > 3) * (x <= 5), x > 5], [1, 2, 3]) assert_array_equal(y, 2) + assert_raises_regex(ValueError, '2 or 3 functions are expected', + piecewise, x, [x <= 3, x > 3], [1]) + assert_raises_regex(ValueError, '2 or 3 functions are expected', + piecewise, x, [x <= 3, x > 3], [1, 1, 1, 1]) + def test_0d_0d_condition(self): x = np.array(3) c = np.array(x > 3)