"""
Ray queries using the embreex package with the
API wrapped to match our native raytracer.
"""
import numpy as np
# `pip install embreex` installs from wheels
from embreex import rtcore_scene
from embreex.mesh_construction import TriangleMesh
from .. import caching, intersections, util
from ..constants import log_time
from ..typed import ArrayLike, Integer
from .ray_util import contains_points
# embree operates on float32 values
_embree_dtype = np.float32
# scale-aware base ray offset to step past hit triangles above f32 ULP
_ray_offset_factor = 1e-6
_ray_offset_floor = 1e-8
class RayMeshIntersector:
def __init__(self, geometry, scale_to_box: bool = True):
"""
Do ray- mesh queries.
Parameters
-------------
geometry : Trimesh object
Mesh to do ray tests on
scale_to_box : bool
If true, will scale mesh to approximate
unit cube to avoid problems with extreme
large or small meshes.
"""
self.mesh = geometry
self._scale_to_box = scale_to_box
self._cache = caching.Cache(id_function=self.mesh.__hash__)
@property
def _scale(self):
"""
Scaling factor for precision.
"""
if self._scale_to_box:
# scale vertices to approximately a cube to help with
# numerical issues at very large/small scales
scale = 100.0 / self.mesh.scale
else:
scale = 1.0
return scale
@caching.cache_decorator
def _scene(self):
"""
A cached version of the embreex scene.
"""
return _EmbreeWrap(
vertices=self.mesh.vertices, faces=self.mesh.faces, scale=self._scale
)
def intersects_location(
self,
ray_origins: ArrayLike,
ray_directions: ArrayLike,
multiple_hits: bool = True,
):
"""
Return the location of where a ray hits a surface.
Parameters
----------
ray_origins : (n, 3) float
Origins of rays
ray_directions : (n, 3) float
Direction (vector) of rays
Returns
---------
locations : (m) sequence of (p, 3) float
Intersection points
index_ray : (m,) int
Indexes of ray
index_tri : (m,) int
Indexes of mesh.faces
"""
(index_tri, index_ray, locations) = self.intersects_id(
ray_origins=ray_origins,
ray_directions=ray_directions,
multiple_hits=multiple_hits,
return_locations=True,
)
return locations, index_ray, index_tri
@log_time
def intersects_id(
self,
ray_origins: ArrayLike,
ray_directions: ArrayLike,
multiple_hits: bool = True,
max_hits: Integer = 100,
return_locations: bool = False,
):
"""
Find the triangles hit by a list of rays, including
optionally multiple hits along a single ray.
Parameters
----------
ray_origins : (n, 3) float
Origins of rays
ray_directions : (n, 3) float
Direction (vector) of rays
multiple_hits : bool
If True will return every hit along the ray
If False will only return first hit
max_hits : int
Maximum number of hits per ray
return_locations : bool
Should we return hit locations or not
Returns
---------
index_tri : (m,) int
Indexes of mesh.faces
index_ray : (m,) int
Indexes of ray
locations : (m) sequence of (p, 3) float
Intersection points, only returned if return_locations
"""
# make sure input is _dtype for embree
ray_origins = np.array(ray_origins, dtype=np.float64)
ray_directions = np.array(ray_directions, dtype=np.float64)
if ray_origins.shape != ray_directions.shape:
raise ValueError("Ray origin and direction don't match!")
ray_directions = util.unitize(ray_directions)
# stack results for multiple hits into a sequence
result_triangle = [np.zeros(0, dtype=np.int64)]
result_ray_idx = [np.zeros(0, dtype=np.int64)]
result_locations = [np.zeros((0, 3), dtype=np.float64)]
if multiple_hits or return_locations:
# how much to offset ray to transport to the other side of face
base_offset = max(_ray_offset_floor, self.mesh.scale * _ray_offset_factor)
ray_offsets = ray_directions * base_offset
# grab the planes from triangles
plane_origins = self.mesh.triangles[:, 0, :]
plane_normals = self.mesh.face_normals
# what each ray hit last iteration; -1 means nothing, used to
# detect a ray stuck on the same face due to precision issues
last_hit = np.full(len(ray_origins), -1, dtype=np.int64)
# absolute indices of rays still being queried; shrinks each
# iteration as rays miss, escape, or get culled
live = np.arange(len(ray_origins))
# use a for loop rather than a while to ensure this exits
# if a ray is offset from a triangle and then is reported
# hitting itself this could get stuck on that one triangle
for _depth in range(max_hits):
# if you set output=1 embreex returns distance along the ray
# which is bizarrely slower than our own plane-line calc
# TODO: switch `run(..., output=True)` once embreex>=4.4.0rc1 is stable
query = self._scene.run(ray_origins[live], ray_directions[live])
hit = query != -1
if not hit.any():
break
# absolute indices and triangle indices for rays that hit
hit_rays = live[hit]
hit_tris = query[hit]
# first-hit-only fast path: no duplicates or locations to track
if not multiple_hits and not return_locations:
result_triangle.append(hit_tris)
result_ray_idx.append(hit_rays)
break
# rays that hit the same triangle as last iteration are stuck
dupe = last_hit[hit_rays] == hit_tris
# store the last hit so we can track duplicates
last_hit[hit_rays] = hit_tris
# subset to separate duplicates and non-duplicates
ok_rays = hit_rays[~dupe]
ok_tris = hit_tris[~dupe]
dupe_rays = hit_rays[dupe]
# compute where clean hits actually land on their triangle;
# planes_lines silently drops near-parallel rays, so `valid`
# shortens new_origins — we must trim ok_rays / ok_tris to match
new_origins, valid = intersections.planes_lines(
plane_origins=plane_origins[ok_tris],
plane_normals=plane_normals[ok_tris],
line_origins=ray_origins[ok_rays],
line_directions=ray_directions[ok_rays],
)
ok_rays, ok_tris = ok_rays[valid], ok_tris[valid]
result_locations.append(new_origins)
result_triangle.append(ok_tris)
result_ray_idx.append(ok_rays)
if not multiple_hits:
break
# clean hits step onto the new face with a fresh base offset
ray_origins[ok_rays] = new_origins + ray_offsets[ok_rays]
ray_offsets[ok_rays] = ray_directions[ok_rays] * base_offset
# stuck rays double their offset and step further to try to clear
ray_offsets[dupe_rays] *= 2.0
ray_origins[dupe_rays] += ray_offsets[dupe_rays]
# carry forward only rays we successfully advanced;
# dropped rays (misses, near-parallel planes) die here
live = np.concatenate([ok_rays, dupe_rays])
index_tri = np.concatenate(result_triangle)
index_ray = np.concatenate(result_ray_idx)
if return_locations:
return index_tri, index_ray, np.concatenate(result_locations)
return index_tri, index_ray
@log_time
def intersects_first(self, ray_origins, ray_directions):
"""
Find the index of the first triangle a ray hits.
Parameters
----------
ray_origins : (n, 3) float
Origins of rays
ray_directions : (n, 3) float
Direction (vector) of rays
Returns
----------
triangle_index : (n,) int
Index of triangle ray hit, or -1 if not hit
"""
# make sure our arrays are in the `embree` dtype
ray_origins = np.array(ray_origins, dtype=_embree_dtype)
ray_directions = np.array(ray_directions, dtype=_embree_dtype)
if ray_origins.shape != ray_directions.shape:
raise ValueError("Ray origin and direction don't match!")
ray_directions = util.unitize(ray_directions)
return self._scene.run(ray_origins, ray_directions)
def intersects_any(self, ray_origins, ray_directions):
"""
Check if a list of rays hits the surface.
Parameters
-----------
ray_origins : (n, 3) float
Origins of rays
ray_directions : (n, 3) float
Direction (vector) of rays
Returns
----------
hit : (n,) bool
Did each ray hit the surface
"""
first = self.intersects_first(
ray_origins=ray_origins, ray_directions=ray_directions
)
hit = first != -1
return hit
def contains_points(self, points):
"""
Check if a mesh contains a list of points, using ray tests.
If the point is on the surface of the mesh, behavior is undefined.
Parameters
---------
points: (n, 3) points in space
Returns
---------
contains: (n,) bool
Whether point is inside mesh or not
"""
return contains_points(self, points)
def __getstate__(self):
state = self.__dict__.copy()
# don't pickle cache
state.pop("_cache", None)
return state
def __setstate__(self, state):
self.__dict__.update(state)
# Add cache back since it doesn't exist in the pickle
self._cache = caching.Cache(id_function=self.mesh.__hash__)
def __deepcopy__(self, *args):
return self.__copy__()
def __copy__(self, *args):
return RayMeshIntersector(geometry=self.mesh, scale_to_box=self._scale_to_box)
class _EmbreeWrap:
"""
A light wrapper for Embreex scene objects which
allows queries to be scaled to help with precision
issues, as well as selecting the correct dtypes.
"""
def __init__(self, vertices, faces, scale):
scaled = np.array(vertices, dtype=np.float64)
self.origin = scaled.min(axis=0)
self.scale = float(scale)
scaled = (scaled - self.origin) * self.scale
self.scene = rtcore_scene.EmbreeScene()
# assign the geometry to the scene
TriangleMesh(
scene=self.scene,
vertices=scaled.astype(_embree_dtype),
indices=faces.view(np.ndarray).astype(np.int32),
)
def run(self, origins, normals, **kwargs):
scaled = (np.array(origins, dtype=np.float64) - self.origin) * self.scale
return self.scene.run(
scaled.astype(_embree_dtype), normals.astype(_embree_dtype), **kwargs
)