A visual guide to multidimensional NumPy array aggregation
By: Marton Meszaros
Tags: Data Science, Numpy, Python, Arrays

A visual guide to multidimensional NumPy array aggregation

If you are just learning to use the aggregate functions in NumPy like sum, mean or median, understanding what the axis parameter does can be difficult. In the following article I provide an intuitive visual guide to axis numbering so you can always make sure that your code works as intended.

Aggregating a 1-dimensional array is straightforward since there is only one axis to consider. With 2 or more dimensions, however, we must think about the direction along which we wish to aggregate our array. In this guide we will discuss one-, two- and three-dimensional problems using the numpy.sum() function. (Note: other aggregate functions like mean, median, amin, amax use the same logic.)

Dimension reduction and axis numbering

We can think of aggregation along a single axis as reducing the dimension of our array by one. If we aggregate a one-dimensional array, we get a single point. If we aggregate a two-dimensional array, we get a one-dimensional array and so on. If we have more than one dimensions, we must decide which dimension (or axis) we want to eliminate. In NumPy this can be done by specifying the axis parameter in the function call.

Let me emphasize: When you set the parameter, you select the dimension you wish to eliminate.

What we must know about axis numbering is that indexing always starts with the highest dimension. With each added dimension, the new dimension gets index 0 and all other indices are increased by one. We will see this in detail in the examples below. The default value for the axis parameter in NumPy is None, in which case the array gets flattened and then aggregated.

Tip: If you are not sure how many dimensions you have but you want to select the lowest, you can use negative indexing. For example, axis=-1 will eliminate the lowest dimension.

Aggregating one dimensional arrays

>>> a = numpy.array([1,2])

>>> numpy.sum(a, axis=0)

3

In this simple case there is only one dimension along which we can aggregate the array. This dimension is the 0th dimensions, or 0th axis. (It is not necessary to specify the axis parameter because a flattened 2-dimensional array will be itself.)

Aggregating two dimensional arrays

>>> b = numpy.array([[1,2],[3,4]])

In this case we have added a second one-dimensional array to our previous array, creating a two-dimensional array. Note that the added dimension is the highest, so that will have index 0, and the previous will have index 1.

Here we can choose whether to eliminate the 0th dimension and collapse the array along its vertical axis or eliminate the 1st dimension and collapse the array along its horizontal axis. Let’s see both:

>>> numpy.sum(b, axis=0)

array([4,6])

With axis=0 we wish to eliminate the 0th dimension. We achieve this by collapsing the array along its vertical axis.

>>> numpy.sum(b, axis=1)

array([3,7])

Following the same logic, here we eliminate the 1st dimension and collapse the array along its horizontal axis.

Aggregating three dimensional arrays

c = numpy.array([[[1,2],[3,4]],[[5,6],[7,8]]])

In this three-dimensional example we simply add one more 2 dimensional array to the previous array, creating a three-dimensional array of shape(2,2,2). This added dimension will be the highest, so it gets index 0, while all other indices are increased by one. Since we have three dimensions, there are three options to consider when aggregating the array.

>>> numpy.sum(c, axis=0)

array([[6,8],[10,12]])

With axis=0 we specify that we wish to eliminate the highest dimension (the one we have just added). We can achieve this by collapsing the two ‘layers’ of the three-dimensional array.

>>> numpy.sum(c, axis=1)

array([[4,6],[12,14]])

With axis=1 we wish to eliminate the 1st dimension. We achieve this by collapsing both of our two-dimensional arrays (the layers) along their vertical axes.

This way we reduced an array of shape(2,2,2) into an array of shape(2,1,2). Note that NumPy will completely eliminate the dimension of size 1 and output a lower dimensional array of shape(2,2) as seen as the third stage of my chart unless you set keepdims=True.

>>> numpy.sum(c, axis=2)

array([[3,7],[11,15]])

With axis=2 we wish to eliminate the 2nd dimension. We achieve this by collapsing both of our 2-dimensional arrays (the layers) along their horizontal axes.

Similar to the previous example, we get an array of shape(2,2,1) under the hood, but NumPy will output a lower dimensional array of shape(2,2) by default.

Key points to remember

You can think of aggregation as reducing the dimensionality of the array When you choose an axis, you actually set the dimension you wish to eliminate Indexing always starts with the highest dimension

Reference

[1] The SciPy community, NumPy v1.20 Manual

(All images are created and owned by the author.)