Fix order argument for quat_to_euler, quat_from_euler & mat_from_eulet (#94)

Korijn · web-flow · commit fc49cdb6aacf · 2024-12-28T18:17:27.000+01:00
* add failing tests

* fix quat_to_euler and quat_from_euler

* self review

* bump version because its a breaking change

* also align mat_from_euler
diff --git a/pylinalg/matrix.py b/pylinalg/matrix.py
@@ -1,5 +1,3 @@
-from math import cos, sin
-
 import numpy as np
 from numpy.lib.stride_tricks import as_strided
 
@@ -129,8 +127,9 @@ def mat_from_euler(angles, /, *, order="xyz", out=None, dtype=None):
     angles : ndarray, [3]
         The euler angles.
     order : string, optional
-        The order in which the rotations should be applied. Default
-        is "xyz".
+        The rotation order as a string. Can include 'X', 'Y', 'Z' for intrinsic
+        rotation (uppercase) or 'x', 'y', 'z' for extrinsic rotation (lowercase).
+        Default is "xyz".
     out : ndarray, optional
         A location into which the result is stored. If provided, it
         must have a shape that the inputs broadcast to. If not provided or
@@ -155,31 +154,42 @@ def mat_from_euler(angles, /, *, order="xyz", out=None, dtype=None):
     ccw around the y-axis, and finally 0° around the x axis.
 
     """
-    angles = np.asarray(angles, dtype=float)
-    order = order.lower()
-
-    if angles.ndim == 0:
-        # add dimension to allow zip
-        angles = angles[None]
-
-    matrices = []
+    fill_out_first = out is not None
+    angles = np.atleast_1d(np.asarray(angles))
+    extrinsic = order.islower()
+    order = order.upper()
+    axis_lookup = {"X": 0, "Y": 1, "Z": 2}
+    if extrinsic:
+        angles = reversed(angles)
+        order = reversed(order)
     for angle, axis in zip(angles, order):
-        axis_idx = {"x": 0, "y": 1, "z": 2}[axis]
-
-        matrix = np.array([[cos(angle), -sin(angle)], [sin(angle), cos(angle)]])
-        if axis_idx == 1:
-            matrix = matrix.T
-        matrix = np.insert(matrix, axis_idx, 0, axis=0)
-        matrix = np.insert(matrix, axis_idx, 0, axis=1)
-        matrix[axis_idx, axis_idx] = 1
-
+        axis_idx = axis_lookup[axis]
         affine_matrix = np.identity(4, dtype=dtype)
-        affine_matrix[:3, :3] = matrix
-
-        matrices.append(affine_matrix)
 
-    # note: combining in the loop would save time and memory usage
-    return mat_combine([x for x in reversed(matrices)], out=out, dtype=dtype)
+        if axis_idx == 0:
+            affine_matrix[1, 1] = np.cos(angle)
+            affine_matrix[1, 2] = -np.sin(angle)
+            affine_matrix[2, 1] = np.sin(angle)
+            affine_matrix[2, 2] = np.cos(angle)
+        elif axis_idx == 1:
+            affine_matrix[0, 0] = np.cos(angle)
+            affine_matrix[0, 2] = np.sin(angle)
+            affine_matrix[2, 0] = -np.sin(angle)
+            affine_matrix[2, 2] = np.cos(angle)
+        elif axis_idx == 2:
+            affine_matrix[0, 0] = np.cos(angle)
+            affine_matrix[0, 1] = -np.sin(angle)
+            affine_matrix[1, 0] = np.sin(angle)
+            affine_matrix[1, 1] = np.cos(angle)
+
+        if fill_out_first:
+            out[:] = affine_matrix
+            fill_out_first = False
+        elif out is None:
+            out = affine_matrix
+        else:
+            out @= affine_matrix
+    return out
 
 
 def mat_from_axis_angle(axis, angle, /, *, out=None, dtype=None):
diff --git a/pylinalg/quaternion.py b/pylinalg/quaternion.py
@@ -260,7 +260,7 @@ def quat_from_axis_angle(axis, angle, /, *, out=None, dtype=None):
     return out
 
 
-def quat_from_euler(angles, /, *, order="XYZ", out=None, dtype=None):
+def quat_from_euler(angles, /, *, order="xyz", out=None, dtype=None):
     """
     Create a quaternion from euler angles.
 
@@ -269,8 +269,9 @@ def quat_from_euler(angles, /, *, order="XYZ", out=None, dtype=None):
     angles : ndarray, [3]
         A set of XYZ euler angles.
     order : string, optional
-        The order in which the rotations should be applied. Default
-        is "xyz".
+        The rotation order as a string. Can include 'X', 'Y', 'Z' for intrinsic
+        rotation (uppercase) or 'x', 'y', 'z' for extrinsic rotation (lowercase).
+        Default is "xyz".
     out : ndarray, optional
         A location into which the result is stored. If provided, it
         must have a shape that the inputs broadcast to. If not provided or
@@ -293,8 +294,9 @@ def quat_from_euler(angles, /, *, order="XYZ", out=None, dtype=None):
         out = np.empty((*batch_shape, 4), dtype=dtype)
 
     # work out the sequence in which to apply the rotations
-    is_extrinsic = [x.isupper() for x in order]
-    order = [{"X": 0, "Y": 1, "Z": 2}[x.upper()] for x in order]
+    is_extrinsic = [x.islower() for x in order]
+    basis_index = {"x": 0, "y": 1, "z": 2}
+    order = [basis_index[x] for x in order.lower()]
 
     # convert each euler matrix into a quaternion
     quaternions = np.zeros((len(order), *batch_shape, 4), dtype=float)
@@ -305,9 +307,9 @@ def quat_from_euler(angles, /, *, order="XYZ", out=None, dtype=None):
     out[:] = quaternions[0]
     for idx in range(1, len(quaternions)):
         if is_extrinsic[idx]:
-            quat_mul(out, quaternions[idx], out=out)
-        else:
             quat_mul(quaternions[idx], out, out=out)
+        else:
+            quat_mul(out, quaternions[idx], out=out)
 
     return out
 
diff --git a/pylinalg/vector.py b/pylinalg/vector.py
@@ -480,13 +480,17 @@ def vec_spherical_safe(vector, /, *, out=None, dtype=None):
     return out
 
 
-def quat_to_euler(quaternion, /, *, out=None, dtype=None):
-    """Convert quaternions to XYZ Euler angles.
+def quat_to_euler(quaternion, /, *, order="xyz", out=None, dtype=None):
+    """Convert quaternions to Euler angles with specified rotation order.
 
     Parameters
     ----------
     quaternion : ndarray, [4]
         The quaternion to convert (in xyzw format).
+    order : str, optional
+        The rotation order as a string. Can include 'X', 'Y', 'Z' for intrinsic
+        rotation (uppercase) or 'x', 'y', 'z' for extrinsic rotation (lowercase).
+        Default is "xyz".
     out : ndarray, optional
         A location into which the result is stored. If provided, it
         must have a shape that the inputs broadcast to. If not provided or
@@ -498,35 +502,130 @@ def quat_to_euler(quaternion, /, *, out=None, dtype=None):
     Returns
     -------
     out : ndarray, [3]
-        The XYZ Euler angles.
+        The Euler angles in the specified order.
+
+    References
+    ----------
+    This implementation is based on the method described in the following paper:
+    Bernardes E, Viollet S (2022) Quaternion to Euler angles conversion: A direct, general and computationally efficient method.
+    PLoS ONE 17(11): e0276302. https://doi.org/10.1371/journal.pone.0276302
     """
     quaternion = np.asarray(quaternion, dtype=float)
+    quat = np.atleast_2d(quaternion)
+    num_rotations = quat.shape[0]
 
     if out is None:
         out = np.empty((*quaternion.shape[:-1], 3), dtype=dtype)
 
-    ysqr = quaternion[..., 1] ** 2
-
-    t0 = 2 * (
-        quaternion[..., 3] * quaternion[..., 0]
-        + quaternion[..., 1] * quaternion[..., 2]
-    )
-    t1 = 1 - 2 * (quaternion[..., 0] ** 2 + ysqr)
-    out[..., 0] = np.arctan2(t0, t1)
-
-    t2 = 2 * (
-        quaternion[..., 3] * quaternion[..., 1]
-        - quaternion[..., 2] * quaternion[..., 0]
-    )
-    t2 = np.clip(t2, a_min=-1, a_max=1)
-    out[..., 1] = np.arcsin(t2)
-
-    t3 = 2 * (
-        quaternion[..., 3] * quaternion[..., 2]
-        + quaternion[..., 0] * quaternion[..., 1]
-    )
-    t4 = 1 - 2 * (ysqr + quaternion[..., 2] ** 2)
-    out[..., 2] = np.arctan2(t3, t4)
+    extrinsic = order.islower()
+    order = order.lower()
+    if not extrinsic:
+        order = order[::-1]
+
+    basis_index = {"x": 0, "y": 1, "z": 2}
+    i = basis_index[order[0]]
+    j = basis_index[order[1]]
+    k = basis_index[order[2]]
+
+    is_proper = i == k
+
+    if is_proper:
+        k = 3 - i - j  # get third axis
+
+    # Step 0
+    # Check if permutation is even (+1) or odd (-1)
+    sign = int((i - j) * (j - k) * (k - i) / 2)
+
+    eps = 1e-7
+    for ind in range(num_rotations):
+        if num_rotations == 1 and out.ndim == 1:
+            _angles = out
+        else:
+            _angles = out[ind, :]
+
+        if is_proper:
+            a = quat[ind, 3]
+            b = quat[ind, i]
+            c = quat[ind, j]
+            d = quat[ind, k] * sign
+        else:
+            a = quat[ind, 3] - quat[ind, j]
+            b = quat[ind, i] + quat[ind, k] * sign
+            c = quat[ind, j] + quat[ind, 3]
+            d = quat[ind, k] * sign - quat[ind, i]
+
+        n2 = a**2 + b**2 + c**2 + d**2
+
+        # Step 3
+        # Compute second angle...
+        # _angles[1] = 2*np.arccos(np.sqrt((a**2 + b**2) / n2))
+        _angles[1] = np.arccos(2 * (a**2 + b**2) / n2 - 1)
+
+        # ... and check if it is equal to 0 or pi, causing a singularity
+        safe1 = np.abs(_angles[1]) >= eps
+        safe2 = np.abs(_angles[1] - np.pi) >= eps
+        safe = safe1 and safe2
+
+        # Step 4
+        # compute first and third angles, according to case
+        if safe:
+            half_sum = np.arctan2(b, a)  # == (alpha+gamma)/2
+            half_diff = np.arctan2(-d, c)  # == (alpha-gamma)/2
+
+            _angles[0] = half_sum + half_diff
+            _angles[2] = half_sum - half_diff
+
+        else:
+            # _angles[0] = 0
+
+            if not extrinsic:
+                # For intrinsic, set first angle to zero so that after reversal we
+                # ensure that third angle is zero
+                # 6a
+                if not safe:
+                    _angles[0] = 0
+                if not safe1:
+                    half_sum = np.arctan2(b, a)
+                    _angles[2] = 2 * half_sum
+                # 6c
+                if not safe2:
+                    half_diff = np.arctan2(-d, c)
+                    _angles[2] = -2 * half_diff
+            else:
+                # For extrinsic, set third angle to zero
+                # 6b
+                if not safe:
+                    _angles[2] = 0
+                if not safe1:
+                    half_sum = np.arctan2(b, a)
+                    _angles[0] = 2 * half_sum
+                # 6c
+                if not safe2:
+                    half_diff = np.arctan2(-d, c)
+                    _angles[0] = 2 * half_diff
+
+        for i_ in range(3):
+            if _angles[i_] < -np.pi:
+                _angles[i_] += 2 * np.pi
+            elif _angles[i_] > np.pi:
+                _angles[i_] -= 2 * np.pi
+
+        # for Tait-Bryan angles
+        if not is_proper:
+            _angles[2] *= sign
+            _angles[1] -= np.pi / 2
+
+        if not extrinsic:
+            # reversal
+            _angles[0], _angles[2] = _angles[2], _angles[0]
+
+        # Step 8
+        if not safe:
+            logger.warning(
+                "Gimbal lock detected. Setting third angle to zero "
+                "since it is not possible to uniquely determine "
+                "all angles."
+            )
 
     return out
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -2,7 +2,7 @@
 
 [project]
 name = "pylinalg"
-version = "0.5.1"
+version = "0.6.0"
 description = "Linear algebra utilities for Python"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -30,6 +30,7 @@ tests = [
     "hypothesis[numpy]~=6.61.0",
     "packaging",
     "twine",
+    "scipy",
 ]
 dev = ["pylinalg[lint,tests,docs]"]
 
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
@@ -356,3 +356,27 @@ def test_mat_look_at(eye, target, up_reference):
     # (map up_reference from target to source space and check if it's in the YZ-plane)
     new_reference = rotation.T @ la.vec_homogeneous(up_reference)
     assert np.abs(new_reference[0]) < 1e-10
+
+
+def test_mat_euler_vs_scipy():
+    """Compare our implementation with scipy's."""
+    from scipy.spatial.transform import Rotation as R  # noqa: N817
+
+    cases = [
+        ("XYZ", [np.pi / 2, np.pi / 180, 0]),
+        ("xyz", [np.pi / 2, np.pi / 180, 0]),
+        ("ZXY", [np.pi, np.pi / 180, -np.pi / 180]),
+        ("zxy", [np.pi, np.pi / 180, -np.pi / 180]),
+        ("ZYX", [0, np.pi / 2, np.pi / 2]),
+        ("zyx", [0, np.pi / 2, np.pi / 2]),
+    ]
+
+    for order, angles in cases:
+        scipy_mat = np.identity(4)
+        scipy_mat[:3, :3] = R.from_euler(order, angles).as_matrix()
+
+        npt.assert_allclose(
+            la.mat_from_euler(angles, order=order),
+            scipy_mat,
+            atol=1e-15,
+        )
diff --git a/tests/test_vectors.py b/tests/test_vectors.py
