Understanding axis in Python

Question

I have seen a couple of codes using numpy.apply_along_axis and I always have to test the codes to see how this works 'cause I didn't understand the axis idea in Python yet.

For example, I tested this simple codes from the reference.

I can see that for the first case it was took the first column of each row of the matrix, and in the second case, the row itself was considered.

So I build an example to test how this works with an array of matrices (the problem that took me to this axis question), which can also be seen as a 3d matrix, where each row is a matrix, right?

a = [[[1,2,3],[2,3,4]],[[4,5,6],[9,8,7]]]

import numpy
data = numpy.array([b for b in a])

def my_func(x):
    return (x[0] + x[-1]) * 0.5

b = numpy.apply_along_axis(my_func, 0, data)
b = numpy.apply_along_axis(my_func, 1, data)

Which gave me:

array([[ 2.5,  3.5,  4.5],
       [ 5.5,  5.5,  5.5]])

And:

array([[ 1.5,  2.5,  3.5],
       [ 6.5,  6.5,  6.5]])

For the first result I got what I expected. But for the second one, I though I would receive:

array([[ 2.,  3.],
       [ 5.,  8.]])

Then I though that maybe should be an axis=2 and I got the previous result testing it. So, I'm wondering how this works to work it properly.

Thank you.

Answer 1

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

Answer 2

Let there be an array of shape (2,2,3). It can be seen that axis 0, axis 1, axis 2 has 2 ,2, 3 data values respectively.

These are the indexes of the elements of the array

[
    [
        [(0,0,0) (0,0,1), (0,0,2)],
        [(0,1,0) (0,1,1), (0,1,2)]
    ],
    [
        [(1,0,0) (1,0,1), (1,0,2)],
        [(1,1,0) (1,1,1), (1,1,2)]
    ]
]

Now if you apply some operation along some axis, then vary the indexes along this axis only keeping the indices along the two other axis constant.

Example: If we apply some operation F along axis 0, then the elements of the result would be

[
    [F((0,0,0),(1,0,0)), F((0,0,1),(1,0,1)), F((0,0,2),(1,0,2))],
    [F((0,1,0),(1,1,0)), F((0,1,1),(1,1,1)), F((0,1,2),(1,1,2))]
]

Along axis 1:

[
    [F((0,0,0),(0,1,0)), F((0,0,1),(0,1,1)), F((0,0,2),(0,1,2))],
    [F((0,1,0),(1,1,0)), F((0,1,1),(1,1,1)), F((0,1,2),(1,1,2))]
]

Along axis 2:

[
    [F((0,0,0),(0,0,1),(0,0,2)), F((0,1,0),(0,1,1),(0,1,2))],
    [F((1,0,0),(1,0,1),(1,0,2)), F((1,1,0),(1,1,1),(1,1,2))]
]

Also the shape of the resulting array can be inferred by omitting the given axis in the shape of given data.

Answer 3

Perhaps checking the shape of your array will help clarify which axis is which;

print data.shape

>> (2,2,3)

This means that calling

numpy.apply_along_axis(my_func, 2, data)

should indeed give a 2x2 matrix, namely

array([[ 2.,  3.],
      [ 5.,  8.]])

because the 3rd axis (index 2) has length 3, while the remaining axes are both length 2.