Implement rotate in elasticapp._rotations

backend/benchmarking/rotate.py

@@ -0,0 +1,57 @@
+from elasticapp._rotations import rotate, rotate_scalar
+from elasticapp._PyArrays import Tensor, Matrix
+from elastica._rotations import _rotate
+import numpy
+import time
+
+# warm up jit for fair comparison
+random_1 = numpy.random.random((3, 3, 1))
+random_2 = numpy.random.random((3, 1))
+out1 = _rotate(random_1, 1, random_2)
+
+
+def benchmark_batchsize(funcs: list, batches: list[int], num_iterations: int = 1000):
+    ret: dict = {}
+    for batch_size in batches:
+        random_a = numpy.random.random((3, 3, batch_size))
+        random_b = numpy.random.random((3, batch_size))
+
+        ret[batch_size] = {}
+        for func_name, func, func_arg_wrap in funcs:
+            tot = 0.0
+            for _ in range(num_iterations):
+                args = func_arg_wrap(random_a, random_b)
+                start = time.perf_counter()
+                func(*args)
+                tot += time.perf_counter() - start
+
+            ret[batch_size][func_name] = tot / num_iterations
+
+    return ret
+
+
+def _pyelastica_arg_wrap(x, y):
+    return x, 1.0, y
+
+
+def _elasticapp_arg_wrap(x, y):
+    return Tensor(x), Matrix(y)
+
+
+results = benchmark_batchsize(
+    [
+        ("pyelastica", _rotate, _pyelastica_arg_wrap),
+        ("elasticapp_simd", rotate, _elasticapp_arg_wrap),
+        ("elasticapp_scalar", rotate_scalar, _elasticapp_arg_wrap),
+    ],
+    [2**i for i in range(14)],
+)
+for size, data in results.items():
+    pyelastica = data["pyelastica"]
+    elasticapp_simd = data["elasticapp_simd"]
+    elasticapp_scalar = data["elasticapp_scalar"]
+    print(f"{size = }")
+    print(f"{pyelastica = }")
+    print(f"{elasticapp_simd = }, ratio: {elasticapp_simd / pyelastica}")
+    print(f"{elasticapp_scalar = }, ratio: {elasticapp_scalar / pyelastica}")
+    print()

backend/src/_rotations.cpp

@@ -1,18 +1,27 @@
 #include <pybind11/pybind11.h>
 #include <blaze/Math.h>
 
+#include "Systems/States/Expressions/backends/blaze/SO3PrimitiveAddAssign/Scalar.hpp"
+#include "Systems/States/Expressions/backends/blaze/SO3PrimitiveAddAssign/SIMD.hpp"
+
 #include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/Scalar.hpp"
 #include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/SIMD.hpp"
-#include "Systems/CosseratRods/Traits/Operations/BlazeBackend/InvRotateDivide/Blaze.hpp"
 
 namespace detail = elastica::cosserat_rod::detail;
+namespace states = elastica::states;
 
 using ElasticaVector = ::blaze::DynamicVector<double, ::blaze::columnVector,
                                               ::blaze::AlignedAllocator<double>>;
 using ElasticaMatrix = ::blaze::DynamicMatrix<double, ::blaze::rowMajor,
                                               ::blaze::AlignedAllocator<double>>;
 using ElasticaTensor = ::blaze::DynamicTensor<double>;
 
+template <states::detail::SO3AddAssignKind T>
+void rotate(ElasticaTensor &director_collection, ElasticaMatrix &axis_collection)
+{
+    states::detail::SO3AddAssign<T>::apply(director_collection, axis_collection);
+}
+
 template <detail::InvRotateDivideKind T>
 ElasticaMatrix inv_rotate_with_span(ElasticaTensor &director_collection, ElasticaVector &span_vector)
 {
@@ -40,9 +49,19 @@ PYBIND11_MODULE(_rotations, m)
         .. autosummary::
            :toctree: _generate
 
+           rotate
+           rotate_scalar
            inv_rotate
+           inv_rotate_scalar
     )pbdoc";
 
+    m.def("rotate", &rotate<states::detail::SO3AddAssignKind::simd>, R"pbdoc(
+        Perform the rotate operation (SIMD Variant).
+    )pbdoc");
+    m.def("rotate_scalar", &rotate<states::detail::SO3AddAssignKind::scalar>, R"pbdoc(
+        Perform the rotate operation (Scalar Variant).
+    )pbdoc");
+
     m.def("inv_rotate", &inv_rotate<detail::InvRotateDivideKind::simd>, R"pbdoc(
         Perform the inverse-rotate operation (SIMD Variant).
     )pbdoc");

backend/tests/test_rotations.py

@@ -6,8 +6,49 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_allclose
-from elasticapp._PyArrays import Tensor
-from elasticapp._rotations import inv_rotate, inv_rotate_scalar
+from elasticapp._PyArrays import Tensor, Matrix
+from elasticapp._rotations import rotate, rotate_scalar, inv_rotate, inv_rotate_scalar
+
+
+@pytest.mark.parametrize("rotate_func", [rotate, rotate_scalar])
+def test_rotate_correctness(rotate_func):
+    blocksize = 16
+
+    def get_aligned_director_collection(theta_collection):
+        sins = np.sin(theta_collection)
+        coss = np.cos(theta_collection)
+        # Get basic director out, then modify it as you like
+        dir = np.tile(np.eye(3).reshape(3, 3, 1), blocksize)
+        dir[0, 0, ...] = coss
+        # Flip signs on [0,1] and [1,0] to go from our row-wise
+        # representation to the more commonly used
+        # columnwise representation, for similar reasons metioned
+        # before
+        dir[0, 1, ...] = sins
+        dir[1, 0, ...] = -sins
+        dir[1, 1, ...] = coss
+
+        return dir
+
+    base_angle = np.deg2rad(np.linspace(0.0, 90.0, blocksize))
+    rotated_by = np.deg2rad(15.0) + 0.0 * base_angle
+    rotated_about = np.array([0.0, 0.0, 1.0]).reshape(-1, 1)
+
+    director_collection = Tensor(get_aligned_director_collection(base_angle))
+    axis_collection = np.tile(rotated_about, blocksize)
+    axis_collection *= rotated_by
+
+    rotate_func(director_collection, Matrix(axis_collection))
+    test_rotated_director_collection = np.asarray(director_collection)
+    correct_rotation = rotated_by + 1.0 * base_angle
+    correct_rotated_director_collection = get_aligned_director_collection(
+        correct_rotation
+    )
+
+    assert test_rotated_director_collection.shape == (3, 3, blocksize)
+    assert_allclose(
+        test_rotated_director_collection, correct_rotated_director_collection
+    )
 
 
 @pytest.mark.parametrize("inv_rotate_func", [inv_rotate, inv_rotate_scalar])