diff --git a/fvdb/_fvdb_cpp.pyi b/fvdb/_fvdb_cpp.pyi index 99d54e6d..3687afb7 100644 --- a/fvdb/_fvdb_cpp.pyi +++ b/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, diff --git a/fvdb/grid.py b/fvdb/grid.py index 6fbfbf47..de537941 100644 --- a/fvdb/grid.py +++ b/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) + ) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5bfa17de..56752bbb 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/fvdb/GridBatch.cpp b/src/fvdb/GridBatch.cpp index 2dd5e779..cb3e2f97 100644 --- a/src/fvdb/GridBatch.cpp +++ b/src/fvdb/GridBatch.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include @@ -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 +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, diff --git a/src/fvdb/GridBatch.h b/src/fvdb/GridBatch.h index e79d316c..f14c91ff 100644 --- a/src/fvdb/GridBatch.h +++ b/src/fvdb/GridBatch.h @@ -305,6 +305,13 @@ struct GridBatch : torch::CustomClassHolder { const std::optional &min_coord = std::nullopt, const std::optional &grid_size = std::nullopt) const; + + std::vector 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) diff --git a/src/fvdb/detail/ops/AcousticRaySamples.cu b/src/fvdb/detail/ops/AcousticRaySamples.cu new file mode 100644 index 00000000..e1af61c8 --- /dev/null +++ b/src/fvdb/detail/ops/AcousticRaySamples.cu @@ -0,0 +1,77 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#include +#include +#include +#include +#include + +#include +#include + +namespace fvdb { +namespace detail { +namespace ops { +namespace { + +template +__global__ void countAcousticRaySamplesKernel( + BatchGridAccessor batchAccessor, + const JaggedAccessor rayOrigins, + const JaggedAccessor rayDirections, + const JaggedAccessor soundSpeeds, + float stepSize, + fvdb::TorchRAcc64 outCounts) { + + // Kernel parallelizes over rays. + // For each ray, we count the number of samples we will take. +} + +template +__global__ void accousticRaySamplesKernel( + BatchGridAccessor batchAccessor, + const JaggedAccessor rayOrigins, + const JaggedAccessor rayDirections, + const JaggedAccessor soundSpeeds, + float stepSize, + JaggedAccessor 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 acousticRaySamples( + const GridBatchImpl &batchHdl, + const JaggedTensor &rayOrigins, + const JaggedTensor &rayDirections, + const JaggedTensor &soundSpeeds, + float stepSize) { + + const auto gridBatchAcc = gridBatchAccessor(batchHdl); + const auto rayOriginsAcc = jaggedAccessor(rayOrigins); + const auto rayDirectionsAcc = jaggedAccessor(rayDirections); + const auto soundSpeedsAcc = jaggedAccessor(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 diff --git a/src/fvdb/detail/ops/AcousticRaySamples.h b/src/fvdb/detail/ops/AcousticRaySamples.h new file mode 100644 index 00000000..4ed15a28 --- /dev/null +++ b/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 +#include + +#include + +#include + +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 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 diff --git a/src/python/GridBatchBinding.cpp b/src/python/GridBatchBinding.cpp index 9b09c1f5..bc75a739 100644 --- a/src/python/GridBatchBinding.cpp +++ b/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) {