Merge pull request #248 from johnwlambert/patch-3

add rotation matrix -> quaternion and yaw -> quaternion conversion utility functions

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]]

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,15 +148,15 @@ 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])
 
     with pytest.raises(ZeroDivisionError) as e_info:
         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_)