jax_sbgeom.jax_utils.raytracing module

class BVH(left_idx: Array, right_idx: Array, aabb: Array, leafs: Array, order: Array, inverse_order: Array)

Bases: object

Dataclass representing a Bounding Volume Hierarchy (BVH).

Build with build_lbvh.

left_idx: Array

right_idx: Array

aabb: Array

leafs: Array

order: Array

inverse_order: Array

build_lbvh(positions, connectivity)

Build a Linear Bounding Volume Hierarchy (LBVH) for a given set of 3D positions and connectivity.

Parameters:

• positions (jnp.ndarray) – Array of shape (N, 3) with point coordinates.

• connectivity (jnp.ndarray) – Array of shape (M, K) with primitive connectivity indices.

Returns:

LBVH structure with child indices, AABBs, leaf mask, and ordering maps.

Return type:

BVH

points_in_aabbs(point: Array, aabb: Array)

Vectorized version of _point_in_aabb. Takes similar arguments as _point_in_aabb but with additional array axes over which _point_in_aabb is mapped.

Original documentation:

Vectorized version of _point_in_aabb. Takes similar arguments as _point_in_aabb but with additional array axes over which _point_in_aabb is mapped.

Original documentation:

Return whether a point lies inside an axis-aligned bounding box.

Return type:

Array

points_in_aabb(point: Array, aabb: Array)

Vectorized version of _point_in_aabb. Takes similar arguments as _point_in_aabb but with additional array axes over which _point_in_aabb is mapped.

Original documentation:

Return whether a point lies inside an axis-aligned bounding box.

Return type:

Array

points_in_aabbvec(point: Array, aabb: Array)

Vectorized version of _point_in_aabb. Takes similar arguments as _point_in_aabb but with additional array axes over which _point_in_aabb is mapped.

Original documentation:

Return whether a point lies inside an axis-aligned bounding box.

Return type:

Array

probe_bvh(bvh: BVH, points: Array, stack_size: int = 64, max_hit_size: int = 64)

Probe a BVH with a set of points to find which AABBs contain the points.

Parameters:

• bvh (BVH) – The BVH to probe.

• points (Array) – The points to probe the BVH with. Shape (N_points, 3)

• stack_size (int) – The size of the stack to use for traversal.

• max_hit_size (int) – The maximum number of hits to record per point.

Returns:

An array of shape (N_points, max_hit_size) containing the indices of the AABBs of the original mesh. (not in the BVH order, but in the original order) In other words, bvh.order is indexed by the resulting output of _probe_bvh_imp. If no hit, -1 is returned.

Return type:

jnp.ndarray

ray_intersects_aabb(origin: Array, direction: Array, aabb: Array)

Vectorized function to check ray-AABB intersections using the slab method.

Parameters:

• origin (Array) – Ray origin of shape (…, 3)

• direction (Array) – Ray direction of shape (…, 3)

• aabb (Array) – Axis-aligned bounding box of shape (…, 2, 3)

Returns:

Boolean array indicating whether the ray intersects the AABB.

Return type:

jnp.ndarray

ray_traversal_bvh(bvh: BVH, points: Array, directions: Array, stack_size: int = 64, max_hit_size: int = 64)

Traverse a BVH with a set of rays defined by points and directions.

Parameters:

• bvh (BVH) – The BVH to traverse.

• points (Array) – The ray origins. Shape (N_points, 3)

• directions (Array) – The ray directions. Shape (N_points, 3)

• stack_size (int) – The size of the stack to use for traversal.

• max_hit_size (int) – The maximum number of hits to record per ray.

Returns:

An array of shape (N_points, max_hit_size) containing the indices of the AABBs that the rays intersect. If a ray hits fewer than max_hit_size AABBs, the remaining entries are filled with -1.

Return type:

jnp.ndarray

ray_traversal_bvh_single(bvh: BVH, point: Array, direction: Array, stack_size: int = 128, max_hit_size: int = 128)

Traverse a BVH with a single ray defined by point and direction.

Use ray_traversal_bvh_vectorized to handle multiple rays. This is only ~10% slower than ray_traversal_bvh for large number of rays.

Parameters:

• bvh (BVH) – The BVH to traverse.

• point (Array) – The ray origin. Shape (3,)

• direction (Array) – The ray direction. Shape (3,)

• stack_size (int) – The size of the stack to use for traversal.

• max_hit_size (int) – The maximum number of hits to record.

Returns:

An array of shape (max_hit_size,) containing the indices of the AABBs that the ray intersects. If the ray hits fewer than max_hit_size AABBs, the remaining entries are filled with -1.

