From f1fdff1b40170231a31ddce7f1357eeb070b5066 Mon Sep 17 00:00:00 2001 From: Matin Ghavami Date: Tue, 18 Nov 2025 13:26:57 -0500 Subject: [PATCH 01/11] initial implementation of geometry methods --- genmetaballs/src/cuda/core/geometry.cu | 34 +++++++++-- genmetaballs/src/cuda/core/geometry.cuh | 80 +++++++++++++++++++++---- 2 files changed, 99 insertions(+), 15 deletions(-) diff --git a/genmetaballs/src/cuda/core/geometry.cu b/genmetaballs/src/cuda/core/geometry.cu index 9e5c886..f78d317 100644 --- a/genmetaballs/src/cuda/core/geometry.cu +++ b/genmetaballs/src/cuda/core/geometry.cu @@ -1,9 +1,35 @@ #include "geometry.cuh" -Vec3D operator+(const Vec3D a, const Vec3D b) { - return {a.x + b.x, a.y + b.y, a.z + b.z}; + +__host__ __device__ 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; +} + +__host__ __device__ 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}; +__host__ __device__ 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..a3eee34 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -4,28 +4,86 @@ using Vec3D = float3; -Vec3D operator+(const Vec3D a, const Vec3D b); -Vec3D operator-(const Vec3D a, const Vec3D b); +inline Vec3D operator-(const Vec3D a) +{ + return {-a.x, -a.y, -a.z}; +} -class Rotation { +inline Vec3D operator+(const Vec3D a, const Vec3D b) +{ + return {a.x + b.x, a.y + b.y, a.z + b.z}; +} + +inline Vec3D operator-(const Vec3D a, const Vec3D b) +{ + return {a.x - b.x, a.y - b.y, a.z - b.z}; +} + +inline Vec3D operator*(const float scalar, const Vec3D a) +{ + return {a.x * scalar, a.y * scalar, a.z * scalar}; +} + +inline Vec3D operator*(const Vec3D a, const float scalar) +{ + return {a.x * scalar, a.y * scalar, a.z * scalar}; +} + +inline Vec3D operator/(const Vec3D a, const float scalar) +{ + return {a.x / scalar, a.y / scalar, a.z / scalar}; +} + +inline float dot(const Vec3D a, const Vec3D b) +{ + return a.x * b.x + a.y * b.y + a.z * b.z; +} + +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_; public: - Vec3D apply(const Vec3D vec) const; - Rotation compose(const Rotation& rot) const; - Rotation inv() const; + Rotation() unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + + __host__ __device__ Vec3D apply(const Vec3D vec) const; + + __host__ __device__ Rotation compose(const Rotation& rot) const; + + __host__ __device__ Rotation inv() const; }; struct Pose { Rotation rot; Vec3D tran; - Vec3D apply(const Vec3D vec) const; - Pose compose(const Pose& pose) const; - Pose inv() const; + __host__ __device__ inline Vec3D apply(const Vec3D vec) const + { + return tran + rot.apply(vec); + } + + __host__ __device__ inline 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}; + } + + __host__ __device__ inline Pose inv() const + { + auto rotinv = rot.inv(); + return {rotinv, -rotinv.apply(tran)}; + } }; struct Ray { From f143eb9419eee2b2cb6353069f6fc87d1c718276 Mon Sep 17 00:00:00 2001 From: Matin Ghavami Date: Tue, 18 Nov 2025 16:07:56 -0500 Subject: [PATCH 02/11] add bindings --- genmetaballs/src/cuda/bindings.cu | 63 ++++++++++++++++++++----- genmetaballs/src/cuda/core/geometry.cuh | 30 ++++++------ genmetaballs/src/cuda/core/utils.cuh | 2 +- tests/python_tests/test_geometry.py | 10 ++-- 4 files changed, 73 insertions(+), 32 deletions(-) diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index f75a322..b74f891 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -20,19 +20,54 @@ NB_MODULE(_genmetaballs_bindings, m) { m.def("gpu_add", &gpu_add, "Add two lists elementwise on the GPU", nb::arg("a"), nb::arg("b")); - // 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(nb::self + nb::self) - .def(nb::self - nb::self) - .def("__repr__", - [](const Vec3D& v) { return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); }); - - // confidence submodule + .def_ro("x", &Vec3D::x) + .def_ro("y", &Vec3D::y) + .def_ro("z", &Vec3D::z) + .def("__add__", static_cast(&operator+)) + .def("__sub__", static_cast(&operator-)) + .def("__neg__", static_cast(&operator-)) + .def("__mul__", static_cast(&operator*)) + .def("__rmul__", static_cast(&operator*)) + .def("__truediv__", static_cast(&operator/)) + .def("__repr__", [](const Vec3D& v) { + return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); + }); + + 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("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_rw("rot", &Pose::rot) + .def_rw("tran", &Pose::tran) + .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<>()) @@ -42,8 +77,12 @@ 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))"); } // NB_MODULE(_genmetaballs_bindings) + diff --git a/genmetaballs/src/cuda/core/geometry.cuh b/genmetaballs/src/cuda/core/geometry.cuh index a3eee34..b890956 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -2,44 +2,46 @@ #include +#include "utils.cuh" + using Vec3D = float3; -inline Vec3D operator-(const Vec3D a) +CUDA_CALLABLE inline Vec3D operator-(const Vec3D a) { return {-a.x, -a.y, -a.z}; } -inline Vec3D operator+(const Vec3D a, const Vec3D b) +CUDA_CALLABLE inline Vec3D operator+(const Vec3D a, const Vec3D b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; } -inline Vec3D operator-(const Vec3D a, const Vec3D b) +CUDA_CALLABLE inline Vec3D operator-(const Vec3D a, const Vec3D b) { return {a.x - b.x, a.y - b.y, a.z - b.z}; } -inline Vec3D operator*(const float scalar, const Vec3D a) +CUDA_CALLABLE inline Vec3D operator*(const float scalar, const Vec3D a) { return {a.x * scalar, a.y * scalar, a.z * scalar}; } -inline Vec3D operator*(const Vec3D a, const float scalar) +CUDA_CALLABLE inline Vec3D operator*(const Vec3D a, const float scalar) { return {a.x * scalar, a.y * scalar, a.z * scalar}; } -inline Vec3D operator/(const Vec3D a, const float scalar) +CUDA_CALLABLE inline Vec3D operator/(const Vec3D a, const float scalar) { return {a.x / scalar, a.y / scalar, a.z / scalar}; } -inline float dot(const Vec3D a, const Vec3D b) +CUDA_CALLABLE inline float dot(const Vec3D a, const Vec3D b) { return a.x * b.x + a.y * b.y + a.z * b.z; } -inline Vec3D cross(const Vec3D a, const Vec3D b) +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, @@ -53,23 +55,23 @@ private: public: Rotation() unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; - __host__ __device__ Vec3D apply(const Vec3D vec) const; + CUDA_CALLABLE Vec3D apply(const Vec3D vec) const; - __host__ __device__ Rotation compose(const Rotation& rot) const; + CUDA_CALLABLE Rotation compose(const Rotation& rot) const; - __host__ __device__ Rotation inv() const; + CUDA_CALLABLE Rotation inv() const; }; struct Pose { Rotation rot; Vec3D tran; - __host__ __device__ inline Vec3D apply(const Vec3D vec) const + CUDA_CALLABLE inline Vec3D apply(const Vec3D vec) const { return tran + rot.apply(vec); } - __host__ __device__ inline Pose compose(const Pose &pose) const + CUDA_CALLABLE inline Pose compose(const Pose &pose) const { /* * If $A_i$ is the matrix corresponding to pose object `p_i`, then @@ -79,7 +81,7 @@ struct Pose { return {rot.compose(pose.rot), rot.apply(pose.tran) + tran}; } - __host__ __device__ inline Pose inv() const + CUDA_CALLABLE inline Pose inv() const { auto rotinv = rot.inv(); return {rotinv, -rotinv.apply(tran)}; diff --git a/genmetaballs/src/cuda/core/utils.cuh b/genmetaballs/src/cuda/core/utils.cuh index 1dcdf36..060bffd 100644 --- a/genmetaballs/src/cuda/core/utils.cuh +++ b/genmetaballs/src/cuda/core/utils.cuh @@ -55,4 +55,4 @@ public: __host__ __device__ constexpr auto size() const noexcept { return data_view_.size(); } -}; +}; // class Array2D diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index 6403ef2..40dc48e 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from genmetaballs import _genmetaballs_bindings as _gmbb +from genmetaballs import _genmetaballs_bindings.geometry as geometry @pytest.fixture @@ -10,11 +10,11 @@ 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)" @@ -22,7 +22,7 @@ def test_vec3d_repr_returns_valid_string() -> None: 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) + a, b = geometry.Vec3D(*_a), geometry.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])]) @@ -30,7 +30,7 @@ def test_vec3d_add(rng: np.random.Generator) -> None: 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) + a, b = geometry.Vec3D(*_a), geometry.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])]) From 827b5b1ae1ddbf06a13eab5840239e9426cd2fab Mon Sep 17 00:00:00 2001 From: mugamma Date: Wed, 19 Nov 2025 01:52:09 +0000 Subject: [PATCH 03/11] build working --- genmetaballs/src/cuda/core/geometry.cu | 14 +++++++++++--- genmetaballs/src/cuda/core/geometry.cuh | 6 +++++- tests/python_tests/test_geometry.py | 2 +- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/genmetaballs/src/cuda/core/geometry.cu b/genmetaballs/src/cuda/core/geometry.cu index f78d317..35d4e50 100644 --- a/genmetaballs/src/cuda/core/geometry.cu +++ b/genmetaballs/src/cuda/core/geometry.cu @@ -1,7 +1,15 @@ +#include + #include "geometry.cuh" +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}}; +} + -__host__ __device__ Vec3D Rotation::apply(const Vec3D vec) const +CUDA_CALLABLE Vec3D Rotation::apply(const Vec3D vec) const { // v' = q * v * q^(-1) for unit quaternions // where q^(-1) = (-x, -y, -z, w) @@ -15,7 +23,7 @@ __host__ __device__ Vec3D Rotation::apply(const Vec3D vec) const return 2.0f * d * q + (w * w - dot(q, q)) * vec + 2.0f * w * c; } -__host__ __device__ Rotation Rotation::compose(const Rotation& rot) const +CUDA_CALLABLE Rotation Rotation::compose(const Rotation& rot) const { // Quaternion multiplication: q1 * q2 float4 q1 = unit_quat_; @@ -28,7 +36,7 @@ __host__ __device__ Rotation Rotation::compose(const Rotation& rot) const q1.w * q2.w - q1.x * q2.x - q1.y * q2.y - q1.z * q2.z}}; } -__host__ __device__ Rotation Rotation::inv() const +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 b890956..d3dd53f 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -52,8 +52,12 @@ class Rotation { private: float4 unit_quat_; + Rotation(float4 unit_quat): unit_quat_{unit_quat} {}; + public: - Rotation() unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + Rotation(): unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + + static Rotation from_quat(float x, float y, float z, float w); CUDA_CALLABLE Vec3D apply(const Vec3D vec) const; diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index 40dc48e..33a1acd 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from genmetaballs import _genmetaballs_bindings.geometry as geometry +from genmetaballs._genmetaballs_bindings import geometry as geometry @pytest.fixture From e5876f6c1a882c2edd31c1825ff0695962eaa703 Mon Sep 17 00:00:00 2001 From: mugamma Date: Wed, 19 Nov 2025 02:39:50 +0000 Subject: [PATCH 04/11] add basic rotation tests --- genmetaballs/src/cuda/bindings.cu | 2 + tests/python_tests/test_geometry.py | 88 ++++++++++++++++++++++++++--- 2 files changed, 81 insertions(+), 9 deletions(-) diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index b74f891..73f636e 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -47,6 +47,8 @@ NB_MODULE(_genmetaballs_bindings, m) { 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"); diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index 33a1acd..7fbd4bb 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -1,5 +1,8 @@ +import operator as op + import numpy as np import pytest +from scipy.spatial.transform import Rotation as Rot from genmetaballs._genmetaballs_bindings import geometry as geometry @@ -20,17 +23,84 @@ def test_vec3d_repr_returns_valid_string() -> None: assert repr_str == "Vec3D(1.0, 2.0, 3.0)" -def test_vec3d_add(rng: np.random.Generator) -> None: +@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=3), rng.uniform(size=3) a, b = geometry.Vec3D(*_a), geometry.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])]) - + c = vec3d_op(a, b) + _c = np_op(_a, _b) + assert np.allclose(_c, np.array([c.x, c.y, c.z])) -def test_vec3d_sub(rng: np.random.Generator) -> None: +def test_vec3d_dot(rng: np.random.Generator) -> None: _a, _b = rng.uniform(size=3), rng.uniform(size=3) a, b = geometry.Vec3D(*_a), geometry.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])]) + assert np.allclose(np.dot(_a, _b), 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, + ) From b6de85f0aeeaf09586ebbdf690c512de3c4d777c Mon Sep 17 00:00:00 2001 From: mugamma Date: Thu, 20 Nov 2025 15:56:32 +0000 Subject: [PATCH 05/11] add proper constructors for pose --- genmetaballs/src/cuda/bindings.cu | 7 +++- genmetaballs/src/cuda/core/geometry.cuh | 49 ++++++++++++++++++------- 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index 73f636e..c999364 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -55,8 +55,11 @@ NB_MODULE(_genmetaballs_bindings, m) { nb::class_(geometry, "Pose") .def(nb::init<>()) - .def_rw("rot", &Pose::rot) - .def_rw("tran", &Pose::tran) + .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"); diff --git a/genmetaballs/src/cuda/core/geometry.cuh b/genmetaballs/src/cuda/core/geometry.cuh index d3dd53f..464d602 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -52,12 +52,12 @@ class Rotation { private: float4 unit_quat_; - Rotation(float4 unit_quat): unit_quat_{unit_quat} {}; + CUDA_CALLABLE Rotation(float4 unit_quat): unit_quat_{unit_quat} {}; public: - Rotation(): unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + CUDA_CALLABLE Rotation(): unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; - static Rotation from_quat(float x, float y, float z, float w); + static CUDA_CALLABLE Rotation from_quat(float x, float y, float z, float w); CUDA_CALLABLE Vec3D apply(const Vec3D vec) const; @@ -66,29 +66,52 @@ public: 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 inline Vec3D apply(const Vec3D vec) const + CUDA_CALLABLE Vec3D apply(const Vec3D vec) const { - return tran + rot.apply(vec); + return tran_ + rot_.apply(vec); } - CUDA_CALLABLE inline Pose compose(const Pose &pose) const + 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}; + return {rot_.compose(pose.rot_), rot_.apply(pose.tran_) + tran_}; } - CUDA_CALLABLE inline Pose inv() const + CUDA_CALLABLE Pose inv() const { - auto rotinv = rot.inv(); - return {rotinv, -rotinv.apply(tran)}; + auto rotinv = rot_.inv(); + return {rotinv, -rotinv.apply(tran_)}; } }; From 17ff36c088eb15f4db8b91029ad7e608d2661142 Mon Sep 17 00:00:00 2001 From: mugamma Date: Thu, 20 Nov 2025 16:00:49 +0000 Subject: [PATCH 06/11] add pose correctness tests --- tests/python_tests/test_geometry.py | 92 ++++++++++++++++++++++++++--- 1 file changed, 84 insertions(+), 8 deletions(-) diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index 7fbd4bb..a1c9161 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -3,6 +3,7 @@ import numpy as np import pytest from scipy.spatial.transform import Rotation as Rot +from scipy.spatial.transform import RigidTransform as Rigid from genmetaballs._genmetaballs_bindings import geometry as geometry @@ -32,16 +33,18 @@ def test_vec3d_repr_returns_valid_string() -> None: ] ) def test_vec3d_ops(rng: np.random.Generator, np_op, vec3d_op) -> None: - _a, _b = rng.uniform(size=3), rng.uniform(size=3) - a, b = geometry.Vec3D(*_a), geometry.Vec3D(*_b) - c = vec3d_op(a, b) - _c = np_op(_a, _b) - assert np.allclose(_c, np.array([c.x, c.y, c.z])) + _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=3), rng.uniform(size=3) - a, b = geometry.Vec3D(*_a), geometry.Vec3D(*_b) - assert np.allclose(np.dot(_a, _b), geometry.dot(a, b)) + _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)) @@ -104,3 +107,76 @@ def test_rotation_inv(rng: np.random.Generator) -> None: 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) + + 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, + ) From 8fb9bcd1b8cd30e6017f4869424fc1d095f92a55 Mon Sep 17 00:00:00 2001 From: mugamma Date: Thu, 20 Nov 2025 18:12:55 +0000 Subject: [PATCH 07/11] format and lint --- .clang-tidy | 1 + genmetaballs/src/cuda/bindings.cu | 24 +++++----- genmetaballs/src/cuda/core/geometry.cu | 28 +++++------- genmetaballs/src/cuda/core/geometry.cuh | 58 +++++++++---------------- tests/python_tests/test_geometry.py | 53 ++++++++++++---------- 5 files changed, 75 insertions(+), 89 deletions(-) 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 c999364..df3e2e1 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -23,7 +23,7 @@ NB_MODULE(_genmetaballs_bindings, m) { /* * Geometry module bindings */ - + nb::module_ geometry = m.def_submodule("geometry", "Geometry helpers for GenMetaballs"); nb::class_(geometry, "Vec3D") @@ -32,15 +32,14 @@ NB_MODULE(_genmetaballs_bindings, m) { .def_ro("x", &Vec3D::x) .def_ro("y", &Vec3D::y) .def_ro("z", &Vec3D::z) - .def("__add__", static_cast(&operator+)) - .def("__sub__", static_cast(&operator-)) - .def("__neg__", static_cast(&operator-)) - .def("__mul__", static_cast(&operator*)) - .def("__rmul__", static_cast(&operator*)) - .def("__truediv__", static_cast(&operator/)) - .def("__repr__", [](const Vec3D& v) { - return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); - }); + .def("__add__", static_cast(&operator+)) + .def("__sub__", static_cast(&operator-)) + .def("__neg__", static_cast(&operator-)) + .def("__mul__", static_cast(&operator*)) + .def("__rmul__", static_cast(&operator*)) + .def("__truediv__", static_cast(&operator/)) + .def("__repr__", + [](const Vec3D& v) { return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); }); 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")); @@ -56,8 +55,8 @@ NB_MODULE(_genmetaballs_bindings, m) { 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")) + "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")) @@ -90,4 +89,3 @@ NB_MODULE(_genmetaballs_bindings, m) { utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))"); } // NB_MODULE(_genmetaballs_bindings) - diff --git a/genmetaballs/src/cuda/core/geometry.cu b/genmetaballs/src/cuda/core/geometry.cu index 35d4e50..c645bdf 100644 --- a/genmetaballs/src/cuda/core/geometry.cu +++ b/genmetaballs/src/cuda/core/geometry.cu @@ -2,15 +2,13 @@ #include "geometry.cuh" -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}}; +// 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 -{ +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}; @@ -23,21 +21,19 @@ CUDA_CALLABLE Vec3D Rotation::apply(const Vec3D vec) const return 2.0f * d * q + (w * w - dot(q, q)) * vec + 2.0f * w * c; } -CUDA_CALLABLE Rotation Rotation::compose(const Rotation& rot) const -{ +// 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}}; + 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}}; } -CUDA_CALLABLE Rotation Rotation::inv() const -{ +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 464d602..f64dbd2 100644 --- a/genmetaballs/src/cuda/core/geometry.cuh +++ b/genmetaballs/src/cuda/core/geometry.cuh @@ -6,56 +6,46 @@ using Vec3D = float3; -CUDA_CALLABLE inline Vec3D operator-(const Vec3D a) -{ +CUDA_CALLABLE inline Vec3D operator-(const Vec3D a) { return {-a.x, -a.y, -a.z}; } -CUDA_CALLABLE inline Vec3D operator+(const Vec3D a, const Vec3D b) -{ +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) -{ +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) -{ +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) -{ +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) -{ +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) -{ +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}; +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: float4 unit_quat_; - CUDA_CALLABLE Rotation(float4 unit_quat): unit_quat_{unit_quat} {}; + CUDA_CALLABLE Rotation(float4 unit_quat) : unit_quat_{unit_quat} {}; public: - CUDA_CALLABLE Rotation(): unit_quat_{0.0f, 0.0f, 0.0f, 1.0f} {}; + 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); @@ -71,35 +61,30 @@ private: Rotation rot_; Vec3D tran_; - CUDA_CALLABLE Pose(const Rotation rot, const Vec3D tran): rot_{rot}, tran_{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} {} + // these member functions are defined in class body to allow for possible inlining - static CUDA_CALLABLE Pose from_components(const Rotation rot, const Vec3D tran) - { + 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 - { + CUDA_CALLABLE Rotation get_rot() const { return rot_; } - CUDA_CALLABLE Vec3D get_tran() const - { + CUDA_CALLABLE Vec3D get_tran() const { return tran_; } - CUDA_CALLABLE Vec3D apply(const Vec3D vec) const - { + CUDA_CALLABLE Vec3D apply(const Vec3D vec) const { return tran_ + rot_.apply(vec); } - CUDA_CALLABLE Pose compose(const Pose &pose) const - { + 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 @@ -108,8 +93,7 @@ public: return {rot_.compose(pose.rot_), rot_.apply(pose.tran_) + tran_}; } - CUDA_CALLABLE Pose inv() const - { + CUDA_CALLABLE Pose inv() const { auto rotinv = rot_.inv(); return {rotinv, -rotinv.apply(tran_)}; } diff --git a/tests/python_tests/test_geometry.py b/tests/python_tests/test_geometry.py index a1c9161..f2e32d0 100644 --- a/tests/python_tests/test_geometry.py +++ b/tests/python_tests/test_geometry.py @@ -2,8 +2,8 @@ import numpy as np import pytest -from scipy.spatial.transform import Rotation as Rot from scipy.spatial.transform import RigidTransform as Rigid +from scipy.spatial.transform import Rotation as Rot from genmetaballs._genmetaballs_bindings import geometry as geometry @@ -30,8 +30,8 @@ def test_vec3d_repr_returns_valid_string() -> None: (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): @@ -40,12 +40,14 @@ def test_vec3d_ops(rng: np.random.Generator, np_op, vec3d_op) -> None: _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) @@ -61,6 +63,7 @@ def test_rotation_apply(rng: np.random.Generator) -> None: 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) @@ -82,11 +85,12 @@ def test_rotation_compose(rng: np.random.Generator) -> None: 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, - ) + 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)) @@ -102,11 +106,12 @@ def test_rotation_inv(rng: np.random.Generator) -> None: 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, - ) + 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)) @@ -126,6 +131,7 @@ def test_pose_apply(rng: np.random.Generator) -> None: 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)) @@ -151,11 +157,12 @@ def test_pose_compose(rng: np.random.Generator) -> None: 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, - ) + 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)) @@ -175,8 +182,8 @@ def test_pose_inv(rng: np.random.Generator) -> None: 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, - ) + np.array([vec.x, vec.y, vec.z]), + np.array([vec_.x, vec_.y, vec_.z]), + rtol=1e-5, + atol=1e-6, + ) From 08f4653f3db2d31723c6095a728a47af052384b4 Mon Sep 17 00:00:00 2001 From: Arijit Dasgupta Date: Sun, 23 Nov 2025 01:04:44 +0000 Subject: [PATCH 08/11] implemented AllGetter, with edits to fmb.h and forward.cu --- genmetaballs/src/cuda/core/fmb.cuh | 8 ++++++++ genmetaballs/src/cuda/core/forward.cu | 2 +- genmetaballs/src/cuda/core/getter.cuh | 21 +++++++++++++++++++++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/genmetaballs/src/cuda/core/fmb.cuh b/genmetaballs/src/cuda/core/fmb.cuh index fe9a765..ad194b4 100644 --- a/genmetaballs/src/cuda/core/fmb.cuh +++ b/genmetaballs/src/cuda/core/fmb.cuh @@ -17,4 +17,12 @@ public: FMBs(uint32_t size) : fmbs_(size), log_weights_(size) { // TODO: set all log_weights_ to 0 } + CUDA_CALLABLE const containter_template& get_all_fmbs() const { + return fmbs_; + } + + // Method to get a single FMB by index + CUDA_CALLABLE const FMB& get_fmb(uint32_t idx) const { + return fmbs_[idx]; + } }; diff --git a/genmetaballs/src/cuda/core/forward.cu b/genmetaballs/src/cuda/core/forward.cu index a71caf6..ec3c185 100644 --- a/genmetaballs/src/cuda/core/forward.cu +++ b/genmetaballs/src/cuda/core/forward.cu @@ -47,7 +47,7 @@ __global__ render_kernel(const Getter fmb_getter, const Blender blender, template void render_fmbs(const FMBs& fmbs, const Intrinsics& intr, const Pose& extr) { // initialize the fmb_getter - typename Getter::Getter fmb_getter(fmbs, intr, extr); + typename Getter::Getter fmb_getter(fmbs, extr); auto kernel = render_kernel; kernel<<>>(fmb_getter, fmbs, intr, extr); } diff --git a/genmetaballs/src/cuda/core/getter.cuh b/genmetaballs/src/cuda/core/getter.cuh index e69de29..dcc86a0 100644 --- a/genmetaballs/src/cuda/core/getter.cuh +++ b/genmetaballs/src/cuda/core/getter.cuh @@ -0,0 +1,21 @@ +#pragma once + +#include +#include + +#include "camera.cuh" +#include "fmb.cuh" +#include "geometry.cuh" +#include "utils.cuh" + +// This is the dummy version of getter, where all FMBs are relevant to any ray +template +struct AllGetter { + FMBs fmbs; + Pose extr; // Current assumption: rays are in camera frame + + // It does not bother using the ray, because it simply returns all FMBs + CUDA_CALLABLE const containter_template& get_metaballs(const Ray& ray) const { + return fmbs.get_all_fmbs(); + } +}; From eb9a24b4f2bc7805be6ad466b457d61e26187361 Mon Sep 17 00:00:00 2001 From: Arijit Dasgupta Date: Sun, 23 Nov 2025 03:57:33 +0000 Subject: [PATCH 09/11] AllGetter with thrust + std::vector testing --- genmetaballs/src/cuda/core/fmb.cuh | 4 +- genmetaballs/src/cuda/core/getter.cuh | 7 +- tests/cpp_tests/test_getter.cu | 147 ++++++++++++++++++++++++++ 3 files changed, 153 insertions(+), 5 deletions(-) create mode 100644 tests/cpp_tests/test_getter.cu diff --git a/genmetaballs/src/cuda/core/fmb.cuh b/genmetaballs/src/cuda/core/fmb.cuh index ad194b4..bf2f162 100644 --- a/genmetaballs/src/cuda/core/fmb.cuh +++ b/genmetaballs/src/cuda/core/fmb.cuh @@ -7,7 +7,7 @@ struct FMB { float3 extent; }; -template +template