BUG: Fixed piecewise function for 0d input

When `x` has more than one element the condlist `[True, False]` is being made equivalent to `[[True, False]]`, which is correct. However, when `x` is zero dimensional the expected condlist is `[[True], [False]]`: this commit addresses the issue. Besides, the documentation stated that there could be undefined values but actually these are 0 by default: using `nan` would be desirable, but for the moment the docs were corrected. Closes #331.
numpy · juliantaylor · Jun 8, 2014 · Jun 7, 2014 · Jun 8, 2014 · 292b9ff538ea4950c7380c76cf65d1a5b108b75c
commit 292b9ff538ea4950c7380c76cf65d1a5b108b75c
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
@@ -651,7 +651,7 @@ def piecewise(x, condlist, funclist, *args, **kw):
         The output is the same shape and type as x and is found by
         calling the functions in `funclist` on the appropriate portions of `x`,
         as defined by the boolean arrays in `condlist`.  Portions not covered
-        by any condition have undefined values.
+        by any condition have a default value of 0.
 
 
     See Also
@@ -693,32 +693,24 @@ def piecewise(x, condlist, funclist, *args, **kw):
     if (isscalar(condlist) or not (isinstance(condlist[0], list) or
                                    isinstance(condlist[0], ndarray))):
         condlist = [condlist]
-    condlist = [asarray(c, dtype=bool) for c in condlist]
+    condlist = array(condlist, dtype=bool)
     n = len(condlist)
-    if n == n2 - 1:  # compute the "otherwise" condition.
-        totlist = condlist[0]
-        for k in range(1, n):
-            totlist |= condlist[k]
-        condlist.append(~totlist)
-        n += 1
-    if (n != n2):
-        raise ValueError(
-            "function list and condition list must be the same")
-    zerod = False
     # 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
-        newcondlist = []
-        for k in range(n):
-            if condlist[k].ndim == 0:
-                condition = condlist[k][None]
-            else:
-                condition = condlist[k]
-            newcondlist.append(condition)
-        condlist = newcondlist
+        if condlist.shape[-1] != 1:
+            condlist = condlist.T
+    if n == n2 - 1:  # compute the "otherwise" condition.
+        totlist = np.logical_or.reduce(condlist, axis=0)
+        condlist = np.vstack([condlist, ~totlist])
+        n += 1
+    if (n != n2):
+        raise ValueError(
+                "function list and condition list must be the same")
 
     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
@@ -1487,6 +1487,7 @@ def test_simple(self):
         x = piecewise([0, 0], [[False, True]], [lambda x:-1])
         assert_array_equal(x, [0, -1])
 
+    def test_two_conditions(self):
         x = piecewise([1, 2], [[True, False], [False, True]], [3, 4])
         assert_array_equal(x, [3, 4])
 
@@ -1505,6 +1506,15 @@ def test_0d(self):
         assert_(y.ndim == 0)
         assert_(y == 0)
 
+        x = 5
+        y = piecewise(x, [[True], [False]], [1, 0])
+        assert_(y.ndim == 0)
+        assert_(y == 1)
+
+    def test_0d_comparison(self):
+        x = 3
+        y = piecewise(x, [x <= 3, x > 3], [4, 0])
+
 
 class TestBincount(TestCase):
     def test_simple(self):