Refactor to use Vec3D instead of Ray (#29)

Part of MET-36

## Summary of Changes

Replace all instances of the `Ray` class with `Vec3D` (primarily in
intersector and getter) to streamline the code and remove unused
components.

This is part of my sequence of PRs for the forward kernels. The hope is
to keep each of them small so that even if we are not going to review it
now, it will be easier to catch up after the final deadline and
understand what changes were maded.


## Test Plan

```bash
pixi run test
```

Author: horizon-blue

genmetaballs/src/cuda/bindings.cu

@@ -110,11 +110,6 @@ NB_MODULE(_genmetaballs_bindings, m) {
         .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<Vec3D, Vec3D>())
-        .def_ro("start", &Ray::start)
-        .def_ro("direction", &Ray::direction);
-
     /*
      * Camera module bindings
      */
@@ -148,7 +143,7 @@ NB_MODULE(_genmetaballs_bindings, m) {
     nb::module_ intersector = m.def_submodule("intersector");
     intersector.def(
         "linear_intersect",
-        [](const FMB& fmb, const Ray& ray, const Pose& cam_pose) {
+        [](const FMB& fmb, const Vec3D& ray, const Pose& cam_pose) {
             auto [t, d] = LinearIntersector::intersect(fmb, ray, cam_pose);
             return std::make_tuple(t, d);
         },

genmetaballs/src/cuda/core/geometry.cuh

@@ -98,10 +98,3 @@ public:
         return {rotinv, -rotinv.apply(tran_)};
     }
 };
-
-struct Ray {
-    Vec3D start;
-    Vec3D direction;
-
-    CUDA_CALLABLE Ray(const Vec3D _start, const Vec3D _dir) : start{_start}, direction{_dir} {}
-};

genmetaballs/src/cuda/core/getter.cuh

@@ -11,13 +11,14 @@
 // This is the dummy version of getter, where all FMBs are relevant to any ray
 template <MemoryLocation location>
 struct AllGetter {
-    FMBScene<location>& scene;
-    Pose& extr; // Current assumption: rays are in camera frame
+    const FMBScene<location>& scene;
+    const Pose& extr; // Current assumption: rays are in camera frame
 
-    CUDA_CALLABLE AllGetter(FMBScene<location>& scene, Pose& extr) : scene(scene), extr(extr) {}
+    CUDA_CALLABLE AllGetter(const FMBScene<location>& scene, const Pose& extr)
+        : scene(scene), extr(extr) {}
 
     // It does not bother using the ray, because it simply returns all FMBs
-    CUDA_CALLABLE FMBScene<location>& get_metaballs(const Ray& ray) const {
+    CUDA_CALLABLE const FMBScene<location>& get_metaballs(const Vec3D& ray) const {
         return scene;
     }
 };

genmetaballs/src/cuda/core/intersector.cuh

@@ -9,11 +9,11 @@
 class LinearIntersector {
 public:
     /*
-     * Ray should be in camera frame
+     * ray should be in camera frame
      */
-    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Ray& ray,
+    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Vec3D ray,
                                                                   const Pose& cam_pose) {
-        const auto v = cam_pose.get_rot().apply(ray.direction);
+        const auto v = cam_pose.get_rot().apply(ray);
         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);

tests/cpp_tests/test_getter.cu

@@ -19,17 +19,11 @@ TEST(AllGetterTest, AllGetterHostTest) {
     AllGetter<MemoryLocation::HOST> getter(scene, extr);
 
     // Create test rays
-    std::vector<Ray> rays = {
-        Ray{Vec3D{0.0f, 0.0f, 0.0f}, Vec3D{1.0f, 0.0f, 0.0f}},
-        Ray{Vec3D{1.0f, 1.0f, 1.0f}, Vec3D{0.0f, 1.0f, 0.0f}},
-        Ray{Vec3D{-1.0f, -1.0f, -1.0f}, Vec3D{0.0f, 0.0f, 1.0f}},
-        Ray{Vec3D{2.5f, -3.1f, 0.2f}, Vec3D{-0.5f, 0.6f, 0.0f}},
-        Ray{Vec3D{4.4f, 0.0f, -0.9f}, Vec3D{0.3f, -0.2f, 1.0f}},
-        Ray{Vec3D{5.0f, 2.2f, 1.1f}, Vec3D{-1.0f, 2.0f, 0.2f}},
-        Ray{Vec3D{0.0f, 7.0f, 6.0f}, Vec3D{0.0f, -1.0f, -1.0f}},
-        Ray{Vec3D{-2.0f, 0.0f, 0.0f}, Vec3D{0.2f, 1.1f, 0.7f}},
-        Ray{Vec3D{9.1f, -0.3f, 2.7f}, Vec3D{-0.3f, 0.1f, 0.0f}},
-        Ray{Vec3D{1.2f, 8.8f, -4.5f}, Vec3D{1.0f, 0.0f, 1.0f}},
+    std::vector<Vec3D> rays = {
+        Vec3D{0.0f, 0.0f, 0.0f},  Vec3D{1.0f, 1.0f, 1.0f},  Vec3D{-1.0f, -1.0f, -1.0f},
+        Vec3D{2.5f, -3.1f, 0.2f}, Vec3D{4.4f, 0.0f, -0.9f}, Vec3D{5.0f, 2.2f, 1.1f},
+        Vec3D{0.0f, 7.0f, 6.0f},  Vec3D{-2.0f, 0.0f, 0.0f}, Vec3D{9.1f, -0.3f, 2.7f},
+        Vec3D{1.2f, 8.8f, -4.5f},
     };
 
     // Get reference to all FMBs from the original FMBs object
@@ -49,7 +43,7 @@ TEST(AllGetterTest, AllGetterHostTest) {
 }
 
 __global__ void test_get_metaballs_kernel_device(const AllGetter<MemoryLocation::DEVICE> fmb_getter,
-                                                 const Ray* rays, int num_rays, int* out_sizes) {
+                                                 const Vec3D* rays, int num_rays, int* out_sizes) {
     int idx = threadIdx.x + blockIdx.x * blockDim.x;
     const auto& fmbs_returned = fmb_getter.get_metaballs(rays[idx]);
     out_sizes[idx] = static_cast<int>(fmbs_returned.size());
@@ -63,17 +57,11 @@ TEST(AllGetterTest, AllGetterDeviceTest) {
     AllGetter<MemoryLocation::DEVICE> getter(device_scene, extr);
 
     // Create test rays
-    std::vector<Ray> rays = {
-        Ray{Vec3D{0.0f, 0.0f, 0.0f}, Vec3D{1.0f, 0.0f, 0.0f}},
-        Ray{Vec3D{1.0f, 1.0f, 1.0f}, Vec3D{0.0f, 1.0f, 0.0f}},
-        Ray{Vec3D{-1.0f, -1.0f, -1.0f}, Vec3D{0.0f, 0.0f, 1.0f}},
-        Ray{Vec3D{2.5f, -3.1f, 0.2f}, Vec3D{-0.5f, 0.6f, 0.0f}},
-        Ray{Vec3D{4.4f, 0.0f, -0.9f}, Vec3D{0.3f, -0.2f, 1.0f}},
-        Ray{Vec3D{5.0f, 2.2f, 1.1f}, Vec3D{-1.0f, 2.0f, 0.2f}},
-        Ray{Vec3D{0.0f, 7.0f, 6.0f}, Vec3D{0.0f, -1.0f, -1.0f}},
-        Ray{Vec3D{-2.0f, 0.0f, 0.0f}, Vec3D{0.2f, 1.1f, 0.7f}},
-        Ray{Vec3D{9.1f, -0.3f, 2.7f}, Vec3D{-0.3f, 0.1f, 0.0f}},
-        Ray{Vec3D{1.2f, 8.8f, -4.5f}, Vec3D{1.0f, 0.0f, 1.0f}},
+    std::vector<Vec3D> rays = {
+        Vec3D{0.0f, 0.0f, 0.0f},  Vec3D{1.0f, 1.0f, 1.0f},  Vec3D{-1.0f, -1.0f, -1.0f},
+        Vec3D{2.5f, -3.1f, 0.2f}, Vec3D{4.4f, 0.0f, -0.9f}, Vec3D{5.0f, 2.2f, 1.1f},
+        Vec3D{0.0f, 7.0f, 6.0f},  Vec3D{-2.0f, 0.0f, 0.0f}, Vec3D{9.1f, -0.3f, 2.7f},
+        Vec3D{1.2f, 8.8f, -4.5f},
     };
 
     // Test on GPU for device containers

tests/python_tests/test_intersector.py

@@ -6,7 +6,7 @@
 from genmetaballs.core import fmb, geometry, intersector
 
 FMB = fmb.FMB
-Pose, Vec3D, Rotation, Ray = geometry.Pose, geometry.Vec3D, geometry.Rotation, geometry.Ray
+Pose, Vec3D, Rotation = geometry.Pose, geometry.Vec3D, geometry.Rotation
 
 
 @pytest.fixture
@@ -20,7 +20,6 @@ def test_linear_intersect(rng):
         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()
@@ -31,6 +30,6 @@ def test_linear_intersect(rng):
         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))
+        ray = Vec3D(*ray_dir)
         t_, d_ = intersector.linear_intersect(fmb, ray, cam_pose)
         assert np.isclose(t, t_) and np.isclose(d, d_)