all tests passing

Author: mugamma

genmetaballs/src/cuda/bindings.cu

@@ -62,6 +62,8 @@ NB_MODULE(_genmetaballs_bindings, m) {
                          auto extent = self.get_extent();
                          return std::tuple{extent.x, extent.y, extent.z};
                      })
+        .def("cov_inv_apply", &FMB::cov_inv_apply,
+             "apply the inverse covariance matrix to the given vector", nb::arg("vec"))
         .def("quadratic_form", &FMB::quadratic_form,
              "Evaluate the associated quadratic form at the given vector", nb::arg("vec"));
 
@@ -146,11 +148,11 @@ NB_MODULE(_genmetaballs_bindings, m) {
     nb::module_ intersector = m.def_submodule("intersector");
     intersector.def(
         "linear_intersect",
-        [](const FMB& fmb, const Ray& ray) {
-            auto [t, d] = LinearIntersector::intersect(fmb, ray);
+        [](const FMB& fmb, const Ray& ray, const Pose& cam_pose) {
+            auto [t, d] = LinearIntersector::intersect(fmb, ray, cam_pose);
             return std::make_tuple(t, d);
         },
-        "Linear intersection of ray and FMB.", nb::arg("fmb"), nb::arg("ray"));
+        "Linear intersection of ray and FMB.", nb::arg("fmb"), nb::arg("ray"), nb::arg("cam_pose"));
 
     /*
      * Utils module bindings

genmetaballs/src/cuda/core/fmb.cu

@@ -2,12 +2,18 @@
 #include "geometry.cuh"
 #include "utils.cuh"
 
+CUDA_CALLABLE __forceinline__ Vec3D vecdiv(const Vec3D u, const Vec3D v) {
+    return {u.x / v.x, u.y / v.y, u.z / v.z};
+}
+
+CUDA_CALLABLE Vec3D FMB::cov_inv_apply(const Vec3D vec) const {
+    const auto rot = pose_.get_rot();
+    return rot.inv().apply(vecdiv(rot.apply(vec), extent_));
+}
+
 CUDA_CALLABLE float FMB::quadratic_form(const Vec3D vec) const {
-    const auto shftd_vec = vec - pose_.get_tran();
-    const auto rot_shftd_vec = pose_.get_rot().apply(shftd_vec);
-    const auto scaled_rot_shftd_vec = Vec3D(
-        rot_shftd_vec.x / extent_.x, rot_shftd_vec.y / extent_.y, rot_shftd_vec.z / extent_.z);
-    return dot(rot_shftd_vec, scaled_rot_shftd_vec);
+    const auto shifted_vec = vec - get_mean();
+    return dot(shifted_vec, cov_inv_apply(shifted_vec));
 }
 
 template <>

genmetaballs/src/cuda/core/fmb.cuh

@@ -34,6 +34,11 @@ public:
     CUDA_CALLABLE float3 get_extent() const {
         return extent_;
     }
+    CUDA_CALLABLE float3 get_mean() const {
+        return pose_.get_tran();
+    }
+
+    CUDA_CALLABLE Vec3D cov_inv_apply(const Vec3D) const;
 
     CUDA_CALLABLE float quadratic_form(const Vec3D) const;
 };

genmetaballs/src/cuda/core/forward.cu

@@ -33,7 +33,7 @@ __global__ render_kernel(const Getter fmb_getter, const Blender blender,
     for (const auto& [pixel_coords, ray] : pixel_coords_and_rays) {
         float w0 = 0.0f, tf = 0.0f, sumexpd = 0.0f;
         for (const auto& fmb : fmb_getter->get_metaballs(ray)) {
-            const auto& [t, d] = Intersector::intersect(fmb, ray);
+            const auto& [t, d] = Intersector::intersect(fmb, ray, extr);
             w = blender->blend(t, d, fmb, ray);
             sumexpd += exp(d); // numerically unstable. use logsumexp
             tf += t;

genmetaballs/src/cuda/core/intersector.cuh

@@ -8,13 +8,15 @@
 // implement equation (6) in the paper
 class LinearIntersector {
 public:
-    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Ray& ray) {
-        auto vecdiv = [](const Vec3D& u, const Vec3D& v) {
-            return Vec3D{u.x / v.x, u.y / v.y, u.z / v.z};
-        };
-        auto rot = fmb.get_pose().get_rot();
-        auto tmp = rot.inv().apply(vecdiv(rot.apply(ray.direction), fmb.get_extent()));
-        auto t = dot(fmb.get_pose().get_tran() - ray.start, tmp) / dot(ray.direction, tmp);
-        return {t, fmb.quadratic_form(ray.start + t * ray.direction)};
+    /*
+     * Ray should be in camera frame
+     */
+    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Ray& ray,
+                                                                  const Pose& cam_pose) {
+        const auto v = cam_pose.get_rot().apply(ray.direction);
+        const auto cov_inv_v = fmb.cov_inv_apply(v);
+        const auto cam_tran = cam_pose.get_tran();
+        const auto t = dot(fmb.get_mean() - cam_tran, cov_inv_v) / dot(v, cov_inv_v);
+        return {t, fmb.quadratic_form(cam_tran + t * v)};
     }
 };

