lsurf 1.0.0__py3-none-any.whl

lsurf/propagation/propagators/surface_interaction.py

@@ -0,0 +1,616 @@
+# 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 Interaction Processor
+
+Handles surface interaction physics based on surface role:
+- DETECTOR: Record ray data, terminate ray
+- ABSORBER: Terminate ray
+- OPTICAL: Compute Fresnel coefficients, generate reflected/refracted rays
+
+This is the second component in the propagator architecture:
+MaterialPropagator -> SurfaceInteractionProcessor -> SimulationOrchestrator
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import numpy as np
+import numpy.typing as npt
+
+from ...geometry.cell_geometry import CellGeometry
+from ...surfaces import Surface, SurfaceRole
+from ...utilities.fresnel import (
+    compute_reflection_direction,
+    compute_refraction_direction,
+    fresnel_coefficients,
+    initialize_polarization_vectors,
+    transform_polarization_reflection,
+    transform_polarization_refraction,
+)
+from ...utilities.ray_data import RayBatch, create_ray_batch, merge_ray_batches
+
+if TYPE_CHECKING:
+    from ...geometry import Geometry
+    from ...materials import MaterialField
+    from ...utilities.recording_sphere import RecordedRays
+    from .surface_propagator import HitData
+
+# Speed of light
+SPEED_OF_LIGHT = 299792458.0
+
+
+@dataclass
+class SurfaceCrossing:
+    """
+    Information about rays crossing a specific surface.
+
+    Attributes
+    ----------
+    surface_idx : int
+        Index of the surface in the geometry.
+    surface : Surface
+        The surface object.
+    ray_indices : ndarray
+        Original indices of rays that hit this surface.
+    hit_positions : ndarray, shape (M, 3)
+        Intersection positions.
+    hit_directions : ndarray, shape (M, 3)
+        Ray directions at intersection.
+    hit_normals : ndarray, shape (M, 3)
+        Surface normals at intersection points.
+    from_front : ndarray, shape (M,)
+        Whether rays hit from the front side (normal direction).
+    """
+
+    surface_idx: int
+    surface: Surface
+    ray_indices: npt.NDArray[np.int32]
+    hit_positions: npt.NDArray[np.float32]
+    hit_directions: npt.NDArray[np.float32]
+    hit_normals: npt.NDArray[np.float32]
+    from_front: npt.NDArray[np.bool_]
+
+    @property
+    def num_hits(self) -> int:
+        """Number of rays that hit this surface."""
+        return len(self.ray_indices)
+
+
+@dataclass
+class OpticalSurfaceHit:
+    """
+    Record of hits on an optical surface.
+
+    Attributes
+    ----------
+    surface_name : str
+        Name of the surface.
+    positions : ndarray, shape (N, 3)
+        Hit positions.
+    directions : ndarray, shape (N, 3)
+        Ray directions at hit.
+    intensities : ndarray, shape (N,)
+        Ray intensities.
+    wavelengths : ndarray, shape (N,)
+        Ray wavelengths.
+    """
+
+    surface_name: str
+    positions: npt.NDArray[np.float32]
+    directions: npt.NDArray[np.float32]
+    intensities: npt.NDArray[np.float32]
+    wavelengths: npt.NDArray[np.float32]
+
+
+@dataclass
+class SurfaceInteractionResult:
+    """
+    Result from processing surface interactions.
+
+    Attributes
+    ----------
+    continuing_rays : RayBatch or None
+        Rays that should continue propagating (from optical surfaces).
+    reflected_rays : RayBatch or None
+        Reflected rays from optical surfaces.
+    refracted_rays : RayBatch or None
+        Refracted rays from optical surfaces.
+    detector_hits : dict
+        Mapping from detector name to RecordedRays.
+    absorbed_count : int
+        Number of rays absorbed.
+    detected_count : int
+        Number of rays detected.
+    optical_hits : list of OpticalSurfaceHit
+        If track_surface_hits was enabled, contains hit data for optical surfaces.
+    """
+
+    continuing_rays: RayBatch | None = None
+    reflected_rays: RayBatch | None = None
+    refracted_rays: RayBatch | None = None
+    detector_hits: dict = field(default_factory=dict)
+    absorbed_count: int = 0
+    detected_count: int = 0
+    optical_hits: list[OpticalSurfaceHit] = field(default_factory=list)
+
+
+class SurfaceInteractionProcessor:
+    """
+    Processes ray-surface interactions based on surface role.
+
+    For each surface type:
+    - DETECTOR: Record ray data, terminate ray
+    - ABSORBER: Terminate ray
+    - OPTICAL: Compute Fresnel coefficients, generate reflected/refracted rays
+
+    Parameters
+    ----------
+    surfaces : list of Surface
+        List of surfaces in the geometry.
+    polarization : str, optional
+        Polarization state for Fresnel calculations: 's', 'p', or 'unpolarized'.
+        Default is 'unpolarized'.
+    track_polarization_vector : bool, optional
+        Whether to track 3D polarization vectors through interactions.
+        Default is False.
+
+    Examples
+    --------
+    >>> from lsurf.surfaces import PlaneSurface, SurfaceRole
+    >>> detector = PlaneSurface(
+    ...     point=(0, 0, 1000),
+    ...     normal=(0, 0, 1),
+    ...     role=SurfaceRole.DETECTOR,
+    ... )
+    >>> processor = SurfaceInteractionProcessor([detector])
+    >>> result = processor.process_hits(rays, hit_data)
+    """
+
+    def __init__(
+        self,
+        surfaces: list[Surface],
+        polarization: str = "unpolarized",
+        track_polarization_vector: bool = False,
+        geometry: "Geometry | None" = None,
+    ):
+        self._surfaces = list(surfaces)
+        self._polarization = polarization
+        self._track_polarization_vector = track_polarization_vector
+        self._geometry = geometry
+        self._is_cell_geometry = isinstance(geometry, CellGeometry)
+
+        # Build indices by role
+        self._detector_indices: list[int] = []
+        self._optical_indices: list[int] = []
+        self._absorber_indices: list[int] = []
+
+        for i, surface in enumerate(surfaces):
+            if surface.role == SurfaceRole.DETECTOR:
+                self._detector_indices.append(i)
+            elif surface.role == SurfaceRole.OPTICAL:
+                self._optical_indices.append(i)
+            elif surface.role == SurfaceRole.ABSORBER:
+                self._absorber_indices.append(i)
+
+    @property
+    def surfaces(self) -> list[Surface]:
+        """List of surfaces."""
+        return self._surfaces
+
+    @property
+    def polarization(self) -> str:
+        """Polarization state for Fresnel calculations."""
+        return self._polarization
+
+    @property
+    def num_detectors(self) -> int:
+        """Number of detector surfaces."""
+        return len(self._detector_indices)
+
+    @property
+    def num_optical(self) -> int:
+        """Number of optical surfaces."""
+        return len(self._optical_indices)
+
+    def process_hits(
+        self,
+        rays: RayBatch,
+        hit_data: "HitData",
+        track_surface_hits: bool = False,
+    ) -> SurfaceInteractionResult:
+        """
+        Process all surface hits from a propagation leg.
+
+        Parameters
+        ----------
+        rays : RayBatch
+            Original ray batch (before hit_data was recorded).
+        hit_data : HitData
+            Hit information from MaterialPropagator.
+        track_surface_hits : bool, optional
+            If True, store hit positions for optical surfaces in result.optical_hits.
+            Default is False.
+
+        Returns
+        -------
+        SurfaceInteractionResult
+            Result containing detector hits, reflected/refracted rays, etc.
+        """
+
+        result = SurfaceInteractionResult()
+        all_reflected: list[RayBatch] = []
+        all_refracted: list[RayBatch] = []
+
+        # Group hits by surface
+        for surface_idx, surface in enumerate(self._surfaces):
+            mask = hit_data.hit_surface_idx == surface_idx
+            if not np.any(mask):
+                continue
+
+            # Extract hit data for this surface
+            crossing = self._extract_crossing(
+                rays, hit_data, surface_idx, surface, mask
+            )
+
+            # Process based on role
+            if surface.role == SurfaceRole.DETECTOR:
+                recorded = self._process_detector_hit(rays, crossing)
+                surface_name = surface.name or f"detector_{surface_idx}"
+                result.detector_hits[surface_name] = recorded
+                result.detected_count += crossing.num_hits
+
+            elif surface.role == SurfaceRole.ABSORBER:
+                result.absorbed_count += crossing.num_hits
+
+            elif surface.role == SurfaceRole.OPTICAL:
+                # Store hit data if tracking enabled
+                if track_surface_hits:
+                    surface_name = surface.name or f"optical_{surface_idx}"
+                    result.optical_hits.append(
+                        OpticalSurfaceHit(
+                            surface_name=surface_name,
+                            positions=crossing.hit_positions.copy(),
+                            directions=crossing.hit_directions.copy(),
+                            intensities=rays.intensities[crossing.ray_indices].copy(),
+                            wavelengths=rays.wavelengths[crossing.ray_indices].copy(),
+                        )
+                    )
+
+                reflected, refracted = self._process_optical_hit(rays, crossing)
+                if reflected is not None and reflected.num_rays > 0:
+                    all_reflected.append(reflected)
+                if refracted is not None and refracted.num_rays > 0:
+                    all_refracted.append(refracted)
+
+        # Merge all reflected/refracted rays
+        if all_reflected:
+            result.reflected_rays = merge_ray_batches(all_reflected)
+        if all_refracted:
+            result.refracted_rays = merge_ray_batches(all_refracted)
+
+        return result
+
+    def _extract_crossing(
+        self,
+        rays: RayBatch,
+        hit_data: "HitData",
+        surface_idx: int,
+        surface: Surface,
+        mask: npt.NDArray[np.bool_],
+    ) -> SurfaceCrossing:
+        """Extract crossing data for a specific surface."""
+        ray_indices = np.where(mask)[0].astype(np.int32)
+        hit_positions = hit_data.hit_positions[mask]
+        hit_directions = hit_data.hit_directions[mask]
+
+        # Compute surface normals at hit positions
+        hit_normals = surface.normal_at(hit_positions, hit_directions)
+
+        # Determine if hits are from front (ray going opposite to normal)
+        # cos(angle) = -direction . normal
+        cos_angle = -np.sum(hit_directions * hit_normals, axis=1)
+        from_front = cos_angle > 0
+
+        # Flip normals for back-side hits to always face the ray
+        hit_normals[~from_front] = -hit_normals[~from_front]
+
+        return SurfaceCrossing(
+            surface_idx=surface_idx,
+            surface=surface,
+            ray_indices=ray_indices,
+            hit_positions=hit_positions,
+            hit_directions=hit_directions,
+            hit_normals=hit_normals,
+            from_front=from_front,
+        )
+
+    def _get_materials_at_crossing(
+        self,
+        crossing: SurfaceCrossing,
+    ) -> tuple["MaterialField | None", "MaterialField | None"]:
+        """
+        Get materials for a surface crossing.
+
+        For standard Geometry, returns (surface.material_front, surface.material_back).
+        For CellGeometry, queries cells to determine materials at hit positions.
+
+        Parameters
+        ----------
+        crossing : SurfaceCrossing
+            The surface crossing information.
+
+        Returns
+        -------
+        tuple[MaterialField | None, MaterialField | None]
+            (material_front, material_back) for this crossing.
+        """
+        surface = crossing.surface
+
+        if not self._is_cell_geometry or self._geometry is None:
+            # Standard geometry: use surface's front/back materials
+            return surface.material_front, surface.material_back
+
+        # CellGeometry: query cells for materials
+        # For cell geometry, we need to determine material on each side
+        # by checking which cell each side belongs to
+        cell_geom: CellGeometry = self._geometry  # type: ignore
+
+        # Get a representative position slightly on each side of the surface
+        # Use the first hit position and offset by a small amount in normal direction
+        if crossing.num_hits == 0:
+            return None, None
+
+        # Use hit normals (which are oriented toward the ray)
+        # For front material: offset in normal direction (toward where ray came from)
+        # For back material: offset opposite to normal (where ray is going)
+        sample_pos = crossing.hit_positions[0:1].astype(np.float64)
+        sample_normal = crossing.hit_normals[0:1].astype(np.float64)
+
+        # Offset to check material on each side
+        offset = 0.001  # Small offset
+        front_pos = sample_pos + offset * sample_normal
+        back_pos = sample_pos - offset * sample_normal
+
+        # Query cell geometry for materials
+        front_results = cell_geom.get_material_at(front_pos)
+        back_results = cell_geom.get_material_at(back_pos)
+
+        # Extract materials (first result that covers this position)
+        mat_front = None
+        mat_back = None
+
+        for material, mask in front_results:
+            if mask[0]:
+                mat_front = material
+                break
+
+        for material, mask in back_results:
+            if mask[0]:
+                mat_back = material
+                break
+
+        return mat_front, mat_back
+
+    def _process_detector_hit(
+        self,
+        rays: RayBatch,
+        crossing: SurfaceCrossing,
+    ) -> "RecordedRays":
+        """Process hits on a detector surface."""
+        from ...utilities.recording_sphere import RecordedRays
+
+        idx = crossing.ray_indices
+
+        # Get polarization vectors if available
+        polarization_vectors = None
+        if self._track_polarization_vector and rays.polarization_vector is not None:
+            polarization_vectors = rays.polarization_vector[idx]
+
+        return RecordedRays(
+            positions=crossing.hit_positions.copy(),
+            directions=crossing.hit_directions.copy(),
+            times=rays.accumulated_time[idx].copy(),
+            intensities=rays.intensities[idx].copy(),
+            wavelengths=rays.wavelengths[idx].copy(),
+            generations=rays.generations[idx].copy(),
+            polarization_vectors=polarization_vectors,
+            ray_indices=idx,  # Track original ray indices for correct mapping
+        )
+
+    def _process_optical_hit(
+        self,
+        rays: RayBatch,
+        crossing: SurfaceCrossing,
+    ) -> tuple[RayBatch | None, RayBatch | None]:
+        """
+        Process hits on an optical surface.
+
+        Computes Fresnel coefficients and generates reflected/refracted rays.
+        """
+        idx = crossing.ray_indices
+        positions = crossing.hit_positions
+        directions = crossing.hit_directions
+        normals = crossing.hit_normals
+        from_front = crossing.from_front
+        num_hits = crossing.num_hits
+
+        if num_hits == 0:
+            return None, None
+
+        # Get ray properties
+        wavelengths = rays.wavelengths[idx]
+        intensities = rays.intensities[idx]
+        accumulated_time = rays.accumulated_time[idx]
+        generations = rays.generations[idx]
+
+        # Handle polarization vectors
+        polarization_vectors = None
+        if self._track_polarization_vector:
+            if rays.polarization_vector is not None:
+                polarization_vectors = rays.polarization_vector[idx]
+            else:
+                polarization_vectors = initialize_polarization_vectors(
+                    directions, polarization=self._polarization
+                )
+
+        # Get materials based on hit direction
+        # from_front: light coming from front material
+        # from_back: light coming from back material
+        mat_front, mat_back = self._get_materials_at_crossing(crossing)
+
+        if mat_front is None or mat_back is None:
+            # Missing materials, treat as pass-through
+            return None, None
+
+        # Compute refractive indices
+        # n1 = material ray is coming FROM
+        # n2 = material ray is going TO
+        n1_values = np.empty(num_hits, dtype=np.float32)
+        n2_values = np.empty(num_hits, dtype=np.float32)
+
+        for i in range(num_hits):
+            pos = positions[i]
+            wl = wavelengths[i]
+            if from_front[i]:
+                n1_values[i] = mat_front.get_refractive_index(
+                    pos[0], pos[1], pos[2], wl
+                )
+                n2_values[i] = mat_back.get_refractive_index(pos[0], pos[1], pos[2], wl)
+            else:
+                n1_values[i] = mat_back.get_refractive_index(pos[0], pos[1], pos[2], wl)
+                n2_values[i] = mat_front.get_refractive_index(
+                    pos[0], pos[1], pos[2], wl
+                )
+
+        # Compute incident angle cosine
+        cos_theta_i = np.abs(np.sum(directions * normals, axis=1))
+
+        # Compute Fresnel coefficients
+        R, T = fresnel_coefficients(
+            n1_values, n2_values, cos_theta_i, self._polarization
+        )
+
+        # Get separate R_s, R_p if tracking polarization
+        R_s = R_p = None
+        if self._track_polarization_vector:
+            R_s, _ = fresnel_coefficients(n1_values, n2_values, cos_theta_i, "s")
+            R_p, _ = fresnel_coefficients(n1_values, n2_values, cos_theta_i, "p")
+
+        # Compute reflected directions
+        reflected_directions = compute_reflection_direction(directions, normals)
+        reflected_intensities = intensities * R
+
+        # Compute refracted directions
+        refracted_directions, tir_mask = compute_refraction_direction(
+            directions, normals, n1_values, n2_values
+        )
+        refracted_intensities = intensities * T
+        refracted_intensities[tir_mask] = 0.0
+
+        # Create reflected ray batch
+        reflected_rays = create_ray_batch(
+            num_rays=num_hits,
+            enable_polarization_vector=self._track_polarization_vector,
+        )
+        # Offset slightly to avoid re-intersection
+        reflected_rays.positions[:] = positions + 0.01 * reflected_directions
+        reflected_rays.directions[:] = reflected_directions
+        reflected_rays.wavelengths[:] = wavelengths
+        reflected_rays.intensities[:] = reflected_intensities
+        reflected_rays.optical_path_lengths[:] = rays.optical_path_lengths[idx]
+        reflected_rays.geometric_path_lengths[:] = rays.geometric_path_lengths[idx]
+        reflected_rays.accumulated_time[:] = accumulated_time
+        reflected_rays.generations[:] = generations + 1
+        reflected_rays.domain_ids[:] = rays.domain_ids[idx]
+        reflected_rays.active[:] = reflected_intensities > 1e-10
+
+        # Transform polarization vectors for reflection
+        if self._track_polarization_vector and polarization_vectors is not None:
+            reflected_pol = transform_polarization_reflection(
+                polarization_vectors,
+                directions,
+                reflected_directions,
+                normals,
+                R_s=R_s,
+                R_p=R_p,
+            )
+            reflected_rays.polarization_vector[:] = reflected_pol
+
+        # Create refracted ray batch
+        refracted_rays = create_ray_batch(
+            num_rays=num_hits,
+            enable_polarization_vector=self._track_polarization_vector,
+        )
+        # Offset slightly to avoid re-intersection
+        refracted_rays.positions[:] = positions + 0.01 * refracted_directions
+        refracted_rays.directions[:] = refracted_directions
+        refracted_rays.wavelengths[:] = wavelengths
+        refracted_rays.intensities[:] = refracted_intensities
+        refracted_rays.optical_path_lengths[:] = rays.optical_path_lengths[idx]
+        refracted_rays.geometric_path_lengths[:] = rays.geometric_path_lengths[idx]
+        refracted_rays.accumulated_time[:] = accumulated_time
+        refracted_rays.generations[:] = generations + 1
+        refracted_rays.domain_ids[:] = rays.domain_ids[idx]
+        refracted_rays.active[:] = (refracted_intensities > 1e-10) & (~tir_mask)
+
+        # Transform polarization vectors for refraction
+        if self._track_polarization_vector and polarization_vectors is not None:
+            refracted_pol = transform_polarization_refraction(
+                polarization_vectors,
+                directions,
+                refracted_directions,
+                normals,
+            )
+            refracted_pol[tir_mask] = 0.0
+            refracted_rays.polarization_vector[:] = refracted_pol
+
+        # Compact to remove inactive rays
+        reflected_rays = reflected_rays.compact()
+        refracted_rays = refracted_rays.compact()
+
+        return reflected_rays, refracted_rays
+
+    def get_surface(self, idx: int) -> Surface:
+        """Get surface by index."""
+        return self._surfaces[idx]
+
+    def get_surface_by_name(self, name: str) -> Surface | None:
+        """Get surface by name."""
+        for surface in self._surfaces:
+            if surface.name == name:
+                return surface
+        return None