Adding all elements of a multi-dimensional array in Python

Question

I have an array A, which is x-by-y-by-z-by-w dimensional. I would like to get a z dimensional vector v, where each element of v is the sum of the elements of A controling for the coordinate of v. Is there a way I can do this w/o loops using numpy?

Here is how I would do it with a loop

for i in range(z):
   v[i] = np.sum(A[:,:,i,:])

I can do it with a loop but was just wondering if there is a more standard way. — mathemagician, Sep 06 '17 at 22:23

score 3 · Accepted Answer · answered Sep 06 '17 at 22:33

You can sum over axis by specifying axis parameter; Here you want to keep the third axis and collapse all the other axises, hence just use axis=(0,1,3):

np.sum(A, axis=(0,1,3))

Example:

A = np.arange(24).reshape((2,2,3,2))

# for loop approach
z = A.shape[2]
v = np.empty(z)
for i in range(z):
    v[i] = np.sum(A[:,:,i,:])

v
# array([  76.,   92.,  108.])

# sum over axis
np.sum(A, axis=(0,1,3))
# array([ 76,  92, 108])

Adding all elements of a multi-dimensional array in Python

1 Answers1