Repository: scipy/scipy
Pull request #20950: ENH: spatial: speed up `Rotation.apply` by replacing `np.einsum` with `np.matmul`
Author: login: cosama
Labels: name: enhancement
description: A new feature or improvement
name: scipy.spatial
name: Cython
description: Issues with the internal Cython code base
Created: 2024-06-12T17:22:15Z | Merged: 2025-04-03T07:34:17Z

The use of `np.einsum` is slow compared to `np.matmul` in `scipy.spatial.transform.Rotation.apply`. It is possible to achieve a speed up of over a factor of two by replacing that function, at least for large vectors and single rotation (the most common use case) but most likely also for other uses of that function (not benchmarked).

#### Reference issue
Closes https://github.com/scipy/scipy/issues/20532

#### What does this implement/fix?

I used the following script to check the execution time of the apply method:

```python
import numpy as np
import scipy
import time

def create_data():
    vector = np.random.default_rng().random((100000, 3), dtype=float)
    q = np.random.default_rng().random(4, dtype=float)
    q /= np.linalg.norm(q)
    rot = scipy.spatial.transform.Rotation.from_quat(q)
    return rot, vector

def rot_apply(**kwargs):
    rot, vector = create_data()
    start = time.time()
    _ = rot.apply(vector, **kwargs)
    return time.time() - start

repeats = 1000
time_forward = 0
time_inverse = 0
for _ in range(repeats):
    time_forward += rot_apply()
    time_inverse += rot_apply(inverse=True)

print(time_forward / repeats, time_inverse / repeats)
```

The result is:

```
# OLD
# 0.0026782317161560057 0.0027620444297790526
# 0.002682605266571045 0.002773979425430298
# 0.002711399793624878 0.0028269896507263186

# NEW
# 0.000678567886352539 0.0015280032157897949
# 0.0011942288875579834 0.0015206401348114014
# 0.001080368757247925 0.0015322771072387696
```

#### Additional information

Unfortunately, `np.matmul` is not a direct replacement and it was necessary to change some calling semantics in the function as well. I carefully tested that all cases of inputs (single Rotation, 1 element Rotation and multi element Rotation) and vectors ( `(3,), (1, 3), (N, 3)`) result in identical output. I used the following code to do so:

```python
from itertools import product
import numpy as np
import scipy

def apply_new(self, vectors, inverse=False):
    vectors = np.asarray(vectors)
    if vectors.ndim > 2 or vectors.shape[-1] != 3:
        raise ValueError("Expected input of shape (3,) or (P, 3), "
                         "got {}.".format(vectors.shape))

    matrix = self.as_matrix()
    if inverse:
        matrix = np.swapaxes(matrix, -1, -2)

    if vectors.shape == (3,):
        return np.matmul(matrix, vectors)
    if self.single:
        matrix = matrix[None, :, :]
    
    n_vectors = vectors.shape[0]
    n_rotations = 1 if self.single else len(self)

    if n_rotations == 1:
        return np.matmul(matrix, vectors, axes=[(-2, -1), (-1, -2), (-1, -2)])[0]
    elif n_vectors == 1 or n_vectors == n_rotations:
        return np.matmul(matrix, vectors[:, None, :], axes=[(-2, -1), (-1, -2), (-1, 0)])[0]

    raise ValueError("Expected equal numbers of rotations and vectors "
                        ", or a single rotation, or a single vector, got "
                        "{} rotations and {} vectors.".format(
                        n_rotations, n_vectors))

q = np.random.default_rng().random(4, dtype=float)
q /= np.linalg.norm(q)

vector0 = np.random.default_rng().random((3,), dtype=float)
rot0 = scipy.spatial.transform.Rotation.from_quat(q)
matrix0 = rot0.as_matrix()

vector1 = np.random.default_rng().random((1, 3), dtype=float)
rot1 = scipy.spatial.transform.Rotation.from_quat([q])
matrix1 = rot1.as_matrix()

q = np.random.default_rng().random((2, 4), dtype=float)
q = (q.T / np.linalg.norm(q, axis=1)).T

vector2 = np.random.default_rng().random((2, 3), dtype=float)
rot2 = scipy.spatial.transform.Rotation.from_quat(q)
matrix2 = rot2.as_matrix()

for i, (r, v) in enumerate(product([rot0, rot1, rot2], [vector0, vector1, vector2])):
    print(i, 0 if r.single else len(r), 0 if v.ndim == 1 else len(v))
    old = r.apply(v)
    new = apply_new(r, v)
    print(np.allclose(old, new), np.all(old.shape == new.shape))
    old = r.apply(v, inverse=True)
    new = apply_new(r, v, inverse=True)
    print(np.allclose(old, new), np.all(old.shape == new.shape))
```



