Natural 3D rotation with mouse

MischaMegens2 · MischaMegens2 · commit ad265179748f · 2024-06-03T11:06:21.000-07:00
- Addresses Issue #28288 - Introduces three-dimensional rotation by mouse using a variation on Ken Shoemake's ARCBALL - Provides a minimal Quaternion class, to avoid an additional dependency on a large package like 'numpy-quaternion'
diff --git a/doc/users/next_whats_new/mouse_rotation.rst b/doc/users/next_whats_new/mouse_rotation.rst
@@ -0,0 +1,12 @@
+Rotating 3d plots with the mouse
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Rotating three-dimensional plots with the mouse has been made more intuitive.
+The plot now reacts the same way to mouse movement, independent of the
+particular orientation at hand; and it is possible to control all 3 rotational
+degrees of freedom (azimuth, elevation, and roll). It uses a variation on
+Ken Shoemake's ARCBALL [Shoemake1992]_.
+
+.. [Shoemake1992] Ken Shoemake, "ARCBALL: A user interface for specifying
+  three-dimensional rotation using a mouse." in Proceedings of Graphics
+  Interface '92, 1992, pp. 151-156, https://doi.org/10.20380/GI1992.18
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -14,6 +14,7 @@
 import itertools
 import math
 import textwrap
+import warnings
 
 import numpy as np
 
@@ -1502,6 +1503,23 @@ def _calc_coord(self, xv, yv, renderer=None):
         p2 = p1 - scale*vec
         return p2, pane_idx
 
