Merge pull request #78 from jcapriot/rotation_check

Refactor rotation function.

geoana/em/static/wholespace.py

@@ -1,6 +1,5 @@
 import numpy as np
 from scipy.special import ellipk, ellipe
-from scipy.spatial.transform import Rotation
 
 from ..base import BaseDipole, BaseMagneticDipole, BaseEM, BaseLineCurrent
 from ... import spatial
@@ -11,9 +10,8 @@
     "MagneticPoleWholeSpace", "PointCurrentWholeSpace"
 ]
 
-from ...kernels import prism_fzx, prism_fzy
-from ...spatial import cartesian_2_cylindrical, cylindrical_2_cartesian, cylindrical_to_cartesian, \
-    cartesian_2_spherical, spherical_to_cartesian, cartesian_to_spherical, cartesian_to_cylindrical
+from ...kernels import prism_fzy
+from ...spatial import  cylindrical_to_cartesian, cartesian_to_cylindrical, rotation_matrix_from_normals
 
 
 class MagneticDipoleWholeSpace(BaseEM, BaseMagneticDipole):
@@ -108,7 +106,7 @@ def vector_potential(self, xyz, coordinates="cartesian"):
 
         # orientation of the dipole
         if coordinates.lower() == "cylindrical":
-            xyz = spatial.cylindrical_2_cartesian(xyz)
+            xyz = spatial.cylindrical_to_cartesian(xyz)
 
         r_vec = xyz - self.location
         r = np.linalg.norm(r_vec, axis=-1, keepdims=True)
@@ -117,7 +115,7 @@ def vector_potential(self, xyz, coordinates="cartesian"):
         a = self.moment * self.mu / (4 * np.pi * r**2) * np.cross(self.orientation, r_hat)
 
         if coordinates.lower() == "cylindrical":
-            a = spatial.cartesian_2_cylindrical(xyz, a)
+            a = spatial.cartesian_to_cylindrical(xyz, a)
 
         return a
 
@@ -205,7 +203,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"):
         xyz = check_xyz_dim(xyz)
 
         if coordinates.lower() == "cylindrical":
-            xyz = spatial.cylindrical_2_cartesian(xyz)
+            xyz = spatial.cylindrical_to_cartesian(xyz)
 
         r_vec = xyz - self.location
         r = np.linalg.norm(r_vec, axis=-1, keepdims=True)
@@ -216,7 +214,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"):
         )
 
         if coordinates.lower() == "cylindrical":
-            b = spatial.cartesian_2_cylindrical(xyz, b)
+            b = spatial.cartesian_to_cylindrical(xyz, b)
 
         return b
 
@@ -349,7 +347,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"):
         xyz = check_xyz_dim(xyz)
 
         if coordinates.lower() == "cylindrical":
-            xyz = spatial.cylindrical_2_cartesian(xyz)
+            xyz = spatial.cylindrical_to_cartesian(xyz)
 
         r_vec = xyz - self.location
         r = np.linalg.norm(r_vec, axis=-1, keepdims=True)
@@ -358,7 +356,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"):
         b = self.moment * self.mu / (4 * np.pi * r ** 2) * r_hat
 
         if coordinates.lower() == "cylindrical":
-            b = spatial.cartesian_2_cylindrical(xyz, b)
+            b = spatial.cartesian_to_cylindrical(xyz, b)
 
         return b
 
@@ -572,15 +570,15 @@ def vector_potential(self, xyz, coordinates="cartesian"):
 
         # convert coordinates if not cartesian
         if coordinates.lower() == "cylindrical":
-            xyz = spatial.cylindrical_2_cartesian(xyz)
+            xyz = spatial.cylindrical_to_cartesian(xyz)
 
         # define a rotation matrix that rotates my orientation to z:
-        rot, _ = Rotation.align_vectors(np.array([0, 0, 1]), self.orientation)
+        rot = rotation_matrix_from_normals(self.orientation, [0, 0, 1], as_matrix=False)
 
         # rotate the points
         r_vec = rot.apply(xyz.reshape(-1, 3) - self.location)
 
-        r_cyl = cartesian_2_cylindrical(r_vec)
+        r_cyl = cartesian_to_cylindrical(r_vec)
 
         rho = r_cyl[..., 0]
         z = r_cyl[..., 2]
@@ -632,7 +630,7 @@ def vector_potential(self, xyz, coordinates="cartesian"):
         A = rot.apply(A, inverse=True).reshape(xyz.shape)
 
         if coordinates.lower() == "cylindrical":
-            A = spatial.cartesian_2_cylindrical(xyz, A)
+            A = spatial.cartesian_to_cylindrical(xyz, A)
 
         return A
 
@@ -704,15 +702,15 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"):
 
         # convert coordinates if not cartesian
         if coordinates.lower() == "cylindrical":
