# Copyright Contributors to the OpenVDB Project
# SPDX-License-Identifier: Apache-2.0
#
"""Functional API for ray operations on sparse grids."""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from .. import _fvdb_cpp
from ..jagged_tensor import JaggedTensor
if TYPE_CHECKING:
from ..grid import Grid
from ..grid_batch import GridBatch
# ---------------------------------------------------------------------------
# Batch variants (GridBatch + JaggedTensor)
# ---------------------------------------------------------------------------
[docs]
def voxels_along_rays_batch(
grid: GridBatch,
ray_origins: JaggedTensor,
ray_directions: JaggedTensor,
max_voxels: int,
eps: float = 0.0,
return_ijk: bool = False,
cumulative: bool = False,
) -> tuple[JaggedTensor, JaggedTensor]:
"""Enumerate voxels intersected by rays using a DDA traversal on a grid batch.
Args:
grid (GridBatch): The grid batch to trace through.
ray_origins (JaggedTensor): Ray origin positions, shape ``(B, -1, 3)``.
ray_directions (JaggedTensor): Ray direction vectors, shape ``(B, -1, 3)``.
max_voxels (int): Maximum number of voxels to return per ray.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
return_ijk (bool): If ``True``, return voxel coordinates instead of linear indices.
cumulative (bool): If ``True``, return cumulative indices across the batch.
Returns:
voxels (JaggedTensor): Voxel coordinates or linear indices per ray hit.
distances (JaggedTensor): ``(t_entry, t_exit)`` pairs per ray hit.
.. seealso:: :func:`voxels_along_rays_single`
"""
grid_data = grid.data
result = _fvdb_cpp.voxels_along_rays(
grid_data, ray_origins._impl, ray_directions._impl, max_voxels, eps, return_ijk, cumulative
)
return JaggedTensor(impl=result[0]), JaggedTensor(impl=result[1])
[docs]
def segments_along_rays_batch(
grid: GridBatch,
ray_origins: JaggedTensor,
ray_directions: JaggedTensor,
max_segments: int,
eps: float = 0.0,
) -> JaggedTensor:
"""Return continuous segments of ray traversal through a grid batch.
Args:
grid (GridBatch): The grid batch to trace through.
ray_origins (JaggedTensor): Ray origin positions, shape ``(B, -1, 3)``.
ray_directions (JaggedTensor): Ray direction vectors, shape ``(B, -1, 3)``.
max_segments (int): Maximum number of segments to return per ray.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
segments (JaggedTensor): ``(t_start, t_end)`` pairs per ray segment.
.. seealso:: :func:`segments_along_rays_single`
"""
grid_data = grid.data
return JaggedTensor(
impl=_fvdb_cpp.segments_along_rays(grid_data, ray_origins._impl, ray_directions._impl, max_segments, eps)
)
[docs]
def ray_implicit_intersection_batch(
grid: GridBatch,
ray_origins: JaggedTensor,
ray_directions: JaggedTensor,
grid_scalars: JaggedTensor,
eps: float = 0.0,
) -> JaggedTensor:
"""Find ray intersections with an implicit surface defined by grid scalars on a grid batch.
The first valid (non-NaN) voxel sampled along each ray seeds the sign reference, and the first
subsequent voxel with the opposite sign is reported as the intersection. Both "positive outside"
and "negative outside" SDF conventions are therefore handled identically, and a ray that enters
the bbox already inside the surface is reported at the *exit* of the surface along the ray.
Args:
grid (GridBatch): The grid batch defining the implicit surface topology.
ray_origins (JaggedTensor): Ray origin positions, shape ``(B, -1, 3)``.
ray_directions (JaggedTensor): Ray direction vectors, shape ``(B, -1, 3)``.
grid_scalars (JaggedTensor): Per-voxel scalar values defining the implicit surface.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
distances (JaggedTensor): Intersection distance per ray, or ``-1`` if no intersection.
.. seealso:: :func:`ray_implicit_intersection_single`
"""
grid_data = grid.data
result_impl = _fvdb_cpp.ray_implicit_intersection(
grid_data, ray_origins._impl, ray_directions._impl, grid_scalars._impl, eps
)
return JaggedTensor(impl=result_impl)
[docs]
def rays_intersect_voxels_batch(
grid: GridBatch,
ray_origins: JaggedTensor,
ray_directions: JaggedTensor,
eps: float = 0.0,
) -> JaggedTensor:
"""Check whether rays hit any voxels in a grid batch.
Args:
grid (GridBatch): The grid batch to test against.
ray_origins (JaggedTensor): Ray origin positions, shape ``(B, -1, 3)``.
ray_directions (JaggedTensor): Ray direction vectors, shape ``(B, -1, 3)``.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
hit (JaggedTensor): Boolean mask indicating whether each ray hit a voxel.
.. seealso:: :func:`rays_intersect_voxels_single`
"""
_, ray_times = voxels_along_rays_batch(
grid,
ray_origins=ray_origins,
ray_directions=ray_directions,
max_voxels=1,
eps=eps,
return_ijk=False,
cumulative=False,
)
did_hit = (ray_times.joffsets[1:] - ray_times.joffsets[:-1]) > 0
return ray_origins.jagged_like(did_hit)
# ---------------------------------------------------------------------------
# Single variants (Grid + torch.Tensor)
# ---------------------------------------------------------------------------
[docs]
def voxels_along_rays_single(
grid: Grid,
ray_origins: torch.Tensor,
ray_directions: torch.Tensor,
max_voxels: int,
eps: float = 0.0,
return_ijk: bool = False,
) -> tuple[JaggedTensor, JaggedTensor]:
"""Enumerate voxels intersected by rays using a DDA traversal on a single grid.
Args:
grid (Grid): The single grid to trace through.
ray_origins (torch.Tensor): Ray origin positions, shape ``(N, 3)``.
ray_directions (torch.Tensor): Ray direction vectors, shape ``(N, 3)``.
max_voxels (int): Maximum number of voxels to return per ray.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
return_ijk (bool): If ``True``, return voxel coordinates instead of linear indices.
Returns:
voxels (JaggedTensor): Voxel coordinates or linear indices per ray hit.
distances (JaggedTensor): ``(t_entry, t_exit)`` pairs per ray hit.
.. seealso:: :func:`voxels_along_rays_batch`
"""
grid_data = grid.data
origins_jt = JaggedTensor(ray_origins)
directions_jt = JaggedTensor(ray_directions)
result = _fvdb_cpp.voxels_along_rays(
grid_data, origins_jt._impl, directions_jt._impl, max_voxels, eps, return_ijk, False
)
return JaggedTensor(impl=result[0]), JaggedTensor(impl=result[1])
[docs]
def segments_along_rays_single(
grid: Grid,
ray_origins: torch.Tensor,
ray_directions: torch.Tensor,
max_segments: int,
eps: float = 0.0,
) -> JaggedTensor:
"""Return continuous segments of ray traversal through a single grid.
Args:
grid (Grid): The single grid to trace through.
ray_origins (torch.Tensor): Ray origin positions, shape ``(N, 3)``.
ray_directions (torch.Tensor): Ray direction vectors, shape ``(N, 3)``.
max_segments (int): Maximum number of segments to return per ray.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
segments (JaggedTensor): ``(t_start, t_end)`` pairs per ray segment.
.. seealso:: :func:`segments_along_rays_batch`
"""
grid_data = grid.data
origins_jt = JaggedTensor(ray_origins)
directions_jt = JaggedTensor(ray_directions)
return JaggedTensor(
impl=_fvdb_cpp.segments_along_rays(grid_data, origins_jt._impl, directions_jt._impl, max_segments, eps)
)
[docs]
def ray_implicit_intersection_single(
grid: Grid,
ray_origins: torch.Tensor,
ray_directions: torch.Tensor,
grid_scalars: torch.Tensor,
eps: float = 0.0,
) -> torch.Tensor:
"""Find ray intersections with an implicit surface defined by grid scalars on a single grid.
The first valid (non-NaN) voxel sampled along each ray seeds the sign reference, and the first
subsequent voxel with the opposite sign is reported as the intersection. Both "positive outside"
and "negative outside" SDF conventions are therefore handled identically, and a ray that enters
the bbox already inside the surface is reported at the *exit* of the surface along the ray.
Args:
grid (Grid): The single grid defining the implicit surface topology.
ray_origins (torch.Tensor): Ray origin positions, shape ``(N, 3)``.
ray_directions (torch.Tensor): Ray direction vectors, shape ``(N, 3)``.
grid_scalars (torch.Tensor): Per-voxel scalar values defining the implicit surface.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
distances (torch.Tensor): Intersection distance per ray, or ``-1`` if no intersection.
.. seealso:: :func:`ray_implicit_intersection_batch`
"""
grid_data = grid.data
origins_jt = JaggedTensor(ray_origins)
directions_jt = JaggedTensor(ray_directions)
scalars_jt = JaggedTensor(grid_scalars)
result_impl = _fvdb_cpp.ray_implicit_intersection(
grid_data, origins_jt._impl, directions_jt._impl, scalars_jt._impl, eps
)
return result_impl.jdata
[docs]
def rays_intersect_voxels_single(
grid: Grid,
ray_origins: torch.Tensor,
ray_directions: torch.Tensor,
eps: float = 0.0,
) -> torch.Tensor:
"""Check whether rays hit any voxels in a single grid.
Args:
grid (Grid): The single grid to test against.
ray_origins (torch.Tensor): Ray origin positions, shape ``(N, 3)``.
ray_directions (torch.Tensor): Ray direction vectors, shape ``(N, 3)``.
eps (float): Small offset to avoid self-intersection. Default ``0.0``.
Returns:
hit (torch.Tensor): Boolean mask indicating whether each ray hit a voxel.
.. seealso:: :func:`rays_intersect_voxels_batch`
"""
_, ray_times = voxels_along_rays_single(
grid,
ray_origins=ray_origins,
ray_directions=ray_directions,
max_voxels=1,
eps=eps,
return_ijk=False,
)
did_hit = (ray_times.joffsets[1:] - ray_times.joffsets[:-1]) > 0
return did_hit