diff --git a/geoana/em/static/wholespace.py b/geoana/em/static/wholespace.py index facccb1e..4d9e828f 100644 --- a/geoana/em/static/wholespace.py +++ b/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,7 +702,7 @@ 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 " @@ -712,7 +710,7 @@ def magnetic_flux_density(self, xyz, coordinates="cartesian"): ) # 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) diff --git a/geoana/spatial.py b/geoana/spatial.py index 5c6d067f..163d0a69 100644 --- a/geoana/spatial.py +++ b/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.]): diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 580870b7..c99c2770 100644 --- a/tests/test_spatial.py +++ b/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()