0

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!

A. White
  • 15
  • 6
  • numpy transposition is in a sense a notational convenience. It doesn't actually reorder an array in memory. So it's not actually slow. It's the correct way to vectorise your code. Have you tried timing your solution? – Denziloe Sep 12 '18 at 00:27
  • I would suggest trying out the operations you want with a 1d vector and a 2d array first and I bet you would probably be able to figure it out. You can most definitely achieve this with some clever usage of `np.einsum()`. I would suggest perusing [this blog post](http://ajcr.net/Basic-guide-to-einsum/) (which inspired [this SO answer](https://stackoverflow.com/a/33641428/5087436)). But you can also transpose as you like for vectorized operations---the loop is what's killing you, not transposition. – alkasm Sep 12 '18 at 00:29
  • `matmul` (the @ operator) does the `dot` on the last 2 dimensions. So if your `Rp` was (701,4,4) shape, you could probably do something like `Q @ Rp @ QH`. I'll have test for details. – hpaulj Sep 12 '18 at 00:32
  • @Denziloe thank you, that is a vital piece of information to have. I will try do it with the transpose now. – A. White Sep 12 '18 at 02:15

1 Answers1

2

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)
hpaulj
  • 175,871
  • 13
  • 170
  • 282
  • Thank you for taking the time to write this example. Would you please be able to explain your use of tensordot? – A. White Sep 12 '18 at 02:18
  • `tensordot` uses transpose and reshape to convert the arrays into ones that it can pass to `dot` - and reshapes back as needed. – hpaulj Sep 12 '18 at 04:15