lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/spectral_gpu_gradient.py

@@ -0,0 +1,553 @@
+# 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-accelerated Spectral Gradient Propagator with per-ray wavelengths.
+
+This module provides a CUDA-based propagator for SpectralInhomogeneousModel
+materials where each ray has its own wavelength value, enabling simulation
+of chromatic dispersion.
+"""
+
+import numpy as np
+import numpy.typing as npt
+
+# 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 jit(*args, **kwargs):
+            """Return a no-op decorator."""
+
+            def decorator(func):
+                return func
+
+            if args and callable(args[0]):
+                return args[0]
+            return decorator
+
+        @staticmethod
+        def is_available():
+            return False
+
+        @staticmethod
+        def grid(n):
+            return 0
+
+        @staticmethod
+        def synchronize():
+            pass
+
+        @staticmethod
+        def to_device(arr):
+            """Fake to_device that returns a wrapper."""
+            return _FakeDeviceArray(arr)
+
+    class _FakeDeviceArray:
+        """Fake device array for when numba is not installed."""
+
+        def __init__(self, arr):
+            self._arr = arr
+
+        def copy_to_host(self):
+            return self._arr
+
+    cuda = _FakeCuda()  # type: ignore[assignment]
+    HAS_CUDA = False
+
+from ..propagator_protocol import GPUMaterialProtocol, GPUMaterialID
+
+
+class SpectralGPUGradientPropagator:
+    """
+    GPU-based propagator for spectral materials with per-ray wavelengths.
+
+    This propagator is designed for SpectralInhomogeneousModel materials
+    where each ray can have a different wavelength, enabling simulation
+    of chromatic dispersion effects (e.g., rainbow formation, prism dispersion).
+
+    Unlike GPUGradientPropagator which uses a single wavelength for all rays,
+    this propagator accepts a wavelengths array and uses specialized kernels
+    that read the wavelength for each ray individually.
+
+    Parameters
+    ----------
+    material : SpectralInhomogeneousModel
+        The spectral material to propagate through. Must have
+        `get_gpu_kernels_perray()` method.
+    method : str
+        Integration method: 'euler' or 'rk4'. Default: 'rk4'
+    threads_per_block : int
+        CUDA threads per block. Default: 256
+
+    Example
+    -------
+    >>> from lsurf.materials import LinsleyAtmosphere
+    >>> atmosphere = LinsleyAtmosphere()
+    >>> propagator = SpectralGPUGradientPropagator(atmosphere, method='rk4')
+    >>>
+    >>> # Create rays with different wavelengths (red, green, blue)
+    >>> wavelengths = np.array([700e-9, 550e-9, 450e-9])  # One per ray
+    >>> propagator.propagate(rays, wavelengths, total_distance=500e3, step_size=100.0)
+    """
+
+    def __init__(
+        self,
+        material: GPUMaterialProtocol,
+        method: str = "rk4",
+        threads_per_block: int = 256,
+        enable_absorption: bool = False,
+    ):
+        # Check if CUDA is available
+        if not HAS_CUDA:
+            raise ImportError(
+                "CUDA support requires numba with CUDA. "
+                "Use GradientPropagator for CPU-based propagation."
+            )
+
+        # Store material reference
+        self._material = material
+
+        # Validate material has GPU support
+        if not hasattr(material, "gpu_material_id"):
+            raise ValueError(
+                f"Material {type(material).__name__} does not support GPU acceleration. "
+                f"Use GradientPropagator.propagate_step_cpu() instead."
+            )
+
+        # Validate material provides per-ray kernels
+        if not hasattr(material, "get_gpu_kernels_perray"):
+            raise ValueError(
+                f"Material {type(material).__name__} does not provide per-ray wavelength kernels. "
+                f"Implement get_gpu_kernels_perray() method or use GPUGradientPropagator."
+            )
+
+        self._material_id = material.gpu_material_id
+
+        if method not in ("euler", "rk4"):
+            raise ValueError(f"method must be 'euler' or 'rk4', got {method}")
+        self.method = method
+        self.threads_per_block = threads_per_block
+        self.enable_absorption = enable_absorption
+
+        # Get per-ray kernels from the material
+        self._kernels = material.get_gpu_kernels_perray()
+        self._kernel = self._kernels.get(method)
+        if self._kernel is None:
+            raise ValueError(
+                f"Material {type(material).__name__} does not provide '{method}' per-ray kernel. "
+                f"Available: {list(self._kernels.keys())}"
+            )
+
+        # Get absorption kernel if enabled
+        self._absorption_kernel = None
+        if enable_absorption:
+            self._absorption_kernel = self._get_absorption_kernel()
+
+        # Check if material uses device arrays (LUTs, grids, etc.)
+        self._gpu_arrays = None
+        self._device_arrays = None
+        if hasattr(material, "get_gpu_arrays"):
+            self._gpu_arrays = material.get_gpu_arrays()
+            # Device arrays will be created on first propagate call
+
+    def _get_absorption_kernel(self):
+        """Get the appropriate absorption kernel for spectral materials."""
+        from ..kernels.absorption.spectral import _kernel_absorption_spectral_perray
+
+        if self._material_id == GPUMaterialID.SPECTRAL_INHOMOGENEOUS:
+            return _kernel_absorption_spectral_perray
+        else:
+            raise ValueError(
+                f"No spectral absorption kernel available for material ID {self._material_id}"
+            )
+
+    def _apply_absorption(
+        self,
+        positions_d,
+        directions_d,
+        wavelengths_d,
+        active_d,
+        intensities_d,
+        optical_depth_d,
+        step_size: float,
+        num_steps: int,
+        blocks: int,
+    ) -> None:
+        """Apply absorption using the spectral per-ray kernel."""
+        material_params = self._material.get_gpu_parameters()
+
+        # Spectral kernel: positions, directions, wavelengths, active, intensities, optical_depth,
+        # step_size, num_steps, center_x, center_y, center_z, ref_radius,
+        # alt_min, alt_delta, n_alt, wl_min, wl_delta, n_wl, lut_alpha
+        self._absorption_kernel[blocks, self.threads_per_block](
+            positions_d,
+            directions_d,
+            wavelengths_d,
+            active_d,
+            intensities_d,
+            optical_depth_d,
+            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
+            material_params[6],  # n_alt
+            material_params[7],  # wl_min
+            material_params[8],  # wl_delta
+            material_params[9],  # n_wl
+            self._device_arrays["lut_alpha"],
+        )
+
+    @property
+    def material(self) -> GPUMaterialProtocol:
+        """The material being propagated through."""
+        return self._material
+
+    def propagate_step(
+        self,
+        rays,  # RayBatch
+        wavelengths: npt.NDArray[np.floating],
+        step_size: float,
+    ) -> None:
+        """
+        Propagate rays by a single integration step.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Ray batch containing positions, directions, and other ray properties.
+            Modified in-place.
+        wavelengths : ndarray, shape (N,)
+            Per-ray wavelengths in meters.
+        step_size : float
+            Integration step size in meters.
+        """
+        self.propagate(
+            rays=rays,
+            wavelengths=wavelengths,
+            total_distance=step_size,
+            step_size=step_size,
+        )
+
+    def propagate(
+        self,
+        rays,  # RayBatch
+        wavelengths: npt.NDArray[np.floating],
+        total_distance: float,
+        step_size: float,
+        steps_per_kernel: int = 100,
+    ) -> None:
+        """
+        Propagate rays through material on GPU with per-ray wavelengths.
+
+        Modifies the rays object in-place.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Ray batch to propagate (modified in-place)
+        wavelengths : ndarray, shape (N,)
+            Per-ray wavelengths in meters. Must have same length as number of rays.
+        total_distance : float
+            Total distance to propagate in meters
+        step_size : float
+            Integration step size in meters
+        steps_per_kernel : int
+            Number of steps per kernel launch. Higher values reduce
+            kernel launch overhead but increase register pressure.
+
+        Raises
+        ------
+        ValueError
+            If wavelengths array length doesn't match number of rays.
+        """
+        num_rays = rays.num_rays
+        total_steps = int(total_distance / step_size)
+
+        # Validate wavelengths array
+        wavelengths = np.asarray(wavelengths)
+        if wavelengths.shape[0] != num_rays:
+            raise ValueError(
+                f"wavelengths array length ({wavelengths.shape[0]}) "
+                f"must match number of rays ({num_rays})"
+            )
+
+        # Prepare GPU arrays for ray state
+        positions_d = cuda.to_device(rays.positions.astype(np.float32))
+        directions_d = cuda.to_device(rays.directions.astype(np.float32))
+        wavelengths_d = cuda.to_device(wavelengths.astype(np.float32))
+        active_d = cuda.to_device(rays.active)
+        geo_path_d = cuda.to_device(rays.geometric_path_lengths.astype(np.float32))
+        opt_path_d = cuda.to_device(rays.optical_path_lengths.astype(np.float32))
+        acc_time_d = cuda.to_device(rays.accumulated_time.astype(np.float32))
+
+        # Prepare absorption arrays if enabled
+        intensities_d = None
+        optical_depth_d = None
+        if self.enable_absorption:
+            intensities_d = cuda.to_device(rays.intensities.astype(np.float32))
+            if rays.optical_depth is not None:
+                optical_depth_d = cuda.to_device(rays.optical_depth.astype(np.float32))
+            else:
+                optical_depth_d = cuda.to_device(np.zeros(num_rays, dtype=np.float32))
+
+        # Transfer material arrays to device if needed (LUTs, grids, etc.)
+        if self._gpu_arrays is not None and self._device_arrays is None:
+            self._device_arrays = {}
+            for name, arr in self._gpu_arrays.items():
+                self._device_arrays[name] = cuda.to_device(arr.astype(np.float32))
+
+        # Get material-specific kernel parameters from the material
+        material_params = self._material.get_gpu_parameters()
+
+        # Build kernel arguments: scalar params + device arrays
+        kernel_args = list(material_params)
+        if self._device_arrays is not None:
+            # Append device arrays in the order the kernel expects
+            for name in ["lut_n", "lut_dn_dh"]:
+                if name in self._device_arrays:
+                    kernel_args.append(self._device_arrays[name])
+
+        # Configure kernel launch
+        blocks = (num_rays + self.threads_per_block - 1) // self.threads_per_block
+
+        # Launch kernels in batches
+        remaining_steps = total_steps
+        while remaining_steps > 0:
+            steps_this_call = min(remaining_steps, steps_per_kernel)
+            # Per-ray kernel signature:
+            # positions, directions, wavelengths, active, geo_path, opt_path, acc_time,
+            # step_size, num_steps, ...material_params, lut_n, lut_dn_dh
+            self._kernel[blocks, self.threads_per_block](
+                positions_d,
+                directions_d,
+                wavelengths_d,
+                active_d,
+                geo_path_d,
+                opt_path_d,
+                acc_time_d,
+                float(step_size),
+                steps_this_call,
+                *kernel_args,
+            )
+
+            # Apply absorption AFTER each batch (for proper path integration)
+            if self.enable_absorption and self._absorption_kernel is not None:
+                self._apply_absorption(
+                    positions_d,
+                    directions_d,
+                    wavelengths_d,
+                    active_d,
+                    intensities_d,
+                    optical_depth_d,
+                    step_size,
+                    steps_this_call,
+                    blocks,
+                )
+
+            remaining_steps -= steps_this_call
+
+        # Synchronize and copy back
+        cuda.synchronize()
+        rays.positions[:] = positions_d.copy_to_host()
+        rays.directions[:] = directions_d.copy_to_host()
+        rays.geometric_path_lengths[:] = geo_path_d.copy_to_host()
+        rays.optical_path_lengths[:] = opt_path_d.copy_to_host()
+        rays.accumulated_time[:] = acc_time_d.copy_to_host()
+
+        # Copy back absorption data if enabled
+        if self.enable_absorption:
+            rays.intensities[:] = intensities_d.copy_to_host()
+            if rays.optical_depth is not None:
+                rays.optical_depth[:] = optical_depth_d.copy_to_host()
+
+    def propagate_with_history(
+        self,
+        rays,  # RayBatch
+        wavelengths: npt.NDArray[np.floating],
+        total_distance: float,
+        step_size: float,
+        history_interval: int = 100,
+    ) -> dict:
+        """
+        Propagate rays and record position history at intervals.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Ray batch to propagate (modified in-place)
+        wavelengths : ndarray, shape (N,)
+            Per-ray wavelengths in meters.
+        total_distance : float
+            Total distance to propagate in meters
+        step_size : float
+            Integration step size in meters
+        history_interval : int
+            Record history every N steps
+
+        Returns
+        -------
+        dict
+            Dictionary with 'positions', 'directions', 'distances' arrays
+        """
+        num_rays = rays.num_rays
+        total_steps = int(total_distance / step_size)
+        num_records = total_steps // history_interval + 1
+
+        # Validate wavelengths array
+        wavelengths = np.asarray(wavelengths)
+        if wavelengths.shape[0] != num_rays:
+            raise ValueError(
+                f"wavelengths array length ({wavelengths.shape[0]}) "
+                f"must match number of rays ({num_rays})"
+            )
+
+        # Prepare GPU arrays for ray state
+        positions_d = cuda.to_device(rays.positions.astype(np.float32))
+        directions_d = cuda.to_device(rays.directions.astype(np.float32))
+        wavelengths_d = cuda.to_device(wavelengths.astype(np.float32))
+        active_d = cuda.to_device(rays.active)
+        geo_path_d = cuda.to_device(rays.geometric_path_lengths.astype(np.float32))
+        opt_path_d = cuda.to_device(rays.optical_path_lengths.astype(np.float32))
+        acc_time_d = cuda.to_device(rays.accumulated_time.astype(np.float32))
+
+        # Transfer material arrays to device if needed (LUTs, grids, etc.)
+        if self._gpu_arrays is not None and self._device_arrays is None:
+            self._device_arrays = {}
+            for name, arr in self._gpu_arrays.items():
+                self._device_arrays[name] = cuda.to_device(arr.astype(np.float32))
+
+        # Get material-specific kernel parameters
+        material_params = self._material.get_gpu_parameters()
+
+        # Build kernel arguments: scalar params + device arrays
+        kernel_args = list(material_params)
+        if self._device_arrays is not None:
+            for name in ["lut_n", "lut_dn_dh"]:
+                if name in self._device_arrays:
+                    kernel_args.append(self._device_arrays[name])
+
+        # Configure kernel launch
+        blocks = (num_rays + self.threads_per_block - 1) // self.threads_per_block
+
+        # Storage for history
+        position_history = np.zeros((num_records, num_rays, 3), dtype=np.float32)
+        direction_history = np.zeros((num_records, num_rays, 3), dtype=np.float32)
+        distance_history = np.zeros(num_records, dtype=np.float32)
+
+        # Record initial state
+        position_history[0] = positions_d.copy_to_host()
+        direction_history[0] = directions_d.copy_to_host()
+        distance_history[0] = 0.0
+
+        # Propagate with recording
+        record_idx = 1
+        steps_done = 0
+
+        while steps_done < total_steps:
+            steps_to_next_record = min(
+                history_interval - (steps_done % history_interval),
+                total_steps - steps_done,
+            )
+
+            self._kernel[blocks, self.threads_per_block](
+                positions_d,
+                directions_d,
+                wavelengths_d,
+                active_d,
+                geo_path_d,
+                opt_path_d,
+                acc_time_d,
+                float(step_size),
+                steps_to_next_record,
+                *kernel_args,
+            )
+            steps_done += steps_to_next_record
+
+            # Record if at interval
+            if steps_done % history_interval == 0 and record_idx < num_records:
+                cuda.synchronize()
+                position_history[record_idx] = positions_d.copy_to_host()
+                direction_history[record_idx] = directions_d.copy_to_host()
+                distance_history[record_idx] = steps_done * step_size
+                record_idx += 1
+
+        # Final copy back
+        cuda.synchronize()
+        rays.positions[:] = positions_d.copy_to_host()
+        rays.directions[:] = directions_d.copy_to_host()
+        rays.geometric_path_lengths[:] = geo_path_d.copy_to_host()
+        rays.optical_path_lengths[:] = opt_path_d.copy_to_host()
+        rays.accumulated_time[:] = acc_time_d.copy_to_host()
+
+        return {
+            "positions": position_history[:record_idx],
+            "directions": direction_history[:record_idx],
+            "distances": distance_history[:record_idx],
+        }
+
+    def get_refractive_index(
+        self,
+        positions: npt.NDArray[np.floating],
+        wavelengths: npt.NDArray[np.floating],
+    ) -> npt.NDArray[np.float64]:
+        """
+        Compute refractive index at positions (CPU, delegates to material).
+
+        Parameters
+        ----------
+        positions : ndarray (N, 3)
+            Position coordinates
+        wavelengths : ndarray (N,)
+            Per-ray wavelengths in meters
+
+        Returns
+        -------
+        ndarray (N,)
+            Refractive index at each position
+        """
+        x = positions[:, 0]
+        y = positions[:, 1]
+        z = positions[:, 2]
+        return self._material.get_refractive_index(x, y, z, wavelengths)