Add n-dimensional support to private rotation conversion functions

.zenodo.json

@@ -25,6 +25,11 @@
       "name": "Niels Cautaerts",
       "orcid": "0000-0002-6402-9879"
     },
+    {
+      "name": "Austin Gerlt",
+      "orcid": "0000-0002-2204-2055",
+      "affiliation": "The Ohio State University"
+    },
     {
       "name": "Simon Høgås"
     }

CONTRIBUTING.rst

@@ -312,6 +312,14 @@ prints a nice report in the terminal. For an even nicer presentation, you can us
 Then, you can open the created ``htmlcov/index.html`` in the browser and inspect the
 coverage in more detail.
 
+Tips for writing tests of Numba decorated functions:
+
+- A Numba decorated function ``numba_func()`` is only covered if it is called in the
+  test as ``numba_func.py_func()``.
+- Always test a Numba decorated function calling ``numba_func()`` directly, in addition
+  to ``numba_func.py_func()``, because the machine code function might give different
+  results on different OS with the same Python code.
+
 .. _adding-data-to-data-module:
 
 Adding data to the data module

orix/__init__.py

@@ -11,5 +11,6 @@
     "Paddy Harrison",
     "Duncan Johnstone",
     "Niels Cautaerts",
+    "Austin Gerlt",
     "Simon Høgås",
 ]

orix/quaternion/_conversions.py

@@ -16,7 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with orix.  If not, see <http://www.gnu.org/licenses/>.
 
-"""Conversions of orientations between many common representations from
+"""Conversions of rotations between many common representations from
 :cite:`rowenhorst2015consistent`, accelerated with Numba.
 
 This module and documentation is only relevant for orix developers, not
@@ -44,7 +44,7 @@ def get_pyramid_single(xyz):
     Returns
     -------
     pyramid : int
-        Which pyramid `xyz` belongs to as a 64-bit integer.
+        Which pyramid ``xyz`` belongs to as a 64-bit integer.
 
     Notes
     -----
@@ -67,6 +67,43 @@ def get_pyramid_single(xyz):
         return 6
 
 
+@nb.jit("int64[:](float64[:, :])", cache=True, nogil=True, nopython=True)
+def get_pyramid_2d(xyz):
+    """Determine to which out of six pyramids in the cube a 2D array of
+    (x, y, z) coordinates belongs.
+
+    Parameters
+    ----------
+    xyz : numpy.ndarray
+        2D array of n (x, y, z) as 64-bit floats.
+
+    Returns
+    -------
+    pyramids : numpy.ndarray
+        1D array of pyramids as 64-bit integers.
+
+    Notes
+    -----
+    This function is optimized with Numba, so care must be taken with
+    array shapes and data types.
+    """
+    n_scalars = xyz.shape[0]
+    pyramids = np.zeros(n_scalars, dtype=np.int64)
+    for i in nb.prange(n_scalars):
+        pyramids[i] = get_pyramid_single(xyz[i])
+    return pyramids
+
+
+def get_pyramid(xyz):
+    """n-dimensional wrapper for get_pyramid_2d, see the docstring of
+    that function.
+    """
+    n_xyz = np.prod(xyz.shape[:-1])
+    xyz2d = xyz.astype(np.float64).reshape(n_xyz, 3)
+    pyramids = get_pyramid_2d(xyz2d).reshape(n_xyz)
+    return pyramids
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def cu2ho_single(cu):
     """Conversion from a single set of cubochoric coordinates to
@@ -147,7 +184,7 @@ def cu2ho_single(cu):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def cu2ho(cu):
+def cu2ho_2d(cu):
     """Conversion from multiple cubochoric coordinates to un-normalized
     homochoric coordinates :cite:`singh2016orientation`.
 
@@ -166,12 +203,22 @@ def cu2ho(cu):
     This function is optimized with Numba, so care must be taken with
     array shapes and data types.
     """
-    ho = np.zeros_like(cu)
+    ho = np.zeros_like(cu, dtype=np.float64)
     for i in nb.prange(cu.shape[0]):
         ho[i] = cu2ho_single(cu[i])
     return ho
 
 
+def cu2ho(cu):
+    """N-dimensional wrapper for cu2ho_2d, see the docstring of that
+    function.
+    """
+    n_cu = np.prod(cu.shape[:-1])
+    cu2d = cu.astype(np.float64).reshape(n_cu, 3)
+    ho = cu2ho_2d(cu2d).reshape(cu.shape)
+    return ho
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def ho2ax_single(ho):
     """Conversion from a single set of homochoric coordinates to an
