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