Summing 3d numpy array over plane, not axis

Viewed 144

I have a 3d numpy array (nx5x5). I want to sum each of the n slices together. So the new shape will be (nx1x1), where each n is just the sum of an individual 5x5 array. Can I do this in numpy without using a loop? np.sum has its axis arguments, but they reshape the array into the wrong shape. I think I may need to call np.sum twice? But I'm having trouble thinking about how to do this. Anybody know the answer?

3 Answers

Here are three different ways of doing it:

  1. Use a tuple for axis:

    a.sum(axis=(1, 2))
    
  2. Reshape properly to merge the axes you want to sum over:

    a.reshape(a.shape[0], -1).sum(axis=1)
    
  3. Use multiple sums:

    a.sum(-1).sum(-1)
    

    OR

    a.sum(1).sum(-1)
    

    etc.

np.sum has its axis arguments, but they reshape the array into the wrong shape

Summing is a reduction operation and it makes sense that after reducing in a specific axis (by summing all elements in that axis) that particular dimension is removed. If you don't want that you can pass the optional keepdims argument.

values = np.random.randn(3,5,5)
values.sum(axis=(1,2), keepdims=True)

Think I figured this out, for anyone who is running into the same issue. It turns out you can select two axes from np.sum, making a plane instead of a line. So:

np.sum(a, (1,2)

Does the trick.

Related