2

I have a list of N unit-normalized 3D vectors p stored in a numpy ndarray with shape (N, 3). I have another such list, q. I want to calculate an ndarray U of shape (N, 3, 3) storing the rotation matrices that rotate each point in p to the corresponding point q.

The list of rotation matrices U should satisfy:

np.all(np.einsum('ijk,ik->ij', U, p) ==  q)

On a point-by-point basis, the problem reduces to being able to compute a rotation matrix for a rotation of some angle about some axis. Code solving the single-point case appears below:

def rotation_matrix(angle, direction):
    direction = np.atleast_1d(direction).astype('f4')
    sina = np.sin(angle)
    cosa = np.cos(angle)
    direction = direction/np.sqrt(np.sum(direction*direction))

    R = np.diag([cosa, cosa, cosa])
    R += np.outer(direction, direction) * (1.0 - cosa)
    direction *= sina
    R += np.array(((0.0, -direction[2], direction[1]),
                (direction[2], 0.0, -direction[0]),
                (-direction[1], direction[0],  0.0)))
    return R

What I need is a function that behaves exactly as the above function, but instead of accepting a single angle and a single direction, it accepts an angles array of shape (npts, ) and a directions array of shape (npts, 3). The code below is only partially finished - the problem is that neither np.diag nor np.outer accept an axis argument 

def rotation_matrices(angles, directions):
    directions = np.atleast_2d(directions)
    angles = np.atleast_1d(angles)
    npts = directions.shape[0]
    directions = directions/np.sqrt(np.sum(directions*directions, axis=1)).reshape((npts, 1))

    sina = np.sin(angles)
    cosa = np.cos(angles)

    #  Lines below require extension to 2d case - np.diag and np.outer do not support axis arguments
    R = np.diag([cosa, cosa, cosa])
    R += np.outer(directions, directions) * (1.0 - cosa)
    directions *= sina
    R += np.array(((0.0, -directions[2], directions[1]),
                (directions[2], 0.0, -directions[0]),
                (-directions[1], directions[0],  0.0)))
    return R

Does either numpy or scipy have a compact vectorized function computing the appropriate rotation matrices in a way that avoids using for loops? The problem is that neither np.diag nor np.outer accept axis as an argument. My application will have N be very large, 1e7 or greater, so a vectorized function that keeps all the relevant axes aligned is necessary for performance reasons.

