将点积依次应用于3D numpy数组的切片

时间:2018-09-12 00:12:21

标签: python numpy multidimensional-array dot-product numpy-slicing

Hi Stack Overflow社区,

我有一个形状为4x4x701的3D numpy数组Rp,其中701个4x4切片中的每个切片在不同的时间点表示一定数量。我正在尝试有效地将Givens旋转矩阵Q及其Hermitian转置QH应用于701个切片中的每个切片,并且目前正在迭代进行,就像这样:

for idx in np.arange(Rp.shape[-1]):
    Rp[[j,k],:,idx] = np.dot(Q, Rp[[j,k],:,idx])
    Rp[:,[j,k],idx] = np.dot(Rp[:,[j,k],idx], QH)

,但是必须有一种方法不是迭代地执行此操作(对于Numpy加速)。我意识到我只能在第一种情况下使用np.dot,但是如果没有某种移调,第二种情况下将无法使用,这似乎会使速度变慢。

任何想法将不胜感激!

1 个答案:

答案 0 :(得分:2)

形状的粗略测试;值确实应该更具诊断性:

In [46]: Q = np.eye(4); QH = np.conj(Q)
In [47]: R = np.ones((10,4,4))

In [48]: (Q @ R @ QH).shape
Out[48]: (10, 4, 4)

In [49]: np.einsum('ij,kjl,lm->kil',Q,R,QH).shape
Out[49]: (10, 4, 4)

如果大尺寸是最后一个:

In [50]: Rp = R.transpose(1,2,0)
In [51]: Rp.shape
Out[51]: (4, 4, 10)

In [53]: np.einsum('ij,jlk,lm->ilk',Q,Rp,QH).shape
Out[53]: (4, 4, 10)

In [55]: (Q @ Rp.transpose(2,1,0) @ QH).transpose(1,2,0).shape
Out[55]: (4, 4, 10)

我们也可以用tensordot来写。

In [58]: np.tensordot(QH,np.tensordot(Q,Rp,(1,1)),(0,1)).shape
Out[58]: (4, 4, 10)