diff --git a/argoverse/utils/transform.py b/argoverse/utils/transform.py index 7193bda7..d2d22c35 100644 --- a/argoverse/utils/transform.py +++ b/argoverse/utils/transform.py @@ -16,6 +16,26 @@ logger = logging.getLogger(__name__) +def yaw_to_quaternion3d(yaw: float) -> np.ndarray: + """Convert a rotation angle in the xy plane (i.e. about the z axis) to a quaternion. + + Args: + yaw: angle to rotate about the z-axis, representing an Euler angle, in radians + + Returns: + array w/ quaternion coefficients (qw,qx,qy,qz) in scalar-first order, per Argoverse convention. + """ + qx, qy, qz, qw = Rotation.from_euler(seq="z", angles=yaw, degrees=False).as_quat() + return np.array([qw, qx, qy, qz]) + + +def rotmat2quat(R: np.ndarray) -> np.ndarray: + """Convert a rotation-matrix to a quaternion in Argo's scalar-first notation (w, x, y, z).""" + quat_xyzw = Rotation.from_matrix(R).as_quat() + quat_wxyz = quat_scipy2argo(quat_xyzw) + return quat_wxyz + + def quat2rotmat(q: np.ndarray) -> np.ndarray: """Normalizes a quaternion to unit-length, then converts it into a rotation matrix. @@ -48,6 +68,13 @@ def quat_argo2scipy(q: np.ndarray) -> np.ndarray: return q_scipy +def quat_scipy2argo(q: np.ndarray) -> np.ndarray: + """Re-order Scipy's scalar-last [x,y,z,w] quaternion order to Argoverse's scalar-first [w,x,y,z].""" + x, y, z, w = q + q_argo = np.array([w, x, y, z]) + return q_argo + + def quat_argo2scipy_vectorized(q: np.ndarray) -> np.ndarray: """Re-order Argoverse's scalar-first [w,x,y,z] quaternion order to Scipy's scalar-last [x,y,z,w]""" return q[..., [1, 2, 3, 0]] diff --git a/tests/test_transform.py b/tests/test_transform.py index c895f046..821c01a1 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -8,7 +8,7 @@ import numpy as np import pytest -from argoverse.utils.transform import quat2rotmat +from argoverse.utils.transform import quat2rotmat, rotmat2quat, yaw_to_quaternion3d EPSILON = 1e-10 @@ -148,7 +148,7 @@ def test_quat2rotmat_5() -> None: assert np.allclose(R_gt, R) -def test_invalid_quaternion_zero_norm(): +def test_invalid_quaternion_zero_norm() -> None: """Ensure that passing a zero-norm quaternion raises an error, as normalization would divide by 0.""" q = np.array([0.0, 0.0, 0.0, 0.0]) @@ -156,7 +156,7 @@ def test_invalid_quaternion_zero_norm(): quat2rotmat(q) -def test_quaternion_renormalized(): +def test_quaternion_renormalized() -> None: """Make sure that a quaternion is correctly re-normalized. Normalized and unnormalized quaternion variants should generate the same 3d rotation matrix. @@ -168,3 +168,61 @@ def test_quaternion_renormalized(): R2 = quat2rotmat(q2) assert np.allclose(R1, R2) + + +def test_rotmat2quat() -> None: + """Ensure `rotmat2quat()` correctly converts rotation matrices to scalar-first quaternions.""" + num_trials = 1000 + for trial in range(num_trials): + + # generate random rotation matrices by sampling quaternion elements from normal distribution + # https://en.wikipedia.org/wiki/Rotation_matrix#Uniform_random_rotation_matrices + + q = np.random.randn(4) + q /= np.linalg.norm(q) + + R = quat2rotmat(q) + q_ = rotmat2quat(R) + + # Note: A unit quaternion multiplied by 1 or -1 represents the same 3d rotation + # https://math.stackexchange.com/questions/2016282/negative-quaternion + # https://math.stackexchange.com/questions/1790521/unit-quaternion-multiplied-by-1 + assert np.allclose(q, q_) or np.allclose(-q, q_) + + +def test_yaw_to_quaternion3d() -> None: + """Ensure yaw_to_quaternion3d() outputs correct values. + + Compare `yaw_to_quaternion3d()` output (which relies upon Scipy) to manual + computation of the 3d rotation matrix, followed by a rotation matrix -> quaternion conversion. + """ + + def rot3d_z(angle_rad: float) -> np.ndarray: + """Generate 3d rotation matrix about the z-axis, from a rotation angle in radians.""" + c = np.cos(angle_rad) + s = np.sin(angle_rad) + R = np.array( + [ + [c, -s, 0], + [s, c, 0], + [0, 0, 1], + ] + ) + return R + + num_trials = 1000 + angle_samples_rad = np.random.rand(num_trials) * 2 * np.pi + for angle_rad in angle_samples_rad: + + R = rot3d_z(angle_rad) + q = rotmat2quat(R) + + q_ = yaw_to_quaternion3d(angle_rad) + + assert isinstance(q_, np.ndarray) + assert q_.size == 4 + + # Note: A unit quaternion multiplied by 1 or -1 represents the same 3d rotation + # https://math.stackexchange.com/questions/2016282/negative-quaternion + # https://math.stackexchange.com/questions/1790521/unit-quaternion-multiplied-by-1 + assert np.allclose(q, q_) or np.allclose(q, -q_)