# 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]
[docs]
@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
[docs]
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))
[docs]
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
[docs]
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
),
)
[docs]
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
),
)
[docs]
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
[docs]
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,
)
[docs]
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)),
)
[docs]
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,
)