diff options
| author | Eric Wieser <wieser.eric@gmail.com> | 2017-10-22 16:45:05 -0700 |
|---|---|---|
| committer | Eric Wieser <wieser.eric@gmail.com> | 2017-10-23 20:31:40 -0700 |
| commit | c875b1391286a982c958d349d58eb720eb081069 (patch) | |
| tree | da383b1c94b6219d76c20978c5ff53a8ed6b4561 /numpy/lib | |
| parent | 303941c929ee2938dc55ad633c9fc063610941b9 (diff) | |
| download | numpy-c875b1391286a982c958d349d58eb720eb081069.tar.gz | |
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
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 7 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 10 |
2 files changed, 16 insertions, 1 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 0f221fe05..498853d32 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 e25962da9..8381c2465 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2514,6 +2514,11 @@ class TestPiecewise(object): 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 @@ class TestPiecewise(object): 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) |