Sequentially applying a dot product to slices of 3D numpy array

Question

Hi Stack Overflow community,

I have a 3D numpy array Rp of shape 4x4x701, where each of the 701 4x4 slices represents a certain quantity at a different point in time. I'm trying to efficiently apply a Givens rotation matrix Q and its Hermitian transpose QH to each of the 701 slices, and am currently doing it iteratively, like so:

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)

but there must be a way to do this NOT iteratively (for the numpy speedup). I realise I can just use np.dot for the first case, but this won't work for the second without some transposition, which seems like it would slow things down.

Any ideas would be greatly appreciated!

Answer 1

A rough test for shapes; values really should be more diagnostic:

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)

---

If the big dimension is last:

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)

We could also write this with tensordot .

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