diff options
author | seberg <sebastian@sipsolutions.net> | 2016-11-09 17:36:24 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2016-11-09 17:36:24 +0100 |
commit | 0268680e5db25caeeb69bbb2e62ead3405c8f6f0 (patch) | |
tree | 9584c8caa137d2b73c8e10fcaba426295a20aac0 /numpy/lib | |
parent | 3fd2fa1d3a1b627b6a13fe963b2c5b14eafb0fbe (diff) | |
parent | 6420f844824cdc327026302870b57dfd8c7481f7 (diff) | |
download | numpy-0268680e5db25caeeb69bbb2e62ead3405c8f6f0.tar.gz |
Merge pull request #8194 from alvarosg/scalar-piecewise
BUG: np.piecewise not working for scalars
Diffstat (limited to 'numpy/lib')
-rw-r--r-- | numpy/lib/function_base.py | 17 | ||||
-rw-r--r-- | numpy/lib/tests/test_function_base.py | 12 |
2 files changed, 23 insertions, 6 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 4172c26b5..df9c7f177 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1225,9 +1225,9 @@ def piecewise(x, condlist, funclist, *args, **kw): Parameters ---------- - x : ndarray + x : ndarray or scalar The input domain. - condlist : list of bool arrays + condlist : list of bool arrays or bool scalars Each boolean array corresponds to a function in `funclist`. Wherever `condlist[i]` is True, `funclist[i](x)` is used as the output value. @@ -1296,12 +1296,21 @@ def piecewise(x, condlist, funclist, *args, **kw): >>> np.piecewise(x, [x < 0, x >= 0], [lambda x: -x, lambda x: x]) array([ 2.5, 1.5, 0.5, 0.5, 1.5, 2.5]) + Apply the same function to a scalar value. + + >>> y = -2 + >>> np.piecewise(y, [y < 0, y >= 0], [lambda x: -x, lambda x: x]) + array(2) + """ x = asanyarray(x) n2 = len(funclist) if (isscalar(condlist) or not (isinstance(condlist[0], list) or isinstance(condlist[0], ndarray))): - condlist = [condlist] + if not isscalar(condlist) and x.size == 1 and x.ndim == 0: + condlist = [[c] for c in condlist] + else: + condlist = [condlist] condlist = array(condlist, dtype=bool) n = len(condlist) # This is a hack to work around problems with NumPy's @@ -1311,8 +1320,6 @@ def piecewise(x, condlist, funclist, *args, **kw): if x.ndim == 0: x = x[None] zerod = True - if condlist.shape[-1] != 1: - condlist = condlist.T if n == n2 - 1: # compute the "otherwise" condition. totlist = np.logical_or.reduce(condlist, axis=0) # Only able to stack vertically if the array is 1d or less diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index f396e036b..5bba088a8 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2324,9 +2324,19 @@ class TestPiecewise(TestCase): assert_(y.ndim == 0) assert_(y == 1) + # With 3 ranges (It was failing, before) + y = piecewise(x, [False, False, True], [1, 2, 3]) + assert_array_equal(y, 3) + def test_0d_comparison(self): x = 3 - piecewise(x, [x <= 3, x > 3], [4, 0]) # Should succeed. + y = piecewise(x, [x <= 3, x > 3], [4, 0]) # Should succeed. + assert_equal(y, 4) + + # With 3 ranges (It was failing, before) + x = 4 + y = piecewise(x, [x <= 3, (x > 3) * (x <= 5), x > 5], [1, 2, 3]) + assert_array_equal(y, 2) def test_multidimensional_extrafunc(self): x = np.array([[-2.5, -1.5, -0.5], |