+    def _arcball(self, p):
+        """
+        Convert a point p = [x, y] to a point on a virtual trackball
+        This is Ken Shoemake's arcball
+        See: Ken Shoemake, "ARCBALL: A user interface for specifying
+        three-dimensional rotation using a mouse." in
+        Proceedings of Graphics Interface '92, 1992, pp. 151-156,
+        https://doi.org/10.20380/GI1992.18
+        """
+        p = 2 * p
+        r = p[0]**2 + p[1]**2
+        if r > 1:
+            p = np.concatenate(([0], p/math.sqrt(r)))
+        else:
+            p = np.concatenate(([math.sqrt(1-r)], p))
+        return p
+
     def _on_move(self, event):
         """
         Mouse moving.
@@ -1537,12 +1555,25 @@ def _on_move(self, event):
             if dx == 0 and dy == 0:
                 return
 
+            # Convert to quaternion
+            elev = np.deg2rad(self.elev)
+            azim = np.deg2rad(self.azim)
             roll = np.deg2rad(self.roll)
-            delev = -(dy/h)*180*np.cos(roll) + (dx/w)*180*np.sin(roll)
-            dazim = -(dy/h)*180*np.sin(roll) - (dx/w)*180*np.cos(roll)
-            elev = self.elev + delev
-            azim = self.azim + dazim
-            roll = self.roll
+            q = _Quaternion.from_cardan_angles(elev, azim, roll)
+
+            # Update quaternion - a variation on Ken Shoemake's ARCBALL
+            current_point = np.array([self._sx/w, self._sy/h])
+            new_point = np.array([x/w, y/h])
+            current_vec = self._arcball(current_point)
+            new_vec = self._arcball(new_point)
+            dq = _Quaternion.rotate_from_to(current_vec, new_vec)
+            q = dq*q
+
+            # Convert to elev, azim, roll
+            elev, azim, roll = q.as_cardan_angles()
+            azim = np.rad2deg(azim)
+            elev = np.rad2deg(elev)
+            roll = np.rad2deg(roll)
             vertical_axis = self._axis_names[self._vertical_axis]
             self.view_init(
                 elev=elev,
@@ -3725,3 +3756,123 @@ def get_test_data(delta=0.05):
     Y = Y * 10
     Z = Z * 500
     return X, Y, Z
+
+
+class _Quaternion:
+    """
+    Quaternions
+    consisting of scalar, along 1, and vector, with components along i, j, k
+    """
+
+    def __init__(self, scalar, vector):
+        self.scalar = scalar
+        self.vector = np.array(vector)
+
+    def __neg__(self):
+        return self.__class__(-self.scalar, -self.vector)
+
+    def __mul__(self, other):
+        """
+        Product of two quaternions
+        i*i = j*j = k*k = i*j*k = -1
+        Quaternion multiplication can be expressed concisely
+        using scalar and vector parts,
+        see <https://en.wikipedia.org/wiki/Quaternion#Scalar_and_vector_parts>
+        """
+        return self.__class__(
+            self.scalar*other.scalar - np.dot(self.vector, other.vector),
+            self.scalar*other.vector + self.vector*other.scalar
+            + np.cross(self.vector, other.vector))
+
+    def conjugate(self):
+        """The conjugate quaternion -(1/2)*(q+i*q*i+j*q*j+k*q*k)"""
+        return self.__class__(self.scalar, -self.vector)
+
+    @property
+    def norm(self):
+        """The 2-norm, q*q', a scalar"""
+        return self.scalar*self.scalar + np.dot(self.vector, self.vector)
+
+    def normalize(self):
+        """Scaling such that norm equals 1"""
+        n = np.sqrt(self.norm)
+        return self.__class__(self.scalar/n, self.vector/n)
+
+    def reciprocal(self):
+        """The reciprocal, 1/q = q'/(q*q') = q' / norm(q)"""
+        n = self.norm
+        return self.__class__(self.scalar/n, -self.vector/n)
+
+    def __div__(self, other):
+        return self*other.reciprocal()
+
+    __truediv__ = __div__
+
+    def rotate(self, v):
+        # Rotate the vector v by the quaternion q, i.e.,
+        # calculate (the vector part of) q*v/q
+        v = self.__class__(0, v)
+        v = self*v/self
+        return v.vector
+
+    def __eq__(self, other):
+        return (self.scalar == other.scalar) and (self.vector == other.vector).all
+
+    def __repr__(self):
+        return "_Quaternion({}, {})".format(repr(self.scalar), repr(self.vector))
+
+    @classmethod
+    def rotate_from_to(cls, r1, r2):
+        """
+        The quaternion for the shortest rotation from vector r1 to vector r2
+        i.e., q = sqrt(r2*r1'), normalized.
+        If r1 and r2 are antiparallel, then the result is ambiguous;
+        a normal vector will be returned, and a warning will be issued.
+        """
+        k = np.cross(r1, r2)
+        nk = np.linalg.norm(k)
+        th = np.arctan2(nk, np.dot(r1, r2))
+        th = th/2
+        if nk == 0:  # r1 and r2 are parallel or anti-parallel
+            if np.dot(r1, r2) < 0:
+                warnings.warn("Rotation defined by anti-parallel vectors is ambiguous")
+                k = np.zeros(3)
+                k[np.argmin(r1*r1)] = 1  # basis vector most perpendicular to r1-r2
+                k = np.cross(r1, k)
+                k = k / np.linalg.norm(k)  # unit vector normal to r1-r2
+                q = cls(0, k)
+            else:
+                q = cls(1, [0, 0, 0])  # = 1, no rotation
+        else:
+            q = cls(math.cos(th), k*math.sin(th)/nk)
+        return q
+
+    @classmethod
+    def from_cardan_angles(cls, elev, azim, roll):
+        """
+        Converts the angles to a quaternion
+            q = exp((roll/2)*e_x)*exp((elev/2)*e_y)*exp((-azim/2)*e_z)
+        i.e., the angles are a kind of Tait-Bryan angles, -z,y',x".
+        The angles should be given in radians, not degrees.
+        """
+        ca, sa = np.cos(azim/2), np.sin(azim/2)
+        ce, se = np.cos(elev/2), np.sin(elev/2)
+        cr, sr = np.cos(roll/2), np.sin(roll/2)
+
+        qw = ca*ce*cr + sa*se*sr
+        qx = ca*ce*sr - sa*se*cr
+        qy = ca*se*cr + sa*ce*sr
+        qz = ca*se*sr - sa*ce*cr
+        return cls(qw, [qx, qy, qz])
+
+    def as_cardan_angles(self):
+        """
+        The inverse of `from_cardan_angles()`.
+        Note that the angles returned are in radians, not degrees.
+        """
+        qw = self.scalar
+        qx, qy, qz = self.vector[..., :]
+        azim = np.arctan2(2*(-qw*qz+qx*qy), qw*qw+qx*qx-qy*qy-qz*qz)
+        elev = np.arcsin( 2*( qw*qy+qz*qx)/(qw*qw+qx*qx+qy*qy+qz*qz))  # noqa E201
+        roll = np.arctan2(2*( qw*qx-qy*qz), qw*qw-qx*qx-qy*qy+qz*qz)   # noqa E201
+        return elev, azim, roll
diff --git a/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py b/lib/mpl_toolkits/mplot3d/tests/test_axes3d.py
@@ -5,6 +5,7 @@
 import pytest
 
 from mpl_toolkits.mplot3d import Axes3D, axes3d, proj3d, art3d