aph
  • 1,765
  • 2
  • 19
  • 34
  • Seems more like a math problem to me. – Divakar Dec 03 '17 at 21:39
  • Also, I'm pretty sure it's not unique. (Once you've mapped p to q you are still free to rotate around q) – Paul Panzer Dec 03 '17 at 21:48
  • You may want to look up homogenous coordinates / transformations, the latter being the matrices used to transform a set of coordinates from one coordinate system to another. I don't think there are any python packages that do this but the calculations should be easy to implement with numpy. – ikom Dec 03 '17 at 21:55
  • Hmm, I see your point @PaulPanzer but I'm still confused because each pair of 3d points defines a unique plane, which defines a unique normal vector, and there is a unique rotation angle about that normal vector that rotates p into q. No? – aph Dec 03 '17 at 22:16
  • Yes, but you needn't rotate around the normal to map p to q. any axis that has the same distance to p and q will do – Paul Panzer Dec 03 '17 at 22:20
  • Thanks to @Divakar for asking for more clarity. I revised my question to contain all of the linear algebra operations, with the question now being how to call specific numpy functions in a way that permits 2d arguments. – aph Dec 04 '17 at 00:00
  • Thanks @PaulPanzer - yes, you are right, the rotation matrix was underdetermined by my question. In the revised version, there is a unique matrix for each specified rotation. – aph Dec 04 '17 at 00:00
  • 1
    The diagonals you can set using `R[:, range(3), range(3)] += ...` (Assuming you do the outer product first. Otherwise you'd have to create `R` first using `zeros`) – Paul Panzer Dec 04 '17 at 01:18
  • Yes, I later realized the diagonals were trivial and so asked a separate question (that you answered perfectly, thanks!) – aph Dec 04 '17 at 01:21
  • Getting the broadcasting correct for the final step where the matrix with the `directions` is actually just as tricky though. – aph Dec 04 '17 at 01:22
  • Something like `R[:, [1, 2, 0], [2, 0, 1]] -= directions`, `R[:, [2, 0, 1], [1, 2, 0]] += directions`? – Paul Panzer Dec 04 '17 at 01:34
  • Wow, your broadcasting kung fu is strong, @PaulPanzer. Sign conventions are correct, as well. Thanks again. – aph Dec 04 '17 at 01:44
  • @aph at the end which solution did you use? – ttsesm Jun 10 '20 at 17:28
  • I ended up adapting what I originally had after implementing the broadcasting tricks suggested by @PaulPanzer in the comments. There's an open-source implementation here: https://github.com/astropy/halotools/blob/master/halotools/utils/rotations3d.py – aph Jun 11 '20 at 18:25
  • @aph thanks a lot for the reply, I'll have a look on the linked source code. – ttsesm Jun 17 '20 at 10:29

2 Answers2

1

Dropping this here for now, will explain later. Using levi-cevita symbols from @jaime's answer here and the matrix form of the Rodrigues formula here and some algebra based on k = (a x b)/sin(theta)

def rotmatx(p, q):
    eijk = np.zeros((3, 3, 3))
    eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
    eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
    d = (p * q).sum(-1)[:, None, None]
    c = (p.dot(eijk) @ q[..., None]).squeeze()   # cross product (optimized)
    cx = c.dot(eijk)
    return np.eye(3) + cx + cx @ cx / (1 + d)

EDIT: dang. question changed. 

def rotation_matrices(angles, directions):
    eijk = np.zeros((3, 3, 3))
    eijk[0, 1, 2] = eijk[1, 2, 0] = eijk[2, 0, 1] = 1
    eijk[0, 2, 1] = eijk[2, 1, 0] = eijk[1, 0, 2] = -1
    theta = angles[:, None, None]
    K = directions.dot(eijk)
    return np.eye(3) + K * np.sin(theta) + K @ K * (1 - np.cos(theta))
Daniel F
  • 13,620
  • 2
  • 29
  • 55
  • Nice answer, @DanielF. And sorry about editing the question on you. I've checked that your rotation_matrices function outputs agree exactly with the method outlined in the comments. I'll check this as the accepted answer since you took the time to write it up. Even though this solution is very fast, I'll mention that the method in the comments is ~2x faster for very large number of points. – aph Dec 14 '17 at 17:32
  • Also, instead of the `K@K` operation in the return statement, which only works for python 3, this can be replaced with `np.matmul(K, K)` to accommodate both python 2 and python 3 users – aph Dec 14 '17 at 17:33
1

Dropping another solution for bulk rotation of a Nx3x3 matrix. Where the 3x3 components represent vector components in

[[11, 12, 13],
 [21, 22, 23],
 [31, 32, 33]]

Now matrix rotation by np.einsum is:

data = np.random.uniform(size=(500, 3, 3))
rotmat = np.random.uniform(size=(3, 3))

data_rot = np.einsum('ij,...jk,lk->...il', rotmat, data, rotmat)

This is equivalent to

for data_mat in data:
    np.dot(np.dot(rotmat, data_mat), rotmat.T)

Speedup over a np.dot-loop is around 250x.

n1nj4
  • 511
  • 5
  • 13
  • I do not really understand how your solution extracts the multiple rotation matrices. Can you elaborate a bit more. From what I understand your Nx3x3 is the extracted rotation matrices, right? – ttsesm Jun 10 '20 at 21:03
  • Yessir, Nx3x3 are the input matrices (`data`) to be rotated by `rotmat` – n1nj4 Nov 08 '20 at 21:42