Commit 640cf641 authored by Ariana Peck's avatar Ariana Peck
Browse files

Batched versions of skopi geometry functions

parent 20789d3b
import numpy as np
import skopi as sk
"""
Batched versions of skopi geometry functions for handling rotations.
"""
def quaternion2rot3d(quat):
"""
Convert a set of quaternions to rotation matrices.
This function was originally adopted from:
https://github.com/duaneloh/Dragonfly/blob/master/src/interp.c
and then from skopi and has been modified to convert in batch.
Parameters
----------
quat : numpy.ndarray, shape (n_quat, 4)
series of quaternions
Returns
-------
rotation : numpy.ndarray, shape (n_quat, 3, 3)
series of rotation matrices
"""
q01 = quat[:,0] * quat[:,1]
q02 = quat[:,0] * quat[:,2]
q03 = quat[:,0] * quat[:,3]
q11 = quat[:,1] * quat[:,1]
q12 = quat[:,1] * quat[:,2]
q13 = quat[:,1] * quat[:,3]
q22 = quat[:,2] * quat[:,2]
q23 = quat[:,2] * quat[:,3]
q33 = quat[:,3] * quat[:,3]
# Obtain the rotation matrix
rotation = np.zeros((quat.shape[0], 3, 3))
rotation[:,0, 0] = (1. - 2. * (q22 + q33))
rotation[:,0, 1] = 2. * (q12 - q03)
rotation[:,0, 2] = 2. * (q13 + q02)
rotation[:,1, 0] = 2. * (q12 + q03)
rotation[:,1, 1] = (1. - 2. * (q11 + q33))
rotation[:,1, 2] = 2. * (q23 - q01)
rotation[:,2, 0] = 2. * (q13 - q02)
rotation[:,2, 1] = 2. * (q23 + q01)
rotation[:,2, 2] = (1. - 2. * (q11 + q22))
return rotation
def axis_angle_to_quaternion(axis, theta):
"""
Convert an angular rotation around an axis series to quaternions.
Parameters
----------
axis : numpy.ndarray, size (num_pts, 3)
axis vector defining rotation
theta : numpy.ndarray, size (num_pts)
angle in radians defining anticlockwise rotation around axis
Returns
-------
quat : numpy.ndarray, size (num_pts, 4)
quaternions corresponding to axis/theta rotations
"""
axis /= np.linalg.norm(axis, axis=1)[:,None]
angle = theta / 2
quat = np.zeros((len(theta), 4))
quat[:,0] = np.cos(angle)
quat[:,1:] = np.sin(angle)[:,None] * axis
return quat
def quaternion_product(q1, q0):
"""
Compute quaternion product, q1 x q0, according to:
https://www.mathworks.com/help/aeroblks/quaternionmultiplication.html.
This should yield the same result as pytorch3d.transforms.quaternion_multiply
(or possibly the -1*q_prod, which represents the same rotation).
Parameters
----------
q0 : numpy.ndarray, shape (n, 4)
first quaternion to rotate by
q1 : numpy.ndarray, shape (n, 4)
second quaternion to rotate by
Returns
-------
q_prod : numpy.ndarray, shape (n, 4)
quaternion product q1 x q0, rotation by q0 followed by q1
"""
p0, p1, p2, p3 = q1[:,0], q1[:,1], q1[:,2], q1[:,3]
r0, r1, r2, r3 = q0[:,0], q0[:,1], q0[:,2], q0[:,3]
q_prod = np.array([r0 * p0 - r1 * p1 - r2 * p2 - r3 * p3,
r0 * p1 + r1 * p0 - r2 * p3 + r3 * p2,
r0 * p2 + r1 * p3 + r2 * p0 - r3 * p1,
r0 * p3 - r1 * p2 + r2 * p1 + r3 * p0]).T
return q_prod
def get_preferred_orientation_quat(num_pts, sigma, base_quat=None):
"""
Sample quaternions distributed around a given or random position in a restricted
range in SO(3), where the spread of the distribution is determined by sigma.
Parameters
----------
num_pts : int
number of quaternions to generate
sigma : float
standard deviation in radians for angular sampling
base_quat : numpy.ndarray, size (4)
quaternion about which to distribute samples, random if None
Returns
-------
quat : numpy.ndarray, size (num_quats, 4)
quaternions with preferred orientations
"""
if base_quat is None:
base_quat = sk.get_random_quat(1) ### need to change to skopi
base_quat = np.tile(base_quat, num_pts).reshape(num_pts, 4)
R_random = quaternion2rot3d(sk.get_random_quat(num_pts)) ### need to change to skopi
unitvec = np.array([0,0,1.0])
rot_axis = np.matmul(unitvec, R_random)
theta = sigma * np.random.randn(num_pts)
rand_axis = theta[:,None] * rot_axis
pref_quat = axis_angle_to_quaternion(rand_axis, theta)
quat = quaternion_product(pref_quat, base_quat)
return quat
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment