numpy piecewise function claiming lists are not same size

I'm trying to write a function to pass a variable into a piecewise function, I have:

def trans(a):
    np.piecewise(a, [(a<.05) & (a>=0),
                    (a<.1) & (a>= .05),
                    (a<.2) & (a>=.1),
                    (a<.4) & (a>=.2),
                    (a<1) & (a>=.4)],
                    [0,1,2,3,4])

However, when I run trans(a), I get:

ValueError: function list and condition list must be the same

The function and condition list I used are both length 5, so I'm not sure what the issue is

EDIT: Apparently numpy.piecewise expects an array, so I needed to pass it one, instead of a plain variable?

Solution

That's the error you'd get if a were a list and not an ndarray:

>>> a = [1,2,7]
>>> np.piecewise(a, [a == 1, a > 1, a > 4], [lambda x: 100*x, lambda x: x**2, lambda x: x])
Traceback (most recent call last):
  File "<ipython-input-18-03e300b14962>", line 1, in <module>
    np.piecewise(a, [a == 1, a > 1, a > 4], [lambda x: 100*x, lambda x: x**2, lambda x: x])
  File "/usr/local/lib/python2.7/dist-packages/numpy/lib/function_base.py", line 693, in piecewise
    "function list and condition list must be the same")
ValueError: function list and condition list must be the same

>>> a = np.array(a)
>>> np.piecewise(a, [a == 1, a > 1, a > 4], [lambda x: 100*x, lambda x: x**2, lambda x: x])
array([100,   4,   7])

This happens because if a is a list, then the vector comparisons don't behave as intended:

>>> a = [1,2,7]
>>> [a == 1, a > 1, a > 4]
[False, True, True]

And these lines lines in np.piecewise

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]

mean that you'll wind up with it thinking the condition list actually looks like

[array([False,  True,  True], dtype=bool)]

which is a list of length 1.

[Oops-- I just noticed my conditions weren't mutually exclusive. Oh, well, it doesn't matter for explaining what's going on anyway.]