使用numpy的矩阵矩阵中的元素总和

Question

我有一个任意形状(N1,N2,k,k) 的矩阵矩阵，意思是N1*N2 形状为k*k 的矩阵。 我希望计算每个矩阵（形状为(k,k)）的总和，并用该总和转换矩阵本身。 结果数组的形状为(N1,N2)，其中位于某个索引i,j 中的每个元素都是该给定索引中相应矩阵的总和。 numpy 操作有没有办法做到这一点？ （也就是说 - 没有循环 range(N1) 和 range(N2)）

这是一个简单的示例（我使用* 与第一个数组和第二个数组转置只是为了创建示例）：

m = np.array([[0, 0, 0, 0]]).reshape(2, 2) # matrix element of size k*k (k=2)
a = np.array([m, m + 1, m + 2, m + 3])
b = np.array([m, m + 1, m + 2, m + 3])
reshaped1 = a[:, np.newaxis]  # (N1,1,k,k) where N1=4
reshaped2 = b[np.newaxis, :]  # (1,N2,k,k) where N2=4
mult = reshaped1 * reshaped2  #  (N1,N2,k,k)=(4,4,2,2)

我希望创建一个新数组res，它将包含所有mult 元素的总和。这可以通过以下伪代码来完成：

for i in range(N1):
    for j in range(N2):
        res[i,j] = sum(mult[i,j])

感谢您的帮助！

Answer 1

如果我理解正确，您可以将np.sum 与多个轴一起使用：

np.sum(mult, axis=(2, 3))

输出：

array([[ 0,  0,  0,  0],
       [ 0,  4,  8, 12],
       [ 0,  8, 16, 24],
       [ 0, 12, 24, 36]])

Answer 2

尝试使用np.sum(np.sum(mult,axis=3),axis=2)

import numpy as np
N1=4
N2=4
m = np.array([[0, 0, 0, 0]]).reshape(2, 2) # matrix element of size k*k (k=2)
a = np.array([m, m + 1, m + 2, m + 3])
b = np.array([m, m + 1, m + 2, m + 3])
reshaped1 = a[:, np.newaxis]  # (N1,1,k,k) where N1=4
reshaped2 = b[np.newaxis, :]  # (1,N2,k,k) where N2=4
mult = reshaped1 * reshaped2  #  (N1,N2,k,k)=(4,4,2,2)
np.sum(mult,axis=3)

res=np.zeros((4,4))
for i in range(N1):
    for j in range(N2):
        res[i,j] = np.sum(mult[i,j])
print(np.array_equal(np.sum(np.sum(mult,axis=3),axis=2),res))
>>> True