lsurf 1.0.0__py3-none-any.whl

lsurf/utilities/ray_data.py

@@ -0,0 +1,532 @@
+# 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.
+
+"""
+Ray Data Structures for GPU Raytracing
+
+This module defines the core data structures for representing rays in the
+simulation, including position, direction, optical properties, and state tracking.
+"""
+
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import numpy as np
+import numpy.typing as npt
+
+# Type aliases for clarity
+Float32Array = npt.NDArray[np.float32]
+BoolArray = npt.NDArray[np.bool_]
+Int32Array = npt.NDArray[np.int32]
+
+
+@dataclass
+class RayBatch:
+    """
+    Structure-of-Arrays (SoA) layout for efficient GPU processing of ray batches.
+
+    All arrays have shape (N,) where N is the number of rays, except positions
+    and directions which have shape (N, 3).
+
+    This layout maximizes memory coalescing on GPU and enables efficient SIMD
+    operations on CPU.
+
+    Attributes
+    ----------
+    positions : Float32Array
+        Ray origin positions, shape (N, 3), units: meters
+        [x0, y0, z0; x1, y1, z1; ...]
+    directions : Float32Array
+        Ray direction unit vectors, shape (N, 3), dimensionless
+        [dx0, dy0, dz0; dx1, dy1, dz1; ...]
+    wavelengths : Float32Array
+        Wavelengths, shape (N,), units: meters
+    intensities : Float32Array
+        Current ray intensities/powers, shape (N,), units: W or arbitrary
+    optical_path_lengths : Float32Array
+        Accumulated optical path length ∫n(s)ds, shape (N,), units: meters
+    geometric_path_lengths : Float32Array
+        Accumulated geometric path length ∫ds, shape (N,), units: meters
+    accumulated_time : Float32Array
+        Accumulated propagation time, shape (N,), units: seconds
+    generations : Int32Array
+        Scattering generation (0=primary, 1=first scatter, etc.), shape (N,)
+    domain_ids : Int32Array
+        Current domain ID for each ray, shape (N,)
+    active : BoolArray
+        Whether ray is still active (not terminated), shape (N,)
+    polarization_s : Optional[Float32Array]
+        S-polarization component (optional), shape (N,)
+    polarization_p : Optional[Float32Array]
+        P-polarization component (optional), shape (N,)
+    polarization_vector : Optional[Float32Array]
+        3D polarization vector (electric field direction), shape (N, 3)
+        Unit vector perpendicular to ray direction representing E-field orientation.
+        Used for tracking polarization state through reflections/refractions.
+    phase : Optional[Float32Array]
+        Phase in radians (for coherent simulations), shape (N,)
+    optical_depth : Optional[Float32Array]
+        Accumulated optical depth τ = ∫α·ds, shape (N,), dimensionless
+        Used for Beer-Lambert absorption tracking. τ = -ln(I/I₀)
+
+    Notes
+    -----
+    The SoA layout enables efficient GPU memory access patterns. For example,
+    all x-coordinates are contiguous in memory, allowing coalesced reads.
+
+    References
+    ----------
+    .. [1] https://en.wikipedia.org/wiki/AoS_and_SoA
+    """
+
+    positions: Float32Array
+    directions: Float32Array
+    wavelengths: Float32Array
+    intensities: Float32Array
+    optical_path_lengths: Float32Array
+    geometric_path_lengths: Float32Array
+    accumulated_time: Float32Array
+    generations: Int32Array
+    domain_ids: Int32Array
+    active: BoolArray
+    polarization_s: Float32Array | None = None
+    polarization_p: Float32Array | None = None
+    polarization_vector: Float32Array | None = None
+    phase: Float32Array | None = None
+    optical_depth: Float32Array | None = None
+
+    def __post_init__(self) -> None:
+        """Validate array shapes after initialization."""
+        n_rays = len(self.positions)
+
+        # Validate shapes
+        assert self.positions.shape == (
+            n_rays,
+            3,
+        ), f"positions must have shape (N, 3), got {self.positions.shape}"
+        assert self.directions.shape == (
+            n_rays,
+            3,
+        ), f"directions must have shape (N, 3), got {self.directions.shape}"
+        assert self.wavelengths.shape == (
+            n_rays,
+        ), f"wavelengths must have shape (N,), got {self.wavelengths.shape}"
+        assert self.intensities.shape == (
+            n_rays,
+        ), f"intensities must have shape (N,), got {self.intensities.shape}"
+        assert self.optical_path_lengths.shape == (
+            n_rays,
+        ), f"optical_path_lengths must have shape (N,), got {self.optical_path_lengths.shape}"
+        assert self.geometric_path_lengths.shape == (
+            n_rays,
+        ), f"geometric_path_lengths must have shape (N,), got {self.geometric_path_lengths.shape}"
+        assert self.accumulated_time.shape == (
+            n_rays,
+        ), f"accumulated_time must have shape (N,), got {self.accumulated_time.shape}"
+        assert self.generations.shape == (
+            n_rays,
+        ), f"generations must have shape (N,), got {self.generations.shape}"
+        assert self.domain_ids.shape == (
+            n_rays,
+        ), f"domain_ids must have shape (N,), got {self.domain_ids.shape}"
+        assert self.active.shape == (
+            n_rays,
+        ), f"active must have shape (N,), got {self.active.shape}"
+
+        # Validate dtypes
+        assert self.positions.dtype == np.float32
+        assert self.directions.dtype == np.float32
+        assert self.wavelengths.dtype == np.float32
+        assert self.intensities.dtype == np.float32
+        assert self.optical_path_lengths.dtype == np.float32
+        assert self.geometric_path_lengths.dtype == np.float32
+        assert self.accumulated_time.dtype == np.float32
+        assert self.generations.dtype == np.int32
+        assert self.domain_ids.dtype == np.int32
+        assert self.active.dtype == np.bool_
+
+        # Validate optional arrays
+        if self.polarization_s is not None:
+            assert self.polarization_s.shape == (n_rays,)
+            assert self.polarization_s.dtype == np.float32
+        if self.polarization_p is not None:
+            assert self.polarization_p.shape == (n_rays,)
+            assert self.polarization_p.dtype == np.float32
+        if self.polarization_vector is not None:
+            assert self.polarization_vector.shape == (
+                n_rays,
+                3,
+            ), f"polarization_vector must have shape (N, 3), got {self.polarization_vector.shape}"
+            assert self.polarization_vector.dtype == np.float32
+        if self.phase is not None:
+            assert self.phase.shape == (n_rays,)
+            assert self.phase.dtype == np.float32
+        if self.optical_depth is not None:
+            assert self.optical_depth.shape == (n_rays,)
+            assert self.optical_depth.dtype == np.float32
+
+    @property
+    def num_rays(self) -> int:
+        """Total number of rays in batch."""
+        return len(self.positions)
+
+    @property
+    def num_active(self) -> int:
+        """Number of currently active rays."""
+        return int(np.sum(self.active))
+
+    def normalize_directions(self) -> None:
+        """Normalize all direction vectors to unit length in-place."""
+        norms = np.linalg.norm(self.directions, axis=1, keepdims=True)
+        norms = np.maximum(norms, 1e-12)  # Avoid division by zero
+        self.directions /= norms
+
+    def compact(self) -> "RayBatch":
+        """
+        Remove inactive rays to reduce memory and computation.
+
+        Returns
+        -------
+        RayBatch
+            New batch containing only active rays
+
+        Notes
+        -----
+        This operation is called "stream compaction" in GPU programming.
+        Useful for maintaining efficiency as rays terminate.
+        """
+        mask = self.active
+
+        return RayBatch(
+            positions=self.positions[mask].copy(),
+            directions=self.directions[mask].copy(),
+            wavelengths=self.wavelengths[mask].copy(),
+            intensities=self.intensities[mask].copy(),
+            optical_path_lengths=self.optical_path_lengths[mask].copy(),
+            geometric_path_lengths=self.geometric_path_lengths[mask].copy(),
+            accumulated_time=self.accumulated_time[mask].copy(),
+            generations=self.generations[mask].copy(),
+            domain_ids=self.domain_ids[mask].copy(),
+            active=self.active[mask].copy(),
+            polarization_s=(
+                self.polarization_s[mask].copy()
+                if self.polarization_s is not None
+                else None
+            ),
+            polarization_p=(
+                self.polarization_p[mask].copy()
+                if self.polarization_p is not None
+                else None
+            ),
+            polarization_vector=(
+                self.polarization_vector[mask].copy()
+                if self.polarization_vector is not None
+                else None
+            ),
+            phase=self.phase[mask].copy() if self.phase is not None else None,
+            optical_depth=(
+                self.optical_depth[mask].copy()
+                if self.optical_depth is not None
+                else None
+            ),
+        )
+
+    def clone(self) -> "RayBatch":
+        """Create a deep copy of the ray batch."""
+        return RayBatch(
+            positions=self.positions.copy(),
+            directions=self.directions.copy(),
+            wavelengths=self.wavelengths.copy(),
+            intensities=self.intensities.copy(),
+            optical_path_lengths=self.optical_path_lengths.copy(),
+            geometric_path_lengths=self.geometric_path_lengths.copy(),
+            accumulated_time=self.accumulated_time.copy(),
+            generations=self.generations.copy(),
+            domain_ids=self.domain_ids.copy(),
+            active=self.active.copy(),
+            polarization_s=(
+                self.polarization_s.copy() if self.polarization_s is not None else None
+            ),
+            polarization_p=(
+                self.polarization_p.copy() if self.polarization_p is not None else None
+            ),
+            polarization_vector=(
+                self.polarization_vector.copy()
+                if self.polarization_vector is not None
+                else None
+            ),
+            phase=self.phase.copy() if self.phase is not None else None,
+            optical_depth=(
+                self.optical_depth.copy() if self.optical_depth is not None else None
+            ),
+        )
+
+
+class RayStatistics(NamedTuple):
+    """
+    Statistical summary of ray batch state.
+
+    Attributes
+    ----------
+    total_rays : int
+        Total number of rays (active + inactive)
+    active_rays : int
+        Number of active rays
+    mean_intensity : float
+        Mean intensity of active rays
+    total_power : float
+        Sum of intensities of all active rays
+    mean_optical_path : float
+        Mean optical path length of active rays
+    mean_generation : float
+        Mean scattering generation of active rays
+    max_generation : int
+        Maximum scattering generation
+    """
+
+    total_rays: int
+    active_rays: int
+    mean_intensity: float
+    total_power: float
+    mean_optical_path: float
+    mean_generation: float
+    max_generation: int
+
+
+def create_ray_batch(
+    num_rays: int,
+    enable_polarization: bool = False,
+    enable_polarization_vector: bool = False,
+    enable_phase: bool = False,
+    enable_optical_depth: bool = False,
+) -> RayBatch:
+    """
+    Create an empty ray batch with zero-initialized arrays.
+
+    Parameters
+    ----------
+    num_rays : int
+        Number of rays to allocate
+    enable_polarization : bool, optional
+        Whether to allocate scalar polarization arrays (s and p components)
+    enable_polarization_vector : bool, optional
+        Whether to allocate 3D polarization vector array
+    enable_phase : bool, optional
+        Whether to allocate phase array
+    enable_optical_depth : bool, optional
+        Whether to allocate optical depth array for absorption tracking
+
+    Returns
+    -------
+    RayBatch
+        Initialized ray batch with all fields set to zero/false
+    """
+    positions = np.zeros((num_rays, 3), dtype=np.float32)
+    directions = np.zeros((num_rays, 3), dtype=np.float32)
+    directions[:, 2] = 1.0  # Default to +z direction
+
+    wavelengths = np.zeros(num_rays, dtype=np.float32)
+    intensities = np.zeros(num_rays, dtype=np.float32)
+    optical_path_lengths = np.zeros(num_rays, dtype=np.float32)
+    geometric_path_lengths = np.zeros(num_rays, dtype=np.float32)
+    accumulated_time = np.zeros(num_rays, dtype=np.float32)
+    generations = np.zeros(num_rays, dtype=np.int32)
+    domain_ids = np.zeros(num_rays, dtype=np.int32)
+    active = np.zeros(num_rays, dtype=np.bool_)
+
+    polarization_s = (
+        np.zeros(num_rays, dtype=np.float32) if enable_polarization else None
+    )
+    polarization_p = (
+        np.zeros(num_rays, dtype=np.float32) if enable_polarization else None
+    )
+    polarization_vector = (
+        np.zeros((num_rays, 3), dtype=np.float32)
+        if enable_polarization_vector
+        else None
+    )
+    phase = np.zeros(num_rays, dtype=np.float32) if enable_phase else None
+    optical_depth = (
+        np.zeros(num_rays, dtype=np.float32) if enable_optical_depth else None
+    )
+
+    return RayBatch(
+        positions=positions,
+        directions=directions,
+        wavelengths=wavelengths,
+        intensities=intensities,
+        optical_path_lengths=optical_path_lengths,
+        geometric_path_lengths=geometric_path_lengths,
+        accumulated_time=accumulated_time,
+        generations=generations,
+        domain_ids=domain_ids,
+        active=active,
+        polarization_s=polarization_s,
+        polarization_p=polarization_p,
+        polarization_vector=polarization_vector,
+        phase=phase,
+        optical_depth=optical_depth,
+    )
+
+
+def compute_statistics(batch: RayBatch) -> RayStatistics:
+    """
+    Compute statistical summary of a ray batch.
+
+    Parameters
+    ----------
+    batch : RayBatch
+        Ray batch to analyze
+
+    Returns
+    -------
+    RayStatistics
+        Statistical summary
+    """
+    active_mask = batch.active
+    n_active = np.sum(active_mask)
+
+    if n_active == 0:
+        return RayStatistics(
+            total_rays=batch.num_rays,
+            active_rays=0,
+            mean_intensity=0.0,
+            total_power=0.0,
+            mean_optical_path=0.0,
+            mean_generation=0.0,
+            max_generation=0,
+        )
+
+    active_intensities = batch.intensities[active_mask]
+    active_opl = batch.optical_path_lengths[active_mask]
+    active_gen = batch.generations[active_mask]
+
+    return RayStatistics(
+        total_rays=batch.num_rays,
+        active_rays=int(n_active),
+        mean_intensity=float(np.mean(active_intensities)),
+        total_power=float(np.sum(active_intensities)),
+        mean_optical_path=float(np.mean(active_opl)),
+        mean_generation=float(np.mean(active_gen)),
+        max_generation=int(np.max(batch.generations)),
+    )
+
+
+def merge_ray_batches(batches: list) -> RayBatch:
+    """
+    Merge multiple ray batches into a single batch.
+
+    Parameters
+    ----------
+    batches : list of RayBatch
+        List of ray batches to merge
+
+    Returns
+    -------
+    RayBatch
+        Combined ray batch containing all rays from input batches
+
+    Notes
+    -----
+    This is useful for combining rays from different sources or for
+    collecting rays after ray splitting at interfaces.
+    """
+    if not batches:
+        return create_ray_batch(num_rays=0)
+
+    # Filter out empty batches
+    non_empty = [b for b in batches if b.num_rays > 0]
+    if not non_empty:
+        return create_ray_batch(num_rays=0)
+
+    if len(non_empty) == 1:
+        return non_empty[0].clone()
+
+    # Check for polarization, phase, and optical_depth consistency
+    has_polarization = all(
+        b.polarization_s is not None and b.polarization_p is not None for b in non_empty
+    )
+    has_polarization_vector = all(b.polarization_vector is not None for b in non_empty)
+    has_phase = all(b.phase is not None for b in non_empty)
+    has_optical_depth = all(b.optical_depth is not None for b in non_empty)
+
+    # Concatenate all arrays
+    positions = np.vstack([b.positions for b in non_empty])
+    directions = np.vstack([b.directions for b in non_empty])
+    wavelengths = np.concatenate([b.wavelengths for b in non_empty])
+    intensities = np.concatenate([b.intensities for b in non_empty])
+    optical_path_lengths = np.concatenate([b.optical_path_lengths for b in non_empty])
+    geometric_path_lengths = np.concatenate(
+        [b.geometric_path_lengths for b in non_empty]
+    )
+    accumulated_time = np.concatenate([b.accumulated_time for b in non_empty])
+    generations = np.concatenate([b.generations for b in non_empty])
+    domain_ids = np.concatenate([b.domain_ids for b in non_empty])
+    active = np.concatenate([b.active for b in non_empty])
+
+    polarization_s = None
+    polarization_p = None
+    polarization_vector = None
+    phase = None
+    optical_depth = None
+
+    if has_polarization:
+        polarization_s = np.concatenate([b.polarization_s for b in non_empty])
+        polarization_p = np.concatenate([b.polarization_p for b in non_empty])
+
+    if has_polarization_vector:
+        polarization_vector = np.vstack([b.polarization_vector for b in non_empty])
+
+    if has_phase:
+        phase = np.concatenate([b.phase for b in non_empty])
+
+    if has_optical_depth:
+        optical_depth = np.concatenate([b.optical_depth for b in non_empty])
+
+    return RayBatch(
+        positions=positions,
+        directions=directions,
+        wavelengths=wavelengths,
+        intensities=intensities,
+        optical_path_lengths=optical_path_lengths,
+        geometric_path_lengths=geometric_path_lengths,
+        accumulated_time=accumulated_time,
+        generations=generations,
+        domain_ids=domain_ids,
+        active=active,
+        polarization_s=polarization_s,
+        polarization_p=polarization_p,
+        polarization_vector=polarization_vector,
+        phase=phase,
+        optical_depth=optical_depth,
+    )