# 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 Source Base Class
Defines the abstract base class for all ray sources in the raytracing
framework. Sources generate initial ray conditions for simulation.
Design Notes
------------
- Follows Interface Segregation Principle: focused on ray generation
- Derived classes implement specific spatial/angular distributions
- Ray intensities are normalized to conserve total power
"""
from abc import ABC, abstractmethod
import numpy as np
from numpy.typing import NDArray
from ..utilities.ray_data import RayBatch, create_ray_batch
[docs]
class RaySource(ABC):
"""
Abstract base class for ray sources.
A ray source defines initial conditions for a ray batch, including
spatial distribution, angular distribution, wavelength spectrum,
and intensity distribution.
Parameters
----------
num_rays : int
Number of rays to generate. Must be positive.
wavelength : float or tuple of float
Single wavelength in meters for monochromatic source, or
(min, max) tuple for polychromatic source.
power : float, optional
Total source power in watts. Default is 1.0.
Attributes
----------
num_rays : int
Number of rays generated by this source.
wavelength : float or tuple
Wavelength specification.
power : float
Total source power.
Notes
-----
Derived classes must implement the `generate()` method which returns
a fully initialized RayBatch.
The ray intensities should sum to the total power (conservation of energy).
When generating rays, use `_allocate_rays()` to create the batch with
proper initialization and `_assign_wavelengths()` to set wavelengths.
Examples
--------
Creating a custom source:
>>> class MySource(RaySource):
... def __init__(self, position, num_rays, wavelength, power=1.0):
... super().__init__(num_rays, wavelength, power)
... self.position = np.array(position)
...
... def generate(self):
... rays = self._allocate_rays()
... rays.positions[:] = self.position
... # Set directions...
... self._assign_wavelengths(rays)
... return rays
"""
[docs]
def __init__(
self,
num_rays: int,
wavelength: float | tuple[float, float],
power: float = 1.0,
):
"""
Initialize ray source.
Parameters
----------
num_rays : int
Number of rays to generate. Must be positive.
wavelength : float or tuple of float
Single wavelength (m) or (min, max) range.
power : float, optional
Total source power in watts. Default is 1.0.
Raises
------
ValueError
If num_rays <= 0, power <= 0, wavelength <= 0, or
wavelength range is invalid.
"""
self.num_rays = num_rays
self.wavelength = wavelength
self.power = power
# Validate parameters
if num_rays <= 0:
raise ValueError("num_rays must be positive")
if power <= 0:
raise ValueError("power must be positive")
if isinstance(wavelength, tuple):
if len(wavelength) != 2:
raise ValueError("wavelength range must be (min, max)")
if wavelength[0] >= wavelength[1]:
raise ValueError("wavelength min must be less than max")
if wavelength[0] <= 0:
raise ValueError("wavelength values must be positive")
elif wavelength <= 0:
raise ValueError("wavelength must be positive")
[docs]
@abstractmethod
def generate(self) -> RayBatch:
"""
Generate ray batch with initial conditions.
Creates and initializes a RayBatch with positions, directions,
wavelengths, and intensities according to the source configuration.
Returns
-------
RayBatch
Initialized rays ready for propagation.
Notes
-----
Implementations should ensure:
- All rays are marked as active
- Directions are normalized
- Intensities sum to total power
- Wavelengths are set appropriately
"""
pass
def _allocate_rays(self) -> RayBatch:
"""
Allocate ray batch with uniform intensity distribution.
Creates a RayBatch with all rays active and intensities set
so they sum to the total power.
Returns
-------
RayBatch
Allocated ray batch with intensities initialized.
Notes
-----
Called by generate() implementations to create the ray batch.
"""
rays = create_ray_batch(num_rays=self.num_rays)
rays.active[:] = True
rays.intensities[:] = self.power / self.num_rays
return rays
def _assign_wavelengths(self, rays: RayBatch) -> None:
"""
Assign wavelengths to rays.
For monochromatic sources, sets all wavelengths to the single value.
For polychromatic sources, samples uniformly from the range.
Parameters
----------
rays : RayBatch
Ray batch to assign wavelengths to.
Notes
-----
Called by generate() implementations after setting positions
and directions.
"""
if isinstance(self.wavelength, tuple):
# Uniform distribution over range
rays.wavelengths[:] = np.random.uniform(
self.wavelength[0], self.wavelength[1], self.num_rays
).astype(np.float32)
else:
# Monochromatic
rays.wavelengths[:] = self.wavelength
def _create_perpendicular_basis(
self, direction: NDArray[np.float32]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
"""
Create two unit vectors perpendicular to a direction.
Parameters
----------
direction : ndarray, shape (3,)
Direction vector (must be normalized).
Returns
-------
v1 : ndarray, shape (3,)
First perpendicular unit vector.
v2 : ndarray, shape (3,)
Second perpendicular unit vector (perpendicular to both
direction and v1).
Notes
-----
Used for generating positions/directions in a plane perpendicular
to the beam axis.
"""
if abs(direction[2]) < 0.9:
v1 = np.cross(direction, np.array([0, 0, 1], dtype=np.float32))
else:
v1 = np.cross(direction, np.array([1, 0, 0], dtype=np.float32))
v1 = v1 / np.linalg.norm(v1)
v2 = np.cross(direction, v1)
return v1.astype(np.float32), v2.astype(np.float32)
[docs]
def __repr__(self) -> str:
"""Return string representation."""
wl_str = (
f"({self.wavelength[0]:.2e}, {self.wavelength[1]:.2e})"
if isinstance(self.wavelength, tuple)
else f"{self.wavelength:.2e}"
)
return (
f"{self.__class__.__name__}(num_rays={self.num_rays}, "
f"wavelength={wl_str}, power={self.power})"
)