我有一个[n×n]矩阵,其中包含属于不同组的值,以及一个[1 x n]向量,用于定义每个元素所属的组。 (通常为~1E4,在本例中n = 4)
我想计算通过将属于同一组的所有元素加在一起而获得的矩阵。
我使用np.where()来计算每个组元素所在的索引。当我使用计算的索引时,我没有获得预期的元素,因为我选择了位置而不是范围(我习惯于Matlab,我可以简单地选择M(idx1,idx2))。
import numpy as np
n=4
M = np.random.rand(n,n)
print(M)
# This vector defines to which group each element belong
belongToGroup = np.array([0, 1, 0, 2])
nGroups=np.max(belongToGroup);
# Calculate a matrix obtained by summing elements belonging to the same group
M_sum = np.zeros((nGroups+1,nGroups+1))
for g1 in range(nGroups+1):
idxG1 = np.where(belongToGroup==g1)
for g2 in range(nGroups+1):
idxG2 = np.where(belongToGroup==g2)
print('g1 = ' + str(g1))
print('g2 = ' + str(g2))
print(idxG1[0])
print(idxG2[0])
print(M[idxG1[0],idxG2[0]])
print(np.sum(M[idxG1[0],idxG2[0]]))
M_sum[g1,g2]=np.sum(M[idxG1[0],idxG2[0]])
print('')
print('Example of the problem:')
print('Elements I would like to sum to obtain M_sum[0,0]')
print(M[0:2,0:2])
print('Elements that are summed instead')
print(M[[0,1],[0,1]])
问题示例:在上面的示例中,元素M_sum [0,0]应该是M [0,0],M [0,1],M [1,0]和M [1,1]之和相反,它被计算为M [0,0]和M [1,1]之和
在MATLAB中,使用2个列表(实际上是矩阵)进行索引会选择一个块。另一方面,numpy
尝试相互广播索引数组,并返回选定的点。它的行为接近于sub2ind
在MATLAB中的作用。
In [971]: arr = np.arange(16).reshape(4,4)
In [972]: arr
Out[972]:
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]])
In [973]: i1, i2 = np.array([0,2,3]), np.array([1,2,0])
使用相同大小的2个1d数组进行索引:
In [974]: arr[i1,i2]
Out[974]: array([ 1, 10, 12])
这实际上返回[arr[0,1], arr[2,2], arr[3,0]]
,每个匹配索引点的一个元素。
但是,如果我将一个索引转换为“列向量”,它将从行中进行选择,而i2
则从列中进行选择。
In [975]: arr[i1[:,None], i2]
Out[975]:
array([[ 1, 2, 0],
[ 9, 10, 8],
[13, 14, 12]])
MATLAB使块索引变得容易,而单个访问更难。在numpy
中,块访问有点困难,尽管底层的机制是相同的。
您的示例i1[0]
和i2[0]
可以是以下数组:
array([0, 2]), array([3])
(2,) (1,)
形状(1,)数组也可以使用(2,)或(2,1)数组进行广播。如果is[0]
是np.array([0,1,2])
,一个(3,)数组不能与(2,)数组配对,那么你的代码就会失败。但是用(2,1)它产生一个(2,3)块。
你可以使用np.ix_
来获得matlab-ish行为:
A = np.arange(9).reshape(3, 3)
A[[1,2],[0,2]]
# array([3, 8])
A[np.ix_([1,2],[0,2])]
# array([[3, 5],
# [6, 8]])
在引擎盖下,np.ix_
做了@hpaulj详细描述的内容:
np.ix_([1,2],[0,2])
# (array([[1],
# [2]]), array([[0, 2]]))
您可以将此应用于您的特定问题,如下所示:
M = np.random.randint(0, 10, (n, n))
M
# array([[6, 2, 7, 1],
# [6, 7, 9, 5],
# [9, 4, 3, 2],
# [3, 1, 7, 9]])
idx = np.array([0, 1, 0, 2])
ng = idx.max() + 1
out = np.zeros((ng, ng), M.dtype)
np.add.at(out, np.ix_(idx, idx), M)
out
# array([[25, 6, 3],
# [15, 7, 5],
# [10, 1, 9]])
暂且不说:有一个更快但不太明显的解决方案依赖于平面索引:
np.bincount(np.ravel_multi_index(np.ix_(idx, idx), (ng, ng)).ravel(), M.ravel(), ng*ng).reshape(ng, ng)
# array([[25., 6., 3.],
# [15., 7., 5.],
# [10., 1., 9.]])