@@ -224,7 +271,7 @@ def ho2ax_single(ho):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def ho2ax(ho):
+def ho2ax_2d(ho):
     """Conversion from multiple homochoric coordinates to un-normalized
     axis-angle pairs :cite:`rowenhorst2015consistent`.
 
@@ -235,8 +282,8 @@ def ho2ax(ho):
 
     Returns
     -------
-    ho : numpy.ndarray
-        2D array of n (x, y, z) as 64-bit floats.
+    ax : numpy.ndarray
+        2D array of n (x, y, z, angle) as 64-bit floats.
 
     Notes
     -----
@@ -250,6 +297,16 @@ def ho2ax(ho):
     return ax
 
 
+def ho2ax(ho):
+    """N-dimensional wrapper for ho2ax_2d, see the docstring of that
+    function.
+    """
+    n_ho = np.prod(ho.shape[:-1])
+    ho2d = ho.astype(np.float64).reshape(n_ho, 3)
+    ho = ho2ax_2d(ho2d).reshape(ho.shape[:-1] + (4,))
+    return ho
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def ax2ro_single(ax):
     """Conversion from a single angle-axis pair to an un-normalized
@@ -285,7 +342,7 @@ def ax2ro_single(ax):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def ax2ro(ax):
+def ax2ro_2d(ax):
     """Conversion from multiple axis-angle pairs to un-normalized
     Rodrigues vectors :cite:`rowenhorst2015consistent`.
 
@@ -304,12 +361,22 @@ def ax2ro(ax):
     This function is optimized with Numba, so care must be taken with
     array shapes and data types.
     """
-    ro = np.zeros_like(ax)
+    ro = np.zeros_like(ax, dtype=np.float64)
     for i in nb.prange(ax.shape[0]):
         ro[i] = ax2ro_single(ax[i])
     return ro
 
 
+def ax2ro(ax):
+    """N-dimensional wrapper for ax2ro_2d, see the docstring of that
+    function.
+    """
+    n_ax = np.prod(ax.shape[:-1])
+    ax2d = ax.astype(np.float64).reshape(n_ax, 4)
+    ro = ax2ro_2d(ax2d).reshape(ax.shape)
+    return ro
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def ro2ax_single(ro):
     """Conversion from a single Rodrigues vector to an un-normalized
@@ -340,7 +407,7 @@ def ro2ax_single(ro):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def ro2ax(ro):
+def ro2ax_2d(ro):
     """Conversion from multiple Rodrigues vectors to un-normalized
     axis-angle pairs :cite:`rowenhorst2015consistent`.
 
@@ -366,6 +433,16 @@ def ro2ax(ro):
     return ax
 
 
+def ro2ax(ro):
+    """N-dimensional wrapper for ro2ax_2d, see the docstring of that
+    function.
+    """
+    n_ro = np.prod(ro.shape[:-1])
+    ro2d = ro.astype(np.float64).reshape(n_ro, 4)
+    ax = ro2ax_2d(ro2d).reshape(ro.shape)
+    return ax
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def ax2qu_single(ax):
     """Conversion from a single axis-angle pair to an un-normalized
@@ -395,7 +472,7 @@ def ax2qu_single(ax):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def ax2qu(ax):
+def ax2qu_2d(ax):
     """Conversion from multiple axis-angle pairs to un-normalized
     quaternions :cite:`rowenhorst2015consistent`.
 
@@ -421,6 +498,16 @@ def ax2qu(ax):
     return qu
 
 
+def ax2qu(ax):
+    """N-dimensional wrapper for ax2qu_2d, see the docstring of that
+    function.
+    """
+    n_ax = np.prod(ax.shape[:-1])
+    ax2d = ax.astype(np.float64).reshape(n_ax, 4)
+    qu = ax2qu_2d(ax2d).reshape(ax.shape)
+    return qu
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def ho2ro_single(ho):
     """Conversion from a single set of homochoric coordinates to an
@@ -445,7 +532,7 @@ def ho2ro_single(ho):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def ho2ro(ho):
+def ho2ro_2d(ho):
     """Conversion from multiple homochoric coordinates to un-normalized
     Rodrigues vectors :cite:`rowenhorst2015consistent`.
 
@@ -471,6 +558,16 @@ def ho2ro(ho):
     return ro
 
 
