Understanding NumPy sum. If you are not clear on what NumPy is... | by Kshitij Bajracharya | Towards Data Science

If you are not clear on what NumPy is or how it works, see this post first.
towardsdatascience.com
In that post on introduction to NumPy, I did a row-wise addition on a NumPy array. Then one of the readers of the post responded by saying that what I had done was a column-wise addition, not row-wise. Indeed, when I was learning it, I felt the same that this is not how it should work. It is the opposite of how it should work. So, I looked at the docs, but all it said was
we can sum each row of an array, in which case we operate along columns, or axis 1
It didn’t help. I was still confused. I kept looking and then I found this post by Aerin Kim and it changed the way I looked at summing in NumPy arrays. So using her post as the base, this is my take on NumPy array sums.
If you’re not clear what the problem is, let’s define it formally.
We have a 5x5 NumPy array like the following
a = array([[ 12., -22., -20., -19.,  -3.],       [-23.,  21., -17., -11.,  -1.],       [ -4.,  -5.,  16.,  -9., -14.],       [-10.,  -6., -18.,  15.,  -8.],       [-25.,  -2., -13.,  -7.,  24.]])
We do a row-wise sum on the above array using
sum_matrix = a.sum(axis=1)
And we end up having the following array in sum_matrix
array([-52., -31., -16., -27., -23.])
And that seems to be correct if axis=1 meant row-wise addition. No issues. Except that, axis=0 should have actually meant rows and not axis=1. This will be clearer as we see how a NumPy array is formed. Let’s take another NumPy array.
arr = np.arange(12).reshape(4,3)
What we get back in arr is
array([[ 0,  1,  2],       [ 3,  4,  5],       [ 6,  7,  8],       [ 9, 10, 11]])
Clearly, it’s 4 rows and 3 columns. So, the first axis, axis=0, should denote rows and the second axis, axis=1, should denote columns.
This will also be confirmed if you look at pandas. Let’s take a DataFrame example.
df = pd.DataFrame(data=np.arange(12).reshape(4,3),index=['row1','row2','row3','row4'],columns=['col1','col2','col3'])
df looks like this
If you now do df.drop(labels=[‘row1’], axis=0), what you’ll get back is
And if you do df.drop(labels=[‘col2’], axis=1), you’ll get
Clearly, axis=0 means rows and axis=1 means columns. Then, why is it that NumPy sum does it differently?
To quote Aerin Kim, in her post, she wrote
The way to understand the “axis” of numpy sum is it collapses the specified axis. So when it collapses the axis 0 (row), it becomes just one row and column-wise sum.
Let’s see what that means. Now, it can get a little confusing in 2D, so let’s understand this first in a higher dimension and then we’ll step it down into 2D; much like what she did in her post.
So, let’s take a 3D array with a shape of (4,3,2).
three_d_array = np.arange(24).reshape(4,3,2)
three_d_array now becomes equal to
array([[[ 0,  1],        [ 2,  3],        [ 4,  5]],       [[ 6,  7],        [ 8,  9],        [10, 11]],       [[12, 13],        [14, 15],        [16, 17]],       [[18, 19],        [20, 21],        [22, 23]]])
Now, let us look at elements along the first axis, axis=0. We have 4 elements along axis=0. These can be viewed withthree_d_array[0], three_d_array[1], three_d_array[2], and three_d_array[3]
Each of these is a 2D array of shape (3,2). Now, if we add all of the above 2D arrays,
three_d_array[0]+three_d_array[1]+three_d_array[2]+three_d_array[3]
after element by element addition, we get back a 3x2 array as follows:
array([[36, 40],       [44, 48],       [52, 56]])
This is exactly what we get when we do three_d_array.sum(axis=0). We are taking all the elements along axis=0. And then we sum them, element by element. Originally, we had a 3D array of shape (4,3,2). After summing, we have a 2D array of shape (3,2). So, we lost the first axis 4 and retained the remaining two (3,2). That is probably what Kim meant when she said “it collapses the axis”.
Now, let’s look at axis=1. This time we keep the first axis fixed, and sum along the second axis, axis=1. Here, we have 12 elements, 3 along each from the first axis 0.
Adding all of these, we get back
Similarly, we add for the remaining three
Combine all 4 and we get an array as below:
array([[ 6,  9],       [24, 27],       [42, 45],       [60, 63]])
This is exactly what we get when we do three_d_array.sum(axis=1); performing element by element addition along axis=1. Again, the shape of the sum matrix is (4,2), which shows that we got rid of the second axis 3 from the original (4,3,2).
For the final axis 2, we do the same thing. This time we have all 24 elements and we sum along axis=2 by keeping the first two axes fixed. Here, we look at only the first 6 elements.
This just gives us one row of the sum array. We do the same for the rest and in the end we have the following array,
array([[ 1,  5,  9],       [13, 17, 21],       [25, 29, 33],       [37, 41, 45]])
which is exactly what we get on three_d_array.sum(axis=2). Again the shape of the sum array is (4,3) where we lose the last axis 2.
The array we have is
a = array([[ 12., -22., -20., -19.,  -3.],       [-23.,  21., -17., -11.,  -1.],       [ -4.,  -5.,  16.,  -9., -14.],       [-10.,  -6., -18.,  15.,  -8.],       [-25.,  -2., -13.,  -7.,  24.]])
Now, when we take elements along axis=0, we get
And element by element addition of these elements give us
And that is equivalent to a.sum(axis=0).
And similarly when we take elements along axis=1,
Now, if we sum all a[0]s, all a[1]s, all a[2]s, all a[3]s and all a[4]s separately, we get
which is what we get with a.sum(axis=1). If we check the shape, we get (5,). This may not be clear if we look at only the 2D array, but as we saw earlier, this just collapses the axis 1 and returns the shape for the remaining axis 0.
Now if we look back at the statement from the docs, “we can sum each row of an array, in which case we operate along columns, or axis 1”, I think it makes a lot more sense. So, although we calculated the sum of each row, technically it is a column-wise addition rather than a row-wise addition as axis=0 is row and axis=1 is column.
Why did NumPy not make it straightforward as in case of pandas? Well, to be honest, I don’t have an answer to that. But that’s the way the cookie crumbles.