Handle parallel rays on AABB faces without NaN intersections

nerfstudio/utils/math.py

@@ -156,19 +156,26 @@ def intersect_aabb(
         t_min, t_max - two tensors of shapes N representing distance of intersection from the origin.
     """
 
-    tx_min = (aabb[:3] - origins) / directions
-    tx_max = (aabb[3:] - origins) / directions
+    parallel = directions == 0
+    safe_directions = torch.where(parallel, torch.ones_like(directions), directions)
+    tx_min = (aabb[:3] - origins) / safe_directions
+    tx_max = (aabb[3:] - origins) / safe_directions
 
     t_min = torch.stack((tx_min, tx_max)).amin(dim=0)
     t_max = torch.stack((tx_min, tx_max)).amax(dim=0)
+    # A parallel ray imposes no parameter bound if it lies inside the slab.
+    # Handle outside parallel slabs as misses without ever dividing by zero.
+    t_min = t_min.masked_fill(parallel, -torch.inf)
+    t_max = t_max.masked_fill(parallel, torch.inf)
+    parallel_miss = (parallel & ((origins < aabb[:3]) | (origins > aabb[3:]))).any(dim=-1)
 
     t_min = t_min.amax(dim=-1)
     t_max = t_max.amin(dim=-1)
 
     t_min = torch.clamp(t_min, min=0, max=max_bound)
     t_max = torch.clamp(t_max, min=0, max=max_bound)
 
-    cond = t_max <= t_min
+    cond = (t_max <= t_min) | parallel_miss
     t_min = torch.where(cond, invalid_value, t_min)
     t_max = torch.where(cond, invalid_value, t_max)
 

tests/utils/test_parallel_aabb_intersection.py

@@ -0,0 +1,131 @@
+# Copyright 2026 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test exact slab geometry when ray directions contain zero components."""
+
+import pytest
+import torch
+
+from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.data.scene_box import OrientedBox, SceneBox
+from nerfstudio.utils.math import intersect_aabb, intersect_obb
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize("negative_zero", [False, True])
+def test_parallel_rays_on_faces_and_edges_have_finite_intervals(dtype, negative_zero):
+    aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=dtype)
+    origins = torch.tensor([[0, 0.5, -1], [1, 0.5, -1], [0, 0, -1], [1, 1, -1]], dtype=dtype)
+    zero = -0.0 if negative_zero else 0.0
+    directions = torch.tensor([[zero, zero, 1]], dtype=dtype).expand_as(origins)
+    near, far = intersect_aabb(origins, directions, aabb)
+    torch.testing.assert_close(near, torch.ones_like(near), rtol=0, atol=0)
+    torch.testing.assert_close(far, torch.full_like(far, 2), rtol=0, atol=0)
+
+
+def test_parallel_rays_outside_slabs_are_misses():
+    aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.float64)
+    origins = torch.tensor([[-0.5, 0.5, -1], [1.5, 0.5, -1], [0.5, -0.1, -1]], dtype=torch.float64)
+    directions = origins.new_tensor([[0, 0, 1]]).expand_as(origins)
+    near, far = intersect_aabb(origins, directions, aabb, max_bound=11, invalid_value=123)
+    torch.testing.assert_close(near, torch.full_like(near, 123), rtol=0, atol=0)
+    torch.testing.assert_close(far, torch.full_like(far, 123), rtol=0, atol=0)
+
+
+def test_ray_starting_on_face_inside_box_has_zero_near():
+    origins = torch.tensor([[0, 1, 0.5]], dtype=torch.float64)
+    directions = origins.new_tensor([[0, 0, 1]])
+    aabb = origins.new_tensor([0, 0, 0, 1, 1, 1])
+    near, far = intersect_aabb(origins, directions, aabb)
+    torch.testing.assert_close(near, near.new_tensor([0]), rtol=0, atol=0)
+    torch.testing.assert_close(far, far.new_tensor([0.5]), rtol=0, atol=0)
+
+
+def test_interior_parallel_intersection_has_finite_analytic_gradients():
+    origins = torch.tensor([[0.5, 0.5, -1]], dtype=torch.float64, requires_grad=True)
+    directions = torch.tensor([[0, 0, 1]], dtype=torch.float64, requires_grad=True)
+    aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.float64, requires_grad=True)
+    near, far = intersect_aabb(origins, directions, aabb)
+    origin_grad, direction_grad, box_grad = torch.autograd.grad((near + far).sum(), (origins, directions, aabb))
+    torch.testing.assert_close(origin_grad, origins.new_tensor([[0, 0, -2]]))
+    torch.testing.assert_close(direction_grad, directions.new_tensor([[0, 0, -3]]))
+    torch.testing.assert_close(box_grad, aabb.new_tensor([0, 0, 1, 0, 0, 1]))
+    assert torch.autograd.gradcheck(lambda o, d, b: intersect_aabb(o, d, b), (origins, directions, aabb))
+
+
+def test_parallel_misses_have_finite_zero_gradients():
+    origins = torch.tensor([[-0.5, 0.5, -1]], dtype=torch.float64, requires_grad=True)
+    directions = torch.tensor([[0, 0, 1]], dtype=torch.float64, requires_grad=True)
+    aabb = origins.new_tensor([0, 0, 0, 1, 1, 1])
+    near, far = intersect_aabb(origins, directions, aabb)
+    for gradient in torch.autograd.grad((near + far).sum(), (origins, directions)):
+        torch.testing.assert_close(gradient, torch.zeros_like(gradient), rtol=0, atol=0)
+
+
+def test_nonzero_small_directions_are_not_treated_as_parallel():
+    origins = torch.tensor([[-1e-12, 0.5, -1]], dtype=torch.float64)
+    directions = origins.new_tensor([[1e-12, 0, 1]])
+    aabb = origins.new_tensor([0, 0, 0, 1, 1, 1])
+    near, far = intersect_aabb(origins, directions, aabb)
+    torch.testing.assert_close(near, near.new_tensor([1]), rtol=0, atol=0)
+    torch.testing.assert_close(far, far.new_tensor([2]), rtol=0, atol=0)
+
+
+def test_oriented_crop_box_handles_ray_on_parallel_face():
+    rotation = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=torch.float64)
+    translation = rotation.new_tensor([2, 3, 4])
+    obb = OrientedBox(R=rotation, T=translation, S=rotation.new_tensor([2, 2, 2]))
+    origins = rotation.new_tensor([[1, 0, -2]]) @ rotation.T + translation
+    directions = rotation.new_tensor([[0, 0, 1]]) @ rotation.T
+    near, far = intersect_obb(origins, directions, obb)
+    torch.testing.assert_close(near, near.new_tensor([1]), rtol=0, atol=0)
+    torch.testing.assert_close(far, far.new_tensor([3]), rtol=0, atol=0)
+
+
+def test_camera_center_ray_has_finite_crop_bounds_on_box_face():
+    camera_to_world = torch.eye(4)[:3].clone()
+    camera_to_world[:, 3] = torch.tensor([0, 0.5, 2])
+    camera = Cameras(camera_to_worlds=camera_to_world, fx=100.0, fy=100.0, cx=0.5, cy=0.5, height=1, width=1)
+    crop = SceneBox(aabb=torch.tensor([[0, 0, 0], [1, 1, 1]], dtype=torch.float32))
+    rays = camera.generate_rays(0, aabb_box=crop)
+    torch.testing.assert_close(rays.directions, rays.directions.new_tensor([[[0, 0, -1]]]), rtol=0, atol=0)
+    assert rays.nears is not None and rays.fars is not None
+    torch.testing.assert_close(rays.nears, rays.nears.new_tensor([[[1]]]), rtol=0, atol=0)
+    torch.testing.assert_close(rays.fars, rays.fars.new_tensor([[[2]]]), rtol=0, atol=0)
+
+
+def test_nonparallel_rays_match_original_formula_and_gradients():
+    generator = torch.Generator().manual_seed(19)
+    origins = torch.randn(12, 3, generator=generator, dtype=torch.float64).requires_grad_()
+    directions = torch.randn(12, 3, generator=generator, dtype=torch.float64).requires_grad_()
+    aabb = origins.new_tensor([-2, -2, -2, 2, 2, 2])
+    slabs = torch.stack(((aabb[:3] - origins) / directions, (aabb[3:] - origins) / directions))
+    near = slabs.amin(0).amax(-1).clamp(0, 1e10)
+    far = slabs.amax(0).amin(-1).clamp(0, 1e10)
+    miss = far <= near
+    expected = (torch.where(miss, 1e10, near), torch.where(miss, 1e10, far))
+    actual = intersect_aabb(origins, directions, aabb)
+    for a, b in zip(actual, expected):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+    actual_grads = torch.autograd.grad(sum(t.sum() for t in actual), (origins, directions))
+    expected_grads = torch.autograd.grad(sum(t.sum() for t in expected), (origins, directions))
+    for a, b in zip(actual_grads, expected_grads):
+        torch.testing.assert_close(a, b, rtol=1e-12, atol=1e-12)
+
+
+def test_empty_ray_batch_preserves_shapes():
+    origins = torch.empty(0, 3)
+    aabb = origins.new_tensor([0, 0, 0, 1, 1, 1])
+    near, far = intersect_aabb(origins, origins, aabb)
+    assert near.shape == far.shape == (0,)