Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add n-dimensional support to private rotation conversion functions #376

Merged
merged 11 commits into from
Aug 23, 2022
160 changes: 143 additions & 17 deletions orix/quaternion/_conversions.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,42 @@ def get_pyramid_single(xyz):
return 6


@nb.jit("int64[:](float64[:, :])", cache=True, nogil=True, nopython=True)
def get_pyramid2d(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
-------
pyramid : int
1D array of length n.
Which pyramid `xyz` belongs to as a 64-bit integer.

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_pyramid2d"""
n_xyz = np.prod(xyz.shape[:-1])
xyz2d = xyz.reshape(n_xyz, 3)
pyramids = get_pyramid2d(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
Expand Down Expand Up @@ -147,7 +183,7 @@ def cu2ho_single(cu):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def cu2ho(cu):
def cu2ho2d(cu):
"""Conversion from multiple cubochoric coordinates to un-normalized
homochoric coordinates :cite:`singh2016orientation`.

Expand All @@ -172,6 +208,14 @@ def cu2ho(cu):
return ho


def cu2ho(cu):
"""n-dimensional wrapper for cu2ho2d"""
n_cu = np.prod(cu.shape[:-1])
cu2d = cu.reshape(n_cu, 3)
ho = cu2ho2d(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
Expand Down Expand Up @@ -224,7 +268,7 @@ def ho2ax_single(ho):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def ho2ax(ho):
def ho2ax2d(ho):
"""Conversion from multiple homochoric coordinates to un-normalized
axis-angle pairs :cite:`rowenhorst2015consistent`.

Expand All @@ -250,6 +294,14 @@ def ho2ax(ho):
return ax


def ho2ax(ho):
"""n-dimensional wrapper for ho2ax2d"""
n_ho = np.prod(ho.shape[:-1])
ho2d = ho.reshape(n_ho, 3)
ho = ho2ax2d(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
Expand Down Expand Up @@ -285,7 +337,7 @@ def ax2ro_single(ax):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def ax2ro(ax):
def ax2ro2d(ax):
"""Conversion from multiple axis-angle pairs to un-normalized
Rodrigues vectors :cite:`rowenhorst2015consistent`.

Expand All @@ -310,6 +362,14 @@ def ax2ro(ax):
return ro


def ax2ro(ax):
"""n-dimensional wrapper for ax2ro2d"""
n_ax = np.prod(ax.shape[:-1])
ax2d = ax.reshape(n_ax, 4)
ro = ax2ro2d(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
Expand Down Expand Up @@ -340,7 +400,7 @@ def ro2ax_single(ro):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def ro2ax(ro):
def ro2ax2d(ro):
"""Conversion from multiple Rodrigues vectors to un-normalized
axis-angle pairs :cite:`rowenhorst2015consistent`.

Expand All @@ -366,6 +426,13 @@ def ro2ax(ro):
return ax


def ro2ax(ro):
n_ro = np.prod(ro.shape[:-1])
ro2d = ro.reshape(n_ro, 4)
ax = ro2ax2d(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
Expand Down Expand Up @@ -395,7 +462,7 @@ def ax2qu_single(ax):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def ax2qu(ax):
def ax2qu2d(ax):
"""Conversion from multiple axis-angle pairs to un-normalized
quaternions :cite:`rowenhorst2015consistent`.

Expand All @@ -421,6 +488,14 @@ def ax2qu(ax):
return qu


def ax2qu(ax):
"""n-dimensional wrapper for ax2qu2d"""
n_ax = np.prod(ax.shape[:-1])
ax2d = ax.reshape(n_ax, 4)
qu = ax2qu2d(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
Expand All @@ -445,7 +520,7 @@ def ho2ro_single(ho):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def ho2ro(ho):
def ho2ro2d(ho):
"""Conversion from multiple homochoric coordinates to un-normalized
Rodrigues vectors :cite:`rowenhorst2015consistent`.

Expand All @@ -471,6 +546,14 @@ def ho2ro(ho):
return ro


def ho2ro(ho):
"""n-dimensional wrapper for ho2ro2d"""
n_ho = np.prod(ho.shape[:-1])
ho2d = ho.reshape(n_ho, 3)
ro = ho2ro2d(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
Expand Down Expand Up @@ -498,7 +581,7 @@ def cu2ro_single(cu):


@nb.jit("float64[:, :](float64[:, :])", cache=True, nogil=True, nopython=True)
def cu2ro(cu):
def cu2ro2d(cu):
"""Conversion from multiple cubochoric coordinates to un-normalized
Rodrigues vectors :cite:`rowenhorst2015consistent`.

Expand All @@ -524,8 +607,16 @@ 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 cu2ro2d"""
n_cu = np.prod(cu.shape[:-1])
cu2d = cu.reshape(n_cu, 3)
ro = cu2ro2d(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.

Expand All @@ -547,18 +638,53 @@ 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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is making these arrays necessary?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, or at least something similar is.

Without these arrays, I was getting a problem where the right hand sides of lines 647-650 would produce float32 values that would then be inserted into a float64 array, resulting in incorrect values. This was a jit-only problem, because the .py_func versions of the functions execute correctly (float64 in, float64 out).

I tried several variations, which all gave the incorrect results:

qu[3] = -c * np.sin(sigma)
qu[3] = -c * np.sin(sigma, dtype = np.float64)
qu[3] = np.multiply(c, np.sin(sigma))

etc.
This FEELS weird and wrong, but it was the only method that worked which didn't involve adding extra variables, which seemed to slow down the computation.
If you or anyone else think of a more elegant solution, feel free to change it.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you do qu[0] = float(c * np.cos(sigma))? If not, try qu[0] = np.float64(c * np.cos(sigma)).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

np.float64 passes correctly

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually, nevermind, np.float64 doesn't work either. I was loading an old pycache file. I think np.array might be the easiest fix to this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, I'll have a look myself.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Huh. Okay, this has be totally confused.
For some reason, every variation of np.float64, float, etc, breaks. However, it ONLY breaks if the value it is inserting into qu is less than zero.

I was going to suggest we just approve this and move on, but the underlying implication here is SOMEWHERE, things are getting cast incorrectly, possibly from float64 to float32 and then back, which would translate to a big loss in accuracy.

I also tried explicitly using numpy multiply, divide, etc, like this:

    sigma = np.multiply(0.5, np.add(eu[0], eu[2]))
    delta = np.multiply(0.5, np.subtract(eu[0], eu[2]))
    c = np.cos(np.divide(eu[1], 2))
    s = np.sin(np.divide(eu[1], 2))
    nc = np.multiply(-1, c)
    ns = np.multiply(-1, s)

    qu = np.zeros(4, dtype=np.float64)
    qu[0] = np.multiply(c, np.cos(sigma))
    qu[1] = np.multiply(ns, np.cos(delta))
    qu[2] = np.multiply(ns, np.sin(delta))
    qu[3] = np.multiply(nc, np.sin(sigma))

    if qu[0] < 0:
        qu = -qu

but that failed as well.

Copy link
Member

@hakonanes hakonanes Aug 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that it's OK to leave as is and move on.

Out of curiosity, however, which OS and Python version are you testing with? All pytest -k test_conversions pass without any changes to the original code on my Ubuntu 22.04 with Python 3.9. Are you running other tests than the ones in test_convers.py? If so, we should add similar tests to catch the behaviour you're describing.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm running on a Windows 10 computer with an AMD CPU. specifically from command prompt, When testing, I was running pytest orix/tests/quaternion/test_conversions.py

Since you brought this up, I tried it again, same code, same computer, but using Windows Subsystem for Linux (WSL) to run pytest. It passed without needing the np.array() additions! So, I assume this is a windows specific problem with numba and jit, which honestly, makes it even weirder.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for clarifying. I've encountered surprising effects with Numba in the past, similar to your experience here. This is just a reminder that we should not trust Numba blindly!

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 eu2qu2d(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 (qo, 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 eu2qu2d"""
n_eu = np.prod(eu.shape[:-1])
eu2d = eu.reshape(n_eu, 3)
qu = eu2qu2d(eu2d).reshape(eu.shape[:-1] + (4,))
return qu
Loading