diff --git a/rtxpy/rtx.py b/rtxpy/rtx.py index 5d57f51..b0466d7 100644 --- a/rtxpy/rtx.py +++ b/rtxpy/rtx.py @@ -8,6 +8,8 @@ import os import atexit import struct +from dataclasses import dataclass, field +from typing import Dict, List, Optional # CRITICAL: cupy must be imported before optix for proper CUDA context sharing import cupy @@ -18,6 +20,24 @@ import numpy as np +# ----------------------------------------------------------------------------- +# Data structures for multi-GAS support +# ----------------------------------------------------------------------------- + +@dataclass +class _GASEntry: + """Storage for a single Geometry Acceleration Structure.""" + 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) + + # ----------------------------------------------------------------------------- # Singleton state management # ----------------------------------------------------------------------------- @@ -37,11 +57,19 @@ def __init__(self): self.hit_pg = None self.sbt = None - # Acceleration structure cache + # Single-GAS mode acceleration structure cache self.gas_handle = 0 self.gas_buffer = None self.current_hash = 0xFFFFFFFFFFFFFFFF # uint64(-1) + # Multi-GAS mode state + self.gas_entries: Dict[str, _GASEntry] = {} # Dict[str, _GASEntry] + self.ias_handle = 0 + self.ias_buffer = None + self.ias_dirty = True + self.instances_buffer = None + self.single_gas_mode = True # False when multi-GAS active + # Device memory for params self.d_params = None @@ -62,11 +90,19 @@ def cleanup(self): self.d_rays_size = 0 self.d_hits_size = 0 - # Free acceleration structure + # Free single-GAS mode acceleration structure self.gas_buffer = None self.gas_handle = 0 self.current_hash = 0xFFFFFFFFFFFFFFFF + # Free multi-GAS mode resources + self.gas_entries = {} + self.ias_handle = 0 + self.ias_buffer = None + self.ias_dirty = True + self.instances_buffer = None + self.single_gas_mode = True + # OptiX objects are automatically cleaned up by Python GC self.sbt = None self.pipeline = None @@ -147,7 +183,7 @@ def _init_optix(): pipeline_options = optix.PipelineCompileOptions( usesMotionBlur=False, - traversableGraphFlags=optix.TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS, + traversableGraphFlags=optix.TRAVERSABLE_GRAPH_FLAG_ALLOW_ANY, numPayloadValues=4, numAttributeValues=2, exceptionFlags=optix.EXCEPTION_FLAG_NONE, @@ -222,7 +258,7 @@ def _init_optix(): dc_from_traversal, dc_from_state, continuation, - 1, # maxTraversableDepth + 2, # maxTraversableDepth (IAS -> GAS = 2 levels) ) # Create shader binding table @@ -278,10 +314,190 @@ def _create_sbt(): # Acceleration structure building # ----------------------------------------------------------------------------- +def _build_gas_for_geometry(vertices, indices): + """ + Build a single GAS (Geometry Acceleration Structure) for the given mesh. + + Args: + vertices: Vertex buffer (Nx3 float32, flattened) + indices: Index buffer (Mx3 int32, flattened) + + Returns: + Tuple of (gas_handle, gas_buffer) or (0, None) on error + """ + global _state + + if not _state.initialized: + _init_optix() + + # Ensure data is on GPU as cupy arrays + if isinstance(vertices, cupy.ndarray): + d_vertices = vertices + else: + d_vertices = cupy.asarray(vertices, dtype=cupy.float32) + + if isinstance(indices, cupy.ndarray): + d_indices = indices + else: + d_indices = cupy.asarray(indices, dtype=cupy.int32) + + # Calculate counts + num_vertices = d_vertices.size // 3 + num_triangles = d_indices.size // 3 + + if num_vertices == 0 or num_triangles == 0: + return 0, None + + # Build input + build_input = optix.BuildInputTriangleArray( + vertexBuffers_=[d_vertices.data.ptr], + vertexFormat=optix.VERTEX_FORMAT_FLOAT3, + vertexStrideInBytes=12, # 3 * sizeof(float) + indexBuffer=d_indices.data.ptr, + numIndexTriplets=num_triangles, + indexFormat=optix.INDICES_FORMAT_UNSIGNED_INT3, + indexStrideInBytes=12, # 3 * sizeof(int) + flags_=[optix.GEOMETRY_FLAG_DISABLE_ANYHIT], + numSbtRecords=1, + ) + build_input.numVertices = num_vertices + + # Acceleration structure options + accel_options = optix.AccelBuildOptions( + buildFlags=optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS, + operation=optix.BUILD_OPERATION_BUILD, + ) + + # Compute memory requirements + buffer_sizes = _state.context.accelComputeMemoryUsage( + [accel_options], + [build_input], + ) + + # Allocate buffers + d_temp = cupy.zeros(buffer_sizes.tempSizeInBytes, dtype=cupy.uint8) + gas_buffer = cupy.zeros(buffer_sizes.outputSizeInBytes, dtype=cupy.uint8) + + # Build acceleration structure + gas_handle = _state.context.accelBuild( + 0, # stream + [accel_options], + [build_input], + d_temp.data.ptr, + buffer_sizes.tempSizeInBytes, + gas_buffer.data.ptr, + buffer_sizes.outputSizeInBytes, + [], # emitted properties + ) + + return gas_handle, gas_buffer + + +def _build_ias(): + """ + 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. + """ + global _state + + if not _state.initialized: + _init_optix() + + if not _state.gas_entries: + _state.ias_handle = 0 + _state.ias_buffer = None + _state.ias_dirty = False + return + + num_instances = len(_state.gas_entries) + + # OptixInstance structure is 80 bytes: + # - transform: float[12] (3x4 row-major) = 48 bytes + # - instanceId: uint32 = 4 bytes + # - sbtOffset: uint32 = 4 bytes + # - visibilityMask: uint32 = 4 bytes + # - flags: uint32 = 4 bytes + # - traversableHandle: uint64 = 8 bytes + # - pad: uint32[2] = 8 bytes + # Total = 80 bytes + + INSTANCE_SIZE = 80 + instances_data = bytearray(num_instances * INSTANCE_SIZE) + + for i, (gas_id, entry) in enumerate(_state.gas_entries.items()): + offset = i * INSTANCE_SIZE + + # Pack transform (12 floats, 48 bytes) + transform_bytes = struct.pack('12f', *entry.transform) + instances_data[offset:offset + 48] = transform_bytes + + # Pack instanceId (4 bytes) + struct.pack_into('I', instances_data, offset + 48, i) + + # 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 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) + + # Padding (8 bytes) - already zeros + + # Copy instances to GPU + _state.instances_buffer = cupy.array( + np.frombuffer(instances_data, dtype=np.uint8) + ) + + # Build input for IAS + build_input = optix.BuildInputInstanceArray( + instances=_state.instances_buffer.data.ptr, + numInstances=num_instances, + ) + + # Acceleration structure options + accel_options = optix.AccelBuildOptions( + buildFlags=optix.BUILD_FLAG_ALLOW_UPDATE, + operation=optix.BUILD_OPERATION_BUILD, + ) + + # Compute memory requirements + buffer_sizes = _state.context.accelComputeMemoryUsage( + [accel_options], + [build_input], + ) + + # Allocate buffers + d_temp = cupy.zeros(buffer_sizes.tempSizeInBytes, dtype=cupy.uint8) + _state.ias_buffer = cupy.zeros(buffer_sizes.outputSizeInBytes, dtype=cupy.uint8) + + # Build IAS + _state.ias_handle = _state.context.accelBuild( + 0, # stream + [accel_options], + [build_input], + d_temp.data.ptr, + buffer_sizes.tempSizeInBytes, + _state.ias_buffer.data.ptr, + buffer_sizes.outputSizeInBytes, + [], # emitted properties + ) + + _state.ias_dirty = False + + def _build_accel(hash_value: int, vertices, indices) -> int: """ Build an OptiX acceleration structure for the given triangle mesh. + This enables single-GAS mode and clears any multi-GAS state. + Args: hash_value: Hash to identify this geometry (for caching) vertices: Vertex buffer (Nx3 float32, flattened) @@ -295,6 +511,15 @@ def _build_accel(hash_value: int, vertices, indices) -> int: if not _state.initialized: _init_optix() + # Clear multi-GAS state when switching to single-GAS mode + if not _state.single_gas_mode: + _state.gas_entries = {} + _state.ias_handle = 0 + _state.ias_buffer = None + _state.ias_dirty = True + _state.instances_buffer = None + _state.single_gas_mode = True + # Check if we already have this acceleration structure cached if _state.current_hash == hash_value: return 0 @@ -374,6 +599,9 @@ def _trace_rays(rays, hits, num_rays: int) -> int: """ Trace rays against the current acceleration structure. + Supports both single-GAS mode (using gas_handle) and multi-GAS mode + (using IAS that references multiple GAS). + Args: rays: Ray buffer (Nx8 float32: ox,oy,oz,tmin,dx,dy,dz,tmax) hits: Hit buffer (Nx4 float32: t,nx,ny,nz) @@ -387,8 +615,18 @@ def _trace_rays(rays, hits, num_rays: int) -> int: if not _state.initialized: return -1 - if _state.gas_handle == 0: - return -1 + # Determine which traversable handle to use + if _state.single_gas_mode: + if _state.gas_handle == 0: + return -1 + trace_handle = _state.gas_handle + else: + # Multi-GAS mode: rebuild IAS if dirty + if _state.ias_dirty: + _build_ias() + if _state.ias_handle == 0: + return -1 + trace_handle = _state.ias_handle # Size check if rays.size != num_rays * 8 or hits.size != num_rays * 4: @@ -424,7 +662,7 @@ def _trace_rays(rays, hits, num_rays: int) -> int: # Pack params: handle(8 bytes) + rays_ptr(8 bytes) + hits_ptr(8 bytes) params_data = struct.pack( 'QQQ', - _state.gas_handle, + trace_handle, d_rays.data.ptr, d_hits.data.ptr, ) @@ -493,6 +731,9 @@ def trace(self, rays, hits, numRays: int) -> int: """ Trace rays against the current acceleration structure. + Works with both single-GAS mode (after build()) and multi-GAS mode + (after add_geometry()). + Args: rays: Ray buffer (8 float32 per ray: ox,oy,oz,tmin,dx,dy,dz,tmax) hits: Hit buffer (4 float32 per hit: t,nx,ny,nz) @@ -503,3 +744,166 @@ def trace(self, rays, hits, numRays: int) -> int: 0 on success, non-zero on error """ return _trace_rays(rays, hits, numRays) + + # ------------------------------------------------------------------------- + # Multi-GAS API + # ------------------------------------------------------------------------- + + def add_geometry(self, geometry_id: str, vertices, indices, + transform: Optional[List[float]] = None) -> int: + """ + Add a geometry (GAS) to the scene with an optional transform. + + This enables multi-GAS mode. If called after build(), the single-GAS + state is cleared. Adding a geometry with an existing ID replaces it. + + 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 + affine transform matrix. Defaults to identity. + Format: [Xx, Xy, Xz, Tx, Yx, Yy, Yz, Ty, Zx, Zy, Zz, Tz] + + Returns: + 0 on success, non-zero on error + """ + global _state + + if not _state.initialized: + _init_optix() + + # Switch to multi-GAS mode if currently in single-GAS mode + if _state.single_gas_mode: + _state.gas_handle = 0 + _state.gas_buffer = None + _state.current_hash = 0xFFFFFFFFFFFFFFFF + _state.single_gas_mode = False + + # Build the GAS for this geometry + gas_handle, gas_buffer = _build_gas_for_geometry(vertices, indices) + if gas_handle == 0: + return -1 + + # Compute a hash for caching purposes + if isinstance(vertices, cupy.ndarray): + vertices_for_hash = vertices.get() + else: + 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: + transform = list(transform) + if len(transform) != 12: + return -1 + + # Create or update the GAS entry + _state.gas_entries[geometry_id] = _GASEntry( + gas_id=geometry_id, + gas_handle=gas_handle, + gas_buffer=gas_buffer, + vertices_hash=vertices_hash, + transform=transform, + ) + + # Mark IAS as needing rebuild + _state.ias_dirty = True + + return 0 + + def remove_geometry(self, geometry_id: str) -> int: + """ + Remove a geometry from the scene. + + Args: + geometry_id: The ID of the geometry to remove + + Returns: + 0 on success, -1 if geometry not found + """ + global _state + + if geometry_id not in _state.gas_entries: + return -1 + + del _state.gas_entries[geometry_id] + _state.ias_dirty = True + + return 0 + + def update_transform(self, geometry_id: str, + transform: List[float]) -> int: + """ + Update the transform of an existing 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] + + Returns: + 0 on success, -1 if geometry not found or invalid transform + """ + global _state + + if geometry_id not in _state.gas_entries: + return -1 + + transform = list(transform) + if len(transform) != 12: + return -1 + + _state.gas_entries[geometry_id].transform = transform + _state.ias_dirty = True + + return 0 + + def list_geometries(self) -> List[str]: + """ + Get a list of all geometry IDs in the scene. + + Returns: + List of geometry ID strings + """ + return list(_state.gas_entries.keys()) + + def get_geometry_count(self) -> int: + """ + Get the number of geometries in the scene. + + Returns: + Number of geometries (0 in single-GAS mode) + """ + return len(_state.gas_entries) + + def clear_scene(self) -> None: + """ + Remove all geometries and reset to single-GAS mode. + + After calling this, you can use either build() for single-GAS mode + or add_geometry() for multi-GAS mode. + """ + global _state + + # Clear multi-GAS state + _state.gas_entries = {} + _state.ias_handle = 0 + _state.ias_buffer = None + _state.ias_dirty = True + _state.instances_buffer = None + + # Clear single-GAS state + _state.gas_handle = 0 + _state.gas_buffer = None + _state.current_hash = 0xFFFFFFFFFFFFFFFF + + # Reset to single-GAS mode + _state.single_gas_mode = True diff --git a/rtxpy/tests/test_simple.py b/rtxpy/tests/test_simple.py index a8cf699..42b4ff0 100644 --- a/rtxpy/tests/test_simple.py +++ b/rtxpy/tests/test_simple.py @@ -414,3 +414,516 @@ def test_nan_in_elevation_data_sparse(test_cupy, dtype): # Near NaN - should not crash t_nan_area = float(hits_nan_area[0]) assert np.isfinite(t_nan_area) or np.isnan(t_nan_area) or t_nan_area == -1.0 + + +# ============================================================================= +# Multi-GAS Tests +# ============================================================================= + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_multi_gas_two_meshes(test_cupy): + """Test tracing against two meshes at different Z heights.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + # Two triangles: one at z=0, one at z=5 + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Add ground mesh at z=0 + res = rtx.add_geometry("ground", verts, tris) + assert res == 0 + + # Add elevated mesh at z=5 using transform + # Transform: identity rotation, translation (0, 0, 5) + transform = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 5] + res = rtx.add_geometry("elevated", verts, tris, transform=transform) + assert res == 0 + + # Ray pointing down from z=10 at the triangle center + rays = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + + res = rtx.trace(rays, hits, 1) + assert res == 0 + + # Should hit the elevated mesh first at z=5 (distance ~5) + t_value = float(hits[0]) + np.testing.assert_almost_equal(t_value, 5.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_multi_gas_with_transform(test_cupy): + """Test geometry with translation transform.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + # Triangle at origin + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Translate by (10, 0, 0) + transform = [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0] + res = rtx.add_geometry("translated", verts, tris, transform=transform) + assert res == 0 + + # Ray pointing down at (10.5, 0.33, 10) - should hit translated mesh + rays = backend.float32([10.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + + res = rtx.trace(rays, hits, 1) + assert res == 0 + + t_value = float(hits[0]) + np.testing.assert_almost_equal(t_value, 10.0, decimal=1) + + # Ray at original position should miss + rays_miss = backend.float32([0.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_multi_gas_many_geometries(test_cupy): + """Stress test with many geometries.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + # Small triangle + verts = backend.float32([0, 0, 0, 0.5, 0, 0, 0.25, 0.5, 0]) + tris = backend.int32([0, 1, 2]) + + # Add 100 geometries in a 10x10 grid + num_geoms = 100 + for i in range(num_geoms): + x = (i % 10) * 2 + y = (i // 10) * 2 + transform = [1, 0, 0, x, 0, 1, 0, y, 0, 0, 1, 0] + res = rtx.add_geometry(f"mesh_{i}", verts, tris, transform=transform) + assert res == 0 + + assert rtx.get_geometry_count() == num_geoms + + # Trace a ray at one of the geometries + rays = backend.float32([4.25, 4.25, 10, 0, 0, 0, -1, 1000]) # Should hit mesh_22 + hits = backend.float32([0, 0, 0, 0]) + + res = rtx.trace(rays, hits, 1) + assert res == 0 + + t_value = float(hits[0]) + np.testing.assert_almost_equal(t_value, 10.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_remove_geometry(test_cupy): + """Test adding and removing 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() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Add two geometries + rtx.add_geometry("mesh1", verts, tris) + rtx.add_geometry("mesh2", verts, tris, transform=[1, 0, 0, 5, 0, 1, 0, 0, 0, 0, 1, 0]) + + assert rtx.get_geometry_count() == 2 + assert "mesh1" in rtx.list_geometries() + assert "mesh2" in rtx.list_geometries() + + # Remove one + res = rtx.remove_geometry("mesh1") + assert res == 0 + + assert rtx.get_geometry_count() == 1 + assert "mesh1" not in rtx.list_geometries() + assert "mesh2" in rtx.list_geometries() + + # Remove non-existent should fail + res = rtx.remove_geometry("nonexistent") + assert res == -1 + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_replace_geometry(test_cupy): + """Test adding geometry with the same ID replaces it.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts1 = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + verts2 = backend.float32([0, 0, 5, 1, 0, 5, 0.5, 1, 5]) # At z=5 + tris = backend.int32([0, 1, 2]) + + # Add initial geometry at z=0 + rtx.add_geometry("mesh", verts1, tris) + + # Ray should hit at z=0 (distance 10) + rays = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + rtx.trace(rays, hits, 1) + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + # Replace with geometry at z=5 + rtx.add_geometry("mesh", verts2, tris) + assert rtx.get_geometry_count() == 1 # Still only one geometry + + # Now should hit at z=5 (distance 5) + hits = backend.float32([0, 0, 0, 0]) + rtx.trace(rays, hits, 1) + np.testing.assert_almost_equal(float(hits[0]), 5.0, decimal=1) + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_update_transform(test_cupy): + """Test updating transform of 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() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Add geometry at origin + rtx.add_geometry("mesh", verts, tris) + + # Ray at origin hits + rays_origin = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + rtx.trace(rays_origin, hits, 1) + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + # Update transform to translate by (10, 0, 0) + res = rtx.update_transform("mesh", [1, 0, 0, 10, 0, 1, 0, 0, 0, 0, 1, 0]) + assert res == 0 + + # Now ray at origin should miss + hits = backend.float32([0, 0, 0, 0]) + rtx.trace(rays_origin, hits, 1) + assert float(hits[0]) == -1.0 + + # Ray at new position should hit + rays_new = backend.float32([10.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + rtx.trace(rays_new, hits, 1) + np.testing.assert_almost_equal(float(hits[0]), 10.0, decimal=1) + + # Update non-existent should fail + res = rtx.update_transform("nonexistent", [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0]) + assert res == -1 + + # Invalid transform length should fail + res = rtx.update_transform("mesh", [1, 0, 0]) + assert res == -1 + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_list_geometries(test_cupy): + """Test list_geometries and get_geometry_count methods.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Initially empty + assert rtx.get_geometry_count() == 0 + assert rtx.list_geometries() == [] + + # Add geometries + rtx.add_geometry("a", verts, tris) + rtx.add_geometry("b", verts, tris) + rtx.add_geometry("c", verts, tris) + + assert rtx.get_geometry_count() == 3 + geoms = rtx.list_geometries() + assert "a" in geoms + assert "b" in geoms + assert "c" in geoms + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_clear_scene(test_cupy): + """Test clear_scene removes all geometry and resets state.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Add geometries + rtx.add_geometry("mesh1", verts, tris) + rtx.add_geometry("mesh2", verts, tris) + assert rtx.get_geometry_count() == 2 + + # Clear scene + rtx.clear_scene() + assert rtx.get_geometry_count() == 0 + assert rtx.list_geometries() == [] + + # Can use build() after clear + res = rtx.build(123, verts, tris) + assert res == 0 + assert rtx.getHash() == 123 + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_backward_compat_single_gas(test_cupy): + """Test that existing single-GAS build() API still works.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Use original API + res = rtx.build(12345, verts, tris) + assert res == 0 + assert rtx.getHash() == 12345 + + # Trace should work + rays = backend.float32([0.5, 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_switch_multi_to_single(test_cupy): + """Test that build() clears multi-GAS state.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Start with multi-GAS + rtx.add_geometry("mesh1", verts, tris) + rtx.add_geometry("mesh2", verts, tris, transform=[1, 0, 0, 5, 0, 1, 0, 0, 0, 0, 1, 0]) + assert rtx.get_geometry_count() == 2 + + # Switch to single-GAS with build() + res = rtx.build(999, verts, tris) + assert res == 0 + + # Multi-GAS state should be cleared + assert rtx.get_geometry_count() == 0 + + # Trace should use single-GAS + rays = backend.float32([0.5, 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_switch_single_to_multi(test_cupy): + """Test that add_geometry() clears single-GAS state.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Start with single-GAS + rtx.build(888, verts, tris) + assert rtx.getHash() == 888 + + # Switch to multi-GAS + rtx.add_geometry("mesh", verts, tris) + + # Single-GAS hash should be cleared + assert rtx.getHash() == 0xFFFFFFFFFFFFFFFF + + # Trace should use multi-GAS (IAS) + rays = backend.float32([0.5, 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_empty_scene(test_cupy): + """Test behavior when tracing after removing all 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() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + # Add and then remove all geometry + rtx.add_geometry("mesh", verts, tris) + rtx.remove_geometry("mesh") + assert rtx.get_geometry_count() == 0 + + # Trace should fail gracefully (return error) + rays = backend.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + res = rtx.trace(rays, hits, 1) + assert res == -1 # No geometry to trace against + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_trace_miss_multi_gas(test_cupy): + """Test ray that misses all geometries in multi-GAS mode.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + backend = cupy + else: + backend = np + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + verts = backend.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = backend.int32([0, 1, 2]) + + rtx.add_geometry("mesh", verts, tris) + + # Ray that misses the geometry + rays = backend.float32([100, 100, 10, 0, 0, 0, -1, 1000]) + hits = backend.float32([0, 0, 0, 0]) + + res = rtx.trace(rays, hits, 1) + assert res == 0 + assert float(hits[0]) == -1.0 # Miss + + +@pytest.mark.parametrize("test_cupy", [False, True]) +def test_cupy_buffers_multi_gas(test_cupy): + """Test multi-GAS mode works with cupy buffers for rays/hits.""" + if test_cupy: + if not has_cupy: + pytest.skip("cupy not available") + import cupy + + rtx = RTX() + rtx.clear_scene() # Clear any state from previous tests + + # Use numpy arrays for vertex/triangle data (works with both backends) + verts = np.float32([0, 0, 0, 1, 0, 0, 0.5, 1, 0]) + tris = np.int32([0, 1, 2]) + + rtx.add_geometry("mesh", verts, tris) + + # Create ray/hit buffers on the appropriate backend + if test_cupy: + rays = cupy.array([0.5, 0.33, 10, 0, 0, 0, -1, 1000], dtype=cupy.float32) + hits = cupy.zeros(4, dtype=cupy.float32) + else: + rays = np.float32([0.5, 0.33, 10, 0, 0, 0, -1, 1000]) + hits = np.float32([0, 0, 0, 0]) + + res = rtx.trace(rays, hits, 1) + assert res == 0 + + # Convert to numpy for comparison + if test_cupy: + hits_np = hits.get() + else: + hits_np = hits + + np.testing.assert_almost_equal(hits_np[0], 10.0, decimal=1)