tests/python_tests/test_fmb.py

@@ -14,6 +14,19 @@ def rng() -> np.random.Generator:
     return np.random.default_rng(0)
 
 
+def test_fmb_cov_inv_apply(rng):
+    for _ in range(100):
+        quat = rng.uniform(size=4).astype(np.float32)
+        tran, extent, vec = rng.uniform(size=(3, 3)).astype(np.float32)
+        pose = Pose.from_components(Rotation.from_quat(*quat), Vec3D(*tran))
+        scipy_rot_mat = Rot.from_quat(quat).as_matrix()
+        cov = scipy_rot_mat.T @ np.diag(extent) @ scipy_rot_mat
+        theirs = np.linalg.solve(cov, vec)
+        ours = FMB(pose, *extent).cov_inv_apply(Vec3D(*vec))
+        ourvec = np.array([ours.x, ours.y, ours.z], dtype=np.float32)
+        assert np.allclose(theirs, ourvec, atol=1e-6)
+
+
 def test_fmb_quadratic_form(rng):
     for _ in range(100):
         quat = rng.uniform(size=4)

tests/python_tests/test_intersector.py

@@ -0,0 +1,36 @@
+import numpy as np
+import pytest
+from scipy.spatial.distance import mahalanobis
+from scipy.spatial.transform import Rotation as Rot
+
+from genmetaballs.core import fmb, geometry, intersector
+
+FMB = fmb.FMB
+Pose, Vec3D, Rotation, Ray = geometry.Pose, geometry.Vec3D, geometry.Rotation, geometry.Ray
+
+
+@pytest.fixture
+def rng() -> np.random.Generator:
+    return np.random.default_rng(0)
+
+
+def test_linear_intersect(rng):
+    for _ in range(100):
+        # sample
+        cam_quat, fmb_quat = rng.uniform(size=(2, 4)).astype(np.float32)
+        fmb_extent, fmb_mu = rng.uniform(size=(2, 3)).astype(np.float32)
+        cam_tran, ray_dir = rng.uniform(size=(2, 3)).astype(np.float32)
+        ray_start = np.zeros(3, dtype=np.float32)  # in camera frame
+        # ground truth computation
+        v = Rot.from_quat(cam_quat).apply(ray_dir)
+        fmb_rotmat = Rot.from_quat(fmb_quat).as_matrix()
+        cov_inv = fmb_rotmat.T @ np.diag(1 / fmb_extent) @ fmb_rotmat
+        t = ((fmb_mu - cam_tran) @ cov_inv @ v) / (v @ cov_inv @ v)
+        d = mahalanobis(fmb_mu, cam_tran + t * v, cov_inv) ** 2
+        # cuda computation
+        fmb_pose = Pose.from_components(Rotation.from_quat(*fmb_quat), Vec3D(*fmb_mu))
+        cam_pose = Pose.from_components(Rotation.from_quat(*cam_quat), Vec3D(*cam_tran))
+        fmb = FMB(fmb_pose, *fmb_extent)
+        ray = Ray(Vec3D(*ray_start), Vec3D(*ray_dir))
+        t_, d_ = intersector.linear_intersect(fmb, ray, cam_pose)
+        assert np.isclose(t, t_) and np.isclose(d, d_)