Diff:
diff --git a/scipy/spatial/transform/_rotation.pyx b/scipy/spatial/transform/_rotation.pyx
index e0f086775a7e..53675734f271 100644
--- a/scipy/spatial/transform/_rotation.pyx
+++ b/scipy/spatial/transform/_rotation.pyx
@@ -2580,17 +2580,9 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
-        matrix = self.as_matrix()
-        if self._single:
-            matrix = matrix[None, :, :]
-
-        n_vectors = vectors.shape[0]
-        n_rotations = len(self._quat)
+        cdef bint single_vector = vectors.ndim == 1
+        cdef Py_ssize_t n_vectors = 1 if single_vector else len(vectors)
+        cdef Py_ssize_t n_rotations = 1 if self.single else len(self)
 
         if n_vectors != 1 and n_rotations != 1 and n_vectors != n_rotations:
             raise ValueError("Expected equal numbers of rotations and vectors "
@@ -2598,15 +2590,25 @@ cdef class Rotation:
                              "{} rotations and {} vectors.".format(
                                 n_rotations, n_vectors))
 
+        cdef np.ndarray matrix = self.as_matrix()
+
         if inverse:
-            result = np.einsum('ikj,ik->ij', matrix, vectors)
-        else:
-            result = np.einsum('ijk,ik->ij', matrix, vectors)
+            matrix = np.swapaxes(matrix, -1, -2)
 
-        if self._single and single_vector:
-            return result[0]
-        else:
-            return result
+        if single_vector:
+            return np.matmul(matrix, vectors)
+
+        if self.single:
+            matrix = matrix[None, :, :]
+
+        if n_rotations == 1:
+            # Single rotation/many vectors, use matmul for speed: The axes argument
+            # is such that the input arguments don't need to be transposed and the
+            # output argument is contineous in memory.
+            return np.matmul(matrix, vectors, axes=[(-2, -1), (-1, -2), (-1, -2)])[0]
+
+        # for stacks of matrices einsum is faster
+        return np.einsum('ijk,ik->ij', matrix, vectors)
 
     @cython.embedsignature(True)
     def __mul__(Rotation self, Rotation other):
diff --git a/scipy/spatial/transform/tests/test_rotation.py b/scipy/spatial/transform/tests/test_rotation.py
index 1c06546e382f..2c91018f7a2b 100644
--- a/scipy/spatial/transform/tests/test_rotation.py
+++ b/scipy/spatial/transform/tests/test_rotation.py
@@ -5,7 +5,7 @@
 from numpy.testing import assert_allclose
 from scipy.spatial.transform import Rotation, Slerp
 from scipy.stats import special_ortho_group
-from itertools import permutations
+from itertools import permutations, product
 
 import pickle
 import copy
@@ -1206,6 +1206,23 @@ def test_apply_multiple_rotations_multiple_points():
     v_inverse = np.array([[2, -1, 3], [4, 6, -5]])
     assert_allclose(r.apply(v, inverse=True), v_inverse)
 
