From 5da6a9eac465a170d48f1f3cc2559abea16c7eae Mon Sep 17 00:00:00 2001 From: Ankith Date: Mon, 15 Jul 2024 01:19:02 +0530 Subject: [PATCH 1/3] Implement inv_rotate in elasticapp._rotations - Also adds corresponding test and benchmarking script --- backend/benchmarking/inv_rotate.py | 46 ++++++++++++++++++++ backend/src/_rotations.cpp | 67 ++++++++++++++++++++++++++++++ backend/src/meson.build | 9 ++++ backend/tests/test_rotations.py | 37 +++++++++++++++++ 4 files changed, 159 insertions(+) create mode 100644 backend/benchmarking/inv_rotate.py create mode 100644 backend/src/_rotations.cpp create mode 100644 backend/tests/test_rotations.py diff --git a/backend/benchmarking/inv_rotate.py b/backend/benchmarking/inv_rotate.py new file mode 100644 index 00000000..8bf7169a --- /dev/null +++ b/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() diff --git a/backend/src/_rotations.cpp b/backend/src/_rotations.cpp new file mode 100644 index 00000000..d12f6bb2 --- /dev/null +++ b/backend/src/_rotations.cpp @@ -0,0 +1,67 @@ +#include +#include + +#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; + +using ElasticaVector = ::blaze::DynamicVector>; +using ElasticaMatrix = ::blaze::DynamicMatrix>; +using ElasticaTensor = ::blaze::DynamicTensor; + +template +ElasticaMatrix inv_rotate_with_span(ElasticaTensor &director_collection, ElasticaVector &span_vector) +{ + ElasticaMatrix vector_collection(director_collection.rows(), director_collection.columns() - 1); + detail::InvRotateDivideOp::apply(vector_collection, director_collection, span_vector); + return vector_collection; +} + +// Overloaded function where span_vector is filled with 1 +template +ElasticaMatrix inv_rotate(ElasticaTensor &director_collection) +{ + ElasticaVector span_vector(director_collection.columns(), 1); + return inv_rotate_with_span(director_collection, span_vector); +} + +PYBIND11_MODULE(_rotations, m) +{ + m.doc() = R"pbdoc( + elasticapp _rotations + --------------- + + .. currentmodule:: _rotations + + .. autosummary:: + :toctree: _generate + + inv_rotate + )pbdoc"; + + m.def("inv_rotate", &inv_rotate, R"pbdoc( + Perform the inverse-rotate operation (SIMD Variant). + )pbdoc"); + m.def("inv_rotate", &inv_rotate_with_span, 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, R"pbdoc( + Perform the inverse-rotate operation (Scalar Variant). + )pbdoc"); + m.def("inv_rotate_scalar", &inv_rotate_with_span, 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 +} diff --git a/backend/src/meson.build b/backend/src/meson.build index 74be201e..3b39caa3 100644 --- a/backend/src/meson.build +++ b/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: [ diff --git a/backend/tests/test_rotations.py b/backend/tests/test_rotations.py new file mode 100644 index 00000000..d97d4f37 --- /dev/null +++ b/backend/tests/test_rotations.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 + +# This file is based on pyelastica tests/test_math/test_rotations.py + +# System imports +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 + + +@pytest.mark.parametrize("inv_rotate_func", [inv_rotate, inv_rotate_scalar]) +def test_inv_rotate(inv_rotate_func): + # A rotation of 120 degrees about x=y=z gives + # the permutation matrix P + # {\begin{bmatrix}0&0&1\\1&0&0\\0&1&0\end{bmatrix}} + # Hence if we pass in I and P*I, the vector returned should be + # along [1.0, 1.0, 1.0] / sqrt(3.0) with a rotation of angle = 120/180 * pi + rotate_from_matrix = np.eye(3).reshape(3, 3, 1) + # Q_new = Q_old . R^T + rotate_to_matrix = np.eye(3) @ np.roll(np.eye(3), -1, axis=1).T + input_director_collection = np.dstack( + (rotate_from_matrix, rotate_to_matrix.reshape(3, 3, -1)) + ) + + correct_axis_collection = np.ones((3, 1)) / np.sqrt(3.0) + test_axis_collection = np.asarray( + inv_rotate_func(Tensor(input_director_collection)) + ) + + correct_angle = np.deg2rad(120) + test_angle = np.linalg.norm(test_axis_collection, axis=0) # (3,1) + test_axis_collection /= test_angle + + assert_allclose(test_axis_collection, correct_axis_collection) + assert_allclose(test_angle, correct_angle) From 25c623706f5726e01ceafe4c3e7609cf7e4b17cd Mon Sep 17 00:00:00 2001 From: Ankith Date: Mon, 15 Jul 2024 02:04:56 +0530 Subject: [PATCH 2/3] Port more inv_rotate tests --- backend/tests/test_rotations.py | 124 ++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/backend/tests/test_rotations.py b/backend/tests/test_rotations.py index d97d4f37..e9646f03 100644 --- a/backend/tests/test_rotations.py +++ b/backend/tests/test_rotations.py @@ -35,3 +35,127 @@ def test_inv_rotate(inv_rotate_func): assert_allclose(test_axis_collection, correct_axis_collection) assert_allclose(test_angle, correct_angle) + + +@pytest.mark.parametrize("inv_rotate_func", [inv_rotate, inv_rotate_scalar]) +@pytest.mark.parametrize("blocksize", [32, 128, 512]) +@pytest.mark.parametrize("point_distribution", ["anticlockwise", "clockwise"]) +def test_inv_rotate_correctness_on_circle_in_two_dimensions( + inv_rotate_func, blocksize, point_distribution +): + """Construct a unit circle, which we know has constant curvature, + and see if inv_rotate gives us the correct axis of rotation and + the angle of change + + Do this when d3 = z and d3= -z to cover both cases + + Parameters + ---------- + blocksize + + Returns + ------- + + """ + # FSAL start at 0. and proceeds counter-clockwise + if point_distribution == "anticlockwise": + theta_collection = np.linspace(0.0, 2.0 * np.pi, blocksize) + elif point_distribution == "clockwise": + theta_collection = np.linspace(2.0 * np.pi, 0.0, blocksize) + else: + raise NotImplementedError + + # rate of change, should correspond to frame rotation angles + dtheta_di = np.abs(theta_collection[1] - theta_collection[0]) + + # +1 because last point should be same as first point + director_collection = np.zeros((3, 3, blocksize)) + + # First fill all d1 components + # normal direction + director_collection[0, 0, ...] = -np.cos(theta_collection) + director_collection[0, 1, ...] = -np.sin(theta_collection) + + # Then all d2 components + # tangential direction + director_collection[1, 0, ...] = -np.sin(theta_collection) + director_collection[1, 1, ...] = np.cos(theta_collection) + + # Then all d3 components + director_collection[2, 2, ...] = -1.0 + + # blocksize - 1 to account for end effects + if point_distribution == "anticlockwise": + axis_of_rotation = np.array([0.0, 0.0, -1.0]) + elif point_distribution == "clockwise": + axis_of_rotation = np.array([0.0, 0.0, 1.0]) + else: + raise NotImplementedError + + correct_axis_collection = np.tile(axis_of_rotation.reshape(3, 1), blocksize - 1) + + test_axis_collection = np.asarray(inv_rotate_func(Tensor(director_collection))) + test_scaling = np.linalg.norm(test_axis_collection, axis=0) + test_axis_collection /= test_scaling + + assert test_axis_collection.shape == (3, blocksize - 1) + assert_allclose(test_axis_collection, correct_axis_collection) + assert_allclose(test_scaling, 0.0 * test_scaling + dtheta_di) + + +@pytest.mark.parametrize("inv_rotate_func", [inv_rotate, inv_rotate_scalar]) +@pytest.mark.parametrize("blocksize", [32, 128]) +def test_inv_rotate_correctness_on_circle_in_two_dimensions_with_different_directors( + inv_rotate_func, + blocksize, +): + """Construct a unit circle, which we know has constant curvature, + and see if inv_rotate gives us the correct axis of rotation and + the angle of change + + Here d3 is not z axis, so the `inv_rotate` formula returns the + curvature but with components in the local axis, i.e. it gives + [K1, K2, K3] in kappa_l = K1 . d1 + K2 . d2 + K3 . d3 + + Parameters + ---------- + blocksize + + Returns + ------- + + """ + # FSAL start at 0. and proceeds counter-clockwise + theta_collection = np.linspace(0.0, 2.0 * np.pi, blocksize) + # rate of change, should correspond to frame rotation angles + dtheta_di = theta_collection[1] - theta_collection[0] + + # +1 because last point should be same as first point + director_collection = np.zeros((3, 3, blocksize)) + + # First fill all d3 components + # tangential direction + director_collection[2, 0, ...] = -np.sin(theta_collection) + director_collection[2, 1, ...] = np.cos(theta_collection) + + # Then all d2 components + # normal direction + director_collection[1, 0, ...] = -np.cos(theta_collection) + director_collection[1, 1, ...] = -np.sin(theta_collection) + + # Then all d1 components + # binormal = d2 x d3 + director_collection[0, 2, ...] = -1.0 + + # blocksize - 1 to account for end effects + # returned curvature is in local coordinates! + correct_axis_collection = np.tile( + np.array([-1.0, 0.0, 0.0]).reshape(3, 1), blocksize - 1 + ) + test_axis_collection = np.asarray(inv_rotate_func(Tensor(director_collection))) + test_scaling = np.linalg.norm(test_axis_collection, axis=0) + test_axis_collection /= test_scaling + + assert test_axis_collection.shape == (3, blocksize - 1) + assert_allclose(test_axis_collection, correct_axis_collection) + assert_allclose(test_scaling, 0.0 * test_scaling + dtheta_di) From 37e5860275c8ef3eacd7e27e471f5e0ce88ff707 Mon Sep 17 00:00:00 2001 From: Ankith Date: Mon, 15 Jul 2024 14:37:47 +0530 Subject: [PATCH 3/3] Implement rotate in elasticapp._rotations --- backend/benchmarking/rotate.py | 57 +++++++++++++++++++++++++++++++++ backend/src/_rotations.cpp | 21 +++++++++++- backend/tests/test_rotations.py | 45 ++++++++++++++++++++++++-- 3 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 backend/benchmarking/rotate.py diff --git a/backend/benchmarking/rotate.py b/backend/benchmarking/rotate.py new file mode 100644 index 00000000..e724c136 --- /dev/null +++ b/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() diff --git a/backend/src/_rotations.cpp b/backend/src/_rotations.cpp index d12f6bb2..451d5fed 100644 --- a/backend/src/_rotations.cpp +++ b/backend/src/_rotations.cpp @@ -1,11 +1,14 @@ #include #include +#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>; @@ -13,6 +16,12 @@ using ElasticaMatrix = ::blaze::DynamicMatrix>; using ElasticaTensor = ::blaze::DynamicTensor; +template +void rotate(ElasticaTensor &director_collection, ElasticaMatrix &axis_collection) +{ + states::detail::SO3AddAssign::apply(director_collection, axis_collection); +} + template 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, R"pbdoc( + Perform the rotate operation (SIMD Variant). + )pbdoc"); + m.def("rotate_scalar", &rotate, R"pbdoc( + Perform the rotate operation (Scalar Variant). + )pbdoc"); + m.def("inv_rotate", &inv_rotate, R"pbdoc( Perform the inverse-rotate operation (SIMD Variant). )pbdoc"); diff --git a/backend/tests/test_rotations.py b/backend/tests/test_rotations.py index e9646f03..07882d94 100644 --- a/backend/tests/test_rotations.py +++ b/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])