[WIP] Custom Ray marcher stub

fvdb/_fvdb_cpp.pyi

@@ -447,6 +447,9 @@ class GridBatch:
     def clipped_grid(self, ijk_min: Vec3iBatch, ijk_max: Vec3iBatch) -> GridBatch: ...
     def coarsened_grid(self, coarsening_factor: Vec3iOrScalar) -> GridBatch: ...
     def contiguous(self) -> GridBatch: ...
+    def acoustic_ray_samples(
+        self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, sound_speeds: JaggedTensor, step_size: float
+    ) -> JaggedTensor: ...
     def integrate_tsdf(
         self,
         voxel_truncation_distance: float,

fvdb/grid.py

@@ -2432,3 +2432,14 @@ def world_to_voxel_matrix(self) -> torch.Tensor:
     def _gridbatch(self):
         # Access underlying GridBatchCpp - use sparingly during migration
         return self._impl
+
+    def acoustic_ray_samples(
+        self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, sound_speeds: JaggedTensor, step_size: float
+    ) -> JaggedTensor:
+        """
+        Generate samples along each ray with a given step size where samples are taken
+        by bending the ray according to Snell's law.
+        """
+        return JaggedTensor(
+            impl=self._impl.acoustic_ray_samples(ray_origins._impl, ray_directions._impl, sound_speeds._impl, step_size)
+        )

src/CMakeLists.txt

@@ -75,6 +75,7 @@ set(FVDB_CPP_FILES
 
 set(FVDB_CU_FILES
     fvdb/detail/GridBatchImpl.cu
+    fvdb/detail/ops/AcousticRaySamples.cu
     fvdb/detail/ops/ActiveGridGoords.cu
     fvdb/detail/ops/ActiveVoxelsInBoundsMask.cu
     fvdb/detail/ops/BuildCoarseGridFromFine.cu

src/fvdb/GridBatch.cpp

@@ -34,6 +34,7 @@
 #include <fvdb/detail/ops/VoxelsAlongRays.h>
 #include <fvdb/detail/utils/Utils.h>
 #include <fvdb/detail/utils/nanovdb/TorchNanoConversions.h>
+#include <fvdb/detail/ops/AcousticRaySamples.h>
 
 #include <torch/types.h>
 
@@ -864,6 +865,14 @@ GridBatch::voxels_along_rays(const JaggedTensor &ray_origins,
         *mImpl, ray_origins, ray_directions, max_vox, eps, return_ijk, cumulative);
 }
 
+std::vector<JaggedTensor>
+GridBatch::acoustic_ray_samples(const JaggedTensor &ray_origins,
+                                const JaggedTensor &ray_directions,
+                                const JaggedTensor &sound_speeds,
+                                float step_size) const {
+    return fvdb::detail::ops::acousticRaySamples(*mImpl, ray_origins, ray_directions, sound_speeds, step_size);
+}
+
 JaggedTensor
 GridBatch::segments_along_rays(const JaggedTensor &ray_origins,
                                const JaggedTensor &ray_directions,

src/fvdb/GridBatch.h

@@ -305,6 +305,13 @@ struct GridBatch : torch::CustomClassHolder {
                                         const std::optional<Vec3iBatch> &min_coord = std::nullopt,
                                         const std::optional<Vec3i> &grid_size = std::nullopt) const;
 
+
+    std::vector<JaggedTensor> acoustic_ray_samples(const JaggedTensor &ray_origins,
+                                                  const JaggedTensor &ray_directions,
+                                                  const JaggedTensor &sound_speeds,
+                                                  float step_size) const;
+
+
     /// @brief Convert grid coordinates to world coordinates
     /// @param ijk A JaggedTensor of grid coordinates with shape [B, -1, 3] (one point set per grid
     /// in the batch)

src/fvdb/detail/ops/AcousticRaySamples.cu

@@ -0,0 +1,77 @@
+// Copyright Contributors to the OpenVDB Project
+// SPDX-License-Identifier: Apache-2.0
+//
+#include <fvdb/detail/ops/VoxelsAlongRays.h>
+#include <fvdb/detail/utils/AccessorHelpers.cuh>
+#include <fvdb/detail/utils/ForEachCPU.h>
+#include <fvdb/detail/utils/Utils.h>
+#include <fvdb/detail/utils/cuda/ForEachCUDA.cuh>
+
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAMathCompat.h>
+
+namespace fvdb {
+namespace detail {
+namespace ops {
+namespace {
+
+template <typename ScalarType>
+__global__ void countAcousticRaySamplesKernel(
+    BatchGridAccessor batchAccessor,
+    const JaggedAccessor<ScalarType, 2> rayOrigins,
+    const JaggedAccessor<ScalarType, 2> rayDirections,
+    const JaggedAccessor<ScalarType, 1> soundSpeeds,
+    float stepSize,
+    fvdb::TorchRAcc64<int32_t, 1> outCounts) {
+
+        // Kernel parallelizes over rays.
+        // For each ray, we count the number of samples we will take.
+}
+
+template <typename ScalarType>
+__global__ void accousticRaySamplesKernel(
+    BatchGridAccessor batchAccessor,
+    const JaggedAccessor<ScalarType, 2> rayOrigins,
+    const JaggedAccessor<ScalarType, 2> rayDirections,
+    const JaggedAccessor<ScalarType, 1> soundSpeeds,
+    float stepSize,
+    JaggedAccessor<ScalarType, 2> outSamples) {
+
+        // Kernel parallelizes over rays.
+        // For each ray, we take a step, query the sound speeds, write a value, and bend the ray according to Snell's law.
+}
+
+} // anonymous namespace
+
+// Generate samples along each ray with a given step size where samples are taken
+// by bending the ray according to Snell's law.
+std::vector<JaggedTensor> acousticRaySamples(
+    const GridBatchImpl &batchHdl,
+    const JaggedTensor &rayOrigins,
+    const JaggedTensor &rayDirections,
+    const JaggedTensor &soundSpeeds,
+    float stepSize) {
+
+    const auto gridBatchAcc = gridBatchAccessor<DeviceTag>(batchHdl);
+    const auto rayOriginsAcc = jaggedAccessor<DeviceTag, scalar_t, 2>(rayOrigins);
+    const auto rayDirectionsAcc = jaggedAccessor<DeviceTag, scalar_t, 2>(rayDirections);
+    const auto soundSpeedsAcc = jaggedAccessor<DeviceTag, scalar_t, 1>(soundSpeeds);
+
+    const int64_t N = rayOrigins.jdata().size(0);
+    auto outSamplesData = torch::empty({N, 2}, rayOrigins.jdata().options());
+    auto outPositionsData = torch::empty({N, 3}, rayOrigins.jdata().options());
+    auto outTimesData = torch::empty({N}, rayOrigins.jdata().options());
+    fvdb::JaggedTensor outSamples = rayOrigins.jagged_like(outSamplesData);
+    fvdb::JaggedTensor outPositions = rayOrigins.jagged_like(outPositionsData);
+    fvdb::JaggedTensor outTimes = rayOrigins.jagged_like(outTimesData);
+
+    // 1. Count the number of samples we will take for each ray.
+
+    // 2. Generate the samples.
+
+    return {outSamples, outPositions, outTimes};
+}
+
+} // namespace ops
+} // namespace detail
+} // namespace fvdb

src/fvdb/detail/ops/AcousticRaySamples.h

@@ -0,0 +1,30 @@
+// Copyright Contributors to the OpenVDB Project
+// SPDX-License-Identifier: Apache-2.0
+//
+#ifndef FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H
+#define FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H
+
+#include <fvdb/JaggedTensor.h>
+#include <fvdb/detail/GridBatchImpl.h>
+
+#include <torch/types.h>
+
+#include <vector>
+
+namespace fvdb {
+namespace detail {
+namespace ops {
+
+// Generate samples along each ray with a given step size where samples are taken
+// by bending the ray according to Snell's law.
+std::vector<JaggedTensor> acousticRaySamples(const GridBatchImpl &batchHdl,
+                                             const JaggedTensor &rayOrigins,
+                                             const JaggedTensor &rayDirections,
+                                             const JaggedTensor &soundSpeeds,
+                                             float stepSize);
+
+} // namespace ops
+} // namespace detail
+} // namespace fvdb
+
+#endif // FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H

src/python/GridBatchBinding.cpp

@@ -154,6 +154,13 @@ bind_grid_batch(py::module &m) {
         .def_property_readonly("dual_bbox", &fvdb::GridBatch::dual_bbox)
         .def_property_readonly("address", &fvdb::GridBatch::address)
 
+
+        .def("acoustic_ray_samples",
+             &fvdb::GridBatch::acoustic_ray_samples,
+             py::arg("ray_origins"),
+             py::arg("ray_directions"),
+             py::arg("sound_speeds"),
+             py::arg("step_size"))
         // Read a property for a single grid in the batch
         .def("voxel_size_at",
              [](const fvdb::GridBatch &self, int64_t bi) {