From 94d8281d4227f3e84bd81da165ace04420571735 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 2 Feb 2026 13:24:28 -0800 Subject: [PATCH 1/2] cleaning up directory --- .gitignore | 1 + examples/generate_playground_gif.py | 334 ++++++++++++++++++++++++++++ 2 files changed, 335 insertions(+) create mode 100644 examples/generate_playground_gif.py diff --git a/.gitignore b/.gitignore index 996ce29..27404e5 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,4 @@ __pycache__ build/ dist/ rtxpy.egg-info/ +examples/.ipynb_checkpoints/ diff --git a/examples/generate_playground_gif.py b/examples/generate_playground_gif.py new file mode 100644 index 0000000..e5d04df --- /dev/null +++ b/examples/generate_playground_gif.py @@ -0,0 +1,334 @@ +"""Generate a GIF from the playground viewshed/hillshade animation. + +This script creates frames from the Crater Lake hiking animation and +combines them into a GIF suitable for the README. +""" + +import numpy as np +import cupy +import xarray as xr +from pathlib import Path +from PIL import Image +import io + +from rtxpy import RTX, viewshed, hillshade + + +def load_terrain(): + """Load Crater Lake terrain data.""" + import rioxarray as rxr + + dem_path = Path(__file__).parent / "crater_lake_national_park.tif" + + if not dem_path.exists(): + raise FileNotFoundError(f"DEM file not found at {dem_path}. Run playground.py first to download it.") + + print(f"Loading DEM: {dem_path}") + terrain = rxr.open_rasterio(str(dem_path), masked=True).squeeze() + + # Subsample aggressively for smaller GIF file size + terrain = terrain[::10, ::10] + + # Crop edges to remove invalid border values + crop = 20 + terrain = terrain[crop:-crop, crop:-crop] + + # Scale down elevation for visualization + terrain.data = terrain.data * 0.2 + + # Ensure contiguous array before GPU transfer + terrain.data = np.ascontiguousarray(terrain.data) + + # Convert to cupy for GPU processing + terrain.data = cupy.asarray(terrain.data) + + print(f"Terrain loaded: {terrain.shape}") + return terrain + + +def generate_hiking_path(x_coords, y_coords, num_points=360): + """Generate a hiking path around Crater Lake (roughly circular).""" + cx = (x_coords.min() + x_coords.max()) / 2 + cy = (y_coords.min() + y_coords.max()) / 2 + + rx = (x_coords.max() - x_coords.min()) * 0.25 + ry = (y_coords.max() - y_coords.min()) * 0.25 + + angles = np.linspace(0, 2 * np.pi, num_points) + wobble = np.sin(angles * 8) * 0.1 + + path_x = cx + (rx + rx * wobble) * np.cos(angles) + path_y = cy + (ry + ry * wobble) * np.sin(angles) + + return path_x, path_y + + +def coords_to_pixel(x, y, x_coords, y_coords): + """Convert data coordinates to pixel coordinates.""" + px = np.searchsorted(x_coords, x) + py = np.searchsorted(-y_coords, -y) + return int(np.clip(px, 0, len(x_coords) - 1)), int(np.clip(py, 0, len(y_coords) - 1)) + + +def draw_legend(colors, x=10, y=10): + """Draw a legend in the corner of the frame.""" + from PIL import Image as PILImage, ImageDraw, ImageFont + + H, W = colors.shape[:2] + + # Create a small PIL image for drawing text + legend_w, legend_h = 90, 52 + legend = PILImage.new('RGBA', (legend_w, legend_h), (0, 0, 0, 180)) + draw = ImageDraw.Draw(legend) + + # Use default font + try: + font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 10) + except (OSError, IOError): + font = ImageFont.load_default() + + # Legend entries: color swatch + label + entries = [ + ((50, 220, 50), "Visible"), + ((80, 80, 85), "Not Visible"), + ((0, 255, 255), "Observer"), + ] + + for i, (color, label) in enumerate(entries): + row_y = 5 + i * 15 + # Draw color swatch + draw.rectangle([5, row_y, 15, row_y + 10], fill=color) + # Draw label + draw.text((20, row_y - 1), label, fill=(255, 255, 255), font=font) + + # Convert legend to numpy and overlay on frame + legend_arr = np.array(legend) + + # Blend legend onto frame + for ly in range(legend_h): + for lx in range(legend_w): + fy, fx = y + ly, x + lx + if 0 <= fy < H and 0 <= fx < W: + alpha = legend_arr[ly, lx, 3] / 255.0 + colors[fy, fx, :3] = ( + colors[fy, fx, :3] * (1 - alpha) + legend_arr[ly, lx, :3] * alpha + ).astype(np.uint8) + + +def draw_observer_marker(colors, px, py, radius=6, glow_radius=18): + """Draw a glowing teal marker with dark outline at the observer's position.""" + H, W = colors.shape[:2] + outline_width = 2 + + # Draw outer glow, dark outline, then bright center + for dy in range(-glow_radius, glow_radius + 1): + for dx in range(-glow_radius, glow_radius + 1): + dist_sq = dx*dx + dy*dy + if dist_sq <= glow_radius * glow_radius: + ny, nx = py + dy, px + dx + if 0 <= ny < H and 0 <= nx < W: + dist = np.sqrt(dist_sq) + if dist <= radius: + # Bright cyan/teal center + colors[ny, nx] = [0, 255, 255, 255] + elif dist <= radius + outline_width: + # Dark outline for contrast + colors[ny, nx] = [0, 40, 40, 255] + elif dist <= glow_radius: + # Glow falloff - blend cyan with existing color + t = (dist - radius - outline_width) / (glow_radius - radius - outline_width) + glow_strength = (1 - t) ** 1.5 # Slightly softer falloff + existing = colors[ny, nx, :3].astype(np.float32) + cyan = np.array([0, 255, 255], dtype=np.float32) + blended = existing + (cyan - existing) * glow_strength * 0.7 + colors[ny, nx, :3] = np.clip(blended, 0, 255).astype(np.uint8) + + +def generate_frames(terrain, num_frames=72): + """Generate animation frames. + + Parameters + ---------- + terrain : xarray.DataArray + The terrain data. + num_frames : int + Number of frames to generate. Both the hillshade and hiker will + complete exactly one full 360° loop in this many frames. + """ + H, W = terrain.data.shape + rtx = RTX() + + x_coords = terrain.indexes.get('x').values + y_coords = terrain.indexes.get('y').values + + path_x, path_y = generate_hiking_path(x_coords, y_coords, num_points=360) + + frames = [] + azimuth = 225 + + print(f"Generating {num_frames} frames...") + + # Calculate rotation per frame for full 360° loop + azimuth_step = 360 / num_frames + hiker_step = 360 / num_frames + + for frame_idx in range(num_frames): + path_idx = int((frame_idx * hiker_step) % 360) + + vsw = path_x[path_idx] + vsh = path_y[path_idx] + azimuth = (225 + frame_idx * azimuth_step) % 360 + + # Compute hillshade and viewshed + hs = hillshade(terrain, + shadows=True, + azimuth=azimuth, + angle_altitude=25, + rtx=rtx) + vs = viewshed(terrain, + x=vsw, + y=vsh, + observer_elev=100.0, + rtx=rtx) + + # Convert to numpy arrays + hs_data = hs.data.get() if hasattr(hs.data, 'get') else hs.data + vs_data = vs.data.get() if hasattr(vs.data, 'get') else vs.data + + # Track NaN and zero pixels before converting - these will be transparent + transparent_mask = np.isnan(hs_data) | np.isnan(vs_data) | (hs_data == 0) + + hs_data = np.nan_to_num(hs_data, nan=0.5) + gray = np.uint8(np.clip(hs_data * 200, 0, 255)) + + # Viewshed returns -1 for invisible, 0-180 for visible (angle) + visible_mask = vs_data >= 0 + not_visible_mask = (vs_data < 0) & ~transparent_mask + + # Compose the final image with alpha channel (RGBA) + colors = np.zeros((H, W, 4), dtype=np.uint8) + colors[:, :, 0] = gray + colors[:, :, 1] = gray + colors[:, :, 2] = gray + colors[:, :, 3] = 255 # Fully opaque by default + + # Tint visible areas bright lime green - make it really pop! + colors[visible_mask, 0] = 50 # Low red + colors[visible_mask, 1] = np.minimum(255, gray[visible_mask].astype(np.int16) + 120).astype(np.uint8) # Bright green + colors[visible_mask, 2] = 50 # Low blue + + # Tint non-visible areas darker gray + colors[not_visible_mask, 0] = (colors[not_visible_mask, 0] * 0.5).astype(np.uint8) + colors[not_visible_mask, 1] = (colors[not_visible_mask, 1] * 0.5).astype(np.uint8) + colors[not_visible_mask, 2] = (colors[not_visible_mask, 2] * 0.55).astype(np.uint8) + + # Make NaN and zero pixels transparent + colors[transparent_mask, 3] = 0 + + # Draw observer marker + px, py = coords_to_pixel(vsw, vsh, x_coords, y_coords) + draw_observer_marker(colors, px, py, radius=4) + + # Draw legend + draw_legend(colors, x=10, y=10) + + frames.append(Image.fromarray(colors, mode='RGBA')) + + if (frame_idx + 1) % 10 == 0: + print(f" Frame {frame_idx + 1}/{num_frames}") + + return frames + + +def create_gif(frames, output_path, fps=12, max_colors=64): + """Create a GIF from frames. + + Parameters + ---------- + frames : list of PIL.Image + The frames to combine. + output_path : Path + Output path for the GIF. + fps : int + Frames per second. + max_colors : int + Maximum colors in palette for smaller file size. + """ + duration = int(1000 / fps) # Duration in milliseconds + + # Use a magenta color as the transparency key (unlikely to appear in terrain) + transparent_color = (255, 0, 255) + + # Convert RGBA frames to RGB, replacing transparent pixels with the key color + print("Converting frames for GIF transparency...") + rgb_frames = [] + for frame in frames: + arr = np.array(frame) + rgb = arr[:, :, :3].copy() + alpha = arr[:, :, 3] + # Set transparent pixels to the key color + rgb[alpha == 0] = transparent_color + rgb_frames.append(Image.fromarray(rgb, mode='RGB')) + + # Create global palette from sampled frames to avoid flickering + print(f"Building global palette from {len(rgb_frames)} frames...") + sample_step = max(1, len(rgb_frames) // 10) + sampled = [np.array(rgb_frames[i]) for i in range(0, len(rgb_frames), sample_step)] + combined = np.concatenate([p.reshape(-1, 3) for p in sampled], axis=0) + + h, w = np.array(rgb_frames[0]).shape[:2] + sample_h = int(np.ceil(len(combined) / w)) + padded = np.zeros((sample_h * w, 3), dtype=np.uint8) + padded[:len(combined)] = combined + palette_img = Image.fromarray(padded.reshape(sample_h, w, 3), mode='RGB') + global_palette = palette_img.quantize(colors=max_colors, method=Image.Quantize.MEDIANCUT) + + # Modify palette to reserve index 0 for transparency + palette_data = list(global_palette.getpalette()) + palette_data[0:3] = transparent_color # Force index 0 to be transparent color + global_palette.putpalette(palette_data) + transparency_index = 0 + + # Quantize all frames using the global palette + print(f"Quantizing frames to {max_colors} colors...") + quantized_frames = [] + for frame in rgb_frames: + q_frame = frame.quantize(palette=global_palette, dither=Image.Dither.FLOYDSTEINBERG) + quantized_frames.append(q_frame) + + print(f"Creating GIF at {output_path}...") + save_kwargs = { + 'save_all': True, + 'append_images': quantized_frames[1:], + 'duration': duration, + 'loop': 0, # Loop forever + 'optimize': True + } + if transparency_index is not None: + save_kwargs['transparency'] = transparency_index + save_kwargs['disposal'] = 2 # Restore to background + print(f" Using transparency index: {transparency_index}") + + quantized_frames[0].save(output_path, **save_kwargs) + + file_size = output_path.stat().st_size / (1024 * 1024) + print(f"GIF created: {output_path} ({file_size:.1f} MB)") + + +def main(): + output_path = Path(__file__).parent / "images" / "playground_demo.gif" + output_path.parent.mkdir(exist_ok=True) + + terrain = load_terrain() + + # Generate 120 frames - both hillshade and hiker complete exactly one 360° loop + # At 15fps this gives an 8 second loop that repeats seamlessly + frames = generate_frames(terrain, num_frames=120) + + create_gif(frames, output_path, fps=6, max_colors=128) + + print(f"\nDone! GIF saved to: {output_path}") + + +if __name__ == "__main__": + main() From a2c2a61e1bb74d2f76649991a66fdad73526497b Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 2 Feb 2026 21:25:34 -0800 Subject: [PATCH 2/2] stubbed out supporting .obj file loading --- rtxpy/__init__.py | 8 +- rtxpy/mesh.py | 244 +++++++++++++++++++++- rtxpy/rtx.py | 212 +++++++++++++++---- rtxpy/tests/test_mesh.py | 415 ++++++++++++++++++++++++++++++++++++- rtxpy/tests/test_simple.py | 268 ++++++++++++++++++++++++ 5 files changed, 1103 insertions(+), 44 deletions(-) diff --git a/rtxpy/__init__.py b/rtxpy/__init__.py index 1da61f1..77f5bdd 100644 --- a/rtxpy/__init__.py +++ b/rtxpy/__init__.py @@ -6,7 +6,13 @@ list_devices, get_current_device, ) -from .mesh import triangulate_terrain, write_stl +from .mesh import ( + triangulate_terrain, + write_stl, + load_obj, + make_transform, + make_transforms_on_terrain, +) from .analysis import viewshed, hillshade __version__ = "0.0.5" diff --git a/rtxpy/mesh.py b/rtxpy/mesh.py index 514a2ac..67a4b48 100644 --- a/rtxpy/mesh.py +++ b/rtxpy/mesh.py @@ -1,12 +1,14 @@ -"""Mesh utilities for terrain triangulation and STL export. +"""Mesh utilities for terrain triangulation, STL export, and OBJ loading. This module provides functions for converting raster terrain data into -triangle meshes suitable for ray tracing, and for exporting meshes to STL format. +triangle meshes suitable for ray tracing, for exporting meshes to STL format, +and for loading external mesh files in OBJ format. """ import numba as nb from numba import cuda import numpy as np +from pathlib import Path from .rtx import has_cupy @@ -204,3 +206,241 @@ def write_stl(filename, verts, triangles): content = np.empty(numTris * 50, np.uint8) _fill_stl_contents(content, vb, ib, numTris) f.write(content) + + +def load_obj(filepath, scale=1.0, swap_yz=False): + """Load a Wavefront OBJ file and return vertices and indices for ray tracing. + + This function parses OBJ files and converts them to the flattened vertex + and index arrays expected by the RTX class. Supports triangular and + quadrilateral faces (quads are automatically triangulated). + + Parameters + ---------- + filepath : str or Path + Path to the OBJ file to load. + scale : float, optional + Scale factor applied to all vertex coordinates. Default is 1.0. + swap_yz : bool, optional + If True, swap Y and Z coordinates. Useful when OBJ uses Y-up convention + but the scene uses Z-up (common for terrain/DEM scenes). Default is False. + + Returns + ------- + vertices : numpy.ndarray + Flattened float32 array of vertex positions with shape (N*3,), + where N is the number of vertices. Layout is [x0, y0, z0, x1, y1, z1, ...]. + indices : numpy.ndarray + Flattened int32 array of triangle indices with shape (M*3,), + where M is the number of triangles. Layout is [i0, i1, i2, i3, i4, i5, ...]. + + Raises + ------ + FileNotFoundError + If the specified file does not exist. + ValueError + If the file contains no valid geometry or has faces with fewer than + 3 vertices. + + Examples + -------- + Load an OBJ file and add it to a scene: + + >>> from rtxpy import RTX, load_obj + >>> verts, indices = load_obj("building.obj", scale=0.1) + >>> rtx = RTX() + >>> rtx.add_geometry("building", verts, indices) + + Load with coordinate swap for Z-up terrain scenes: + + >>> verts, indices = load_obj("model.obj", swap_yz=True) + + Notes + ----- + - OBJ files use 1-based indexing; this function converts to 0-based. + - Only vertex positions (v) and faces (f) are parsed. Texture coordinates (vt), + normals (vn), materials, and other OBJ features are ignored. + - Faces with more than 4 vertices are triangulated using a fan pattern. + - Negative face indices (relative references) are supported. + """ + filepath = Path(filepath) + if not filepath.exists(): + raise FileNotFoundError(f"OBJ file not found: {filepath}") + + vertices = [] + faces = [] + + with open(filepath, 'r') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#'): + continue + + parts = line.split() + if not parts: + continue + + if parts[0] == 'v' and len(parts) >= 4: + # Vertex: v x y z [w] + x, y, z = float(parts[1]), float(parts[2]), float(parts[3]) + if swap_yz: + y, z = z, y + vertices.append([x * scale, y * scale, z * scale]) + + elif parts[0] == 'f' and len(parts) >= 4: + # Face: f v1 v2 v3 ... or f v1/vt1/vn1 v2/vt2/vn2 ... + face_indices = [] + for p in parts[1:]: + # Handle v, v/vt, v/vt/vn, or v//vn formats + idx_str = p.split('/')[0] + idx = int(idx_str) + # OBJ uses 1-based indexing, convert to 0-based + # Negative indices are relative to current vertex count + if idx < 0: + idx = len(vertices) + idx + else: + idx = idx - 1 + face_indices.append(idx) + + if len(face_indices) < 3: + continue + + # Triangulate: fan triangulation for polygons + # Triangle: [0, 1, 2] + # Quad: [0, 1, 2], [0, 2, 3] + # Pentagon: [0, 1, 2], [0, 2, 3], [0, 3, 4] + for i in range(1, len(face_indices) - 1): + faces.append([face_indices[0], face_indices[i], face_indices[i + 1]]) + + if not vertices: + raise ValueError(f"No vertices found in OBJ file: {filepath}") + if not faces: + raise ValueError(f"No faces found in OBJ file: {filepath}") + + vertices_array = np.array(vertices, dtype=np.float32).flatten() + indices_array = np.array(faces, dtype=np.int32).flatten() + + return vertices_array, indices_array + + +def make_transform(x=0.0, y=0.0, z=0.0, scale=1.0, rotation_z=0.0): + """Create a 3x4 affine transform matrix for positioning geometry. + + This is a convenience function for creating transform matrices to use + with RTX.add_geometry(). The transform applies scale, then rotation + around the Z axis, then translation. + + Parameters + ---------- + x : float, optional + X translation. Default is 0.0. + y : float, optional + Y translation. Default is 0.0. + z : float, optional + Z translation. Default is 0.0. + scale : float, optional + Uniform scale factor. Default is 1.0. + rotation_z : float, optional + Rotation around Z axis in radians. Default is 0.0. + + Returns + ------- + list + 12-float list representing a 3x4 row-major affine transform matrix. + Format: [Xx, Xy, Xz, Tx, Yx, Yy, Yz, Ty, Zx, Zy, Zz, Tz] + + Examples + -------- + Simple translation: + + >>> transform = make_transform(x=100, y=200, z=50) + >>> rtx.add_geometry("tower", verts, indices, transform=transform) + + Scale and translate: + + >>> transform = make_transform(x=100, y=200, z=50, scale=0.1) + + Rotate 90 degrees and translate: + + >>> import math + >>> transform = make_transform(x=100, y=200, rotation_z=math.pi/2) + """ + import math + c = math.cos(rotation_z) + s = math.sin(rotation_z) + + # Scale * Rotation * Translation (applied right to left) + # Rotation matrix around Z: [[c, -s, 0], [s, c, 0], [0, 0, 1]] + return [ + scale * c, -scale * s, 0.0, x, + scale * s, scale * c, 0.0, y, + 0.0, 0.0, scale, z, + ] + + +def make_transforms_on_terrain(positions, terrain, scale=1.0, rotation_z=0.0): + """Create transforms for placing objects at multiple positions on terrain. + + This convenience function samples terrain elevation at each (x, y) position + and creates transform matrices suitable for RTX.add_geometry(transforms=...). + + Parameters + ---------- + positions : array-like + Sequence of (x, y) coordinate pairs where objects should be placed. + Can be a list of tuples, numpy array of shape (N, 2), etc. + terrain : array-like + 2D array of elevation values with shape (H, W). The terrain uses + pixel coordinates where position (x, y) samples terrain[int(y), int(x)]. + scale : float, optional + Uniform scale factor applied to all transforms. Default is 1.0. + rotation_z : float or array-like, optional + Rotation around Z axis in radians. Can be a single value applied to + all instances, or an array of rotations (one per position). Default is 0.0. + + Returns + ------- + list + List of 12-float transform matrices, one per position. + + Examples + -------- + Place cell towers at multiple locations: + + >>> tower_verts, tower_indices = load_obj("cell_tower.obj") + >>> positions = [(100, 200), (300, 400), (500, 150)] + >>> transforms = make_transforms_on_terrain(positions, dem, scale=0.1) + >>> rtx.add_geometry("towers", tower_verts, tower_indices, transforms=transforms) + + With random rotations: + + >>> import numpy as np + >>> rotations = np.random.uniform(0, 2*np.pi, len(positions)) + >>> transforms = make_transforms_on_terrain(positions, dem, rotation_z=rotations) + """ + positions = np.asarray(positions) + if positions.ndim == 1: + positions = positions.reshape(-1, 2) + + n = len(positions) + + # Handle rotation_z as scalar or array + if np.isscalar(rotation_z): + rotations = np.full(n, rotation_z) + else: + rotations = np.asarray(rotation_z) + if len(rotations) != n: + raise ValueError(f"rotation_z length ({len(rotations)}) must match " + f"positions length ({n})") + + transforms = [] + for i, (x, y) in enumerate(positions): + # Sample terrain elevation at this position + # Terrain uses row, col indexing: terrain[row, col] = terrain[y, x] + row = int(np.clip(y, 0, terrain.shape[0] - 1)) + col = int(np.clip(x, 0, terrain.shape[1] - 1)) + z = float(terrain[row, col]) + + transforms.append(make_transform(x, y, z, scale, rotations[i])) + + return transforms diff --git a/rtxpy/rtx.py b/rtxpy/rtx.py index e229a5b..b242732 100644 --- a/rtxpy/rtx.py +++ b/rtxpy/rtx.py @@ -24,18 +24,25 @@ # Data structures for multi-GAS support # ----------------------------------------------------------------------------- +def _identity_transform(): + """Return an identity 3x4 transform matrix.""" + return [ + 1.0, 0.0, 0.0, 0.0, # Row 0: [Xx, Xy, Xz, Tx] + 0.0, 1.0, 0.0, 0.0, # Row 1: [Yx, Yy, Yz, Ty] + 0.0, 0.0, 1.0, 0.0, # Row 2: [Zx, Zy, Zz, Tz] + ] + + @dataclass class _GASEntry: - """Storage for a single Geometry Acceleration Structure.""" + """Storage for a single Geometry Acceleration Structure with multiple instances.""" gas_id: str gas_handle: int gas_buffer: cupy.ndarray # Must keep reference to prevent GC vertices_hash: int - transform: List[float] = field(default_factory=lambda: [ - 1.0, 0.0, 0.0, 0.0, # Row 0: [Xx, Xy, Xz, Tx] - 0.0, 1.0, 0.0, 0.0, # Row 1: [Yx, Yy, Yz, Ty] - 0.0, 0.0, 1.0, 0.0, # Row 2: [Zx, Zy, Zz, Tz] - ]) # 12 floats (3x4 row-major affine transform) + # List of transforms - each transform creates one instance of this geometry + # Each transform is 12 floats (3x4 row-major affine matrix) + transforms: List[List[float]] = field(default_factory=lambda: [_identity_transform()]) # ----------------------------------------------------------------------------- @@ -569,7 +576,9 @@ def _build_ias(geom_state: _GeometryState): Build an Instance Acceleration Structure (IAS) from all GAS entries. This creates a top-level acceleration structure that references all - geometry acceleration structures with their transforms. + geometry acceleration structures with their transforms. Each GAS entry + can have multiple transforms, creating multiple instances of the same + geometry (GPU instancing). Args: geom_state: The geometry state containing GAS entries to build IAS from. @@ -585,7 +594,8 @@ def _build_ias(geom_state: _GeometryState): geom_state.ias_dirty = False return - num_instances = len(geom_state.gas_entries) + # Count total instances across all GAS entries (each transform = one instance) + num_instances = sum(len(entry.transforms) for entry in geom_state.gas_entries.values()) # OptixInstance structure is 80 bytes: # - transform: float[12] (3x4 row-major) = 48 bytes @@ -600,29 +610,35 @@ def _build_ias(geom_state: _GeometryState): INSTANCE_SIZE = 80 instances_data = bytearray(num_instances * INSTANCE_SIZE) - for i, (gas_id, entry) in enumerate(geom_state.gas_entries.items()): - offset = i * INSTANCE_SIZE + instance_idx = 0 + for geometry_idx, (gas_id, entry) in enumerate(geom_state.gas_entries.items()): + # Create one instance for each transform of this geometry + for transform in entry.transforms: + offset = instance_idx * INSTANCE_SIZE - # Pack transform (12 floats, 48 bytes) - transform_bytes = struct.pack('12f', *entry.transform) - instances_data[offset:offset + 48] = transform_bytes + # Pack transform (12 floats, 48 bytes) + transform_bytes = struct.pack('12f', *transform) + instances_data[offset:offset + 48] = transform_bytes - # Pack instanceId (4 bytes) - struct.pack_into('I', instances_data, offset + 48, i) + # Pack instanceId (4 bytes) - use geometry index so we can identify + # which geometry was hit (all instances of same geometry share this ID) + struct.pack_into('I', instances_data, offset + 48, geometry_idx) - # Pack sbtOffset (4 bytes) - all use same hit group (SBT index 0) - struct.pack_into('I', instances_data, offset + 52, 0) + # Pack sbtOffset (4 bytes) - all use same hit group (SBT index 0) + struct.pack_into('I', instances_data, offset + 52, 0) - # Pack visibilityMask (4 bytes) - 0xFF = visible to all rays - struct.pack_into('I', instances_data, offset + 56, 0xFF) + # Pack visibilityMask (4 bytes) - 0xFF = visible to all rays + struct.pack_into('I', instances_data, offset + 56, 0xFF) - # Pack flags (4 bytes) - OPTIX_INSTANCE_FLAG_NONE = 0 - struct.pack_into('I', instances_data, offset + 60, 0) + # Pack flags (4 bytes) - OPTIX_INSTANCE_FLAG_NONE = 0 + struct.pack_into('I', instances_data, offset + 60, 0) - # Pack traversableHandle (8 bytes) - struct.pack_into('Q', instances_data, offset + 64, entry.gas_handle) + # Pack traversableHandle (8 bytes) - same GAS for all instances + struct.pack_into('Q', instances_data, offset + 64, entry.gas_handle) - # Padding (8 bytes) - already zeros + # Padding (8 bytes) - already zeros + + instance_idx += 1 # Copy instances to GPU geom_state.instances_buffer = cupy.array( @@ -1039,29 +1055,56 @@ def trace(self, rays, hits, numRays: int, primitive_ids=None, instance_ids=None) # ------------------------------------------------------------------------- def add_geometry(self, geometry_id: str, vertices, indices, - transform: Optional[List[float]] = None) -> int: + transform: Optional[List[float]] = None, + transforms: Optional[List[List[float]]] = None) -> int: """ - Add a geometry (GAS) to the scene with an optional transform. + Add a geometry (GAS) to the scene with optional transform(s). This enables multi-GAS mode. If called after build(), the single-GAS state is cleared. Adding a geometry with an existing ID replaces it. + Supports GPU instancing: pass multiple transforms to create multiple + instances of the same geometry efficiently (geometry data is stored + once, but rendered at multiple locations). + Args: geometry_id: Unique identifier for this geometry vertices: Vertex buffer (flattened float32 array, 3 floats per vertex) indices: Index buffer (flattened int32 array, 3 ints per triangle) - transform: Optional 12-float list representing a 3x4 row-major + transform: Optional single 12-float list representing a 3x4 row-major affine transform matrix. Defaults to identity. Format: [Xx, Xy, Xz, Tx, Yx, Yy, Yz, Ty, Zx, Zy, Zz, Tz] + transforms: Optional list of transforms for GPU instancing. Each + transform is a 12-float list. Use this to efficiently + render the same geometry at multiple locations. + Cannot be used together with `transform` parameter. Returns: 0 on success, non-zero on error + + Examples: + # Single instance (default) + rtx.add_geometry("terrain", verts, indices) + + # Single instance with transform + rtx.add_geometry("tower", verts, indices, transform=[1,0,0,x, 0,1,0,y, 0,0,1,z]) + + # Multiple instances (GPU instancing) - efficient for many copies + transforms = [[1,0,0,x1, 0,1,0,y1, 0,0,1,z1], + [1,0,0,x2, 0,1,0,y2, 0,0,1,z2], + ...] + rtx.add_geometry("towers", verts, indices, transforms=transforms) """ global _state if not _state.initialized: _init_optix() + # Validate transform arguments + if transform is not None and transforms is not None: + raise ValueError("Cannot specify both 'transform' and 'transforms'. " + "Use 'transform' for a single instance or 'transforms' for multiple.") + # Switch to multi-GAS mode if currently in single-GAS mode if self._geom_state.single_gas_mode: self._geom_state.gas_handle = 0 @@ -1081,17 +1124,26 @@ def add_geometry(self, geometry_id: str, vertices, indices, vertices_for_hash = np.asarray(vertices) vertices_hash = hash(vertices_for_hash.tobytes()) - # Set transform (identity if not provided) - if transform is None: - transform = [ - 1.0, 0.0, 0.0, 0.0, - 0.0, 1.0, 0.0, 0.0, - 0.0, 0.0, 1.0, 0.0, - ] - else: + # Build transforms list + if transforms is not None: + # Multiple transforms provided (GPU instancing) + transforms_list = [] + for t in transforms: + t = list(t) + if len(t) != 12: + raise ValueError(f"Each transform must have 12 floats, got {len(t)}") + transforms_list.append(t) + if not transforms_list: + raise ValueError("transforms list cannot be empty") + elif transform is not None: + # Single transform provided transform = list(transform) if len(transform) != 12: return -1 + transforms_list = [transform] + else: + # No transform - use identity + transforms_list = [_identity_transform()] # Create or update the GAS entry self._geom_state.gas_entries[geometry_id] = _GASEntry( @@ -1099,7 +1151,7 @@ def add_geometry(self, geometry_id: str, vertices, indices, gas_handle=gas_handle, gas_buffer=gas_buffer, vertices_hash=vertices_hash, - transform=transform, + transforms=transforms_list, ) # Mark IAS as needing rebuild @@ -1126,31 +1178,111 @@ def remove_geometry(self, geometry_id: str) -> int: return 0 def update_transform(self, geometry_id: str, - transform: List[float]) -> int: + transform: List[float], + instance_index: int = 0) -> int: """ - Update the transform of an existing geometry. + Update the transform of a specific instance of a geometry. Args: geometry_id: The ID of the geometry to update transform: 12-float list representing a 3x4 row-major affine transform matrix. Format: [Xx, Xy, Xz, Tx, Yx, Yy, Yz, Ty, Zx, Zy, Zz, Tz] + instance_index: Which instance to update (default 0, the first instance) Returns: - 0 on success, -1 if geometry not found or invalid transform + 0 on success, -1 if geometry not found, invalid transform, or invalid index """ if geometry_id not in self._geom_state.gas_entries: return -1 + entry = self._geom_state.gas_entries[geometry_id] + if instance_index < 0 or instance_index >= len(entry.transforms): + return -1 + transform = list(transform) if len(transform) != 12: return -1 - self._geom_state.gas_entries[geometry_id].transform = transform + entry.transforms[instance_index] = transform + self._geom_state.ias_dirty = True + + return 0 + + def set_transforms(self, geometry_id: str, + transforms: List[List[float]]) -> int: + """ + Replace all transforms for a geometry (change instance count and positions). + + Args: + geometry_id: The ID of the geometry to update + transforms: List of 12-float transform matrices + + Returns: + 0 on success, -1 if geometry not found or invalid transforms + """ + if geometry_id not in self._geom_state.gas_entries: + return -1 + + transforms_list = [] + for t in transforms: + t = list(t) + if len(t) != 12: + return -1 + transforms_list.append(t) + + if not transforms_list: + return -1 + + self._geom_state.gas_entries[geometry_id].transforms = transforms_list self._geom_state.ias_dirty = True return 0 + def add_instances(self, geometry_id: str, + transforms: List[List[float]]) -> int: + """ + Add additional instances to an existing geometry. + + This is useful for incrementally adding more instances without + replacing the existing ones. + + Args: + geometry_id: The ID of the geometry to add instances to + transforms: List of 12-float transform matrices for new instances + + Returns: + 0 on success, -1 if geometry not found or invalid transforms + """ + if geometry_id not in self._geom_state.gas_entries: + return -1 + + entry = self._geom_state.gas_entries[geometry_id] + + for t in transforms: + t = list(t) + if len(t) != 12: + return -1 + entry.transforms.append(t) + + self._geom_state.ias_dirty = True + + return 0 + + def get_instance_count(self, geometry_id: str) -> int: + """ + Get the number of instances for a geometry. + + Args: + geometry_id: The ID of the geometry + + Returns: + Number of instances, or -1 if geometry not found + """ + if geometry_id not in self._geom_state.gas_entries: + return -1 + return len(self._geom_state.gas_entries[geometry_id].transforms) + def list_geometries(self) -> List[str]: """ Get a list of all geometry IDs in the scene. diff --git a/rtxpy/tests/test_mesh.py b/rtxpy/tests/test_mesh.py index 3ee0f24..97af78b 100644 --- a/rtxpy/tests/test_mesh.py +++ b/rtxpy/tests/test_mesh.py @@ -3,7 +3,13 @@ import numpy as np import pytest -from rtxpy import triangulate_terrain, write_stl +from rtxpy import ( + triangulate_terrain, + write_stl, + load_obj, + make_transform, + make_transforms_on_terrain, +) from rtxpy.rtx import has_cupy @@ -180,3 +186,410 @@ def test_write_from_cupy_arrays(self, tmp_path): assert filepath.exists() expected_size = 80 + 4 + 1 * 50 assert filepath.stat().st_size == expected_size + + +class TestLoadObj: + """Tests for load_obj function.""" + + def test_load_simple_triangle(self, tmp_path): + """Test loading OBJ with a single triangle.""" + obj_content = """# Simple triangle +v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +f 1 2 3 +""" + filepath = tmp_path / "triangle.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + # 3 vertices * 3 components = 9 + assert len(verts) == 9 + assert verts.dtype == np.float32 + # 1 triangle * 3 indices = 3 + assert len(indices) == 3 + assert indices.dtype == np.int32 + + # Check vertex values + np.testing.assert_array_almost_equal( + verts, [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.5, 1.0, 0.0] + ) + # Check indices (0-based) + np.testing.assert_array_equal(indices, [0, 1, 2]) + + def test_load_quad_triangulation(self, tmp_path): + """Test that quads are automatically triangulated.""" + obj_content = """# Quad (unit square) +v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 1.0 1.0 0.0 +v 0.0 1.0 0.0 +f 1 2 3 4 +""" + filepath = tmp_path / "quad.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + # 4 vertices + assert len(verts) == 12 + # Quad becomes 2 triangles: 6 indices + assert len(indices) == 6 + + # Fan triangulation: [0,1,2] and [0,2,3] + np.testing.assert_array_equal(indices, [0, 1, 2, 0, 2, 3]) + + def test_load_with_scale(self, tmp_path): + """Test that scale parameter is applied to vertices.""" + obj_content = """v 1.0 2.0 3.0 +v 4.0 5.0 6.0 +v 7.0 8.0 9.0 +f 1 2 3 +""" + filepath = tmp_path / "scaled.obj" + filepath.write_text(obj_content) + + verts, _ = load_obj(filepath, scale=0.1) + + np.testing.assert_array_almost_equal( + verts, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + ) + + def test_load_with_swap_yz(self, tmp_path): + """Test that swap_yz swaps Y and Z coordinates.""" + obj_content = """v 1.0 2.0 3.0 +v 4.0 5.0 6.0 +v 7.0 8.0 9.0 +f 1 2 3 +""" + filepath = tmp_path / "swapped.obj" + filepath.write_text(obj_content) + + verts, _ = load_obj(filepath, swap_yz=True) + + # Y and Z should be swapped: (x, y, z) -> (x, z, y) + np.testing.assert_array_almost_equal( + verts, [1.0, 3.0, 2.0, 4.0, 6.0, 5.0, 7.0, 9.0, 8.0] + ) + + def test_load_face_with_texture_and_normal_indices(self, tmp_path): + """Test parsing faces with v/vt/vn format.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +vt 0.0 0.0 +vt 1.0 0.0 +vt 0.5 1.0 +vn 0.0 0.0 1.0 +f 1/1/1 2/2/1 3/3/1 +""" + filepath = tmp_path / "with_texnorm.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + # Should correctly parse vertex indices, ignoring vt and vn + assert len(verts) == 9 + np.testing.assert_array_equal(indices, [0, 1, 2]) + + def test_load_face_with_vertex_texture_only(self, tmp_path): + """Test parsing faces with v/vt format.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +vt 0.0 0.0 +vt 1.0 0.0 +vt 0.5 1.0 +f 1/1 2/2 3/3 +""" + filepath = tmp_path / "with_tex.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + assert len(verts) == 9 + np.testing.assert_array_equal(indices, [0, 1, 2]) + + def test_load_face_with_vertex_normal_only(self, tmp_path): + """Test parsing faces with v//vn format.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +vn 0.0 0.0 1.0 +f 1//1 2//1 3//1 +""" + filepath = tmp_path / "with_norm.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + assert len(verts) == 9 + np.testing.assert_array_equal(indices, [0, 1, 2]) + + def test_load_negative_indices(self, tmp_path): + """Test parsing faces with negative (relative) indices.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +f -3 -2 -1 +""" + filepath = tmp_path / "negative.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + # -3, -2, -1 should resolve to 0, 1, 2 + np.testing.assert_array_equal(indices, [0, 1, 2]) + + def test_load_multiple_faces(self, tmp_path): + """Test loading OBJ with multiple faces.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 1.0 1.0 0.0 +v 0.0 1.0 0.0 +f 1 2 3 +f 1 3 4 +""" + filepath = tmp_path / "two_triangles.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + assert len(verts) == 12 + assert len(indices) == 6 + np.testing.assert_array_equal(indices, [0, 1, 2, 0, 2, 3]) + + def test_load_with_comments_and_blank_lines(self, tmp_path): + """Test that comments and blank lines are ignored.""" + obj_content = """# This is a comment +# Another comment + +v 0.0 0.0 0.0 + +v 1.0 0.0 0.0 +# Comment between vertices +v 0.5 1.0 0.0 + +f 1 2 3 +# Final comment +""" + filepath = tmp_path / "with_comments.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + assert len(verts) == 9 + assert len(indices) == 3 + + def test_load_file_not_found(self): + """Test that FileNotFoundError is raised for missing files.""" + with pytest.raises(FileNotFoundError): + load_obj("/nonexistent/path/to/file.obj") + + def test_load_no_vertices(self, tmp_path): + """Test that ValueError is raised when file has no vertices.""" + obj_content = """# No vertices +f 1 2 3 +""" + filepath = tmp_path / "no_verts.obj" + filepath.write_text(obj_content) + + with pytest.raises(ValueError, match="No vertices found"): + load_obj(filepath) + + def test_load_no_faces(self, tmp_path): + """Test that ValueError is raised when file has no faces.""" + obj_content = """v 0.0 0.0 0.0 +v 1.0 0.0 0.0 +v 0.5 1.0 0.0 +""" + filepath = tmp_path / "no_faces.obj" + filepath.write_text(obj_content) + + with pytest.raises(ValueError, match="No faces found"): + load_obj(filepath) + + def test_load_pentagon_triangulation(self, tmp_path): + """Test that polygons with more than 4 vertices are triangulated.""" + obj_content = """# Pentagon +v 0.0 1.0 0.0 +v 0.951 0.309 0.0 +v 0.588 -0.809 0.0 +v -0.588 -0.809 0.0 +v -0.951 0.309 0.0 +f 1 2 3 4 5 +""" + filepath = tmp_path / "pentagon.obj" + filepath.write_text(obj_content) + + verts, indices = load_obj(filepath) + + # 5 vertices + assert len(verts) == 15 + # Pentagon becomes 3 triangles: [0,1,2], [0,2,3], [0,3,4] + assert len(indices) == 9 + np.testing.assert_array_equal(indices, [0, 1, 2, 0, 2, 3, 0, 3, 4]) + + +class TestMakeTransform: + """Tests for make_transform function.""" + + def test_identity_transform(self): + """Test that default parameters give identity transform.""" + t = make_transform() + expected = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0] + np.testing.assert_array_almost_equal(t, expected) + + def test_translation_only(self): + """Test translation without rotation or scale.""" + t = make_transform(x=10, y=20, z=30) + expected = [1, 0, 0, 10, 0, 1, 0, 20, 0, 0, 1, 30] + np.testing.assert_array_almost_equal(t, expected) + + def test_scale_only(self): + """Test uniform scaling.""" + t = make_transform(scale=2.0) + expected = [2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0] + np.testing.assert_array_almost_equal(t, expected) + + def test_rotation_90_degrees(self): + """Test 90 degree rotation around Z axis.""" + import math + t = make_transform(rotation_z=math.pi / 2) + # cos(90) = 0, sin(90) = 1 + # Rotation matrix: [[0, -1, 0], [1, 0, 0], [0, 0, 1]] + expected = [0, -1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0] + np.testing.assert_array_almost_equal(t, expected, decimal=10) + + def test_combined_transform(self): + """Test scale + rotation + translation.""" + import math + t = make_transform(x=5, y=10, z=15, scale=2.0, rotation_z=math.pi / 2) + # Scale 2 * Rotation 90deg + expected = [0, -2, 0, 5, 2, 0, 0, 10, 0, 0, 2, 15] + np.testing.assert_array_almost_equal(t, expected, decimal=10) + + def test_returns_list(self): + """Test that result is a list of 12 floats.""" + t = make_transform(x=1, y=2, z=3) + assert isinstance(t, list) + assert len(t) == 12 + + +class TestMakeTransformsOnTerrain: + """Tests for make_transforms_on_terrain function.""" + + def test_single_position(self): + """Test placing one object on terrain.""" + terrain = np.array([ + [0, 0, 0], + [0, 5, 0], + [0, 0, 0], + ], dtype=np.float32) + + positions = [(1, 1)] # Center of terrain, elevation = 5 + transforms = make_transforms_on_terrain(positions, terrain) + + assert len(transforms) == 1 + t = transforms[0] + assert len(t) == 12 + # Check translation components (indices 3, 7, 11) + assert t[3] == 1.0 # x + assert t[7] == 1.0 # y + assert t[11] == 5.0 # z (sampled from terrain) + + def test_multiple_positions(self): + """Test placing multiple objects on terrain.""" + terrain = np.array([ + [1, 2, 3], + [4, 5, 6], + [7, 8, 9], + ], dtype=np.float32) + + positions = [(0, 0), (1, 1), (2, 2)] + transforms = make_transforms_on_terrain(positions, terrain) + + assert len(transforms) == 3 + # Check z values match terrain + assert transforms[0][11] == 1.0 # terrain[0, 0] + assert transforms[1][11] == 5.0 # terrain[1, 1] + assert transforms[2][11] == 9.0 # terrain[2, 2] + + def test_with_scale(self): + """Test scale parameter is applied.""" + terrain = np.zeros((3, 3), dtype=np.float32) + positions = [(1, 1)] + transforms = make_transforms_on_terrain(positions, terrain, scale=0.5) + + t = transforms[0] + # Scale should appear in diagonal elements + assert t[0] == 0.5 # Xx + assert t[5] == 0.5 # Yy + assert t[10] == 0.5 # Zz + + def test_with_scalar_rotation(self): + """Test single rotation value applied to all.""" + import math + terrain = np.zeros((3, 3), dtype=np.float32) + positions = [(0, 0), (1, 1)] + transforms = make_transforms_on_terrain(positions, terrain, + rotation_z=math.pi / 2) + + # Both should have same rotation + for t in transforms: + np.testing.assert_almost_equal(t[0], 0.0) # cos(90) = 0 + np.testing.assert_almost_equal(t[4], 1.0) # sin(90) = 1 + + def test_with_array_rotation(self): + """Test different rotation for each position.""" + import math + terrain = np.zeros((3, 3), dtype=np.float32) + positions = [(0, 0), (1, 1)] + rotations = [0.0, math.pi / 2] + transforms = make_transforms_on_terrain(positions, terrain, + rotation_z=rotations) + + # First: no rotation + np.testing.assert_almost_equal(transforms[0][0], 1.0) # cos(0) + np.testing.assert_almost_equal(transforms[0][4], 0.0) # sin(0) + # Second: 90 degree rotation + np.testing.assert_almost_equal(transforms[1][0], 0.0) # cos(90) + np.testing.assert_almost_equal(transforms[1][4], 1.0) # sin(90) + + def test_position_clipping(self): + """Test that out-of-bounds positions are clipped.""" + terrain = np.array([ + [1, 2], + [3, 4], + ], dtype=np.float32) + + # Position outside terrain bounds + positions = [(10, 10)] # Should be clipped to (1, 1) + transforms = make_transforms_on_terrain(positions, terrain) + + # Should sample terrain[1, 1] = 4 + assert transforms[0][11] == 4.0 + + def test_rotation_length_mismatch(self): + """Test error when rotation array length doesn't match positions.""" + terrain = np.zeros((3, 3), dtype=np.float32) + positions = [(0, 0), (1, 1), (2, 2)] + rotations = [0.0, 0.0] # Only 2, but 3 positions + + with pytest.raises(ValueError, match="rotation_z length"): + make_transforms_on_terrain(positions, terrain, rotation_z=rotations) + + def test_numpy_array_positions(self): + """Test that numpy array positions work.""" + terrain = np.array([ + [1, 2], + [3, 4], + ], dtype=np.float32) + + positions = np.array([[0, 0], [1, 1]]) + transforms = make_transforms_on_terrain(positions, terrain) + + assert len(transforms) == 2 + assert transforms[0][11] == 1.0 # terrain[0, 0] + assert transforms[1][11] == 4.0 # terrain[1, 1] diff --git a/rtxpy/tests/test_simple.py b/rtxpy/tests/test_simple.py index 446dcdd..440433d 100644 --- a/rtxpy/tests/test_simple.py +++ b/rtxpy/tests/test_simple.py @@ -1205,3 +1205,271 @@ def test_primitive_and_instance_ids_optional(test_cupy): inst_ids = backend.int32([0]) res = rtx.trace(rays, hits, 1, instance_ids=inst_ids) assert res == 0 + + +# ----------------------------------------------------------------------------- +# GPU Instancing Tests +# ----------------------------------------------------------------------------- + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_gpu_instancing_multiple_transforms(test_cupy): + """Test adding one geometry with multiple transforms (GPU instancing).""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() + + # Small triangle at origin + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Create 3 instances at different positions along x-axis + transforms = [ + [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], # At origin + [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0], # At x=10 + [1, 0, 0, 20, 0, 1, 0, 0, 0, 0, 1, 0], # At x=20 + ] + + res = rtx.add_geometry("triangles", verts, tris, transforms=transforms) + assert res == 0 + assert rtx.get_instance_count("triangles") == 3 + + # Ray hitting first instance (at origin) + rays0 = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits0 = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays0, hits0, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits0[0]), 10.0, decimal=1) + + # Ray hitting second instance (at x=10) + rays1 = backend.float32([10.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits1 = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays1, hits1, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits1[0]), 10.0, decimal=1) + + # Ray hitting third instance (at x=20) + rays2 = backend.float32([20.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits2 = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays2, hits2, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits2[0]), 10.0, decimal=1) + + # Ray at x=5 should miss (between instances) + rays_miss = backend.float32([5.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits_miss = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays_miss, hits_miss, 1) + assert res == 0 + assert float(hits_miss[0]) == -1.0 # Miss + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_gpu_instancing_many_instances(test_cupy): + """Stress test with many instances of the same geometry.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() + + # Small triangle + verts = backend.float32([0, 0, 0, 0.5, 0, 0, 0.25, 0.5, 0]) + tris = backend.int32([0, 1, 2]) + + # Create 100 instances in a 10x10 grid + transforms = [] + for row in range(10): + for col in range(10): + x, y = col * 2.0, row * 2.0 + transforms.append([1, 0, 0, x, 0, 1, 0, y, 0, 0, 1, 0]) + + res = rtx.add_geometry("grid", verts, tris, transforms=transforms) + assert res == 0 + assert rtx.get_instance_count("grid") == 100 + + # Test a few positions + for row in [0, 5, 9]: + for col in [0, 5, 9]: + x, y = col * 2.0 + 0.25, row * 2.0 + 0.25 + rays = backend.float32([x, y, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays, hits, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_add_instances_to_existing_geometry(test_cupy): + """Test adding more instances to an existing geometry.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Start with one instance + res = rtx.add_geometry("mesh", verts, tris, transform=[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0]) + assert res == 0 + assert rtx.get_instance_count("mesh") == 1 + + # Add two more instances + new_transforms = [ + [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 20, 0, 1, 0, 0, 0, 0, 1, 0], + ] + res = rtx.add_instances("mesh", new_transforms) + assert res == 0 + assert rtx.get_instance_count("mesh") == 3 + + # Verify all three positions are hittable + for x in [0.5, 10.5, 20.5]: + rays = backend.float32([x, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays, hits, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_set_transforms_replace_all(test_cupy): + """Test replacing all transforms for a geometry.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Start with 3 instances + transforms = [ + [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 20, 0, 1, 0, 0, 0, 0, 1, 0], + ] + res = rtx.add_geometry("mesh", verts, tris, transforms=transforms) + assert res == 0 + assert rtx.get_instance_count("mesh") == 3 + + # Replace with just 2 instances at different positions + new_transforms = [ + [1, 0, 0, 5, 0, 1, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 15, 0, 1, 0, 0, 0, 0, 1, 0], + ] + res = rtx.set_transforms("mesh", new_transforms) + assert res == 0 + assert rtx.get_instance_count("mesh") == 2 + + # Old positions should miss + rays_old = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits_old = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays_old, hits_old, 1) + assert res == 0 + assert float(hits_old[0]) == -1.0 # Miss + + # New positions should hit + for x in [5.5, 15.5]: + rays = backend.float32([x, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays, hits, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_update_transform_specific_instance(test_cupy): + """Test updating transform of a specific instance.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Two instances + transforms = [ + [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0], + ] + res = rtx.add_geometry("mesh", verts, tris, transforms=transforms) + assert res == 0 + + # Move second instance (index 1) to x=20 + res = rtx.update_transform("mesh", [1, 0, 0, 20, 0, 1, 0, 0, 0, 0, 1, 0], instance_index=1) + assert res == 0 + + # First instance still at origin + rays0 = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits0 = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays0, hits0, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits0[0]), 10.0, decimal=1) + + # Old position (x=10) should miss + rays_old = backend.float32([10.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits_old = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays_old, hits_old, 1) + assert res == 0 + assert float(hits_old[0]) == -1.0 # Miss + + # New position (x=20) should hit + rays_new = backend.float32([20.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits_new = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays_new, hits_new, 1) + assert res == 0 + np.testing.assert_almost_equal(float(hits_new[0]), 10.0, decimal=1) + + +def test_transforms_validation(): + """Test error handling for invalid transforms.""" + rtx = RTX() + rtx.clear_scene() + + verts = np.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = np.int32([0, 1, 2]) + + # Cannot specify both transform and transforms + with pytest.raises(ValueError, match="Cannot specify both"): + rtx.add_geometry("mesh", verts, tris, + transform=[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], + transforms=[[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0]]) + + # Empty transforms list + with pytest.raises(ValueError, match="cannot be empty"): + rtx.add_geometry("mesh", verts, tris, transforms=[]) + + # Wrong transform length + with pytest.raises(ValueError, match="must have 12 floats"): + rtx.add_geometry("mesh", verts, tris, transforms=[[1, 0, 0]]) # Too short