-            xyz = spatial.cylindrical_2_cartesian(xyz)
+            xyz = spatial.cylindrical_to_cartesian(xyz)
         elif coordinates.lower() != "cartesian":
             raise TypeError(
                 f"coordinates must be 'cartesian' or 'cylindrical', the coordinate "
                 f"system you provided, {coordinates}, is not yet supported."
             )
 
         # define a rotation matrix that rotates my orientation to z:
-        rot, _ = Rotation.align_vectors(np.array([0, 0, 1]), self.orientation)
+        rot = rotation_matrix_from_normals(self.orientation, [0, 0, 1], as_matrix=False)
 
         # rotate the points
         r_vec = rot.apply(xyz.reshape(-1, 3) - self.location)
@@ -873,7 +871,7 @@ def vector_potential(self, xyz):
 
             # find the rotation from the line segments orientation
             # to the x_hat direction.
-            rot, _ = Rotation.align_vectors([1, 0, 0], l_hat.reshape(-1, 3))
+            rot = rotation_matrix_from_normals(l_hat, [1, 0, 0], as_matrix=False)
 
             # shift and rotate the grid points
             r0_vec = rot.apply(xyz.reshape(-1, 3) - p0).reshape(xyz.shape)
@@ -1006,7 +1004,7 @@ def magnetic_flux_density(self, xyz):
 
             # find the rotation from the line segments orientation
             # to the x_hat direction.
-            rot, _ = Rotation.align_vectors([1, 0, 0], l_hat)
+            rot = rotation_matrix_from_normals(l_hat, [1, 0, 0], as_matrix=False)
 
             # shift and rotate the grid points
             r0_vec = rot.apply(xyz.reshape(-1, 3) - p0).reshape(xyz.shape)
@@ -1018,10 +1016,10 @@ def magnetic_flux_density(self, xyz):
             r0_hat = r0_vec / r0
             r1_hat = r1_vec / r1
 
-            cyl_points = cartesian_2_cylindrical(r0_vec[..., 1:])
+            cyl_points = cartesian_to_cylindrical(r0_vec[..., 1:])
 
             temp_storage[..., 1] = (r1_hat[..., 0] - r0_hat[..., 0])/cyl_points[..., 0]
-            temp_storage[..., 1:] = cylindrical_2_cartesian(cyl_points, temp_storage)
+            temp_storage[..., 1:] = cylindrical_to_cartesian(cyl_points, temp_storage)
 
             # the undo the local rotation...
             out += rot.apply(temp_storage.reshape(-1, 3), inverse=True).reshape(xyz.shape)

geoana/spatial.py

@@ -23,6 +23,8 @@
 
 """
 import numpy as np
+from scipy.spatial.transform import Rotation
+
 from .utils import mkvc
 
 
@@ -531,20 +533,17 @@ def repeat_scalar(scalar, dim=3):
 
     return np.kron(np.ones((1, dim)), np.atleast_2d(scalar).T)
 
-def rotation_matrix_from_normals(v0, v1, tol=1e-20):
+
+def rotation_matrix_from_normals(v0, v1, tol=1e-20, as_matrix=True):
     """
     Generate a 3x3 rotation matrix defining the rotation from vector v0 to v1.
 
-    This function uses Rodrigues' rotation formula to generate the rotation
-    matrix :math:`\\mathbf{A}` going from vector :math:`\\mathbf{v_0}` to
-    vector :math:`\\mathbf{v_1}`. Thus:
+    This function builds a quaternion representing the rotation, then constructs
+    the rotation matrix.
 
     .. math::
         \\mathbf{Av_0} = \\mathbf{v_1}
 
-    For detailed desciption of the algorithm, see
-    https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
-
     Parameters
     ----------
     v0 : (3) numpy.ndarray
@@ -554,46 +553,61 @@ def rotation_matrix_from_normals(v0, v1, tol=1e-20):
     tol : float, optional
         Numerical tolerance. If the length of the rotation axis is below this value,
         it is assumed to be no rotation, and an identity matrix is returned.
+    as_matrix : bool, optional
+        If ``True``, a ``(3,3) numpy.ndarray`` is returned.
+        If ``False``, the respective ``scipy.spatial.transform.Rotation`` object is returned.
 
     Returns
     -------
