numpy apply_along_axis on a 1d array

Question

What happens when numpy.apply_along_axis takes a 1d array as input? When I use it on 1d array, I see something strange:

y=array([1,2,3,4])

First try:

apply_along_axis(lambda x: x > 2, 0, y)
apply_along_axis(lambda x: x - 2, 0, y)

returns:

array([False, False,  True,  True], dtype=bool)
array([-1,  0,  1,  2])

However when I try:

apply_along_axis(lambda x: x - 2 if x > 2 else x, 0, y)

I get an error:

The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I could of course use list comprehension then convert back to array instead, but that seems convoluted and I feel like I'm missing something about apply_along_axis when applied to a 1d array.

UPDATE : as per Jeff G's answer, my confusion stems from the fact that for 1d array with only one axis, what is being passed to the function is in fact the 1d array itself rather than the individual elements.

"numpy.where" is clearly better for my chosen example (and no need for apply_along_axis), but my question is really about the proper idiom for applying a general function (that takes one scalar and returns one scalar) to each element of an array (other than list comprehension), something akin to pandas.Series.apply (or map). I know of 'vectorize' but it seems no less unwieldy than list comprehension.

Answer 1

I'm unclear whether you're asking if y must be 1-D (answer is no, it can be multidimensional) or if you're asking about the function passed into apply_along_axis . To that, the answer is yes: the function you pass must take a 1-D array. (This is stated clearly in the function's documentation ).

In your three examples, the type of x is always a 1-D array. The reason your first two examples work is because Python is implicitly broadcasting the > and - operators along that array.

Your third example fails because there is no such broadcasting along an array for if / else . For this to work with apply_along_axis you need to pass a function that takes a 1-D array. numpy.where would work for this:

>>> apply_along_axis(lambda x: numpy.where(x > 2, x - 2, x), 0, y)
array([1, 2, 1, 2])

PS In all these examples, apply_along_axis is unnecessary, thanks to broadcasting. You could achieve the same results with these:

>>> y > 2
>>> y - 2
>>> numpy.where(y > 2, y - 2, y)

Answer 2

This answer addresses the updated addendum to your original question:

numpy.vectorize will take an elementwise function and return a new function. The new function can be applied to an entire array. It's like map , but it uses the broadcasting rules of numpy.

f = lambda x: x - 2 if x > 2 else x # your elementwise fn
fv = np.vectorize(f)
fv(np.array([1,2,3,4]))
# Out[5]: array([1, 2, 1, 2])