# 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.
"""
Protocol definitions for ray propagators and GPU materials.
This module defines:
- RayPropagatorProtocol: Common interface for CPU/GPU propagators
- MaterialFieldProtocol: Interface for materials providing refractive index
- GPUMaterialProtocol: Interface for GPU-accelerated materials
- GPUMaterialID: Enum for GPU material types
"""
from enum import IntEnum
from typing import Protocol, runtime_checkable
import numpy.typing as npt
# =============================================================================
# GPU Material Registry
# =============================================================================
[docs]
class GPUMaterialID(IntEnum):
"""
Registry of GPU-compatible material types.
Each material type has specialized device functions compiled into the
CUDA kernels. This enum identifies which material's functions to use.
Tiers:
- EXPONENTIAL_ATMOSPHERE (1): Legacy specialized kernels
- SIMPLE_INHOMOGENEOUS (2): 1D radial profile via LUT interpolation
- GRID_INHOMOGENEOUS (3): 3D grid data via trilinear interpolation
- SPECTRAL_INHOMOGENEOUS (4): 2D LUT-based (altitude × wavelength)
"""
EXPONENTIAL_ATMOSPHERE = 1 # Legacy: specialized hard-coded kernels
SIMPLE_INHOMOGENEOUS = 2 # Tier 1: 1D LUT-based (SimpleInhomogeneousModel)
GRID_INHOMOGENEOUS = 3 # Tier 2: 3D grid interpolation (GridInhomogeneousModel)
SPECTRAL_INHOMOGENEOUS = 4 # 2D LUT-based (SpectralInhomogeneousModel)
DUCT_ATMOSPHERE = 5 # Analytical duct atmosphere (no LUT)
US_DUCT_ATMOSPHERE = 6 # Analytical US Standard duct atmosphere (no LUT)
@runtime_checkable
class RayPropagatorProtocol(Protocol):
"""
Protocol for ray propagators in inhomogeneous media.
All propagators (CPU and GPU) must implement this interface to ensure
they are interchangeable drop-in replacements.
The propagator operates on RayBatch objects and modifies them in-place.
"""
def propagate_step(
self,
rays, # RayBatch - avoiding import to prevent circular dependency
step_size: float,
wavelength: float = 532e-9,
) -> None:
"""
Propagate rays by a single integration step.
This is the primary interface method that all propagators must implement.
It takes a RayBatch object and advances all active rays by one step.
Parameters
----------
rays : RayBatch
Ray batch containing positions, directions, and other ray properties.
Modified in-place.
step_size : float
Integration step size in meters. For gradient-based propagators,
this is the nominal step size that may be adaptively adjusted.
wavelength : float, optional
Wavelength in meters, default 532 nm (green light).
Used for wavelength-dependent effects like dispersion.
Notes
-----
- Modifies `rays.positions`, `rays.directions` in-place
- Updates path length accumulators if present
- Only propagates rays where `rays.active` is True
- Implementations may use different integration schemes (Euler, RK4, etc.)
"""
...
def propagate(
self,
rays, # RayBatch
total_distance: float,
step_size: float,
wavelength: float = 532e-9,
) -> None:
"""
Propagate rays through a total distance.
This is a convenience method that calls propagate_step repeatedly.
Implementations may optimize this for batch processing.
Parameters
----------
rays : RayBatch
Ray batch to propagate (modified in-place)
total_distance : float
Total distance to propagate in meters
step_size : float
Integration step size in meters
wavelength : float, optional
Wavelength in meters, default 532 nm
"""
...
@runtime_checkable
class MaterialFieldProtocol(Protocol):
"""
Protocol for material fields that provide refractive index.
Materials must implement these methods to work with gradient propagators.
"""
def get_refractive_index(
self,
x: npt.NDArray,
y: npt.NDArray,
z: npt.NDArray,
wavelength: float = 532e-9,
) -> npt.NDArray:
"""
Get refractive index at given positions.
Parameters
----------
x, y, z : ndarray
Positions in meters
wavelength : float
Wavelength in meters
Returns
-------
ndarray
Refractive index at each position
"""
...
def get_refractive_index_gradient(
self,
x: npt.NDArray,
y: npt.NDArray,
z: npt.NDArray,
wavelength: float = 532e-9,
) -> tuple[npt.NDArray, npt.NDArray, npt.NDArray]:
"""
Get gradient of refractive index at given positions.
Parameters
----------
x, y, z : ndarray
Positions in meters
wavelength : float
Wavelength in meters
Returns
-------
tuple of (dn_dx, dn_dy, dn_dz)
Gradient components in m^-1
"""
...
@runtime_checkable
class GPUMaterialProtocol(Protocol):
"""
Protocol for GPU-accelerated materials.
Materials implementing this protocol can be used with GPU propagators.
"""
@property
def gpu_material_id(self) -> GPUMaterialID:
"""Return GPU material ID enum value."""
...
def get_gpu_kernels(self) -> dict:
"""Return dict of GPU kernel functions keyed by method name."""
...
def get_gpu_parameters(self) -> tuple:
"""Return tuple of material-specific parameters for GPU kernels."""
...