First, data=numpy.array(a)
is already enough, no need to use numpy.array([b for b in a])
.
data
is now a 3D ndarray
with the shape (2,2,3)
, and has 3 axes 0, 1, 2
. The first axis has a length of 2, the second axis's length is also 2 and the third axis's length is 3.
Therefore both numpy.apply_along_axis(my_func, 0, data)
and numpy.apply_along_axis(my_func, 1, data)
will result in a 2D array of shape (2,3)
. In both cases the shape is (2,3)
, those of the remaining axes, 2nd and 3rd or 1st and 3rd.
numpy.apply_along_axis(my_func, 2, data)
returns the (2,2)
shape array you showed, where (2,2)
is the shape of the first 2 axes, as you apply
along the 3rd axis (by giving index 2
).
The way to understand it is whichever axis you apply along will be 'collapsed' into the shape of your my_func
, which in this case returns a single value. The order and shape of the remaining axis will remain unchanged.
The alternative way to think of it is: apply_along_axis
means apply that function to the values on that axis, for each combination of the remaining axis/axes. Fetch the result, and organize them back into the shape of the remaining axis/axes. So, if my_func
returns a tuple
of 4 values:
def my_func(x):
return (x[0] + x[-1]) * 2,1,1,1
we will expect numpy.apply_along_axis(my_func, 0, data).shape
to be (4,2,3)
.
- See also
numpy.apply_over_axes
for applying a function repeatedly over multiple axes