MET-30: Implement Geometry Primitives

.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,

genmetaballs/src/cuda/bindings.cu

@@ -11,19 +11,58 @@ namespace nb = nanobind;
 
 NB_MODULE(_genmetaballs_bindings, m) {
 
-    // exposing Vec3D
-    nb::class_<Vec3D>(m, "Vec3D")
+    /*
+     * Geometry module bindings
+     */
+
+    nb::module_ geometry = m.def_submodule("geometry", "Geometry helpers for GenMetaballs");
+
+    nb::class_<Vec3D>(geometry, "Vec3D")
         .def(nb::init<>())
         .def(nb::init<float, float, float>())
-        .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_<Rotation>(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_<Pose>(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_<Ray>(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_<ZeroParameterConfidence>(confidence, "ZeroParameterConfidence")
         .def(nb::init<>())
@@ -33,7 +72,10 @@ NB_MODULE(_genmetaballs_bindings, m) {
         .def(nb::init<float, float>())
         .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))");
 

genmetaballs/src/cuda/core/geometry.cu

@@ -1,9 +1,39 @@
+#include <cmath>
+
 #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}};
 }

genmetaballs/src/cuda/core/geometry.cuh

@@ -2,30 +2,101 @@
 
 #include <cuda_runtime.h>
 
+#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 {

genmetaballs/src/cuda/core/utils.cuh

@@ -52,4 +52,4 @@ public:
     __host__ __device__ constexpr auto size() const noexcept {
         return data_view_.size();
     }
-};
+}; // class Array2D

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",
 ]