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.

Sign up for GitHub

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
Merged

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

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
bb820d8
add get_pyramid and eu2qu
argerlt Aug 17, 2022
853853a
added n-dimensional support to _conversions.py
argerlt Aug 18, 2022
5d8a24a
ran pre-commit
argerlt Aug 18, 2022
b5faf26
Apply suggestions from code review
argerlt Aug 21, 2022
30422a0
Additional suggestions from code review
argerlt Aug 22, 2022
d340eac
ensure wrappers pass correct data type, formatting
argerlt Aug 22, 2022
6496d6d
black compliance
argerlt Aug 22, 2022
ce37ac1
Set dtype explicitly in a few arrays
hakonanes Aug 23, 2022
df1d57a
Add clarifying comments, remove redundant parentheses
hakonanes Aug 23, 2022
b89e940
Add advice on testing Numba functions in contributing guide
hakonanes Aug 23, 2022
7fe5aa1
Add new contributor Austin Gerlt to credits and Zenodo file
hakonanes Aug 23, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 165 additions & 21 deletions orix/quaternion/_conversions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
-----
Expand All @@ -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 2hich pyramids ``xyz`` belong to 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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -172,6 +209,16 @@ def cu2ho(cu):
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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -310,6 +367,16 @@ def ax2ro(ax):
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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -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.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
Expand All @@ -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`.

Expand All @@ -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.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
Expand Down Expand Up @@ -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`.

Expand All @@ -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.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
-------
Expand All @@ -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)
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 •
edited
Loading

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