diff --git a/.clang-tidy b/.clang-tidy index db4c130..1ad12ab 100644 --- a/.clang-tidy +++ b/.clang-tidy @@ -19,6 +19,7 @@ Checks: > -readability-avoid-const-params-in-decls, -readability-braces-around-statements, -readability-isolate-declaration, + -readability-math-missing-parentheses, -cppcoreguidelines-avoid-magic-numbers, -cppcoreguidelines-pro-bounds-array-to-pointer-decay, -cppcoreguidelines-pro-bounds-pointer-arithmetic, diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index 5a48a8a..b07e8ad 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -11,19 +11,58 @@ namespace nb = nanobind; NB_MODULE(_genmetaballs_bindings, m) { - // exposing Vec3D - nb::class_(m, "Vec3D") + /* + * Geometry module bindings + */ + + nb::module_ geometry = m.def_submodule("geometry", "Geometry helpers for GenMetaballs"); + + nb::class_(geometry, "Vec3D") .def(nb::init<>()) .def(nb::init()) - .def_rw("x", &Vec3D::x) - .def_rw("y", &Vec3D::y) - .def_rw("z", &Vec3D::z) + .def_ro("x", &Vec3D::x) + .def_ro("y", &Vec3D::y) + .def_ro("z", &Vec3D::z) .def(nb::self + nb::self) .def(nb::self - nb::self) + .def(-nb::self) + .def(nb::self * float()) + .def(float() * nb::self) + .def(nb::self / float()) .def("__repr__", [](const Vec3D& v) { return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); }); - // confidence submodule + geometry.def("dot", &dot, "Dot product of two `Vec3D`s", nb::arg("a"), nb::arg("b")); + geometry.def("cross", &cross, "Cross product of two `Vec3D`s", nb::arg("a"), nb::arg("b")); + + nb::class_(geometry, "Rotation") + .def(nb::init<>()) + .def_static("from_quat", &Rotation::from_quat, "Create rotation from quaternion", + nb::arg("x"), nb::arg("y"), nb::arg("z"), nb::arg("w")) + .def("apply", &Rotation::apply, "Apply rotation to vector", nb::arg("vec")) + .def("compose", &Rotation::compose, "Compose with another rotation", nb::arg("rot")) + .def("inv", &Rotation::inv, "Inverse rotation"); + + nb::class_(geometry, "Pose") + .def(nb::init<>()) + .def_static("from_components", &Pose::from_components, + "Create rotation from a rotation and a translation", nb::arg("rot"), + nb::arg("tran")) + .def_prop_ro("rot", &Pose::get_rot, "get the rotation component") + .def_prop_ro("tran", &Pose::get_tran, "get the translation component") + .def("apply", &Pose::apply, "Apply pose to vector", nb::arg("vec")) + .def("compose", &Pose::compose, "Compose with another pose", nb::arg("pose")) + .def("inv", &Pose::inv, "Inverse pose"); + + nb::class_(geometry, "Ray") + .def(nb::init<>()) + .def_rw("start", &Ray::start) + .def_rw("direction", &Ray::direction); + + /* + * Confidence module bindings + */ + nb::module_ confidence = m.def_submodule("confidence"); nb::class_(confidence, "ZeroParameterConfidence") .def(nb::init<>()) @@ -33,7 +72,10 @@ NB_MODULE(_genmetaballs_bindings, m) { .def(nb::init()) .def("get_confidence", &TwoParameterConfidence::get_confidence); - // utils submodule + /* + * Utils module bindings + */ + nb::module_ utils = m.def_submodule("utils"); utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))"); diff --git a/genmetaballs/src/cuda/core/geometry.cu b/genmetaballs/src/cuda/core/geometry.cu index 9e5c886..c645bdf 100644 --- a/genmetaballs/src/cuda/core/geometry.cu +++ b/genmetaballs/src/cuda/core/geometry.cu @@ -1,9 +1,39 @@ +#include + #include "geometry.cuh" -Vec3D operator+(const Vec3D a, const Vec3D b) { - return {a.x + b.x, a.y + b.y, a.z + b.z}; +// NOLINTNEXTLINE(readability-convert-member-functions-to-static) +CUDA_CALLABLE Rotation Rotation::from_quat(float x, float y, float z, float w) { + auto modulus = std::sqrt(x * x + y * y + z * z + w * w); + return Rotation{{x / modulus, y / modulus, z / modulus, w / modulus}}; +} + +CUDA_CALLABLE Vec3D Rotation::apply(const Vec3D vec) const { + // v' = q * v * q^(-1) for unit quaternions + // where q^(-1) = (-x, -y, -z, w) + Vec3D q = {unit_quat_.x, unit_quat_.y, unit_quat_.z}; + float w = unit_quat_.w; + + // v' = 2*(q·v)*q + (w²-|q|²)*v + 2*w*(q×v) + float d = dot(q, vec); + Vec3D c = cross(q, vec); + + return 2.0f * d * q + (w * w - dot(q, q)) * vec + 2.0f * w * c; +} + +// NOLINTNEXTLINE(readability-convert-member-functions-to-static) +CUDA_CALLABLE Rotation Rotation::compose(const Rotation& rot) const { + // Quaternion multiplication: q1 * q2 + float4 q1 = unit_quat_; + float4 q2 = rot.unit_quat_; + + return Rotation{{q1.w * q2.x + q1.x * q2.w + q1.y * q2.z - q1.z * q2.y, + q1.w * q2.y - q1.x * q2.z + q1.y * q2.w + q1.z * q2.x, + q1.w * q2.z + q1.x * q2.y - q1.y * q2.x + q1.z * q2.w, + q1.w * q2.w - q1.x * q2.x - q1.y * q2.y - q1.z * q2.z}}; } -Vec3D operator-(const Vec3D a, const Vec3D b) { - return {a.x - b.x, a.y - b.y, a.z - b.z}; +CUDA_CALLABLE Rotation Rotation::inv() const { + // For unit quaternions, inverse = conjugate: (-x, -y, -z, w) + return Rotation{{-unit_quat_.x, -unit_quat_.y, -unit_quat_.z, unit_quat_.w}}; } diff --git a/genmetaballs/src/cuda/core/geometry.cuh b/genmetaballs/src/cuda/core/geometry.cuh index e1f6334..f64dbd2 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -2,30 +2,101 @@ #include +#include "utils.cuh" + using Vec3D = float3; -Vec3D operator+(const Vec3D a, const Vec3D b); -Vec3D operator-(const Vec3D a, const Vec3D b); +CUDA_CALLABLE inline Vec3D operator-(const Vec3D a) { + return {-a.x, -a.y, -a.z}; +} -class Rotation { +CUDA_CALLABLE inline Vec3D operator+(const Vec3D a, const Vec3D b) { + return {a.x + b.x, a.y + b.y, a.z + b.z}; +} + +CUDA_CALLABLE inline Vec3D operator-(const Vec3D a, const Vec3D b) { + return {a.x - b.x, a.y - b.y, a.z - b.z}; +} + +CUDA_CALLABLE inline Vec3D operator*(const float scalar, const Vec3D a) { + return {a.x * scalar, a.y * scalar, a.z * scalar}; +} + +CUDA_CALLABLE inline Vec3D operator*(const Vec3D a, const float scalar) { + return {a.x * scalar, a.y * scalar, a.z * scalar}; +} + +CUDA_CALLABLE inline Vec3D operator/(const Vec3D a, const float scalar) { + return {a.x / scalar, a.y / scalar, a.z / scalar}; +} +CUDA_CALLABLE inline float dot(const Vec3D a, const Vec3D b) { + return a.x * b.x + a.y * b.y + a.z * b.z; +} + +CUDA_CALLABLE inline Vec3D cross(const Vec3D a, const Vec3D b) { + return {a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x}; +} + +class Rotation { private: - // ... - float rotmat_[9]; + float4 unit_quat_; + + CUDA_CALLABLE Rotation(float4 unit_quat) : unit_quat_{unit_quat} {}; public: - Vec3D apply(const Vec3D vec) const; - Rotation compose(const Rotation& rot) const; - Rotation inv() const; + CUDA_CALLABLE Rotation() : unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + + static CUDA_CALLABLE Rotation from_quat(float x, float y, float z, float w); + + CUDA_CALLABLE Vec3D apply(const Vec3D vec) const; + + CUDA_CALLABLE Rotation compose(const Rotation& rot) const; + + CUDA_CALLABLE Rotation inv() const; }; -struct Pose { - Rotation rot; - Vec3D tran; +class Pose { +private: + Rotation rot_; + Vec3D tran_; + + CUDA_CALLABLE Pose(const Rotation rot, const Vec3D tran) : rot_{rot}, tran_{tran} {} + +public: + // these member functions are defined in class body to allow for possible inlining + + CUDA_CALLABLE Pose() : rot_{Rotation()}, tran_{0.0f, 0.0f, 0.0f} {} + + static CUDA_CALLABLE Pose from_components(const Rotation rot, const Vec3D tran) { + return {rot, tran}; + } + + CUDA_CALLABLE Rotation get_rot() const { + return rot_; + } + + CUDA_CALLABLE Vec3D get_tran() const { + return tran_; + } + + CUDA_CALLABLE Vec3D apply(const Vec3D vec) const { + return tran_ + rot_.apply(vec); + } + + CUDA_CALLABLE Pose compose(const Pose& pose) const { + /* + * If $A_i$ is the matrix corresponding to pose object `p_i`, then + * $A_1A_2$ is the matrix corresponding to the pose object + * `p_1.compose(p2)`. + */ + return {rot_.compose(pose.rot_), rot_.apply(pose.tran_) + tran_}; + } - Vec3D apply(const Vec3D vec) const; - Pose compose(const Pose& pose) const; - Pose inv() const; + CUDA_CALLABLE Pose inv() const { + auto rotinv = rot_.inv(); + return {rotinv, -rotinv.apply(tran_)}; + } }; struct Ray { diff --git a/genmetaballs/src/cuda/core/utils.cuh b/genmetaballs/src/cuda/core/utils.cuh index d9b8f28..20f5a66 100644 --- a/genmetaballs/src/cuda/core/utils.cuh +++ b/genmetaballs/src/cuda/core/utils.cuh @@ -52,4 +52,4 @@ public: __host__ __device__ constexpr auto size() const noexcept { return data_view_.size(); } -}; +}; // class Array2D diff --git a/genmetaballs/src/genmetaballs/core/__init__.py b/genmetaballs/src/genmetaballs/core/__init__.py index 341a681..23c7077 100644 --- a/genmetaballs/src/genmetaballs/core/__init__.py +++ b/genmetaballs/src/genmetaballs/core/__init__.py @@ -1,3 +1,4 @@ +from genmetaballs._genmetaballs_bindings import geometry from genmetaballs._genmetaballs_bindings.confidence import ( TwoParameterConfidence, ZeroParameterConfidence, @@ -7,5 +8,6 @@ __all__ = [ "ZeroParameterConfidence", "TwoParameterConfidence", + "geometry", "sigmoid", ] diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index 6403ef2..a150bc5 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -1,7 +1,11 @@ +import operator as op + import numpy as np import pytest +from scipy.spatial.transform import RigidTransform as Rigid +from scipy.spatial.transform import Rotation as Rot -from genmetaballs import _genmetaballs_bindings as _gmbb +from genmetaballs.core import geometry @pytest.fixture @@ -10,27 +14,176 @@ def rng() -> np.random.Generator: def test_vec3d_smoke() -> None: - _gmbb.Vec3D(0, 0, 0) + geometry.Vec3D(0, 0, 0) def test_vec3d_repr_returns_valid_string() -> None: - v = _gmbb.Vec3D(1.0, 2.0, 3.0) + v = geometry.Vec3D(1.0, 2.0, 3.0) repr_str = repr(v) assert isinstance(repr_str, str) assert repr_str == "Vec3D(1.0, 2.0, 3.0)" -def test_vec3d_add(rng: np.random.Generator) -> None: - _a, _b = rng.uniform(size=3), rng.uniform(size=3) - a, b = _gmbb.Vec3D(*_a), _gmbb.Vec3D(*_b) - c = a + b - _c = _a + _b - assert all([np.isclose(c.x, _c[0]), np.isclose(c.y, _c[1]), np.isclose(c.z, _c[2])]) +@pytest.mark.parametrize( + "np_op,vec3d_op", + [ + (np.add, op.add), + (np.subtract, op.sub), + (np.cross, geometry.cross), + ], +) +def test_vec3d_ops(rng: np.random.Generator, np_op, vec3d_op) -> None: + _a, _b = rng.uniform(size=(100, 3)), rng.uniform(size=(100, 3)) + for i in range(100): + a, b = geometry.Vec3D(*_a[i]), geometry.Vec3D(*_b[i]) + c = vec3d_op(a, b) + _c = np_op(_a[i], _b[i]) + assert np.allclose(_c, np.array([c.x, c.y, c.z])) + + +def test_vec3d_dot(rng: np.random.Generator) -> None: + _a, _b = rng.uniform(size=(100, 3)), rng.uniform(size=(100, 3)) + for i in range(100): + a, b = geometry.Vec3D(*_a[i]), geometry.Vec3D(*_b[i]) + assert np.allclose(np.dot(_a[i], _b[i]), geometry.dot(a, b)) + + +def test_rotation_apply(rng: np.random.Generator) -> None: + quats = rng.uniform(-1, 1, size=(100, 4)) + quats /= np.linalg.norm(quats, axis=1, keepdims=True) + vecs = rng.uniform(size=(100, 3)) + + for i in range(100): + rot_scipy = Rot.from_quat(quats[i]) + rot_geom = geometry.Rotation.from_quat(*quats[i]) + vec_geom = geometry.Vec3D(*vecs[i]) + + result_scipy = rot_scipy.apply(vecs[i]) + result_geom = rot_geom.apply(vec_geom) + + assert np.allclose(result_scipy, np.array([result_geom.x, result_geom.y, result_geom.z])) + + +def test_rotation_compose(rng: np.random.Generator) -> None: + quats1 = rng.uniform(-1, 1, size=(100, 4)) + quats1 /= np.linalg.norm(quats1, axis=1, keepdims=True) + quats2 = rng.uniform(-1, 1, size=(100, 4)) + quats2 /= np.linalg.norm(quats2, axis=1, keepdims=True) + vecs = rng.uniform(size=(50, 3)) + + for i in range(100): + rot1_scipy = Rot.from_quat(quats1[i]) + rot2_scipy = Rot.from_quat(quats2[i]) + composed_scipy = rot1_scipy * rot2_scipy + + rot1_geom = geometry.Rotation.from_quat(*quats1[i]) + rot2_geom = geometry.Rotation.from_quat(*quats2[i]) + composed_geom = rot1_geom.compose(rot2_geom) + + for j in range(50): + vec_geom = geometry.Vec3D(*vecs[j]) + result_scipy = composed_scipy.apply(vecs[j]) + result_geom = composed_geom.apply(vec_geom) + assert np.allclose( + result_scipy, + np.array([result_geom.x, result_geom.y, result_geom.z]), + rtol=1e-5, + atol=1e-6, + ) + + +def test_rotation_inv(rng: np.random.Generator) -> None: + quats = rng.uniform(-1, 1, size=(100, 4)) + quats /= np.linalg.norm(quats, axis=1, keepdims=True) + vecs = rng.uniform(size=(50, 3)) + + for i in range(100): + rot = geometry.Rotation.from_quat(*quats[i]) + rotinv = rot.inv() + composed = rot.compose(rotinv) + + for j in range(50): + vec = geometry.Vec3D(*vecs[j]) + vec_ = composed.apply(vec) + assert np.allclose( + np.array([vec.x, vec.y, vec.z]), + np.array([vec_.x, vec_.y, vec_.z]), + rtol=1e-5, + atol=1e-6, + ) + + +def test_pose_apply(rng: np.random.Generator) -> None: + exp_coords = rng.uniform(size=(100, 6)) + vecs = rng.uniform(size=(100, 3)) + + for i in range(100): + pose_scipy = Rigid.from_exp_coords(exp_coords[i]) + + pose_geom = geometry.Pose.from_components( + rot=geometry.Rotation.from_quat(*pose_scipy.rotation.as_quat()), + tran=geometry.Vec3D(*pose_scipy.translation), + ) + vec_geom = geometry.Vec3D(*vecs[i]) + + result_scipy = pose_scipy.apply(vecs[i]) + result_geom = pose_geom.apply(vec_geom) + + assert np.allclose(result_scipy, np.array([result_geom.x, result_geom.y, result_geom.z])) + + +def test_pose_compose(rng: np.random.Generator) -> None: + exp_coords1 = rng.uniform(size=(100, 6)) + exp_coords2 = rng.uniform(size=(100, 6)) + vecs = rng.uniform(size=(50, 3)) + + for i in range(100): + pose1_scipy = Rigid.from_exp_coords(exp_coords1[i]) + pose2_scipy = Rigid.from_exp_coords(exp_coords2[i]) + composed_scipy = pose1_scipy * pose2_scipy + + pose1_geom = geometry.Pose.from_components( + rot=geometry.Rotation.from_quat(*pose1_scipy.rotation.as_quat()), + tran=geometry.Vec3D(*pose1_scipy.translation), + ) + pose2_geom = geometry.Pose.from_components( + rot=geometry.Rotation.from_quat(*pose2_scipy.rotation.as_quat()), + tran=geometry.Vec3D(*pose2_scipy.translation), + ) + composed_geom = pose1_geom.compose(pose2_geom) + + for j in range(50): + vec_geom = geometry.Vec3D(*vecs[j]) + result_scipy = composed_scipy.apply(vecs[j]) + result_geom = composed_geom.apply(vec_geom) + assert np.allclose( + result_scipy, + np.array([result_geom.x, result_geom.y, result_geom.z]), + rtol=1e-5, + atol=1e-6, + ) + + +def test_pose_inv(rng: np.random.Generator) -> None: + exp_coords = rng.uniform(size=(100, 6)) + vecs = rng.uniform(size=(50, 3)) + + for i in range(100): + pose_scipy = Rigid.from_exp_coords(exp_coords[i]) + pose = geometry.Pose.from_components( + rot=geometry.Rotation.from_quat(*pose_scipy.rotation.as_quat()), + tran=geometry.Vec3D(*pose_scipy.translation), + ) + poseinv = pose.inv() + composed = pose.compose(poseinv) -def test_vec3d_sub(rng: np.random.Generator) -> None: - _a, _b = rng.uniform(size=3), rng.uniform(size=3) - a, b = _gmbb.Vec3D(*_a), _gmbb.Vec3D(*_b) - c = a - b - _c = _a - _b - assert all([np.isclose(c.x, _c[0]), np.isclose(c.y, _c[1]), np.isclose(c.z, _c[2])]) + for j in range(50): + vec = geometry.Vec3D(*vecs[j]) + vec_ = composed.apply(vec) + assert np.allclose( + np.array([vec.x, vec.y, vec.z]), + np.array([vec_.x, vec_.y, vec_.z]), + rtol=1e-5, + atol=1e-6, + )