summaryrefslogtreecommitdiff
path: root/numpy/lib
diff options
context:
space:
mode:
authorseberg <sebastian@sipsolutions.net>2016-11-09 17:36:24 +0100
committerGitHub <noreply@github.com>2016-11-09 17:36:24 +0100
commit0268680e5db25caeeb69bbb2e62ead3405c8f6f0 (patch)
tree9584c8caa137d2b73c8e10fcaba426295a20aac0 /numpy/lib
parent3fd2fa1d3a1b627b6a13fe963b2c5b14eafb0fbe (diff)
parent6420f844824cdc327026302870b57dfd8c7481f7 (diff)
downloadnumpy-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.py17
-rw-r--r--numpy/lib/tests/test_function_base.py12
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],