Broadcast rotation matrices multiplication

Question

How to do the line marked with # <---- in a more direct way?

In the program, each row of x is coordinates of a point, rot_mat[0] and rot_mat[1] are two rotation matrices. The program rotates x by each rotation matrix.

Changing the order of multiplication between each rotation matrix and the coordinates is fine, if it makes things simpler. I want to have each row of x or the result representing coordinate of a point.

The result should match the checks.

Program:

# Rotation of coordinates of 4 points by 
# each of the 2 rotation matrices.
import numpy as np
from scipy.stats import special_ortho_group
rot_mats = special_ortho_group.rvs(dim=3, size=2)  # 2 x 3 x 3
x = np.arange(12).reshape(4, 3)
result = np.dot(rot_mats, x.T).transpose((0, 2, 1))  # <----
print("---- result ----")
print(result)
print("---- check ----")
print(np.dot(x, rot_mats[0].T))
print(np.dot(x, rot_mats[1].T))

Result:

---- result ----
[[[  0.20382264   1.15744672   1.90230739]
  [ -2.68064533   3.71537598   5.38610452]
  [ -5.56511329   6.27330525   8.86990165]
  [ -8.44958126   8.83123451  12.35369878]]

 [[  1.86544623   0.53905202  -1.10884323]
  [  5.59236544  -1.62845022  -4.00918928]
  [  9.31928465  -3.79595246  -6.90953533]
  [ 13.04620386  -5.9634547   -9.80988139]]]
---- check ----
[[  0.20382264   1.15744672   1.90230739]
 [ -2.68064533   3.71537598   5.38610452]
 [ -5.56511329   6.27330525   8.86990165]
 [ -8.44958126   8.83123451  12.35369878]]
[[  1.86544623   0.53905202  -1.10884323]
 [  5.59236544  -1.62845022  -4.00918928]
 [  9.31928465  -3.79595246  -6.90953533]
 [ 13.04620386  -5.9634547   -9.80988139]]

I hope I can get rid of `transpose((0, 2, 1))` . – R zu Dec 19 '17 at 17:19 — R zu, Dec 19 '17 at 17:19
That transpose isn't the bottleneck, but `np.dot` is. – Divakar Dec 19 '17 at 17:40 — Divakar, Dec 19 '17 at 17:40

Divakar · Accepted Answer · 2017-12-19T17:27:05.583

4

Use np.tensordot for multiplication involving such tensors -

np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)

Here's some timings to convince ourselves why tensordot works better with tensors -

In [163]: rot_mats = np.random.rand(20,30,30)
     ...: x = np.random.rand(40,30)

# With numpy.dot
In [164]: %timeit np.dot(rot_mats, x.T).transpose((0, 2, 1))
1000 loops, best of 3: 670 µs per loop

# With numpy.tensordot
In [165]: %timeit np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)
10000 loops, best of 3: 75.7 µs per loop

In [166]: rot_mats = np.random.rand(200,300,300)
     ...: x = np.random.rand(400,300)

# With numpy.dot
In [167]: %timeit np.dot(rot_mats, x.T).transpose((0, 2, 1))
1 loop, best of 3: 1.82 s per loop

# With numpy.tensordot
In [168]: %timeit np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)
10 loops, best of 3: 185 ms per loop

edited Dec 19 '17 at 17:27

answered Dec 19 '17 at 17:21

Divakar

218,885
19
262
358

does `swapaxes` give a new array instead of a view? – R zu Dec 19 '17 at 17:22
@Rzu Gives a view. – Divakar Dec 19 '17 at 17:22
@Rzu For NumPy >= 1.10.0, if a is an ndarray, then a view of a is returned; otherwise a new array is created. For earlier NumPy versions a view of a is returned only if the order of the axes is changed, otherwise the input array is returned. - [source](https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.swapaxes.html) – FHTMitchell Dec 19 '17 at 17:23
Don't think it can be a view, because order of axes is changed from 0, 1, 2 to 0, 2, 1. Simple test would be to ravel everything before and after swapaxes, and see if they give the same result. – R zu Dec 19 '17 at 17:24
1

@Rzu Maybe you are thinking of some other view. I am talking about whether its a copy or not. More info - http://scipy-cookbook.readthedocs.io/items/ViewsVsCopies.html – Divakar Dec 19 '17 at 17:25
@Rzu I'm pretty sure it's the `ravel` that makes a copy. You can check with `np.shares_memory`. – Paul Panzer Dec 19 '17 at 17:28
Kind of amazing that tensordot is faster than dot, or swapaxes is faster than transpose, or both. – R zu Dec 19 '17 at 17:34
@Rzu Related posts on more comparisons between dot and tensordot : https://stackoverflow.com/a/47737721/, https://stackoverflow.com/a/47646944/. – Divakar Dec 19 '17 at 17:37

Broadcast rotation matrices multiplication

1 Answers1