Return type:

jnp.ndarray

ray_traversal_bvh_vectorized(bvh: BVH, point: Array, direction: Array, stack_size: int = 128, max_hit_size: int = 128)

Traverse a BVH with a single ray defined by point and direction.

Use ray_traversal_bvh_vectorized to handle multiple rays. This is only ~10% slower than ray_traversal_bvh for large number of rays.

Parameters:

• bvh (BVH) – The BVH to traverse.

• point (Array) – The ray origin. Shape (3,)

• direction (Array) – The ray direction. Shape (3,)

• stack_size (int) – The size of the stack to use for traversal.

• max_hit_size (int) – The maximum number of hits to record.

Returns:

An array of shape (max_hit_size,) containing the indices of the AABBs that the ray intersects. If the ray hits fewer than max_hit_size AABBs, the remaining entries are filled with -1.

Return type:

jnp.ndarray

ray_triangle_intersection_single(point: Array, direction: Array, triangle: Array, eps=1e-08, eps_bary_center=1e-08)

Compute intersections of one ray with one trianlge

Parameters:

• point (Array) – Ray origin. Shape (3,)

• direction (Array) – Ray direction. Shape (3,)

• triangle (Array) – Triangle vertices. Shape (3, 3)

Returns:

distance along ray (jnp.inf if no hit)

Return type:

t

ray_triangle_intersection_vectorized(point: Array, direction: Array, triangle: Array, eps=1e-08, eps_bary_center=1e-08)

Compute intersections of one ray with one trianlge

Parameters:

• point (Array) – Ray origin. Shape (3,)

• direction (Array) – Ray direction. Shape (3,)

• triangle (Array) – Triangle vertices. Shape (3, 3)

Returns:

distance along ray (jnp.inf if no hit)

Return type:

t

find_minimum_distance_to_mesh(points: Array, directions: Array, mesh)

Convenience function to find the minimum distance from rays to a triangle mesh.

Parameters:

• points (Array) – Ray origins. Shape (N_points, 3)

• directions (Array) – Ray directions. Shape (N_points, 3)

• mesh (tuple of (positions, connectivity)) – positions: jnp.ndarray of shape (M, 3) representing the 3D coordinates of mesh vertices. connectivity: jnp.ndarray of shape (K, 3) representing the triangle connectivity of the mesh.

Returns:

An array of shape (N_points,) containing the minimum distance from each ray to the mesh. If no intersection occurs, the distance is jnp.inf.

Return type:

jnp.ndarray

closest_point_on_triangle(p: Array, a: Array, b: Array, c: Array)

Compute the closest point on triangle abc to point p. Not vectorized itself. Reimplements the Embree reference algorithm without explicit if conditionals.

Parameters:

• p (Array) – (3,) array of point coordinates

• a (Array) – (3,) array of triangle vertex a coordinates

• b (Array) – (3,) array of triangle vertex b coordinates

• c (Array) – (3,) array of triangle vertex c coordinates

Returns:

closest_point – (3,) array of closest point coordinates

Return type:

jnp.ndarray

bvh_closest_point_vectorized(bvh: BVH, point: Array, mesh, stack_size: int = 64, max_hit_size: int = 64)

Get closest point on triangle mesh for a single point using BVH traversal.

For multiple points, use bvh_closest_point_vectorized (handles arbitrary point shapes)

Parameters:

• bvh (BVH) – BVH structure.

• point (Array) – Query point with shape (3,).

• mesh (tuple) – Mesh tuple (positions, connectivity) with triangle connectivity.

Returns:

• closest_point (jnp.ndarray) – Closest point on the mesh.

• distance (jnp.ndarray) – Distance from query point to closest point.

• triangle_index (jnp.ndarray) – Triangle index of the closest point.

get_closest_points_on_mesh(points: Array, mesh)

Get closest points on triangle mesh for a set of points.

Parameters:

• points (Array) – Array of points

• mesh (tuple of (positions, connectivity)) – positions: jnp.ndarray of shape (M, 3) representing the 3D coordinates of mesh vertices. connectivity: jnp.ndarray of shape (K, 3) representing the triangle connectivity of the mesh.

Return type:

Array

Returns:

closest_points: jnp.ndarray [N,3]

Array of closest points on the mesh for each input point.

distances: jnp.ndarray [N,]

Squared distances from each input point to its closest point on the mesh.

triangle_indices: jnp.ndarray [N,]

Indices of the triangles on the mesh corresponding to the closest points.