Merge branch 'ankith26-gsoc' of https://github.com/ankith26/PyElastica …

…into ankith26-gsoc

backend/benchmarking/inv_rotate.py

@@ -0,0 +1,46 @@
+from elasticapp._rotations import inv_rotate, inv_rotate_scalar
+from elasticapp._PyArrays import Tensor
+from elastica._rotations import _inv_rotate
+import numpy
+import time
+
+# warm up jit for fair comparison
+random_1 = numpy.random.random((3, 3, 1))
+out1 = _inv_rotate(random_1)
+
+
+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))
+
+        ret[batch_size] = {}
+        for func_name, func, func_wrap in funcs:
+            random_a_w = func_wrap(random_a) if func_wrap else random_a
+
+            start = time.perf_counter()
+            for _ in range(num_iterations):
+                func(random_a_w)
+
+            ret[batch_size][func_name] = (time.perf_counter() - start) / num_iterations
+
+    return ret
+
+
+results = benchmark_batchsize(
+    [
+        ("pyelastica", _inv_rotate, None),
+        ("elasticapp_simd", inv_rotate, Tensor),
+        ("elasticapp_scalar", inv_rotate_scalar, Tensor),
+    ],
+    [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/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

@@ -0,0 +1,86 @@
+#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"
+
+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)
+{
+    ElasticaMatrix vector_collection(director_collection.rows(), director_collection.columns() - 1);
+    detail::InvRotateDivideOp<T>::apply(vector_collection, director_collection, span_vector);
+    return vector_collection;
+}
+
+// Overloaded function where span_vector is filled with 1
+template <detail::InvRotateDivideKind T>
+ElasticaMatrix inv_rotate(ElasticaTensor &director_collection)
+{
+    ElasticaVector span_vector(director_collection.columns(), 1);
+    return inv_rotate_with_span<T>(director_collection, span_vector);
+}
+
+PYBIND11_MODULE(_rotations, m)
+{
+    m.doc() = R"pbdoc(
+        elasticapp _rotations
+        ---------------
+
+        .. currentmodule:: _rotations
+
+        .. 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");
+    m.def("inv_rotate", &inv_rotate_with_span<detail::InvRotateDivideKind::simd>, R"pbdoc(
+        Perform the inverse-rotate operation (SIMD Variant).
+        This overload also accepts a vector (as the second vector) to perform elementwise division.
+    )pbdoc");
+
+    m.def("inv_rotate_scalar", &inv_rotate<detail::InvRotateDivideKind::scalar>, R"pbdoc(
+        Perform the inverse-rotate operation (Scalar Variant).
+    )pbdoc");
+    m.def("inv_rotate_scalar", &inv_rotate_with_span<detail::InvRotateDivideKind::scalar>, R"pbdoc(
+        Perform the inverse-rotate operation (Scalar Variant).
+        This overload also accepts a vector (as the second vector) to perform elementwise division.
+    )pbdoc");
+
+#ifdef VERSION_INFO
+    m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
+#else
+    m.attr("__version__") = "dev";
+#endif
+}

backend/src/meson.build

@@ -25,6 +25,15 @@ _linalg_numpy = py.extension_module(
     subdir: package,
 )
 
+_rotations = py.extension_module(
+    '_rotations',
+    sources: ['_rotations.cpp'],
+    dependencies: py_deps + blaze_deps,
+    link_language: 'cpp',
+    install: true,
+    subdir: package,
+)
+
 _PyArrays = py.extension_module(
     '_PyArrays',
     sources: [