lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/surface_propagator.py

@@ -0,0 +1,719 @@
+# The Clear BSD License
+#
+# Copyright (c) 2026 Tobias Heibges
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+#
+#      * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#
+#      * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#
+#      * Neither the name of the copyright holder nor the names of its
+#      contributors may be used to endorse or promote products derived from this
+#      software without specific prior written permission.
+#
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
+# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+Surface Propagator for GPU Ray Tracing
+
+Orchestrates GPU kernel launches for ray propagation with multi-surface
+detection. Uses separate kernels for propagation and intersection,
+with GPU memory management handled in the propagator layer.
+
+Architecture
+------------
+This propagator follows the clean kernel architecture:
+1. Kernels = Pure CUDA device functions + parallel kernels (no memory management)
+2. Propagators = Orchestration + GPU memory management
+
+The propagation loop:
+1. Take one propagation step for all active rays
+2. Check intersection with all surfaces
+3. Record hits and deactivate rays that hit
+4. Repeat until all rays hit or max_steps reached
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+from .signed_distance_handler import (
+    compute_signed_distance_gpu,
+)
+
+# GPU support is optional
+try:
+    from numba import cuda
+
+    HAS_CUDA = True
+except ImportError:
+
+    class _FakeCuda:
+        """Fake cuda module for when numba is not installed."""
+
+        @staticmethod
+        def is_available():
+            return False
+
+        @staticmethod
+        def synchronize():
+            pass
+
+        @staticmethod
+        def to_device(arr):
+            return _FakeDeviceArray(arr)
+
+    class _FakeDeviceArray:
+        def __init__(self, arr):
+            self._arr = arr
+
+        def copy_to_host(self):
+            return self._arr
+
+    cuda = _FakeCuda()  # type: ignore[assignment]
+    HAS_CUDA = False
+
+
+if TYPE_CHECKING:
+    from ...utilities.ray_data import RayBatch
+
+# 64 params allows up to 8 wave components for multi-wave surfaces
+MAX_SURFACE_PARAMS = 64
+
+# Speed of light
+SPEED_OF_LIGHT = 299792458.0
+
+
+@dataclass
+class HitData:
+    """
+    Data about ray-surface intersections.
+
+    Attributes
+    ----------
+    hit_surface_idx : ndarray, shape (N,)
+        Index of surface hit for each ray (-1 = no hit)
+    hit_positions : ndarray, shape (N, 3)
+        Position of each intersection
+    hit_directions : ndarray, shape (N, 3)
+        Direction at each intersection
+    num_rays : int
+        Total number of rays
+    """
+
+    hit_surface_idx: npt.NDArray[np.int32]
+    hit_positions: npt.NDArray[np.float32]
+    hit_directions: npt.NDArray[np.float32]
+    num_rays: int
+
+    def get_hit_mask(self, surface_idx: int) -> npt.NDArray[np.bool_]:
+        """Return mask of rays that hit a specific surface."""
+        return self.hit_surface_idx == surface_idx
+
+    def get_no_hit_mask(self) -> npt.NDArray[np.bool_]:
+        """Return mask of rays that didn't hit any surface."""
+        return self.hit_surface_idx == -1
+
+    def count_hits(self, surface_idx: int) -> int:
+        """Count rays that hit a specific surface."""
+        return int(np.sum(self.hit_surface_idx == surface_idx))
+
+    def count_no_hits(self) -> int:
+        """Count rays that didn't hit any surface."""
+        return int(np.sum(self.hit_surface_idx == -1))
+
+
+class SurfacePropagator:
+    """
+    GPU propagator with multi-surface intersection detection.
+
+    This propagator advances rays through a material while checking for
+    intersections with multiple surfaces. When a ray crosses any surface,
+    the intersection point is recorded and the ray is deactivated.
+    The hit data is returned for CPU post-processing based on surface role.
+
+    Parameters
+    ----------
+    material : GPUMaterialProtocol
+        The material to propagate through. Must support get_refractive_index()
+        and get_refractive_index_gradient() methods.
+    surfaces : list of Surface
+        List of surfaces to check for intersections.
+    method : str
+        Integration method: 'euler' (default).
+    threads_per_block : int
+        CUDA threads per block. Default: 256.
+    use_gpu : bool, optional
+        Whether to use GPU acceleration. If True (default), uses CUDA if
+        available, otherwise falls back to CPU.
+
+    Notes
+    -----
+    This propagator uses a step-by-step approach:
+    1. Propagate rays one step
+    2. Check all surfaces for intersection
+    3. Record hits and deactivate hitting rays
+    4. Repeat until done
+    """
+
+    def __init__(
+        self,
+        material,
+        surfaces: list,
+        method: str = "euler",
+        threads_per_block: int = 256,
+        use_gpu: bool = True,
+        apply_absorption: bool = False,
+    ):
+        if method not in ("euler",):
+            raise ValueError(f"method must be 'euler', got {method}")
+
+        self._material = material
+        self._surfaces = list(surfaces)
+        self.method = method
+        self.threads_per_block = threads_per_block
+        self._use_gpu = use_gpu and HAS_CUDA and cuda.is_available()
+        self._apply_absorption = apply_absorption
+
+    @property
+    def material(self):
+        """The material being propagated through."""
+        return self._material
+
+    @property
+    def surfaces(self) -> list:
+        """List of surfaces being checked for intersection."""
+        return self._surfaces
+
+    @property
+    def num_surfaces(self) -> int:
+        """Number of surfaces."""
+        return len(self._surfaces)
+
+    def _compute_min_surface_distance(
+        self,
+        positions: np.ndarray,
+        active_mask: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Compute minimum absolute signed distance to any surface for each ray.
+
+        Parameters
+        ----------
+        positions : ndarray, shape (N, 3)
+            Ray positions.
+        active_mask : ndarray, shape (N,)
+            Boolean mask for active rays.
+
+        Returns
+        -------
+        ndarray, shape (N,)
+            Minimum distance to any surface for each ray. Inactive rays get inf.
+        """
+        num_rays = len(positions)
+        min_distances = np.full(num_rays, np.inf, dtype=np.float32)
+
+        if not np.any(active_mask):
+            return min_distances
+
+        active_indices = np.where(active_mask)[0]
+        active_positions = positions[active_indices].astype(np.float32)
+
+        for surface in self._surfaces:
+            if hasattr(surface, "signed_distance"):
+                # Use GPU if available and surface supports it
+                if (
+                    self._use_gpu
+                    and hasattr(surface, "gpu_capable")
+                    and surface.gpu_capable
+                    and hasattr(surface, "get_gpu_parameters")
+                ):
+                    geometry_id = surface.geometry_id
+                    params = surface.get_gpu_parameters()
+                    sd = compute_signed_distance_gpu(
+                        active_positions, geometry_id, params
+                    )
+                else:
+                    sd = surface.signed_distance(active_positions)
+
+                # Take absolute value (we care about distance, not which side)
+                abs_sd = np.abs(sd)
+                min_distances[active_indices] = np.minimum(
+                    min_distances[active_indices], abs_sd
+                )
+
+        return min_distances
+
+    def propagate_to_surface(
+        self,
+        rays: "RayBatch",
+        step_size: float,
+        max_steps: int,
+        adaptive_stepping: bool = False,
+        min_step_size: float = 3e-4,
+        surface_proximity_factor: float = 0.5,
+        surface_proximity_threshold: float = 10.0,
+    ) -> HitData:
+        """
+        Propagate rays until any surface is hit (or max_steps reached).
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Ray batch to propagate. Modified in-place.
+        step_size : float
+            Maximum integration step size in meters.
+        max_steps : int
+            Maximum number of propagation steps.
+        adaptive_stepping : bool, optional
+            Whether to use adaptive step sizing near surfaces (default False).
+        min_step_size : float, optional
+            Minimum step size in meters (default 3e-4 = 0.3mm → ~1ps resolution).
+        surface_proximity_factor : float, optional
+            Step = distance * factor when within threshold (default 0.5).
+        surface_proximity_threshold : float, optional
+            Distance within which to start adaptive stepping (default 10.0 m).
+
+        Returns
+        -------
+        HitData
+            Information about which rays hit which surfaces.
+        """
+        if self._use_gpu and self._all_surfaces_gpu_capable():
+            return self._propagate_gpu(
+                rays,
+                step_size,
+                max_steps,
+                adaptive_stepping=adaptive_stepping,
+                min_step_size=min_step_size,
+                surface_proximity_factor=surface_proximity_factor,
+                surface_proximity_threshold=surface_proximity_threshold,
+            )
+        return self._propagate_cpu(
+            rays,
+            step_size,
+            max_steps,
+            adaptive_stepping=adaptive_stepping,
+            min_step_size=min_step_size,
+            surface_proximity_factor=surface_proximity_factor,
+            surface_proximity_threshold=surface_proximity_threshold,
+        )
+
+    def _all_surfaces_gpu_capable(self) -> bool:
+        """Check if all surfaces support GPU acceleration."""
+        return all(getattr(s, "gpu_capable", False) for s in self._surfaces)
+
+    def _propagate_gpu(
+        self,
+        rays: "RayBatch",
+        step_size: float,
+        max_steps: int,
+        adaptive_stepping: bool = False,
+        min_step_size: float = 3e-4,
+        surface_proximity_factor: float = 0.5,
+        surface_proximity_threshold: float = 10.0,
+    ) -> HitData:
+        """GPU-accelerated propagation using GPUSurfacePropagator."""
+        from .gpu_surface_propagator import GPUSurfacePropagator
+
+        gpu_prop = GPUSurfacePropagator(
+            material=self._material,
+            surfaces=self._surfaces,
+            apply_absorption=self._apply_absorption,
+        )
+        return gpu_prop.propagate_to_surface(
+            rays,
+            step_size,
+            max_steps,
+            adaptive_stepping=adaptive_stepping,
+            min_step_size=min_step_size,
+            surface_proximity_factor=surface_proximity_factor,
+            surface_proximity_threshold=surface_proximity_threshold,
+        )
+
+    def _propagate_cpu(
+        self,
+        rays: "RayBatch",
+        step_size: float,
+        max_steps: int,
+        adaptive_stepping: bool = False,
+        min_step_size: float = 3e-4,
+        surface_proximity_factor: float = 0.5,
+        surface_proximity_threshold: float = 10.0,
+    ) -> HitData:
+        """
+        CPU implementation of propagation with surface detection.
+
+        Uses Euler integration for ray propagation through the material,
+        checking for surface intersections at each step.
+
+        When adaptive_stepping is enabled, step sizes are reduced as rays
+        approach surfaces to achieve sub-nanosecond timing precision.
+        """
+        num_rays = rays.num_rays
+        c = SPEED_OF_LIGHT
+
+        # Initialize hit tracking
+        hit_surface_idx = np.full(num_rays, -1, dtype=np.int32)
+        hit_positions = np.zeros((num_rays, 3), dtype=np.float32)
+        hit_directions = np.zeros((num_rays, 3), dtype=np.float32)
+
+        # Get working arrays
+        positions = rays.positions.copy()
+        directions = rays.directions.copy()
+        active = rays.active.copy()
+        geo_path = rays.geometric_path_lengths.copy()
+        opt_path = rays.optical_path_lengths.copy()
+        acc_time = rays.accumulated_time.copy()
+
+        # Per-ray step sizes (for adaptive stepping)
+        ray_step_sizes = np.full(num_rays, step_size, dtype=np.float32)
+
+        # Default wavelength for material queries
+        default_wavelength = 500e-9
+
+        for step in range(max_steps):
+            # Check if any rays are still active
+            if not np.any(active):
+                break
+
+            # Compute adaptive step sizes if enabled
+            if adaptive_stepping:
+                min_distances = self._compute_min_surface_distance(positions, active)
+
+                # Vectorized adaptive step computation
+                # For rays within threshold: step = distance * factor, clamped
+                # For rays outside threshold: use max step_size
+                within_threshold = min_distances < surface_proximity_threshold
+                adaptive_steps = np.where(
+                    within_threshold,
+                    np.clip(
+                        min_distances * surface_proximity_factor,
+                        min_step_size,
+                        step_size,
+                    ),
+                    step_size,
+                )
+                ray_step_sizes = adaptive_steps.astype(np.float32)
+
+            # Store previous positions for intersection detection
+            prev_positions = positions.copy()
+
+            # Propagate each active ray one step
+            for i in range(num_rays):
+                if not active[i]:
+                    continue
+
+                x, y, z = positions[i]
+                dx, dy, dz = directions[i]
+
+                # Get per-ray step size
+                current_step = ray_step_sizes[i]
+
+                # Get material properties
+                n = self._material.get_refractive_index(x, y, z, default_wavelength)
+                gx, gy, gz = self._material.get_refractive_index_gradient(
+                    x, y, z, default_wavelength
+                )
+
+                # Euler step: update direction (ray curvature in gradient field)
+                dot = dx * gx + dy * gy + dz * gz
+                kx = (gx - dot * dx) / n
+                ky = (gy - dot * dy) / n
+                kz = (gz - dot * dz) / n
+
+                # Update position using adaptive step
+                positions[i, 0] = x + dx * current_step
+                positions[i, 1] = y + dy * current_step
+                positions[i, 2] = z + dz * current_step
+
+                # Update direction
+                dx_new = dx + kx * current_step
+                dy_new = dy + ky * current_step
+                dz_new = dz + kz * current_step
+
+                # Normalize direction
+                norm = np.sqrt(dx_new**2 + dy_new**2 + dz_new**2)
+                if norm > 1e-12:
+                    directions[i] = [dx_new / norm, dy_new / norm, dz_new / norm]
+
+                # Update path lengths with adaptive step
+                geo_path[i] += current_step
+                opt_path[i] += n * current_step
+                acc_time[i] += n * current_step / c
+
+            # Check each surface for intersections
+            for surf_idx, surface in enumerate(self._surfaces):
+                if not np.any(active):
+                    break
+
+                # Get active ray indices
+                active_indices = np.where(active)[0]
+                if len(active_indices) == 0:
+                    continue
+
+                # Compute signed distances for all active rays (vectorized)
+                active_prev_pos = prev_positions[active_indices].astype(np.float32)
+                active_curr_pos = positions[active_indices].astype(np.float32)
+
+                # Use GPU signed distance if surface supports it and GPU is enabled
+                if (
+                    self._use_gpu
+                    and hasattr(surface, "gpu_capable")
+                    and surface.gpu_capable
+                    and hasattr(surface, "get_gpu_parameters")
+                ):
+                    geometry_id = surface.geometry_id
+                    params = surface.get_gpu_parameters()
+                    prev_sd = compute_signed_distance_gpu(
+                        active_prev_pos, geometry_id, params
+                    )
+                    curr_sd = compute_signed_distance_gpu(
+                        active_curr_pos, geometry_id, params
+                    )
+                else:
+                    # CPU fallback using surface's own method
+                    prev_sd = surface.signed_distance(active_prev_pos)
+                    curr_sd = surface.signed_distance(active_curr_pos)
+
+                # Check for sign change (crossing) - vectorized
+                crossing = ((prev_sd >= 0) & (curr_sd < 0)) | (
+                    (prev_sd < 0) & (curr_sd >= 0)
+                )
+
+                if not np.any(crossing):
+                    continue
+
+                # Process rays that crossed this surface
+                crossing_indices = active_indices[crossing]
+
+                for i in crossing_indices:
+                    local_idx = np.where(active_indices == i)[0][0]
+
+                    # Bisect to find exact intersection point
+                    p0 = prev_positions[i].copy()
+                    p1 = positions[i].copy()
+                    sd0 = prev_sd[local_idx]
+
+                    for _ in range(
+                        20
+                    ):  # 20 iterations gives ~nm precision from 10m step
+                        mid = (p0 + p1) / 2
+
+                        # Get signed distance at midpoint
+                        if (
+                            self._use_gpu
+                            and hasattr(surface, "gpu_capable")
+                            and surface.gpu_capable
+                        ):
+                            geometry_id = surface.geometry_id
+                            params = surface.get_gpu_parameters()
+                            mid_sd = compute_signed_distance_gpu(
+                                mid[np.newaxis, :].astype(np.float32),
+                                geometry_id,
+                                params,
+                            )[0]
+                        else:
+                            mid_sd = surface.signed_distance(
+                                mid[np.newaxis, :].astype(np.float32)
+                            )[0]
+
+                        if abs(mid_sd) < 1e-6:
+                            break
+
+                        # Determine which half contains the crossing
+                        if (sd0 >= 0 and mid_sd < 0) or (sd0 < 0 and mid_sd >= 0):
+                            p1 = mid
+                        else:
+                            p0 = mid
+                            sd0 = mid_sd
+
+                    # Record hit
+                    hit_pos = (p0 + p1) / 2
+                    hit_surface_idx[i] = surf_idx
+                    hit_positions[i] = hit_pos
+                    hit_directions[i] = directions[i]
+
+                    # Correct accumulated time/path for exact intersection distance
+                    # The step already accumulated time for current_step, but
+                    # actual distance traveled is less (to hit_pos, not full step)
+                    actual_distance = np.linalg.norm(hit_pos - prev_positions[i])
+                    excess_distance = ray_step_sizes[i] - actual_distance
+
+                    # Get refractive index at intersection for correction
+                    # (use hit_pos which is close enough for correction purposes)
+                    n_hit = self._material.get_refractive_index(
+                        hit_pos[0], hit_pos[1], hit_pos[2], default_wavelength
+                    )
+
+                    # Remove excess accumulated values
+                    geo_path[i] -= excess_distance
+                    opt_path[i] -= n_hit * excess_distance
+                    acc_time[i] -= n_hit * excess_distance / c
+
+                    # Update ray state to intersection point
+                    positions[i] = hit_pos
+                    active[i] = False
+
+        # Copy final state back to rays
+        rays.positions[:] = positions
+        rays.directions[:] = directions
+        rays.active[:] = active
+        rays.geometric_path_lengths[:] = geo_path
+        rays.optical_path_lengths[:] = opt_path
+        rays.accumulated_time[:] = acc_time
+
+        return HitData(
+            hit_surface_idx=hit_surface_idx,
+            hit_positions=hit_positions,
+            hit_directions=hit_directions,
+            num_rays=num_rays,
+        )
+
+    def get_surface_by_index(self, idx: int):
+        """Get surface by index."""
+        return self._surfaces[idx]
+
+    def extract_hits_for_surface(
+        self,
+        rays: "RayBatch",
+        hit_data: HitData,
+        surface_idx: int,
+    ) -> "RayBatch":
+        """
+        Extract rays that hit a specific surface as a new RayBatch.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Original ray batch
+        hit_data : HitData
+            Hit data from propagation
+        surface_idx : int
+            Index of surface to extract hits for
+
+        Returns
+        -------
+        RayBatch
+            New batch containing only rays that hit the specified surface.
+            Positions and directions are set to intersection values.
+        """
+        from ...utilities.ray_data import RayBatch
+
+        mask = hit_data.get_hit_mask(surface_idx)
+        num_hits = np.sum(mask)
+
+        if num_hits == 0:
+            # Return empty batch
+            from ...utilities.ray_data import create_ray_batch
+
+            return create_ray_batch(0)
+
+        # Create new batch with intersection data
+        return RayBatch(
+            positions=hit_data.hit_positions[mask].copy(),
+            directions=hit_data.hit_directions[mask].copy(),
+            wavelengths=rays.wavelengths[mask].copy(),
+            intensities=rays.intensities[mask].copy(),
+            optical_path_lengths=rays.optical_path_lengths[mask].copy(),
+            geometric_path_lengths=rays.geometric_path_lengths[mask].copy(),
+            accumulated_time=rays.accumulated_time[mask].copy(),
+            generations=rays.generations[mask].copy(),
+            domain_ids=rays.domain_ids[mask].copy(),
+            active=np.ones(num_hits, dtype=np.bool_),
+            polarization_s=(
+                rays.polarization_s[mask].copy()
+                if rays.polarization_s is not None
+                else None
+            ),
+            polarization_p=(
+                rays.polarization_p[mask].copy()
+                if rays.polarization_p is not None
+                else None
+            ),
+            polarization_vector=(
+                rays.polarization_vector[mask].copy()
+                if rays.polarization_vector is not None
+                else None
+            ),
+            phase=rays.phase[mask].copy() if rays.phase is not None else None,
+        )
+
+    def extract_no_hits(
+        self,
+        rays: "RayBatch",
+        hit_data: HitData,
+    ) -> "RayBatch":
+        """
+        Extract rays that didn't hit any surface.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Original ray batch
+        hit_data : HitData
+            Hit data from propagation
+
+        Returns
+        -------
+        RayBatch
+            New batch containing only rays that didn't hit any surface.
+        """
+        from ...utilities.ray_data import RayBatch
+
+        mask = hit_data.get_no_hit_mask()
+        num_no_hits = np.sum(mask)
+
+        if num_no_hits == 0:
+            from ...utilities.ray_data import create_ray_batch
+
+            return create_ray_batch(0)
+
+        return RayBatch(
+            positions=rays.positions[mask].copy(),
+            directions=rays.directions[mask].copy(),
+            wavelengths=rays.wavelengths[mask].copy(),
+            intensities=rays.intensities[mask].copy(),
+            optical_path_lengths=rays.optical_path_lengths[mask].copy(),
+            geometric_path_lengths=rays.geometric_path_lengths[mask].copy(),
+            accumulated_time=rays.accumulated_time[mask].copy(),
+            generations=rays.generations[mask].copy(),
+            domain_ids=rays.domain_ids[mask].copy(),
+            active=rays.active[mask].copy(),
+            polarization_s=(
+                rays.polarization_s[mask].copy()
+                if rays.polarization_s is not None
+                else None
+            ),
+            polarization_p=(
+                rays.polarization_p[mask].copy()
+                if rays.polarization_p is not None
+                else None
+            ),
+            polarization_vector=(
+                rays.polarization_vector[mask].copy()
+                if rays.polarization_vector is not None
+                else None
+            ),
+            phase=rays.phase[mask].copy() if rays.phase is not None else None,
+        )