Make numpy-quaternion optional

[hakonanes]

Description of the change

This PR makes numpy-quaternion an optional dependency. It can be installed using pip install orix[all] or pip install orix numpy-quaternion.

Fixes #514.

I've made a private orix.constants.installed dictionary with (package, installed bool) available. Whether a dependency is installed or not is then checked with installed["numpy-quaternion"].

Functionality previously done by numpy-quaternion can now be done using Numba's generalized universal functions, which handles broadcasting in a smart way. Reduces need for us to check shapes.

How should we handle the optional dependencies on anaconda? I suggest to make an orix-base available, without numpy-quaternion. orix will still have it. What do people think?

Progress of the PR

• Docstrings for all functions
• Unit tests with pytest for all lines
• Clean code style by running black via pre-commit
• Improve installation guide now that we have optional dependencies
• Fix strange drawing of lines between vectors in "combining rotations" doc example

For reviewers

• The PR title is short, concise, and will make sense 1 year later.
• New functions are imported in corresponding __init__.py.
• New features, API changes, and deprecations are mentioned in the unreleased
  section in CHANGELOG.rst.
• Contributor(s) are listed correctly in __credits__ in orix/__init__.py and in
  .zenodo.json.

[hakonanes]

tomllib requires Python 3.11... Will see if we don't need it. Guess we only have to list optional dependencies in orix.constants.installed.

[hakonanes]

The failing tests for Python 3.8 are expected and will be addressed in a separate PR addressing #515.

[hakonanes]

Wanted to benchmark the re-introduced functionality against numpy-quaternion. In an environment with numpy-quaternion available, I think this is a fair benchmarking of the code used:

import numpy as np
import numba as nb
import quaternion

from orix.quaternion import Quaternion
from orix.vector import Vector3d

n_v = 10
n_qu = 1000
v = Vector3d(np.random.random(n_v * 3).reshape((n_v, 3)))
Q = Quaternion.random(n_qu)


# Called when numpy-quaternion is not available
@nb.guvectorize("(n),(m)->(m)", cache=True)
def qu_rotate_vec_gufunc(
    qu: np.ndarray, v1: np.ndarray, v2: np.ndarray
) -> None:  # pragma: no cover
    a, b, c, d = qu
    x, y, z = v1
    tx = 2 * (c * z - d * y)
    ty = 2 * (d * x - b * z)
    tz = 2 * (b * y - c * x)
    v2[0] = x + a * tx - d * ty + c * tz
    v2[1] = y + d * tx + a * ty - b * tz
    v2[2] = z - c * tx + b * ty + a * tz


@nb.guvectorize("(n),(n)->(n)", cache=True)
def qu_multiply_gufunc(
    qu1: np.ndarray, qu2: np.ndarray, qu12: np.ndarray
) -> None:  # pragma: no cover
    qu12[0] = qu1[0] * qu2[0] - qu1[1] * qu2[1] - qu1[2] * qu2[2] - qu1[3] * qu2[3]
    qu12[1] = qu1[1] * qu2[0] + qu1[0] * qu2[1] - qu1[3] * qu2[2] + qu1[2] * qu2[3]
    qu12[2] = qu1[2] * qu2[0] + qu1[3] * qu2[1] + qu1[0] * qu2[2] - qu1[1] * qu2[3]
    qu12[3] = qu1[3] * qu2[0] - qu1[2] * qu2[1] + qu1[1] * qu2[2] + qu1[0] * qu2[3]


# numpy-quaternion implementation
%timeit \
    qu = quaternion.from_float_array(Q.data);\
    quaternion.rotate_vectors(qu, v.data)
# 69.8 μs ± 571 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit \
    qu = quaternion.from_float_array(Q.data);\
    qu = quaternion.from_float_array(Q.data);\
    quaternion.as_float_array(qu * qu)
# 4.05 μs ± 35.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

# Alternative implementation
%timeit \
    v2 = np.empty(Q.shape + v.shape + (3,));\
    qu_rotate_vec_gufunc(Q.unit.reshape(-1, 1).data, v.reshape(1, -1).data, v2);\
    v2 = v2.reshape(*Q.shape, *v.shape, 3)
# 92.2 μs ± 988 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
%timeit \
    qu12 = np.empty_like(Q.data);\
    qu_multiply_gufunc(Q.data, Q.data, qu12)
# 37.4 μs ± 170 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

So, numpy-quaternion beats orix with Numba 0.60.0 and NumPy 1.26.4 by:

• quaternion-vector outer: 1.3x
• quaternion-quaternion multiplication: 9x

[pc494]

Sorry, realise I'm blocking here, plan is to get to it tomorrow/Friday :)

[hakonanes]

That would be great, thanks! Even though all tests pass w/o numpy-quaternion, I get different behavior from Vector3d.get_circle() that I have to fix.

[hakonanes]

Fixed the issue with Vector3d.get_circle(). Just need to update tests and the changelog.

[hakonanes]

Hope to finish testing and changelog updates within Monday. Should be OK to review as is, I don't expect any code changes.

[pc494]

This LGTM, couple of tiny things for you to look at and action/ignore @hakonanes.

[pc494]

extremely boring typo, should be orix's

[hakonanes]

I like typos, means I get to know the language better! I thought words ending with x shouldn't have the s, thank you for pointing this out (:

[pc494]

The guidance I have found simply says "However, if it sounds better, it is acceptable to just add the apostrophe." so that didn't help, although I would say orix's if I was speaking.

[pc494]

maybe not a typo then?

[hakonanes]

An intentional typo, if you will

[pc494]

Would this be obvious if I thought about it?

[pc494]

(as in the stuff below this line)

[hakonanes]

Perhaps not. I've added a comment-block to help with the understanding!

[pc494]

maybe 1 line here about why we have/want to use two of them (I notice that we do use both)

[hakonanes]

I've added an explanation. In general, I'd like for us to use the highest precision possible. That isn't always practical, so we have to lower it, but typically should use either of the two. I guess you could call them "high" and "lower" precisions.

[pc494]

Perfect, thanks @hakonanes I'm merging :)

UPDATE: I notice we fail on py3.8 here but I also do not mind (see #515) and think getting this in is more important. I know merging on reds is bad, but I'm making an exception.