Make numpy.sum() return a sum of matrices instead of a single number

Question

I am doing a fairly complicated summation using a matrix with numpy. The shape of the matrix is matrix.shape = (500, 500) and the shape of the array is arr.shape = (25,). The operation is as follows:

totalsum = np.sum([i * matrix for i in arr])

Here is what I don't understand:

np.sum() is very slow and returns a single float, float64. Doing the same operation with Python's sum.(), i.e.

totalsum2 = sum([i*matrix for i in arr])

Preserves the shape of the matrix. That is, the resulting shape is totalsum2.shape() = (500, 500). Huh?

I also think it is strange that np.sum() takes longer than sum(), particularly when we are working with numpy ndarrays.

What exactly is going on here? How is np.sum() summing the above values in comparison to sum()?

I would like np.sum() to preserve the matrix shape. How can I set the dimensions such that np.sum() preserves the matrix size and does not return a single float?

Answer 1

You must call np.sum with the optional axis parameter set to 0 (summation over the axis 0, i.e the one created by your list comprehension)

totalsum = np.sum([i * matrix for i in arr], 0)

Alternatively, you can omit the brackets so np.sum evaluate a generator.

totalsum = np.sum(i * matrix for i in arr)