+def ho2ro(ho):
+    """N-dimensional wrapper for ho2ro_2d, see the docstring of that
+    function.
+    """
+    n_ho = np.prod(ho.shape[:-1])
+    ho2d = ho.astype(np.float64).reshape(n_ho, 3)
+    ro = ho2ro_2d(ho2d).reshape(ho.shape[:-1] + (4,))
+    return ro
+
+
 @nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
 def cu2ro_single(cu):
     """Conversion from a single set of cubochoric coordinates to an
@@ -498,7 +595,7 @@ def cu2ro_single(cu):
 
 
 @nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
-def cu2ro(cu):
+def cu2ro_2d(cu):
     """Conversion from multiple cubochoric coordinates to un-normalized
     Rodrigues vectors :cite:`rowenhorst2015consistent`.
 
@@ -524,16 +621,26 @@ def cu2ro(cu):
     return ro
 
 
-@nb.jit("float64[:](float64, float64, float64)", cache=True, nogil=True, nopython=True)
-def eu2qu_single(alpha, beta, gamma):
+def cu2ro(cu):
+    """N-dimensional wrapper for cu2ro_2d, see the docstring of that
+    function.
+    """
+    n_cu = np.prod(cu.shape[:-1])
+    cu2d = cu.astype(np.float64).reshape(n_cu, 3)
+    ro = cu2ro_2d(cu2d).reshape(cu.shape[:-1] + (4,))
+    return ro
+
+
+@nb.jit("float64[:](float64[:])", cache=True, nogil=True, nopython=True)
+def eu2qu_single(eu):
     """Convert three Euler angles (alpha, beta, gamma) to a unit
     quaternion.
 
     Parameters
     ----------
-    alpha, beta, gamma : float
-        Euler angles in the Bunge convention in radians as 64-bit
-        floats.
+    eu : numpy.ndarray
+        1D array of (alpha, beta, gamma) Euler angles given in radians
+        in the Bunge convention (ie, passive Z-X-Z) as 64-bit floats.
 
     Returns
     -------
@@ -547,18 +654,55 @@ def eu2qu_single(alpha, beta, gamma):
     This function is optimized with Numba, so care must be taken with
     array shapes and data types.
     """
-    sigma = 0.5 * np.add(alpha, gamma)
-    delta = 0.5 * np.subtract(alpha, gamma)
-    c = np.cos(beta / 2)
-    s = np.sin(beta / 2)
+    sigma = 0.5 * np.add(eu[0], eu[2])
+    delta = 0.5 * np.subtract(eu[0], eu[2])
+    c = np.cos(eu[1] / 2)
+    s = np.sin(eu[1] / 2)
 
     qu = np.zeros(4, dtype=np.float64)
-    qu[0] = c * np.cos(sigma)
-    qu[1] = -s * np.cos(delta)
-    qu[2] = -s * np.sin(delta)
-    qu[3] = -c * np.sin(sigma)
+    qu[0] = np.array(c * np.cos(sigma), dtype=np.float64)
+    qu[1] = np.array(-s * np.cos(delta), dtype=np.float64)
+    qu[2] = np.array(-s * np.sin(delta), dtype=np.float64)
+    qu[3] = np.array(-c * np.sin(sigma), dtype=np.float64)
 
     if qu[0] < 0:
         qu = -qu
 
     return qu
+
+
+@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
+def eu2qu_2d(eu):
+    """Conversion from multiple Euler angles (alpha, beta, gamma) to unit
+    quaternions
+
+    Parameters
+    ----------
+    eu : numpy.ndarray
+        2D array of n (alpha, beta, gamma) as 64-bit floats.
+
+    Returns
+    -------
+    qu : numpy.ndarray
+        2D array of n (q0, q1, q2, q3) quaternions as 64-bit floats.
+
+    Notes
+    -----
+    This function is optimized with Numba, so care must be taken with
+    array shapes and data types.
+    """
+    n_vectors = eu.shape[0]
+    qu = np.zeros((n_vectors, 4), dtype=np.float64)
+    for i in nb.prange(n_vectors):
+        qu[i] = eu2qu_single(eu[i])
+    return qu
+
+
+def eu2qu(eu):
+    """N-dimensional wrapper for eu2qu_2d, see the docstring of that
+    function.
+    """
+    n_eu = np.prod(eu.shape[:-1])
+    eu2d = eu.astype(np.float64).reshape(n_eu, 3)
+    qu = eu2qu_2d(eu2d).reshape(eu.shape[:-1] + (4,))
+    return qu