+from mpl_toolkits.mplot3d.axes3d import _Quaternion as Quaternion
 import matplotlib as mpl
 from matplotlib.backend_bases import (MouseButton, MouseEvent,
                                       NavigationToolbar2)
@@ -1766,29 +1767,127 @@ def test_shared_axes_retick():
     assert ax2.get_zlim() == (-0.5, 2.5)
 
 
+def test_quaternion():
+    # 1:
+    q1 = Quaternion(1, [0, 0, 0])
+    assert q1.scalar == 1
+    assert (q1.vector == [0, 0, 0]).all
+    # __neg__:
+    assert (-q1).scalar == -1
+    assert ((-q1).vector == [0, 0, 0]).all
+    # i, j, k:
+    qi = Quaternion(0, [1, 0, 0])
+    assert qi.scalar == 0
+    assert (qi.vector == [1, 0, 0]).all
+    qj = Quaternion(0, [0, 1, 0])
+    assert qj.scalar == 0
+    assert (qj.vector == [0, 1, 0]).all
+    qk = Quaternion(0, [0, 0, 1])
+    assert qk.scalar == 0
+    assert (qk.vector == [0, 0, 1]).all
+    # i^2 = j^2 = k^2 = -1:
+    assert qi*qi == -q1
+    assert qj*qj == -q1
+    assert qk*qk == -q1
+    # identity:
+    assert q1*qi == qi
+    assert q1*qj == qj
+    assert q1*qk == qk
+    # i*j=k, j*k=i, k*i=j:
+    assert qi*qj == qk
+    assert qj*qk == qi
+    assert qk*qi == qj
+    assert qj*qi == -qk
+    assert qk*qj == -qi
+    assert qi*qk == -qj
+    # __mul__:
+    assert (Quaternion(2, [3, 4, 5]) * Quaternion(6, [7, 8, 9])
+            == Quaternion(-86, [28, 48, 44]))
+    # conjugate():
+    for q in [q1, qi, qj, qk]:
+        assert q.conjugate().scalar == q.scalar
+        assert (q.conjugate().vector == -q.vector).all
+        assert q.conjugate().conjugate() == q
+        assert ((q*q.conjugate()).vector == 0).all
+    # norm:
+    q0 = Quaternion(0, [0, 0, 0])
+    assert q0.norm == 0
+    assert q1.norm == 1
+    assert qi.norm == 1
+    assert qj.norm == 1
+    assert qk.norm == 1
+    for q in [q0, q1, qi, qj, qk]:
+        assert q.norm == (q*q.conjugate()).scalar
+    # normalize():
+    for q in [
+        Quaternion(2, [0, 0, 0]),
+        Quaternion(0, [3, 0, 0]),
+        Quaternion(0, [0, 4, 0]),
+        Quaternion(0, [0, 0, 5]),
+        Quaternion(6, [7, 8, 9])
+    ]:
+        assert q.normalize().norm == 1
+    # reciprocal():
+    for q in [q1, qi, qj, qk]:
+        assert q*q.reciprocal() == q1
+        assert q.reciprocal()*q == q1
+    # rotate():
+    assert (qi.rotate([1, 2, 3]) == np.array([1, -2, -3])).all
+    # rotate_from_to():
+    for r1, r2, q in [
+        ([1, 0, 0], [0, 1, 0], Quaternion(np.sqrt(1/2), [0, 0, np.sqrt(1/2)])),
+        ([1, 0, 0], [0, 0, 1], Quaternion(np.sqrt(1/2), [0, -np.sqrt(1/2), 0])),
+        ([1, 0, 0], [1, 0, 0], Quaternion(1, [0, 0, 0]))
+    ]:
+        assert Quaternion.rotate_from_to(r1, r2) == q
+    # rotate_from_to(), special case:
+    for r1 in [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1]]:
+        r1 = np.array(r1)
+        with pytest.warns(UserWarning):
+            q = Quaternion.rotate_from_to(r1, -r1)
+        assert np.isclose(q.norm, 1)
+        assert np.dot(q.vector, r1) == 0
+    # from_cardan_angles(), as_cardan_angles():
+    for elev, azim, roll in [(0, 0, 0),
+                             (90, 0, 0), (0, 90, 0), (0, 0, 90),
+                             (0, 30, 30), (30, 0, 30), (30, 30, 0),
+                             (47, 11, -24)]:
+        for mag in [1, 2]:
+            q = Quaternion.from_cardan_angles(
+                np.deg2rad(elev), np.deg2rad(azim), np.deg2rad(roll))
+            assert np.isclose(q.norm, 1)
+            q = Quaternion(mag * q.scalar, mag * q.vector)
+            e, a, r = np.rad2deg(Quaternion.as_cardan_angles(q))
+            assert np.isclose(e, elev)
+            assert np.isclose(a, azim)
+            assert np.isclose(r, roll)
+
+
 def test_rotate():
     """Test rotating using the left mouse button."""