+def test_apply_shapes():
+    vector0 = np.array([1.0, 2.0, 3.0])
+    vector1 = np.array([vector0])
+    vector2 = np.array([vector0, vector0])
+    matrix0 = np.identity(3)
+    matrix1 = np.array([matrix0])
+    matrix2 = np.array([matrix0, matrix0])
+
+    for m, v in product([matrix0, matrix1, matrix2], [vector0, vector1, vector2]):
+        r = Rotation.from_matrix(m)
+        shape = v.shape
+        if not r.single and (v.shape == (3,) or v.shape == (1, 3)):
+            shape = (len(r), 3)
+        x = r.apply(v)
+        assert x.shape == shape
+        x = r.apply(v, inverse=True)
+        assert x.shape == shape
 
 def test_getitem():
     mat = np.empty((2, 3, 3))


Reviews:
user: login: scottshambaugh
body: This all looks generally good to me, and a nice speedup. Doing a search for `.shape` in [test_rotation.py](https://github.com/scipy/scipy/blob/main/scipy/spatial/transform/tests/test_rotation.py), I'm not seeing any unit tests for checking that vector shapes are preserved through `.apply()`. Could you please add some?

Also, on the topic of performance enhancements, would you be interested in taking a look at https://github.com/scipy/scipy/issues/19300? Might even be able to roll it into this current MR, as a small change that's performance focused.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2024-06-12T19:00:11Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2024-06-12T22:40:24Z
user: login: nmayorov
body: Minor suggestions regarding code organization.
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2024-09-21T20:03:50Z
user: login: nmayorov
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T21:16:12Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T13:33:55Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T13:34:14Z
user: login: nmayorov
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T20:07:51Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T20:29:51Z
user: login: nmayorov
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T21:23:58Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T22:11:45Z
user: login: cosama
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T19:11:06Z

Review comments:
diff_hunk: @@ -2515,33 +2515,29 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
+            return np.matmul(matrix, vectors)
         if self._single:
             matrix = matrix[None, :, :]
 
         n_vectors = vectors.shape[0]
         n_rotations = len(self._quat)
 
-        if n_vectors != 1 and n_rotations != 1 and n_vectors != n_rotations:
-            raise ValueError("Expected equal numbers of rotations and vectors "
-                             ", or a single rotation, or a single vector, got "
-                             "{} rotations and {} vectors.".format(
-                                n_rotations, n_vectors))
-
-        if inverse:
-            result = np.einsum('ikj,ik->ij', matrix, vectors)
-        else:
-            result = np.einsum('ijk,ik->ij', matrix, vectors)
-
-        if self._single and single_vector:
-            return result[0]
-        else:
-            return result
+        if n_rotations == 1:
+            return np.matmul(matrix, vectors, axes=[(-2, -1), (-1, -2), (-1, -2)])[0]
+        elif n_vectors == 1 or n_vectors == n_rotations:
+            return np.matmul(
+                matrix, vectors[:, None, :], axes=[(-2, -1), (-1, -2), (-1, 0)]
path: scipy/spatial/transform/_rotation.pyx
user: login: scottshambaugh
body: Could you add a comment here and above describing what's happening? The axes indices are kind of hard to interpret easily.
created_at: 2024-06-12T18:57:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,29 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
+            return np.matmul(matrix, vectors)
         if self._single:
             matrix = matrix[None, :, :]
 
         n_vectors = vectors.shape[0]
         n_rotations = len(self._quat)
 
-        if n_vectors != 1 and n_rotations != 1 and n_vectors != n_rotations:
-            raise ValueError("Expected equal numbers of rotations and vectors "
-                             ", or a single rotation, or a single vector, got "
-                             "{} rotations and {} vectors.".format(
-                                n_rotations, n_vectors))
-
-        if inverse:
-            result = np.einsum('ikj,ik->ij', matrix, vectors)
-        else:
-            result = np.einsum('ijk,ik->ij', matrix, vectors)
-
-        if self._single and single_vector:
-            return result[0]
-        else:
-            return result
+        if n_rotations == 1:
+            return np.matmul(matrix, vectors, axes=[(-2, -1), (-1, -2), (-1, -2)])[0]
+        elif n_vectors == 1 or n_vectors == n_rotations:
+            return np.matmul(
+                matrix, vectors[:, None, :], axes=[(-2, -1), (-1, -2), (-1, 0)]
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: Agreed, took me forever to understand how to set them properly. I will add some comment. I wasn't sure if they were typically used.
created_at: 2024-06-12T22:40:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,31 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
+            return np.matmul(matrix, vectors)
         if self._single:
             matrix = matrix[None, :, :]
 
         n_vectors = vectors.shape[0]
         n_rotations = len(self._quat)
 
-        if n_vectors != 1 and n_rotations != 1 and n_vectors != n_rotations:
path: scipy/spatial/transform/_rotation.pyx
user: login: nmayorov
body: I think it's preferable to keep the verification block before the computations.
created_at: 2024-09-21T19:58:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,31 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
path: scipy/spatial/transform/_rotation.pyx
user: login: nmayorov
body: Is it possible to move it somehow into the main computation block?

Will it work when `matrix` has 3 dimensions?
created_at: 2024-09-21T20:00:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,31 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
path: scipy/spatial/transform/_rotation.pyx
user: login: nmayorov
body: `vectors.ndim == 1` would be cleaner.
created_at: 2025-03-21T21:16:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,31 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: Yes this works with 3 dimensions and I have tested it (also I added test to test this). I fixed the second issue.
created_at: 2025-03-31T13:33:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2515,33 +2515,31 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
-        single_vector = False
-        if vectors.shape == (3,):
-            single_vector = True
-            vectors = vectors[None, :]
-
         matrix = self.as_matrix()
+        if inverse:
+            matrix = np.swapaxes(matrix, -1, -2)
+
+        if vectors.shape == (3,):
+            return np.matmul(matrix, vectors)
         if self._single:
             matrix = matrix[None, :, :]
 
         n_vectors = vectors.shape[0]
         n_rotations = len(self._quat)
 
-        if n_vectors != 1 and n_rotations != 1 and n_vectors != n_rotations:
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: I reorganized the code.
created_at: 2025-03-31T13:34:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2580,31 +2580,39 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
+        n_vectors = vectors.shape[0]
path: scipy/spatial/transform/_rotation.pyx
user: login: nmayorov
body: I would put

```
single_vector = vectors.ndim == 1
n_vectors = 1 if single_vector else len(vectors)
```

Also some reasons to use types for all variables, can matter if `apply` is called often I suppose.
created_at: 2025-04-01T20:07:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2580,31 +2580,39 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
+        n_vectors = vectors.shape[0]
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: I can add those changes.

I am not entirely sure what you mean by types for all variables. Could you elaborate?

However, trying to understand what you mean I actually realized that if you apply a rotation to a `np.float32` vector, you will get a `float` vector back. This is how it has been done so far, but it could be unanticipated and lead to memory issues for very large vectors. If you like me to change that, let me know, both `np.matmul` and `np.einsum` take a `dtype` kwarg, however it would affect the precision of the result, as both the rotation matrix and the vector are "cast" to the new dtype before performing the operation (at least that is my limited understanding of what is happening).
created_at: 2025-04-01T20:29:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2580,31 +2580,39 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
+        n_vectors = vectors.shape[0]
path: scipy/spatial/transform/_rotation.pyx
user: login: nmayorov
body: I meant a small thing, that `n_vectors`, `single_vector` probably should be defined with `cdef Py_ssize_t`, etc..

But as there are many small non-optimal things involving Python interaction it doesn't look necessary.


created_at: 2025-04-01T21:23:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2580,31 +2580,39 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
+        n_vectors = vectors.shape[0]
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: I see, I didn't even consider cython types because there were none before. I added them, let me know if there is anything wrong with them, I am a absolute cython beginner.
created_at: 2025-04-01T22:11:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2580,31 +2580,39 @@ cdef class Rotation:
             raise ValueError("Expected input of shape (3,) or (P, 3), "
                              "got {}.".format(vectors.shape))
 
+        n_vectors = vectors.shape[0]
path: scipy/spatial/transform/_rotation.pyx
user: login: cosama
body: I implemented the types and the requested changes. Let me know if there is anything else that I should work on.
created_at: 2025-04-02T19:11:06Z
author_association: CONTRIBUTOR

PR comments:
user: login: lucascolley
created_at: 2024-06-12T18:07:31Z
body: cc @scottshambaugh 
author_association: MEMBER
user: login: dschmitz89
created_at: 2024-06-13T06:14:40Z
body: In the [Von Mises-Fisher distribution](https://github.com/scipy/scipy/blob/44e4ebaac992fde33f04638b99629d23973cb9b2/scipy/stats/_multivariate.py#L6747), I used `np.einsum` to apply the same rotation to an arbitrary number of vectors as well. Would `np.matmul` also be faster for that use case? I am not familiar with rotations in scipy. @scottshambaugh 
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2024-06-13T19:21:10Z
body: I think all shapes are correctly tested already, but I agree it would be good to have a dedicated test, I thought about that already while putting this merge request together. The above code has most of the code for a good test already included, one of the reason I put it here. 
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2024-06-13T19:23:09Z
body: @dschmitz89 I am not sure about the use of np.einsum elsewhere, here it makes sense because applying a rotation to a vector is inherently a matrix multiplication. I am certain there would be a few places where einsum could be replaced by matmul but it might also be worthwhile looking into why matmul is faster.
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2024-06-14T18:37:04Z
body: First comment: The old implementation returned a view whenever rotation was single and `vectors.shape == (, 3)`. The new implementation returns a view for most cases.

Second comment: I run some more, complete benchmarks and `np.einsum` is faster for the case where multiple rotations were multiplied with vectors. So I returned that part of the code to use `np.einsum` again. There is still one case, where the current implementation is slower, for `len(Rotation) > 1000` of rotations with a vector of shape `(,3)`. This seems to be a very unusual case and I don't want to switch code paths just because it is faster on my particular machine. The same code is faster for `len(Rotation) < 1000` in general and much faster (factor of 2) for `Rotation.single == True`.

Here the full benchmark code and the result. @scottshambaugh Would probably make sense to run it on a few other computers.

```python
import numpy as np
import scipy
import time
from itertools import product

def apply_new(self, vectors, inverse=False):
    vectors = np.asarray(vectors)
    if vectors.ndim > 2 or vectors.shape[-1] != 3:
        raise ValueError("Expected input of shape (3,) or (P, 3), "
                         "got {}.".format(vectors.shape))

    matrix = self.as_matrix()
    if inverse:
        matrix = np.swapaxes(matrix, -1, -2)

    if vectors.shape == (3,):
        return np.matmul(matrix, vectors)
    if self.single:
        matrix = matrix[None, :, :]
    
    n_vectors = vectors.shape[0]
    n_rotations = 1 if self.single else len(self)

    if n_rotations == 1:
        return (np.matmul(matrix, vectors, axes=[(-2, -1), (-1, -2), (-1, -2)])[0])
    elif n_vectors == 1 or n_vectors == n_rotations:
        return np.einsum('ijk,ik->ij', matrix, vectors)

    raise ValueError("Expected equal numbers of rotations and vectors "
                     ", or a single rotation, or a single vector, got "
                     "{} rotations and {} vectors.".format(
                     n_rotations, n_vectors))

def create_data(qlength, vlength):
    if vlength == 0:
        vector = np.random.default_rng().random(3, dtype=float)
    else:
        vector = np.random.default_rng().random((vlength, 3), dtype=float)
    if qlength == 0:
        q = np.random.default_rng().random(4, dtype=float)
        q /= np.linalg.norm(q)
    else:
        q = np.random.default_rng().random((qlength, 4), dtype=float)
        q = (q.T / np.linalg.norm(q, axis=1)).T
    rot = scipy.spatial.transform.Rotation.from_quat(q)
    return rot, vector

def rot_apply(qlength, vlength, **kwargs):
    rot, vector = create_data(qlength, vlength)
    start = time.time()
    vold = rot.apply(vector, **kwargs)
    old = time.time() - start
    start = time.time()
    vnew = apply_new(rot, vector, **kwargs)
    new = time.time() - start
    assert np.allclose(vold, vnew)
    assert np.all(vold.shape == vnew.shape)
    # print(f"{qlength}, {vlength}, {vold.base}, {vnew.base}")
    return old, new

repeats = 1000

for qlength, vlength in product([0, 1, 100, 10000], [0, 1, 100, 10000]):
    time_old = 0
    time_new = 0
    try:
        for _ in range(repeats):
            old, new = rot_apply(qlength, vlength)
            time_old += old
            time_new += new
            old, new = rot_apply(qlength, vlength, inverse=True)
            time_old += old
            time_new += new
        print(qlength, vlength, time_old / time_new)
    except ValueError as e:
        print(qlength, vlength, e)
```

```
# RotationLength VectorLength Old/New 
0 0 1.6128810708528978
0 1 1.1584153250628633
0 100 1.3335232488209414
0 10000 5.467099769567047
1 0 2.0145978222690553
1 1 1.2266894155082901
1 100 1.3443168367510157
1 10000 4.561584595210014
100 0 1.2184889499841387
100 1 1.202783661119516
100 100 1.2001426406347508
100 10000 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 100 rotations and 10000 vectors.
10000 0 0.5991654889963401
10000 1 1.0108455611809377
10000 100 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 10000 rotations and 100 vectors.
10000 10000 1.049880155989495
```

Third I added some tests now and just pushed the new code.

Edits: Fixed the missing `from itertools import product`

Additional information: This was run with numpy==1.24.4 and python3.8.
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2024-06-14T18:42:24Z
body: By the way:
 * There is a open pull request that tries to replace einsum internals with matmul in certain cases where it performs better: https://github.com/numpy/numpy/pull/23513
 * Here a more general discussion of einsum vs matmul : https://github.com/numpy/numpy/issues/22604
author_association: CONTRIBUTOR
user: login: scottshambaugh
created_at: 2024-06-16T04:08:18Z
body: Here's my benchmark results, follows a similar pattern to yours. 

```
0 0 0.8921083151669145
0 1 1.1040477565037288
0 100 0.6159805957873901
0 10000 3.77882823527405
1 0 1.3328017417713094
1 1 1.117239321695207
1 100 1.0407421398109065
1 10000 3.718852274239589
100 0 0.9686032244871602
100 1 1.2368293831024262
100 100 1.224578172008924
100 10000 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 100 rotations and 10000 vectors.
10000 0 0.3426780849381298
10000 1 1.054243055281351
10000 100 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 10000 rotations and 100 vectors.
10000 10000 1.22044644462435
```

I think the added tests and code change looks good!
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2024-06-16T04:14:49Z
body: @scottshambaugh Thanks. I am honestly quite bafflet that the cases with `vector.shape = (3,)` seem to be general faster with einsum compared to matmul. This is a prim example of what matmul is for (particular the `(3, 3) @ (3,)` case), so I would have expected it being extremely optimized. I can allow this case to fall back to the einsum version as well if there is agreement here about this. Might be also worth filing an issue with numpy for this.
author_association: CONTRIBUTOR
user: login: scottshambaugh
created_at: 2024-09-05T14:10:41Z
body: Pinging @nmayorov for review!
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-01-29T21:31:38Z
body: As far as I understand this PR is worthwhile to get back to.
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-03-21T21:12:11Z
body: @cosama will you have time to finish it?

Only few code issues to consider and then it should go in.
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-03-21T22:12:16Z
body: If necessary, I can open another PR with minor adjustments from this branch I'd like to see. 

It's definitely minor, but why not keep the code slightly tidier.

I'll wait for @cosama for few days.
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2025-03-22T04:21:52Z
body: I should be able to add those changes. I haven't been pushing this since I am still unsure if some of the speed ups are really not just pure hardware dependent. I have a new laptop and can also run some of the initial tests. Give me a week or so to wrap it up.
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2025-03-31T13:30:59Z
body: @nmayorov I tried to address your comments and moved code around a bit, hopeful this addresses all your concerns.

 I also broke up the benchmark into a forward and inverse part and ran it again on my new computer. Here the results:
```
0 0 1.8452974770476398 1.4531733854094016
0 1 1.336271584847663 1.1087501844473957
0 100 1.3964632509670798 1.2113899613899615
0 10000 2.0539749043415063 2.1320597355197313
1 0 1.7097630454486599 1.3387907439661608
1 1 1.3591692538272133 1.1176571551169276
1 100 1.4486333209371514 1.222575637093415
1 10000 2.0582821588865303 2.1438766617546863
100 0 1.3246859004187994 1.165620834281627
100 1 1.4224954823923488 1.1653989766284054
100 100 1.4106442577030813 1.1664950929929312
100 10000 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 100 rotations and 10000 vectors.
10000 0 0.7846246732697737 0.8442266003171767
10000 1 1.280833456895538 1.235395807497396
10000 100 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 10000 rotations and 100 vectors.
10000 10000 1.2580216429958289 1.2240184030114867
```

Except for the case where one has a ten thousand rotations and a single vector, it is always faster with this change.

author_association: CONTRIBUTOR
user: login: cosama
created_at: 2025-04-01T07:20:13Z
body: @nmayorov @tylerjereddy @peterbell10 The CI passed. I think this is ready for a second round of reviews.
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-04-02T22:31:22Z
body: Let's wait for all pipelines to succeed and I'll merge.
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-04-03T07:36:29Z
body: @cosama thank you for the persistency.
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2025-04-03T08:18:21Z
body: I just tried to benchmark the new code and weirdly the new implementation is still about 50% slower than my pure python numpy implementation that I used for benchmarking, the changes might have made things even worse?

```
0 0 1.7305148679762896 1.7025304079158357
0 1 1.519727144962808 1.507895557508576
0 100 1.4967968461254157 1.4849362538467505
0 10000 1.0530161855863172 1.0624768
1 0 1.6433901655547318 1.633105005059547
1 1 1.627465098603944 1.6023462125889578
1 100 1.5625677423866298 1.5682991934079964
1 10000 1.0497874766742692 1.058891571389028
100 0 1.352968725607535 1.3904833151976785
100 1 1.4841569694221954 1.4395166365963161
100 100 1.4489400808724422 1.4298251125772332
100 10000 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 100 rotations and 10000 vectors.
10000 0 1.1797042145792036 1.1827059732094514
10000 1 1.2949531725896601 1.2676253213715802
10000 100 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 10000 rotations and 100 vectors.
10000 10000 1.2400124445219305 1.2287987675006589
```

Any ideas why this might be?

Should I open a new issue? I will try to look into it some but it is a bit puzzling to me, must be some weird `cython` slow down. I just assumed the `cython` code will be as good as pure python, but that seems not to be the case. If I benchmark the old and new version in pure python I see the previously observed patterns, so maybe some of the changes I introduced occur some unanticipated overhead.
author_association: CONTRIBUTOR
user: login: nmayorov
created_at: 2025-04-03T10:40:56Z
body: Have you accounted for quaternion to matrix conversion?
author_association: CONTRIBUTOR
user: login: cosama
created_at: 2025-04-03T12:38:48Z
body: Yes I did. It turns out that the differences I saw were again due to some CPU memory optimization where the second call somehow managed to be faster due to some memory catching. It is really difficult to make sure these effects are not messing with your benchmark. Here my final result for the old implementation (run in cython), vs the new implementation (also run in cython), >1 means new is better:

```
0 0 1.4253966475753035
0 1 1.1266725832878208
0 100 1.2162649361351463
0 10000 5.553617506311027
1 0 1.4105507920360412
1 1 1.1402151788310555
1 100 1.229784653864009
1 10000 5.664163661230228
100 0 1.1767079740353683
100 1 0.969317042823074
100 100 0.9640595918991731
100 10000 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 100 rotations and 10000 vectors.
10000 0 0.6642819117004968
10000 1 0.8861196881814045
10000 100 Expected equal numbers of rotations and vectors , or a single rotation, or a single vector, got 10000 rotations and 100 vectors.
10000 10000 1.1155628016131678
```

I think all is good. Also thanks for working through everything with me.
author_association: CONTRIBUTOR