-    (3, 3) numpy.ndarray
-        The rotation matrix from v0 to v1.
+    (3, 3) numpy.ndarray or scipy.spatial.transform.Rotation
+        The rotation matrix from v0 to v1, whose type depends on the `as_matrix` parameter.
     """
 
-    v0 = mkvc(v0)
-    v1 = mkvc(v1)
+    v0 = np.asarray(v0, dtype=float, copy=True).squeeze()
+    v1 = np.asarray(v1, dtype=float, copy=True).squeeze()
 
-    # ensure both n0, n1 are vectors of length 1
-    assert len(v0) == 3, "Length of n0 should be 3"
-    assert len(v1) == 3, "Length of n1 should be 3"
+    if v0.shape != (3, ):
+        raise ValueError(f"v0 shape should be (3,), got {v0.shape}")
+
+    if v1.shape != (3, ):
+        raise ValueError(f"v1 shape should be (3,), got {v1.shape}")
 
     # ensure both are true normals
-    n0 = v0*1./np.linalg.norm(v0)
-    n1 = v1*1./np.linalg.norm(v1)
+    v0 /= np.linalg.norm(v0)
+    v1 /= np.linalg.norm(v1)
 
-    n0dotn1 = n0.dot(n1)
+    v0dotv1 = v0.dot(v1)
 
     # define the rotation axis, which is the cross product of the two vectors
-    rotAx = np.cross(n0, n1)
-
-    if np.linalg.norm(rotAx) < tol:
-        return np.eye(3, dtype=float)
-
-    rotAx *= 1./np.linalg.norm(rotAx)
-
-    cosT = n0dotn1/(np.linalg.norm(n0)*np.linalg.norm(n1))
-    sinT = np.sqrt(1.-n0dotn1**2)
-
-    ux = np.array(
-        [
-            [0., -rotAx[2], rotAx[1]],
-            [rotAx[2], 0., -rotAx[0]],
-            [-rotAx[1], rotAx[0], 0.]
-        ], dtype=float
-    )
+    rotation_axis = np.cross(v0, v1)
+
+    if np.linalg.norm(rotation_axis) < tol:
+        # check if vectors were anti-parallel.
+        if v0dotv1 < 0:
+            # find another vector that is perpendicular to v1,
+            # by crossing with z_hat or y_hat (One of them must
+            # work because it can't be parallel to both of them)
+            trial = np.cross(v0, np.array([0., 0., 1.]))
+            if np.linalg.norm(trial) > tol:
+                rotation_axis = trial
+            else:
+                rotation_axis = np.cross(v0, np.array([0., 1., 0.]))
+        else:
+            mat = np.eye(3, dtype=float)
+            if as_matrix:
+                return mat
+            else:
+                return Rotation.from_matrix(mat)
+
+
+    w = 1 + v0dotv1
+    rot = Rotation.from_quat(np.r_[rotation_axis, w])
+
+    if as_matrix:
+        return rot.as_matrix()
+    else:
+        return rot
 
-    return np.eye(3, dtype=float) + sinT*ux + (1.-cosT)*(ux.dot(ux))
 
 
 def rotate_points_from_normals(xyz, n0, n1, x0=np.r_[0., 0., 0.]):

tests/test_spatial.py

@@ -5,6 +5,7 @@
 import numpy.testing as npt
 from geoana import utils
 from geoana import spatial
+from geoana.spatial import rotation_matrix_from_normals
 
 
 def test_errors():
@@ -208,6 +209,51 @@ def test_aliases(self):
         np.testing.assert_equal(s2c, spatial.spherical_to_cartesian(grid, vec))
         np.testing.assert_equal(c2s, spatial.cartesian_to_spherical(grid, vec))
 
+@pytest.mark.parametrize(
+    'source_vector',
+    [
+        [0, 0, 1],
+        [0, 1, 0],
+        [1, 0, 0],
+        [0, 0, -1],
+        [0, -1, 0],
+        [-1, 0, 0],
+        [2/np.sqrt(5), 0, 1/np.sqrt(5)],
+        [2/np.sqrt(6), 1/np.sqrt(6), 1/np.sqrt(6)],
+        [-2/np.sqrt(6), -1/np.sqrt(6), -1/np.sqrt(6)],
+     ],
+)
+@pytest.mark.parametrize(
+    'target_vector',
+    [
+        [0, 0, 1],
+        [0, 1, 0],
+        [1, 0, 0],
+        [0, 0, -1],
+        [0, -1, 0],
+        [-1, 0, 0],
+        [2/np.sqrt(5), 0, 1/np.sqrt(5)],
+        [2/np.sqrt(6), 1/np.sqrt(6), 1/np.sqrt(6)],
+        [-2/np.sqrt(6), -1/np.sqrt(6), -1/np.sqrt(6)],
+     ],
+)
+@pytest.mark.parametrize('as_matrix', [True, False])
+def test_rotation(source_vector, target_vector, as_matrix):
+
+    rot = rotation_matrix_from_normals(source_vector, target_vector, as_matrix=as_matrix)
+    atol = 1E-15
+    if as_matrix:
+        npt.assert_allclose(rot @ source_vector, target_vector, atol=atol)
+        npt.assert_allclose(rot.T @ target_vector, source_vector, atol=atol)
+    else:
+        npt.assert_allclose(rot.apply(source_vector), target_vector, atol=atol)
+        npt.assert_allclose(rot.apply(target_vector, inverse=True), source_vector, atol=atol)
+
+def test_rotation_errors():
+    with pytest.raises(ValueError, match="v0 shape should be.*"):
+        rotation_matrix_from_normals([0, 1, 2, 3], [0, 1, 3])
+    with pytest.raises(ValueError, match="v1 shape should be.*"):
+        rotation_matrix_from_normals([0, 1, 2], [0, 1, 3, 3])
 
 if __name__ == '__main__':
     unittest.main()