-    for roll in [0, 30]:
+    for roll, dx, dy, new_elev, new_azim, new_roll in [
+            [0, 0.5, 0, 0, -90, 0],
+            [30, 0.5, 0, 30, -90, 0],
+            [0, 0, 0.5, -90, 0, 0],
+            [30, 0, 0.5, -60, -90, 90],
+            [0, 0.5, 0.5, -45, -90, 45],
+            [30, 0.5, 0.5, -15, -90, 45]]:
         fig = plt.figure()
         ax = fig.add_subplot(1, 1, 1, projection='3d')
         ax.view_init(0, 0, roll)
         ax.figure.canvas.draw()
 
-        # drag mouse horizontally to change azimuth
-        dx = 0.1
-        dy = 0.2
+        # drag mouse to change orientation
         ax._button_press(
             mock_event(ax, button=MouseButton.LEFT, xdata=0, ydata=0))
         ax._on_move(
             mock_event(ax, button=MouseButton.LEFT,
                            xdata=dx*ax._pseudo_w, ydata=dy*ax._pseudo_h))
         ax.figure.canvas.draw()
-        roll_radians = np.deg2rad(ax.roll)
-        cs = np.cos(roll_radians)
-        sn = np.sin(roll_radians)
-        assert ax.elev == (-dy*180*cs + dx*180*sn)
-        assert ax.azim == (-dy*180*sn - dx*180*cs)
-        assert ax.roll == roll
+
+        assert np.isclose(ax.elev, new_elev)
+        assert np.isclose(ax.azim, new_azim)
+        assert np.isclose(ax.roll, new_roll)
 
 
 def test_pan():
