lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/gpu_surface_propagator.py

@@ -0,0 +1,707 @@
+# 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.
+
+"""
+GPU Surface Propagator
+
+Fully GPU-resident propagator with surface detection.
+Uses shared GPUDeviceRays structure for both material propagation
+and surface intersection - no duplicate memory allocation.
+
+Architecture
+------------
+This propagator minimizes CPU-GPU transfers:
+1. One-time upload of ray data at start
+2. GPU-resident loop: propagation + detection + bisection
+3. Small reduction transfer per step (8 bytes) for loop termination
+4. One-time download of results at end
+
+For N=1M rays, 10k steps:
+- Previous: ~240 GB transferred
+- This: ~48 MB transferred (5000x reduction)
+"""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from ..gpu_device_rays import GPUDeviceRays
+from .surface_propagator import HitData
+
+# Import MAX_SURFACE_PARAMS
+MAX_SURFACE_PARAMS = 64
+
+# Speed of light
+SPEED_OF_LIGHT = 299792458.0
+
+# GPU support
+try:
+    from numba import cuda
+
+    _HAS_CUDA = cuda.is_available()
+
+    if _HAS_CUDA:
+
+        @cuda.jit
+        def kernel_homogeneous_step(
+            positions, directions, active, geo_path, opt_path, acc_time, step_size, n, c
+        ):
+            """
+            Simple straight-line propagation for homogeneous materials.
+
+            For homogeneous materials (constant n, zero gradient):
+            - Position advances in straight line
+            - Direction stays constant
+            - Path lengths and time accumulate
+            """
+            idx = cuda.grid(1)
+            if idx >= positions.shape[0]:
+                return
+            if not active[idx]:
+                return
+
+            # Update position (straight line)
+            positions[idx, 0] += directions[idx, 0] * step_size
+            positions[idx, 1] += directions[idx, 1] * step_size
+            positions[idx, 2] += directions[idx, 2] * step_size
+
+            # Update path lengths
+            geo_path[idx] += step_size
+            opt_path[idx] += n * step_size
+            acc_time[idx] += n * step_size / c
+
+        @cuda.jit
+        def kernel_homogeneous_step_adaptive(
+            positions,
+            directions,
+            active,
+            geo_path,
+            opt_path,
+            acc_time,
+            step_sizes,
+            n,
+            c,
+        ):
+            """
+            Straight-line propagation with per-ray adaptive step sizes.
+
+            Same as kernel_homogeneous_step but uses per-ray step_sizes array
+            instead of a scalar step size. This enables adaptive stepping
+            where rays near surfaces use smaller steps for precision.
+            """
+            idx = cuda.grid(1)
+            if idx >= positions.shape[0]:
+                return
+            if not active[idx]:
+                return
+
+            step_size = step_sizes[idx]
+
+            # Update position (straight line)
+            positions[idx, 0] += directions[idx, 0] * step_size
+            positions[idx, 1] += directions[idx, 1] * step_size
+            positions[idx, 2] += directions[idx, 2] * step_size
+
+            # Update path lengths
+            geo_path[idx] += step_size
+            opt_path[idx] += n * step_size
+            acc_time[idx] += n * step_size / c
+
+except ImportError:
+    _HAS_CUDA = False
+
+    class _FakeCuda:
+        @staticmethod
+        def is_available():
+            return False
+
+        @staticmethod
+        def synchronize():
+            pass
+
+        @staticmethod
+        def to_device(arr):
+            return arr
+
+    cuda = _FakeCuda()  # type: ignore[assignment]
+    kernel_homogeneous_step = None  # type: ignore[assignment]
+    kernel_homogeneous_step_adaptive = None  # type: ignore[assignment]
+
+
+if TYPE_CHECKING:
+    from ...utilities.ray_data import RayBatch
+
+
+class GPUSurfacePropagator:
+    """
+    Fully GPU-resident propagator with surface detection.
+
+    Uses shared GPUDeviceRays structure for both material propagation
+    and surface intersection - no duplicate memory allocation.
+
+    Parameters
+    ----------
+    material : MaterialField
+        Material to propagate through. Must support get_gpu_kernels().
+    surfaces : list of Surface
+        List of surfaces to check for intersections. All must be gpu_capable.
+    threads_per_block : int
+        CUDA threads per block. Default: 256.
+    apply_absorption : bool
+        Whether to apply Beer-Lambert absorption during propagation. Default: False.
+
+    Notes
+    -----
+    Falls back to CPU SurfacePropagator if:
+    - CUDA is not available
+    - Material doesn't support GPU propagation
+    - Any surface is not gpu_capable
+
+    Example
+    -------
+    >>> propagator = GPUSurfacePropagator(material, surfaces)
+    >>> hit_data = propagator.propagate_to_surface(rays, step_size=100.0, max_steps=10000)
+    """
+
+    def __init__(
+        self,
+        material,
+        surfaces: list,
+        threads_per_block: int = 256,
+        apply_absorption: bool = False,
+    ):
+        self._material = material
+        self._surfaces = list(surfaces)
+        self._threads_per_block = threads_per_block
+        self._apply_absorption = apply_absorption
+        self._use_gpu = self._check_gpu_support()
+
+        # Initialize absorption kernel (will be set in _setup_gpu_resources if enabled)
+        self._absorption_kernel = None
+
+        if self._use_gpu:
+            self._setup_gpu_resources()
+
+    def _check_gpu_support(self) -> bool:
+        """Check if GPU propagation is possible."""
+        if not _HAS_CUDA:
+            return False
+
+        # Check material: either has GPU kernels OR is homogeneous
+        self._is_homogeneous_material = getattr(
+            self._material, "_is_homogeneous", False
+        )
+        has_gpu_kernels = hasattr(self._material, "get_gpu_kernels")
+
+        if not (has_gpu_kernels or self._is_homogeneous_material):
+            return False
+
+        # Check all surfaces are GPU-capable
+        for surface in self._surfaces:
+            if not getattr(surface, "gpu_capable", False):
+                return False
+
+        return True
+
+    def _setup_gpu_resources(self) -> None:
+        """Pre-upload surface parameters to GPU (one-time setup)."""
+        num_surfaces = len(self._surfaces)
+
+        # Collect geometry IDs
+        geometry_ids = np.array([s.geometry_id for s in self._surfaces], dtype=np.int32)
+        self._d_geometry_ids = cuda.to_device(geometry_ids)
+
+        # Concatenate all surface parameters
+        all_params = np.zeros(num_surfaces * MAX_SURFACE_PARAMS, dtype=np.float32)
+        for i, surface in enumerate(self._surfaces):
+            params = surface.get_gpu_parameters()
+            offset = i * MAX_SURFACE_PARAMS
+            all_params[offset : offset + len(params)] = params
+
+        self._d_surface_params = cuda.to_device(all_params)
+
+        # Upload material LUT arrays to GPU if available (for spectral/grid materials)
+        self._d_material_arrays = {}
+        if hasattr(self._material, "get_gpu_arrays"):
+            gpu_arrays = self._material.get_gpu_arrays()
+            for name, arr in gpu_arrays.items():
+                if arr is not None:
+                    self._d_material_arrays[name] = cuda.to_device(
+                        np.ascontiguousarray(arr, dtype=np.float32)
+                    )
+
+        # Import kernels
+        from ..kernels.surface import (
+            kernel_save_prev_positions,
+            kernel_detect_crossing,
+            kernel_init_signed_distances,
+            kernel_bisect_crossing,
+            kernel_reduce_status,
+            kernel_compute_min_surface_distance,
+            kernel_compute_adaptive_steps,
+        )
+
+        self._kernel_save_prev = kernel_save_prev_positions
+        self._kernel_detect = kernel_detect_crossing
+        self._kernel_init_sd = kernel_init_signed_distances
+        self._kernel_bisect = kernel_bisect_crossing
+        self._kernel_reduce = kernel_reduce_status
+        self._kernel_min_dist = kernel_compute_min_surface_distance
+        self._kernel_adaptive_steps = kernel_compute_adaptive_steps
+
+        # Set up absorption kernel if enabled
+        if self._apply_absorption:
+            self._absorption_kernel = self._get_absorption_kernel()
+
+    def _get_absorption_kernel(self):
+        """Get the appropriate absorption kernel based on material type."""
+        from ..propagator_protocol import GPUMaterialID
+        from ..kernels.absorption.simple import _kernel_absorption_simple
+        from ..kernels.absorption.grid import _kernel_absorption_grid
+        from ..kernels.absorption.spectral import _kernel_absorption_spectral_perray
+
+        material_id = getattr(self._material, "gpu_material_id", None)
+
+        if material_id == GPUMaterialID.SIMPLE_INHOMOGENEOUS:
+            return _kernel_absorption_simple
+        elif material_id == GPUMaterialID.EXPONENTIAL_ATMOSPHERE:
+            return _kernel_absorption_simple  # Uses same LUT structure
+        elif material_id == GPUMaterialID.GRID_INHOMOGENEOUS:
+            return _kernel_absorption_grid
+        elif material_id == GPUMaterialID.SPECTRAL_INHOMOGENEOUS:
+            return _kernel_absorption_spectral_perray
+        else:
+            # No absorption kernel available for this material type
+            return None
+
+    def _apply_absorption_step(
+        self,
+        gpu_rays: "GPUDeviceRays",
+        step_size: float,
+        num_steps: int,
+        blocks: int,
+    ) -> None:
+        """Apply absorption using the appropriate kernel."""
+        from ..propagator_protocol import GPUMaterialID
+
+        if self._absorption_kernel is None:
+            return
+
+        material_params = self._material.get_gpu_parameters()
+        material_id = getattr(self._material, "gpu_material_id", None)
+
+        if material_id in (
+            GPUMaterialID.SIMPLE_INHOMOGENEOUS,
+            GPUMaterialID.EXPONENTIAL_ATMOSPHERE,
+        ):
+            # Simple kernel: positions, directions, active, intensities, optical_depth,
+            # step_size, num_steps, center_x, center_y, center_z, ref_radius,
+            # min_alt, delta_h, lut_size, lut_alpha
+            if "lut_alpha" not in self._d_material_arrays:
+                return  # No absorption LUT available
+            self._absorption_kernel[blocks, self._threads_per_block](
+                gpu_rays.d_positions,
+                gpu_rays.d_directions,
+                gpu_rays.d_active,
+                gpu_rays.d_intensities,
+                gpu_rays.d_optical_depth,
+                float(step_size),
+                num_steps,
+                material_params[0],  # center_x
+                material_params[1],  # center_y
+                material_params[2],  # center_z
+                material_params[3],  # ref_radius
+                material_params[4],  # min_alt
+                material_params[5],  # delta_h
+                int(material_params[6]),  # lut_size
+                self._d_material_arrays["lut_alpha"],
+            )
+        elif material_id == GPUMaterialID.GRID_INHOMOGENEOUS:
+            # Grid kernel: positions, directions, active, intensities, optical_depth,
+            # step_size, num_steps, x_min, y_min, z_min, grid_dx, grid_dy, grid_dz,
+            # nx, ny, nz, alpha_grid
+            if "alpha_grid" not in self._d_material_arrays:
+                return  # No absorption grid available
+            self._absorption_kernel[blocks, self._threads_per_block](
+                gpu_rays.d_positions,
+                gpu_rays.d_directions,
+                gpu_rays.d_active,
+                gpu_rays.d_intensities,
+                gpu_rays.d_optical_depth,
+                float(step_size),
+                num_steps,
+                material_params[0],  # x_min
+                material_params[1],  # y_min
+                material_params[2],  # z_min
+                material_params[3],  # grid_dx
+                material_params[4],  # grid_dy
+                material_params[5],  # grid_dz
+                int(material_params[6]),  # nx
+                int(material_params[7]),  # ny
+                int(material_params[8]),  # nz
+                self._d_material_arrays["alpha_grid"],
+            )
+        elif material_id == GPUMaterialID.SPECTRAL_INHOMOGENEOUS:
+            # Spectral kernel with per-ray wavelengths
+            if "lut_alpha" not in self._d_material_arrays:
+                return  # No absorption LUT available
+            self._absorption_kernel[blocks, self._threads_per_block](
+                gpu_rays.d_positions,
+                gpu_rays.d_directions,
+                gpu_rays.d_wavelengths,
+                gpu_rays.d_active,
+                gpu_rays.d_intensities,
+                gpu_rays.d_optical_depth,
+                float(step_size),
+                num_steps,
+                material_params[0],  # center_x
+                material_params[1],  # center_y
+                material_params[2],  # center_z
+                material_params[3],  # ref_radius
+                material_params[4],  # alt_min
+                material_params[5],  # alt_delta
+                int(material_params[6]),  # n_alt
+                material_params[7],  # wl_min
+                material_params[8],  # wl_delta
+                int(material_params[9]),  # n_wl
+                self._d_material_arrays["lut_alpha"],
+            )
+
+    @property
+    def use_gpu(self) -> bool:
+        """Whether GPU acceleration is active."""
+        return self._use_gpu
+
+    @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
+
+    def propagate_to_surface(
+        self,
+        rays: "RayBatch",
+        step_size: float,
+        max_steps: int,
+        wavelength: float = 532e-9,
+        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).
+
+        Uses GPU-resident loop with minimal CPU-GPU transfers.
+
+        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.
+        wavelength : float, optional
+            Default wavelength for material queries.
+        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:
+            return self._propagate_gpu(
+                rays,
+                step_size,
+                max_steps,
+                wavelength,
+                adaptive_stepping=adaptive_stepping,
+                min_step_size=min_step_size,
+                surface_proximity_factor=surface_proximity_factor,
+                surface_proximity_threshold=surface_proximity_threshold,
+            )
+        else:
+            # Fall back to CPU propagator
+            from .surface_propagator import SurfacePropagator
+
+            cpu_prop = SurfacePropagator(
+                material=self._material,
+                surfaces=self._surfaces,
+                use_gpu=False,
+            )
+            return cpu_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_gpu(
+        self,
+        rays: "RayBatch",
+        step_size: float,
+        max_steps: int,
+        wavelength: float,
+        adaptive_stepping: bool = False,
+        min_step_size: float = 3e-4,
+        surface_proximity_factor: float = 0.5,
+        surface_proximity_threshold: float = 10.0,
+    ) -> HitData:
+        """GPU-resident propagation loop with optional adaptive stepping."""
+        num_rays = rays.num_rays
+        num_surfaces = len(self._surfaces)
+        blocks = (num_rays + self._threads_per_block - 1) // self._threads_per_block
+
+        # === ONE-TIME UPLOAD ===
+        gpu_rays = GPUDeviceRays.from_ray_batch(rays, num_surfaces)
+
+        # Initialize signed distances for all surfaces
+        self._kernel_init_sd[blocks, self._threads_per_block](
+            gpu_rays.d_positions,
+            gpu_rays.d_active,
+            self._d_geometry_ids,
+            self._d_surface_params,
+            num_surfaces,
+            MAX_SURFACE_PARAMS,
+            gpu_rays.d_prev_signed_dist,
+        )
+        cuda.synchronize()
+
+        # Determine propagation mode: homogeneous (straight-line) or gradient-based
+        use_homogeneous = getattr(self, "_is_homogeneous_material", False)
+
+        # Allocate GPU arrays for adaptive stepping if needed
+        if adaptive_stepping:
+            d_min_distances = cuda.to_device(np.zeros(num_rays, dtype=np.float32))
+            d_adaptive_steps = cuda.to_device(
+                np.full(num_rays, step_size, dtype=np.float32)
+            )
+
+        if use_homogeneous:
+            # For homogeneous materials, use simple straight-line kernel
+            # Get refractive index (constant for homogeneous materials)
+            n = self._material.get_refractive_index(0, 0, 0, wavelength)
+        else:
+            # Get material propagation kernel for inhomogeneous materials
+            material_kernels = self._material.get_gpu_kernels()
+            prop_kernel = material_kernels.get("euler") or material_kernels.get("rk4")
+            if prop_kernel is None:
+                raise RuntimeError("Material does not provide GPU propagation kernel")
+            material_params = self._material.get_gpu_parameters()
+            # Get reference refractive index for time correction in bisection
+            # Use value at reference position (e.g., sea level for atmosphere)
+            n = self._material.get_refractive_index(0, 0, 0, wavelength)
+
+        # === GPU-RESIDENT LOOP ===
+        for step in range(max_steps):
+            # 1. Save positions before propagation (for bisection)
+            self._kernel_save_prev[blocks, self._threads_per_block](
+                gpu_rays.d_positions,
+                gpu_rays.d_prev_positions,
+                gpu_rays.d_active,
+            )
+
+            # 1b. Compute adaptive step sizes if enabled
+            if adaptive_stepping:
+                # Compute minimum distance to any surface
+                self._kernel_min_dist[blocks, self._threads_per_block](
+                    gpu_rays.d_positions,
+                    gpu_rays.d_active,
+                    self._d_geometry_ids,
+                    self._d_surface_params,
+                    num_surfaces,
+                    MAX_SURFACE_PARAMS,
+                    d_min_distances,
+                )
+
+                # Compute adaptive step sizes based on distances
+                self._kernel_adaptive_steps[blocks, self._threads_per_block](
+                    d_min_distances,
+                    gpu_rays.d_active,
+                    step_size,
+                    min_step_size,
+                    surface_proximity_factor,
+                    surface_proximity_threshold,
+                    d_adaptive_steps,
+                )
+
+            # 2. Propagate one step using appropriate kernel
+            if use_homogeneous:
+                if adaptive_stepping:
+                    # Use per-ray step sizes
+                    kernel_homogeneous_step_adaptive[blocks, self._threads_per_block](
+                        gpu_rays.d_positions,
+                        gpu_rays.d_directions,
+                        gpu_rays.d_active,
+                        gpu_rays.d_geo_path,
+                        gpu_rays.d_opt_path,
+                        gpu_rays.d_acc_time,
+                        d_adaptive_steps,
+                        float(n),
+                        SPEED_OF_LIGHT,
+                    )
+                else:
+                    # Use fixed step size
+                    kernel_homogeneous_step[blocks, self._threads_per_block](
+                        gpu_rays.d_positions,
+                        gpu_rays.d_directions,
+                        gpu_rays.d_active,
+                        gpu_rays.d_geo_path,
+                        gpu_rays.d_opt_path,
+                        gpu_rays.d_acc_time,
+                        step_size,
+                        float(n),
+                        SPEED_OF_LIGHT,
+                    )
+            else:
+                # Inhomogeneous: use material kernel with gradient-based bending
+                # Note: Adaptive stepping for inhomogeneous materials would require
+                # per-ray step size support in material kernels (future enhancement)
+                #
+                # Build kernel arguments: scalar params + device arrays (LUTs)
+                # Convention matches GPUGradientPropagator:
+                #   - lut_n, lut_dn_dh for SimpleAtmosphere/SpectralAtmosphere
+                #   - n_grid, grad_grid for GridAtmosphere
+                kernel_args = list(material_params)
+                for name in ["lut_n", "lut_dn_dh", "n_grid", "grad_grid"]:
+                    if name in self._d_material_arrays:
+                        kernel_args.append(self._d_material_arrays[name])
+
+                prop_kernel[blocks, self._threads_per_block](
+                    gpu_rays.d_positions,
+                    gpu_rays.d_directions,
+                    gpu_rays.d_active,
+                    gpu_rays.d_geo_path,
+                    gpu_rays.d_opt_path,
+                    gpu_rays.d_acc_time,
+                    step_size,
+                    1,  # num_steps
+                    *kernel_args,
+                    wavelength,
+                )
+
+            # 2b. Apply absorption if enabled
+            if self._apply_absorption and self._absorption_kernel is not None:
+                self._apply_absorption_step(gpu_rays, step_size, 1, blocks)
+
+            # 3. Detect surface crossings
+            self._kernel_detect[blocks, self._threads_per_block](
+                gpu_rays.d_positions,
+                gpu_rays.d_prev_signed_dist,
+                gpu_rays.d_active,
+                self._d_geometry_ids,
+                self._d_surface_params,
+                num_surfaces,
+                MAX_SURFACE_PARAMS,
+                gpu_rays.d_crossing_mask,
+                gpu_rays.d_hit_surface_idx,
+            )
+
+            # 4. Reduce to check status (ONLY transfer: 8 bytes)
+            gpu_rays.d_reduction_result.copy_to_device(np.zeros(2, dtype=np.int32))
+            self._kernel_reduce[blocks, self._threads_per_block](
+                gpu_rays.d_active,
+                gpu_rays.d_crossing_mask,
+                gpu_rays.d_reduction_result,
+            )
+            num_active, num_crossing = gpu_rays.get_reduction_result()
+
+            # 5. Refine crossings if any
+            if num_crossing > 0:
+                # Get step sizes for time correction
+                if adaptive_stepping:
+                    d_step_sizes_for_correction = d_adaptive_steps
+                else:
+                    # Create array with fixed step size for correction
+                    d_step_sizes_for_correction = cuda.to_device(
+                        np.full(num_rays, step_size, dtype=np.float32)
+                    )
+
+                self._kernel_bisect[blocks, self._threads_per_block](
+                    gpu_rays.d_prev_positions,
+                    gpu_rays.d_positions,
+                    gpu_rays.d_crossing_mask,
+                    gpu_rays.d_hit_surface_idx,
+                    self._d_geometry_ids,
+                    self._d_surface_params,
+                    MAX_SURFACE_PARAMS,
+                    gpu_rays.d_hit_positions,
+                    gpu_rays.d_active,
+                    20,  # num_iterations (20 gives ~nm precision from 10m step)
+                    gpu_rays.d_geo_path,  # for time correction
+                    gpu_rays.d_opt_path,
+                    gpu_rays.d_acc_time,
+                    d_step_sizes_for_correction,
+                    float(n),
+                    SPEED_OF_LIGHT,
+                )
+
+            # 6. Early exit if all rays done
+            if num_active == 0:
+                break
+
+            # Reset crossing mask for next iteration
+            gpu_rays.d_crossing_mask.copy_to_device(np.zeros(num_rays, dtype=np.bool_))
+
+        # === ONE-TIME DOWNLOAD ===
+        gpu_rays.to_ray_batch(rays)
+
+        # Build hit data
+        hit_surface_idx, hit_positions, _ = gpu_rays.get_hit_data()
+        hit_directions = gpu_rays.d_directions.copy_to_host()
+
+        return HitData(
+            hit_surface_idx=hit_surface_idx,
+            hit_positions=hit_positions.astype(np.float32),
+            hit_directions=hit_directions.astype(np.float32),
+            num_rays=num_rays,
+        )
+
+
+__all__ = ["GPUSurfacePropagator"]