Package¶

Config¶

Configuration module for diffinytrace.

This module provides global configuration options for controlling ray intersection solver behavior, such as tolerance, maximum iterations, damping factor, and whether to display iteration counts. These settings can be adjusted at runtime to tune performance and accuracy.

Example

>>> import diffinytrace.config as config
>>> config.set_tolerance(1e-8)
>>> config.set_show_iteration_count(True)
>>> config.restore_default_settings()

diffinytrace.config.set_show_iteration_count(flag)[source]¶

Set the option to show the number of iterations for each intersection.

Parameters:: flag (bool) – True to show the number of iterations, False otherwise.

diffinytrace.config.get_show_iteration_count()[source]¶

Check if the number of iterations should be shown.

Returns:: True if the number of iterations should be shown, False otherwise.
Return type:: bool

diffinytrace.config.set_tolerance(new_tolerance)[source]¶

Set the tolerance for ray intersection calculations.

Parameters:: new_tolerance (float) – The new tolerance value (must be > 0).
Raises:: ValueError – If new_tolerance is not greater than 0.

diffinytrace.config.get_tolerance()[source]¶

Get the current tolerance for ray intersection calculations.

Returns:: The current tolerance value.
Return type:: float

diffinytrace.config.set_max_iterations(new_max_iterations)[source]¶

Set the maximum number of iterations for the ray intersection solver.

Parameters:: new_max_iterations (int) – The new maximum number of iterations (must be > 0).
Raises:: ValueError – If new_max_iterations is not greater than 0.

diffinytrace.config.get_max_iterations()[source]¶

Get the current maximum number of iterations for the ray intersection solver.

Returns:: The current maximum number of iterations.
Return type:: int

diffinytrace.config.set_damping_factor(new_damping_factor)[source]¶

Set the damping factor for the Newton method used in ray intersections.

Parameters:: new_damping_factor (float) – The new damping factor (0 < new_damping_factor <= 1).
Raises:: ValueError – If new_damping_factor is not between 0 and 1 (exclusive of 0, inclusive of 1).

diffinytrace.config.get_damping_factor()[source]¶

Get the current damping factor for the Newton method.

Returns:: The current damping factor.
Return type:: float

diffinytrace.config.restore_default_settings()[source]¶

Reset to the default configuration settings for the ray tracer.

This will reset all configuration parameters to their default values.

Constraints¶

class diffinytrace.constraints.Constraint(fun, type)[source]¶

Bases: object

Base class for optimization constraints.

fun¶

Function defining the constraint.

Type:: Callable

type¶

Type of constraint (‘eq’ or ‘ineq’).

Type:: str

class diffinytrace.constraints.EqualZero(fun)[source]¶

Bases: Constraint

Equality constraint enforcing fun() == 0.

Parameters:: fun (Callable) – The constraint function.

class diffinytrace.constraints.GEQZero(fun)[source]¶

Bases: Constraint

Inequality constraint enforcing fun() >= 0.

Parameters:: fun (Callable) – The constraint function.

class diffinytrace.constraints.LEQZero(fun)[source]¶

Bases: Constraint

Inequality constraint enforcing fun() <= 0.

Parameters:: fun (Callable) – The constraint function.

Element¶

diffinytrace.element.is_valid_square_circle(transform: Transform, O: Tensor, aperture_radius: float, is_square: bool) → Tensor[source]¶

Checks whether points lie within a circular or square aperture after transformation.

Parameters:

transform (Transform) – Transformation object to convert global to local coordinates.
O (torch.Tensor) – Points in global coordinates of shape (N, 3).
aperture_radius (float or torch.Tensor) – Radius of the circular or square aperture.
is_square (bool) – If True, aperture is square; if False, circular.

Returns:

Boolean tensor of shape (N,) indicating whether each point lies within the aperture.

Return type:

torch.Tensor

Note

For a square, checks if ( |x| < r ) and ( |y| < r ). For a circle, checks if ( sqrt{x^2 + y^2} < r ).

class diffinytrace.element.OpticalSystem(modules_dict: Dict)[source]¶

Bases: Module, Plotable

Base class for optical systems composed of multiple optical modules.

This class serves as a container for modules such as lenses, mirrors, and detectors. It supports visualization and modular organization.

modules_dict¶

Dictionary of named optical modules.

Type:: nn.ModuleDict

forward()[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_plotable_childs()[source]¶: Returns a list of all plotable child objects of this object. Each child is represented as a list containing the child object and its name.

get_plot_points_2D(resolution: int)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution: int)[source]¶

Returns a list of 3D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 3D plot points.
Return type:: list

class diffinytrace.element.SequentialOpticalSystem(modules_dict: Dict, n_func_enviroment=RefractiveIndex())[source]¶

Bases: OpticalSystem

Optical system that processes rays in a defined sequence.

Useful for simulating light propagation through a sequence of elements, e.g., source → lens → detector.

n_func_enviroment¶

Function returning refractive index of the surrounding medium.

Type:: Callable

forward(x, mapping_sequence: List[str])[source]¶

Propagates rays through the defined sequence of modules.

Parameters:

x (Any) – Input rays or sampling data.
mapping_sequence (list[str]) – Ordered list of module names defining propagation sequence.

Returns:

Output after final module in the sequence.

Return type:

Any

class diffinytrace.element.OpticalElement(fill_color='white', outline_color='black', is_volume=False)[source]¶

Bases: PhysicalObject, Plotable

Abstract base class for optical elements like lenses, mirrors, and detectors.

Provides interface for geometric transformation and ray propagation.

forward(O2: Tensor, D2: Tensor, wl: Tensor, n_func_enviroment, meta_data)[source]¶

Propagates rays through the optical element.

Parameters:

O2 (torch.Tensor) – Ray origins.
D2 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment (Callable) – Function returning environmental refractive index.
meta_data (dict) – Dictionary with path length and validity information.

Raises:

NotImplementedError – Must be overridden by subclasses.

get_transform()[source]¶

Returns the transformation associated with the surface.

Returns:: The local-to-global transformation object.
Return type:: Transform

class diffinytrace.element.OpticalSurface(transform: Transform, surface, aperture_radius: float, is_square: bool = False, fill_color: str = 'white', outline_color: str = 'black')[source]¶

Bases: OpticalElement, PhysicalSurface

Represents a surface in 3D space with a defined aperture and transformation.

Supports both square and circular apertures, and provides methods for parametric sampling, CAD conversion, ray intersection, and plotting.

surface¶

Object with a method explicit(parametric_pos) returning z-values.

Type:: object

aperture_radius¶

Radius of the circular or square aperture.

Type:: float

is_square¶

Whether the aperture is square-shaped.

Type:: bool

transform¶

Local-to-global transformation.

Type:: Transform

integrator¶

Integration object (Disc or Cube) for parametric sampling.

Type:: Integrator

get_constraint_funs_leq_zero()[source]¶

Returns constraint functions used for integration and optimization over the surface.

Returns:: List of functions f(param_pos) <= 0 indicating valid parametric regions.
Return type:: list[Callable]
Raises:: RuntimeError – If is_square is True (not yet implemented).

get_plot_points_2D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 2D slices through the surface (z-y plane) for plotting.

Parameters:: resolution (int) – Number of sample points along the y-axis.
Returns:: List of (z, y) coordinate tuples.
Return type:: List[Tuple[torch.Tensor]]

get_plot_points_3D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 3D grid of surface points for visualization.

Parameters:: resolution (int) – Grid resolution in x and y.
Returns:: List of (x, y, z) meshgrids as torch tensors.
Return type:: List[Tuple[torch.Tensor]]

get_CAD_points(resolution: int) → List[Tuple[Tensor]][source]¶

Generates a 3D surface point grid for CAD conversion.

Parameters:: resolution (int) – Sampling resolution.
Returns:: (x, y, z) coordinate grids for CAD modeling.
Return type:: Tuple[torch.Tensor]

get_CAD_face(resolution: int, tol: float = 0.001, smoothing=None, minDeg: int = 1, maxDeg: int = 3)[source]¶

Converts the surface into a CAD face using B-spline approximation.

Parameters:

resolution (int) – Sampling resolution.
tol (float, optional) – Approximation tolerance. Defaults to 0.001.
smoothing (Optional[int]) – Smoothing value for fitting.
minDeg (int) – Minimum degree of the spline.
maxDeg (int) – Maximum degree of the spline.

Returns:

CAD face object.

Return type:

cadquery.Face

parametric_sample(num_points: int, method: str = 'sobol') → tuple[Tensor, Tensor][source]¶

Samples parametric positions on the aperture using the integrator.

Parameters:

num_points (int) – Number of sample points.
method (str) – Sampling method. Options: “sobol”, “monte_carlo”, “midpoint”, etc.

Returns:

Sampled positions and integration weights.

Return type:

Tuple[torch.Tensor, torch.Tensor]

parametric_surface(parametric_pos: Tensor) → Tensor[source]¶

Maps 2D parametric coordinates to 3D global coordinates using the surface height and transform.

Parameters:: parametric_pos (torch.Tensor) – 2D parametric positions of shape (N, 2).
Returns:: 3D positions of shape (N, 3) in global space.
Return type:: torch.Tensor
Raises:: RuntimeError – If input does not have shape […, 2].

get_surface_and_normal_func_with_params()[source]¶

Constructs a callable for surface position and normal computation with parameter tracking.

Returns:: Callable computes (position, normal), and the list contains parameters to be optimized.
Return type:: Tuple[Callable, List]

get_ray_intersect_length(O, D) → Tensor[source]¶

Computes intersection length along ray until hitting the surface.

Parameters:

O (torch.Tensor) – Ray origins of shape (N, 3).
D (torch.Tensor) – Ray directions of shape (N, 3).

Returns:

Intersection distances t such that O + t*D lies on the surface.

Return type:

torch.Tensor

get_new_is_valid(O, valid) → Tensor[source]¶

Updates a boolean mask indicating which rays are still valid after hitting the aperture.

Parameters:

O (torch.Tensor) – Ray intersection points.
valid (torch.Tensor) – Previous boolean validity mask.

Returns:

Updated validity mask.

Return type:

torch.Tensor

get_transform() → Transform[source]¶

Returns the transformation associated with the surface.

Returns:: The local-to-global transformation object.
Return type:: Transform

class diffinytrace.element.LensSurfaceTransmissionEnter(transform: Transform, surface, aperture_radius: float, n_func, is_square: bool = False)[source]¶

Bases: OpticalSurface

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the lens entry surface.

Parameters:

O1 (torch.Tensor) – Ray origins.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.element.LensSurfaceTransmissionLeave(transform: Transform, surface, aperture_radius: float, n_func, is_square: bool = False)[source]¶

Bases: OpticalSurface

forward(O2, D2, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the lens exit surface.

Parameters:

O2 (torch.Tensor) – Ray origins.
D2 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.element.LensSurfaceSide(surface1: PhysicalSurface, surface2: PhysicalSurface, aperture_radius, is_square: bool)[source]¶

Bases: PhysicalSurface, Plotable

Non-optical surface connecting two curved lens surfaces for visualization.

Used to render the full 3D body of the lens.

surface1¶

First lens surface.

Type:: PhysicalSurface

surface2¶

Second lens surface.

Type:: PhysicalSurface

aperture_radius¶

Radius or half-width of aperture.

Type:: float

is_square¶

Whether aperture is square.

Type:: bool

parametric_sample(num_points: int, method: str = 'sobol')[source]¶

Samples parametric positions on the lens side surface.

Parameters:

num_points (int) – Number of sample points.
method (str) – Sampling method (“sobol”, “monte_carlo”, “midpoint”).

Returns:

Sampled positions and integration weights.

Return type:

Tuple[torch.Tensor, torch.Tensor]

Raises:

RuntimeError – If unsupported method is provided.

parametric_surface(parametric_pos: Tensor) → Tensor[source]¶

Maps parametric coordinates to 3D global coordinates for the lens side.

Parameters:: parametric_pos (torch.Tensor) – Parametric positions of shape (N, 2).
Returns:: 3D positions of shape (N, 3).
Return type:: torch.Tensor

get_plot_points_2D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 2D slices through the surface (z-y plane) for plotting.

Parameters:: resolution (int) – Number of sample points along the y-axis.
Returns:: List of (z, y) coordinate tuples.
Return type:: List[Tuple[torch.Tensor]]

get_plot_points_3D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 3D grid of surface points for visualization.

Parameters:: resolution (int) – Grid resolution in x and y.
Returns:: List of (x, y, z) meshgrids as torch tensors.
Return type:: List[Tuple[torch.Tensor]]

get_plotly_color_scale()[source]¶

Returns color scale for plotly visualization.

Returns:: Color scale values.
Return type:: List

class diffinytrace.element.Lens(transform: Transform, lens_thickness: float, surface1, surface2, n_func, aperture_radius: float, is_square=False)[source]¶

Bases: OpticalElement

Represents a transmissive lens consisting of two refractive surfaces.

The lens is modeled as a sequence of: - Entry surface (refraction from external medium into the lens) - Exit surface (refraction from lens into external medium) - Side surface (purely for visualization)

In our implementation, lenses consist of two explicit surfaces, a transformation matrix \(M\), a lens thickness, an aperture radius, and a material. When the lens is initialized, one can also optionally specify whether the lens is round or square. If the keyword is_square is not specified, the lens will default to being round.

Example

Below is an example of initializing a square lens:

>>> import diffinytrace as dit
>>> aperture_half = 30.
>>> lens_thickness = 8.
>>> material = dit.materials["NBK7"]
>>> transform = dit.transforms.Identity()
>>> bspline = dit.Bspline(aperture_half, [3, 3], [8, 8])
>>> plane = dit.Plane()
>>> lens = dit.Lens(transform, lens_thickness,
>>>          bspline, plane,
>>>          material, aperture_half, is_square=True)

n_func¶

Function mapping wavelength to refractive index of the lens material.

Type:: Callable

_transform1¶

Transform for the first surface.

Type:: Transform

_transform2¶

Transform for the second surface.

Type:: Transform

lens_thickness¶

Learnable thickness of the lens.

Type:: torch.nn.Parameter

surface1¶

Entry surface.

Type:: LensSurfaceTransmissionEnter

surface2¶

Exit surface.

Type:: LensSurfaceTransmissionLeave

lens_surface_side¶

Side surface (for 3D rendering).

Type:: LensSurfaceSide

aperture_radius¶

Radius (or half-width) of aperture.

Type:: float

is_square¶

Whether the aperture is square.

Type:: bool

get_plot_points_2D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 2D slices through the lens for plotting.

Parameters:: resolution (int) – Number of sample points.
Returns:: List of (z, y) coordinate tuples.
Return type:: List[Tuple[torch.Tensor]]

get_plot_points_3D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 3D grid of lens surface points for visualization.

Parameters:: resolution (int) – Grid resolution.
Returns:: List of (x, y, z) meshgrids.
Return type:: List[Tuple[torch.Tensor]]

get_plotly_color_scale() → List[source]¶

Returns color scale for plotly visualization.

Returns:: Color scale values.
Return type:: List

get_plotable_childs() → List[source]¶

Returns plotable child elements.

Returns:: List of child elements.
Return type:: List

forward(O1: Tensor, D1: Tensor, wl: Tensor, n_func_enviroment, meta_data)[source]¶

Simulates light passing through the lens.

Parameters:

O1 (torch.Tensor) – Ray origin positions.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment (Callable) – Function returning external medium refractive index.
meta_data (dict) – Ray metadata (PL, OPL, paths, valid).

Returns:

Updated ray origins, directions, etc.

Return type:

Tuple[torch.Tensor]

get_transform()[source]¶

Returns the transformation of the lens exit surface.

Returns:: The transformation object.
Return type:: Transform

class diffinytrace.element.Mirror(transform, surface, aperture_radius, is_square=False)[source]¶

Bases: OpticalSurface

Reflective optical element that reflects rays according to the law of reflection.

Visualization is colored in a warm gold tone.

Inherits:: OpticalSurface: Full support for surface transformation and intersection.

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the mirror surface.

Parameters:

O1 (torch.Tensor) – Ray origins.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.element.Detector(transform, surface, aperture_radius, is_square=True)[source]¶

Bases: OpticalSurface

Represents a terminal optical element that collects ray data.

Detectors consist of an explicit surface, a transformation matrix \(M\), and an aperture radius. The detector class represents a target surface used to track the rays that hit it. When the detector is initialized, one can also optionally specify whether the detector is round or square. If the keyword is_square is not specified, the detector defaults to being square.

Example

Below is an example of how to initialize a detector:

>>> import diffinytrace as dit
>>> aperture_half = 30.
>>> transform = dit.transforms.Identity()
>>> plane = dit.Plane()
>>> detector = dit.Detector(transform, plane,
>>>                         aperture_half, is_square=False)

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Captures the final ray interaction without altering its direction.

Parameters:

O1 (torch.Tensor) – Ray origin.
D1 (torch.Tensor) – Ray direction.
wl (torch.Tensor) – Wavelength.
n_func_enviroment (Callable) – Function for surrounding medium.
meta_data (dict) – Ray tracing metadata.

Returns:

Final ray data.

Return type:

Tuple[torch.Tensor]

diffinytrace.element.trace_to_detector(optical_system: SequentialOpticalSystem, sequence: List, source, detector: Detector, num_rays: int, device=device(type='cpu'), method_ray_tracing: str = 'sobol_pow2')[source]¶

Traces rays through a system to a detector and returns the impact coordinates.

Parameters:

optical_system (SequentialOpticalSystem) – Ray-tracing pipeline.
sequence (list[str]) – Ordered names of system modules.
source – Source object with .sample() method.
detector (Detector) – Final surface to collect rays.
num_rays (int) – Number of rays to simulate.
device – Torch device (CPU/GPU).
method_ray_tracing (str) – Sampling method for source rays.

Returns:

(input samples, weights, detector plane hits, wavelengths)

Return type:

Tuple[torch.Tensor]

diffinytrace.element.set_unused_params_to_zero(optical_system: SequentialOpticalSystem, sequence, source, params, num_rays=200000, method_ray_tracing='sobol')[source]¶

Sets unused parameters (those with zero gradient across ray paths) to zero.

Parameters:

optical_system (SequentialOpticalSystem) – Full system.
sequence (list) – Ordered module names.
source – Ray source.
params (list[torch.nn.Parameter] or torch.nn.Parameter) – Parameters to clean.
num_rays (int) – Ray sample count.
method_ray_tracing (str) – Sampling method.

diffinytrace.element.get_unused_params_mask(optical_system: SequentialOpticalSystem, sequence: List[str], source, params, num_rays: int = 100000, method_ray_tracing='sobol') → List[BoolTensor][source]¶

Returns a boolean mask identifying which parameters are unused in the ray tracing process.

Parameters:

optical_system (SequentialOpticalSystem) – Full system.
sequence (list) – Ordered module names.
source – Ray source.
params (list[torch.nn.Parameter]) – Parameter list.
num_rays (int) – Number of rays to test.
method_ray_tracing (str) – Sampling method.

Returns:

Masks of the same shape as each parameter.

Return type:

list[torch.BoolTensor]

diffinytrace.element.set_used_params_bounds_to_constant(optical_system, sequence, source, params, bounds_attr_name_new, bounds_attr_name_old='bounds', num_rays=100000, method_ray_tracing='sobol')[source]¶

Locks unused parameters by copying their current value as bounds, making them constant.

Parameters:

bounds_attr_name_new (str) – Name of the new bounds attribute to write.
bounds_attr_name_old (str) – Name of the original bounds attribute.

class diffinytrace.element.FresnelOpticalSurface(transform, surface, aperture_radius, surface_derivative_x, surface_derivative_y, is_square=False)[source]¶

Bases: OpticalSurface

get_virtual_normals(O)[source]¶

Computes virtual surface normals for Fresnel surfaces.

Parameters:: O (torch.Tensor) – Positions.
Returns:: Virtual normals.
Return type:: torch.Tensor

class diffinytrace.element.FresnelVirtualLensSurfaceTransmissionEnter(transform, surface, aperture_radius, n_func, surface_derivative_x, surface_derivative_y, is_square=False)[source]¶

Bases: FresnelOpticalSurface

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the Fresnel lens entry surface.

Parameters:

O1 (torch.Tensor) – Ray origins.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.element.FresnelVirtualLensSurfaceTransmissionLeave(transform, surface, aperture_radius, n_func, surface_derivative_x, surface_derivative_y, is_square=False)[source]¶

Bases: FresnelOpticalSurface

forward(O2, D2, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the Fresnel lens exit surface.

Parameters:

O2 (torch.Tensor) – Ray origins.
D2 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.element.FresnelVirtualLens(transform, lens_thickness, surface1, surface2, n_func, aperture_radius, surface1_derivative_x=None, surface1_derivative_y=None, surface2_derivative_x=None, surface2_derivative_y=None, is_square=False)[source]¶

Bases: OpticalElement

get_plot_points_2D(resolution: int)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution)[source]¶

Returns 3D grid of Fresnel lens surface points for visualization.

Parameters:: resolution (int) – Grid resolution.
Returns:: List of (x, y, z) meshgrids.
Return type:: List[Tuple[torch.Tensor]]

get_plotly_color_scale()[source]¶

Returns color scale for plotly visualization.

Returns:: Color scale values.
Return type:: List

get_plotable_childs() → List[source]¶

Returns plotable child elements.

Returns:: List of child elements.
Return type:: List

forward(O1: Tensor, D1: Tensor, wl: Tensor, n_func_enviroment, meta_data) → Tensor[source]¶

Simulates light passing through the Fresnel lens.

Parameters:

O1 (torch.Tensor) – Ray origin positions.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning external medium refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, etc.

Return type:

Tuple

get_transform() → Transform[source]¶

Returns the transformation of the Fresnel lens exit surface.

Returns:: The transformation object.
Return type:: Transform

diffinytrace.element.compute_reflected_directions(D: Tensor, N: Tensor) → Tensor[source]¶

Computes reflected ray directions using the reflection law.

Parameters:

D (torch.Tensor) – Incident directions of shape (M, 3), normalized.
N (torch.Tensor) – Surface normals at points of incidence, shape (M, 3).

Returns:

Reflected directions of shape (M, 3).

Return type:

torch.Tensor

diffinytrace.element.get_refracted_directions(D: Tensor, N: Tensor, n1: Tensor | float, n2: Tensor | float) → Tensor[source]¶

Computes refracted ray directions using Snell’s law.

At material interfaces, the transmitted direction \(\mathbf{D'}\) is computed based on the surface normal \(\mathbf{N} = \nabla s / \|\nabla s\|\) and the incident direction \(\mathbf{D}\), using Snell’s law (see [WCH22]):

\[\mathbf{D'} = \mathbf{N} \sqrt{1 - (1 - \cos^2 \psi_i) \eta^2} + \eta (\mathbf{D} - \mathbf{N} \cos \psi_i),\]

where \(\cos \psi_i = \mathbf{D} \cdot \mathbf{N}\) and \(\eta = n / n'\) is the ratio of the refractive indices of the two materials.

Parameters:

D (torch.Tensor) – Incident directions of shape (M, 3), normalized.
N (torch.Tensor) – Surface normals at points of incidence, shape (M, 3).
n1 (float or torch.Tensor) – Refractive index of the incident medium.
n2 (float or torch.Tensor) – Refractive index of the transmission medium.

Returns:

Refracted directions of shape (M, 3).

Return type:

torch.Tensor

diffinytrace.element.set_unused_bspline_coeff_to_nearest(optical_system, sequence: list[str], source, bspline_surface, num_rays=100000, method_ray_tracing='sobol')[source]¶

Fills only the unused B-spline coefficients with the nearest used value.

This function identifies B-spline coefficients that have no influence on the ray paths (i.e., gradients are zero), and updates only those by copying the value from the closest neighboring coefficient that is used. Used coefficients remain unchanged.

This is useful for having geometry that is simple to manifacture while not tempering with the overall performance.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system used for tracing.
sequence (list[str]) – Ordered list of module names for ray propagation.
source – Ray source with a .sample() method.
bspline_surface – Surface object with a .coeff tensor.
num_rays (int, optional) – Number of rays used to detect unused coefficients. Default is 100000.
method_ray_tracing (str, optional) – Sampling method (e.g., “sobol”). Default is “sobol”.

Raises:

RuntimeError – If all coefficients are unused — likely due to insufficient ray coverage.

Gaussian Smoother¶

diffinytrace.gaussian_smoother.gaussian_func_1D(eval_points: Tensor, x_range, num_gauss_points: int, sigma: float, include_boundary=False) → Tensor[source]¶

Gaussian function for 1D convolution.

Parameters:

eval_points (torch.Tensor) – Points where the Gaussian function is evaluated.
x_range (tuple) – Range of the target plane.
num_gauss_points (int) – Number of Gaussian points.
sigma (float) – Standard deviation of the Gaussian function.
include_boundary (bool) – Whether to include the boundary points.

Returns:

Evaluated Gaussian function.

Return type:

torch.Tensor

diffinytrace.gaussian_smoother.gaussian_func_2D(eval_points: Tensor, x_range, y_range, x_grid_size: int, y_grid_size: int, sigma: float | Tensor, val_multi: Tensor | None = None, summed: bool = True, include_boundary=False) → Tensor[source]¶

Gaussian function for 2D convolution.

Parameters:

eval_points (torch.Tensor) – Points where the Gaussian function is evaluated.
y_range (tuple) – Range of the target plane in the vertical direction.
x_range (tuple) – Range of the target plane in the horizontal direction.
y_grid_size (int) – Number of Gaussian points in the vertical direction.
x_grid_size (int) – Number of Gaussian points in the horizontal direction.
sigma (float) – Standard deviation of the Gaussian function.
val_multi (torch.Tensor|None) – Optional multiplier for the Gaussian function.
summed (bool) – Whether to sum the Gaussian function.
include_boundary (bool) – Whether to include the boundary points.

Returns:

Evaluated Gaussian function.

Return type:

torch.Tensor

diffinytrace.gaussian_smoother.calc_smooth_desired_irradiance(desired_irradiance_fun: Callable, x_range: List[float], y_range: List[float], x_grid_size: int, y_grid_size: int, sigma: float, num_integration_points: int, num_splits=5, dtype=torch.float64, device=device(type='cpu')) → Tensor[source]¶

Calculates the smoothed desired irradiance using Gaussian convolution.

Parameters:

desired_irradiance_fun (Callable) – Function that computes the desired irradiance at given points.
x_range (List[float]) – Range of the target plane in the x direction [min, max].
y_range (List[float]) – Range of the target plane in the y direction [min, max].
x_grid_size (int) – Number of pixels in the x direction.
y_grid_size (int) – Number of pixels in the y direction.
sigma (float) – Standard deviation of the Gaussian kernel.
num_integration_points (int) – Number of integration points for numerical integration.
num_splits (int, optional) – Number of splits for integration to reduce memory usage. Defaults to 5.
dtype (torch.dtype, optional) – Data type for tensors. Defaults to torch.get_default_dtype().
device (torch.device, optional) – Device for computation. Defaults to torch.get_default_device().

Returns:

Smoothed desired irradiance map.

Return type:

torch.Tensor

class diffinytrace.gaussian_smoother.GaussianSmoother(x_range: list, y_range: list, x_grid_size: int, y_grid_size: int, sigma: float, desired_irradiance_fun: Callable, smoothed_num_integration_points: int, smoothed_num_splits: int, dtype=torch.float64, device=device(type='cpu'))[source]¶

Bases: object

The GaussianSmoother class implements gaussian measurement functions but also computes smoothed desired irradiance distributions. For more information on this class please refer to the examples.

Parameters:

x_range (list) – Range of the target plane in the x direction [min, max].
y_range (list) – Range of the target plane in the y direction [min, max].
x_grid_size (int) – Number of pixels in the x direction.
y_grid_size (int) – Number of pixels in the y direction.
sigma (float) – Standard deviation of the Gaussian kernel.
desired_irradiance_fun (Callable) – Function that computes the desired irradiance at given points.
smoothed_num_integration_points (int) – Number of integration points for smoothing.
smoothed_num_splits (int) – Number of splits for integration to reduce memory usage.
dtype (torch.dtype, optional) – Data type for tensors. Defaults to torch.get_default_dtype().
device (torch.device, optional) – Device for computation. Defaults to torch.get_default_device().

x_grid_size¶

Number of pixels in the x direction.

Type:: int

y_grid_size¶

Number of pixels in the y direction.

Type:: int

sigma¶

Standard deviation of the Gaussian kernel.

Type:: float

include_boundary¶

Whether to include boundary points in the grid.

Type:: bool

x_range¶

Range of the target plane in the x direction.

Type:: list

y_range¶

Range of the target plane in the y direction.

Type:: list

grid¶

Grid object for pixel centers.

Type:: Grid

discrete_desired_irradiance¶

Desired irradiance at pixel centers.

Type:: torch.Tensor

smoothed_desired_irradiance¶

Smoothed desired irradiance map.

Type:: torch.Tensor

smoothed_irradiance(points: Tensor, ray_multi: Tensor, x_range=None, y_range=None) → Tensor[source]¶

Computes the smoothed irradiance at given points using a Gaussian kernel.

Parameters:

points (torch.Tensor) – Array of points where the irradiance is evaluated, shape [N, 2].
ray_multi (torch.Tensor) – Multiplicative weights for each point, e.g., ray flux.
x_range (tuple, optional) – Range of the target plane in the x direction. Defaults to None.
y_range (tuple, optional) – Range of the target plane in the y direction. Defaults to None.

Returns:

Smoothed irradiance values at the specified points.

Return type:

torch.Tensor

none_smoothed_irradiance(points: Tensor, ray_multi: Tensor) → Tensor[source]¶

Computes the non-smoothed irradiance at given points by summing ray contributions in each grid cell.

Parameters:

points (torch.Tensor) – Array of points where the irradiance is evaluated, shape [N, 2].
ray_multi (torch.Tensor) – Multiplicative weights for each point, e.g., ray flux.

Returns:

Non-smoothed irradiance values at the specified grid cells.

Return type:

torch.Tensor

integrate_values(vals: Tensor, x_range=None, y_range=None) → Tensor[source]¶

Integrates the provided values over the grid using midpoint rule.

Parameters:: vals (torch.Tensor) – Values to integrate, typically irradiance or residuals, shape matching the grid.
Returns:: The integrated sum over the grid.
Return type:: torch.Tensor

diffinytrace.gaussian_smoother.make_evaluation_function(optical_system: SequentialOpticalSystem, sequence: List, source: LightSource, detector, smoother: GaussianSmoother, num_splits: int = 10, num_rays_per_split: int = 100000, method_ray_tracing='monte_carlo', device=device(type='cpu')) → Callable[source]¶

Creates an evaluation function for comparing simulated and desired irradiance.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system to be used for ray tracing.
sequence – The sequence of optical elements.
source (LightSource) – The light source for the simulation.
detector – The detector object.
smoother (GaussianSmoother) – Smoother object for irradiance comparison.
num_splits (int, optional) – Number of splits for ray tracing to reduce memory usage. Defaults to 10.
num_rays_per_split (int, optional) – Number of rays per split. Defaults to 1,000,000.
method_ray_tracing (str, optional) – Ray tracing method (‘monte_carlo’, etc.). Defaults to “monte_carlo”.
device (torch.device, optional) – Device for computation. Defaults to torch.get_default_device().

Returns:

A function that computes the L2 error between simulated and desired irradiance.

Return type:

Callable

diffinytrace.gaussian_smoother.make_merit_function(optical_system: SequentialOpticalSystem, sequence: List, source: LightSource, detector, smoother: GaussianSmoother, num_rays: int, method_ray_tracing='sobol_pow2', use_desired_irradiance_smoothing=True, device=device(type='cpu'), T_margin=None) → Callable[source]¶

Creates a merit function to obtain a desired irradiance distribution for the given optical system, source, and detector.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system to be used.
sequence – The sequence of elements in the optical system.
source (LightSource) – The light source to be used.
detector – The detector to be used.
num_rays (int) – Number of rays to be traced.
smoother (Smoother) – The smoother object for merit function calculation.
device – The device to be used for calculations.
method_ray_tracing (str) – Method for ray tracing (‘sobol’ or ‘midpoint’).
use_desired_irradiance_smoothing (bool) – Whether to use desired irradiance smoothing.
use_power_correction (bool) – Whether to use power correction.
save_last_eval (bool) – Whether to save the last evaluation.
T_margin (float|None) – Optional margin for integration domain if it is None the integration domain will not be adjusted on the fly.

Returns:

A function that computes the merit value.

Return type:

Callable

class diffinytrace.gaussian_smoother.GaussianSmootherSquare(aperture_radius: list, grid_size: int, sigma: float, desired_irradiance_fun: Callable, smoothed_num_integration_points: int, smoothed_num_splits: int, dtype=torch.float64, device=device(type='cpu'))[source]¶

Bases: GaussianSmoother

This class is a specialized version of GaussianSmoother for cases where the x and y ranges are identical, and the grid is square (same number of pixels in both directions).

Parameters:

x_range (list) – Range of the target plane in both x and y directions [min, max].
x_grid_size (int) – Number of pixels in both x and y directions.
sigma (float) – Standard deviation of the Gaussian kernel.
desired_irradiance_fun (Callable) – Function that computes the desired irradiance at given points.
smoothed_num_integration_points (int) – Number of integration points for smoothing.
smoothed_num_splits (int) – Number of splits for integration to reduce memory usage.
dtype (torch.dtype, optional) – Data type for tensors. Defaults to torch.get_default_dtype().
device (torch.device, optional) – Device for computation. Defaults to torch.get_default_device().

Hat Smoother¶

class diffinytrace.hat_smoother.HatSmoother(x_range: list, y_range: list, x_grid_size: int, y_grid_size: int, desired_irradiance_fun: Callable, dtype=torch.float64, device=device(type='cpu'))[source]¶: Bases: object

Integrators¶

class diffinytrace.integrators.Integrator[source]¶

Bases: object

sample(num_points: int | list[int], method: IntegrationMethod) → tuple[Tensor, Tensor][source]¶

Sample points and weights using the specified method. :param num_points: Number of points in each dimension. :type num_points: int or list :param method: The integration method to use. Options are ‘simpson’, ‘midpoint’, ‘monte_carlo’, ‘sobol’, and ‘sobol_pow2’. :type method: str

Returns:: A tuple containing the sampled points and their corresponding weights.
Return type:: tuple

in_bounds(x: Tensor) → Tensor[source]¶

get_volume() → float[source]¶

class diffinytrace.integrators.Cube(bounds)[source]¶

Bases: Integrator

Integrator for a multi-dimensional cube (hyperrectangle).

Parameters:: bounds (array-like) – The bounds for each dimension of the cube. Should be a list or array of shape (n_dim, 2), where each row specifies [lower_bound, upper_bound] for a dimension.

Example

>>> cube = dit.integrators.Cube([[0, 1], [0, 1]])
>>> points, weights = cube.sample([10, 10], method=IntegrationMethod.MIDPOINT)
>>> volume = cube.get_volume()
>>> all_in_bounds = cube.in_bounds(points)
>>> print("Sampled points:", points)
>>> print("Integration weights:", weights)
>>> print("Cube volume:", volume)
>>> print("All points in bounds:", all_in_bounds)

sample(num_points: int | list[int], method: IntegrationMethod = IntegrationMethod.MIDPOINT) → tuple[Tensor, Tensor][source]¶

Sample points and weights using the specified method.

Parameters:

num_points (int or list) – Number of points in each dimension.
method (str) – The integration method to use. Options are ‘simpson’, ‘midpoint’, ‘monte_carlo’, ‘sobol’, and ‘sobol_pow2’.

Returns:

A tuple containing the sampled points and their corresponding weights.

Return type:

tuple

in_bounds(x: Tensor) → Tensor[source]¶

get_volume() → float[source]¶

Returns:: Volume of the Cube.
Return type:: float

class diffinytrace.integrators.Disc(radius)[source]¶

Bases: Integrator

Integrator for a 2D disc (circle).

Parameters:: radius (float) – The radius of the disc.

Example

>>> disc = dit.integrators.Disc(1.0)
>>> points, weights = disc.sample(2**4, method="sobol_pow2")
>>> volume = disc.get_volume()
>>> all_in_bounds = disc.in_bounds(points)
>>> print("Sampled points:", points)
>>> print("Integration weights:", weights)
>>> print("Disc area:", volume)
>>> print("All points in bounds:", all_in_bounds)

sample(num_points: int | list[int], method: IntegrationMethod = IntegrationMethod.SOBOL) → tuple[Tensor, Tensor][source]¶

Sample points and weights using the specified method.

Parameters:

num_points (int or list) – Number of points in each dimension.
method (str) – The integration method to use. Options are ‘simpson’, ‘midpoint’, ‘monte_carlo’, ‘sobol’, and ‘sobol_pow2’.

Returns:

A tuple containing the sampled points and their corresponding weights.

Return type:

tuple

in_bounds(x)[source]¶

Check if points are within the disc.

Parameters:: x (torch.Tensor) – Points to check.
Returns:: Boolean tensor indicating if points are within the disc.
Return type:: torch.Tensor

get_volume()[source]¶

Calculate the volume of the disc.

Returns:: Volume of the disc.
Return type:: float

class diffinytrace.integrators.IntegrationMethod(*values)[source]¶

Bases: Enum

SIMPSON = 'simpson'¶

MIDPOINT = 'midpoint'¶

MONTE_CARLO = 'monte_carlo'¶

SOBOL = 'sobol'¶

SOBOL_POW2 = 'sobol_pow2'¶

Intersection¶

class diffinytrace.intersection.SemiFunctionalModule[source]¶

Bases: Module

Abstract base class for semi-functional surface modules.

These modules define a static method functional that computes a functional transformation on inputs and parameters, and a method to list their functional parameters for optimization purposes.

static functional(O: Tensor, *params)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

diffinytrace.intersection.cat_semi_functionals(functional_modules: List[SemiFunctionalModule]) → Callable[source]¶

Recursively chains a list of `SemiFunctionalModule`s into a single composite function.

Each module’s functional() method is applied in sequence using the respective slice of the parameter list.

Parameters:: functional_modules (list[SemiFunctionalModule]) – List of functional modules.
Returns:: A function f(O, *params) that applies all modules in sequence.
Return type:: Callable

diffinytrace.intersection.get_functional_param_args(semi_functional_module_list: List[SemiFunctionalModule]) → List[source]¶

Collects all functional parameters from a list of semi-functional modules.

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: Flattened list of all parameters.
Return type:: list[torch.nn.Parameter]

diffinytrace.intersection.construct_surface_and_normal_func(semi_functional_module_list: List[SemiFunctionalModule]) → Callable[source]¶

Constructs a function to evaluate both the surface value and its gradient (normal direction) with respect to the ray origin O.

The surface is defined by composing the provided semi-functional modules.

Returns a callable:

\[(O, p_1, ..., p_n) \mapsto ( s(O), \frac{\partial s}{\partial O} )\]

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: A function s_dsd(O, *params, only_s=False) returning surface value s and optionally gradient ds/dO.
Return type:: Callable

diffinytrace.intersection.construct_surface_and_normal_func_with_params(semi_functional_module_list: List[SemiFunctionalModule]) → Tuple[Callable, List][source]¶

Constructs both the surface function and a list of its functional parameters.

Useful for optimization workflows that require parameter tracking.

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: Callable: A function computing surface and its gradient. list[torch.nn.Parameter]: The list of parameters for the surface.
Return type:: tuple

class diffinytrace.intersection.CustomAutogradRule_t(*args, **kwargs)[source]¶

Bases: Function

Custom PyTorch autograd rule for ray-surface intersection.

Computes a differentiable intersection length t such that:

\[s(O + t D) = 0\]

where O is the ray origin, D is the direction, and s is the surface function.

This rule enables backpropagation through t with respect to O, D, and surface parameters.

static forward(ctx, O: Tensor, D: Tensor, surface_and_normal_func: Callable, t_detached: Tensor, *param_args) → Tensor[source]¶

Stores inputs for backward pass and returns precomputed t.

Parameters:

O (torch.Tensor) – Ray origin of shape (N, 3).
D (torch.Tensor) – Ray direction of shape (N, 3).
surface_and_normal_func (Callable) – Surface function returning (s, ds/dR).
t_detached (torch.Tensor) – Estimated intersection length (detached).
*param_args – Surface parameters.

Returns:

Intersection length t.

Return type:

torch.Tensor

static backward(ctx, grad_outputs: Tensor) → Tuple[source]¶

Computes gradients of intersection length t with respect to: - ray origin O - ray direction D - surface parameters

Parameters:: grad_outputs (torch.Tensor) – Gradient of the loss w.r.t. output t.
Returns:: Gradients with respect to inputs (O, D, None, None, *param_args).
Return type:: tuple

diffinytrace.intersection.get_ray_intersection_length(O: Tensor, D: Tensor, surface_and_normal_func: Callable, param_args: List, t_init: Tensor | None = None) → Tensor[source]¶

Solves for the intersection length t such that:

\[s(O + t D) = 0\]

using a Newton-style iteration method with damping.

This function finds the length t where a ray intersects a parametric surface, given by a composed function with normal information.

Parameters:

O (torch.Tensor) – Ray origins of shape (N, 3).
D (torch.Tensor) – Ray directions of shape (N, 3).
surface_and_normal_func (Callable) – A function returning (s, ds/dR).
param_args (list) – List of surface parameters.
t_init (torch.Tensor, optional) – Initial guess for t. If None, starts from zero.

Returns:

Estimated intersection lengths t with autograd support.

Return type:

torch.Tensor

Raises:

Warning is printed (not exception) if convergence fails within max_iter. –

Optimize¶

Optimization Utilities for PyTorch-SciPy Integration¶

This submodule provides a set of tools for constrained and unconstrained optimization of PyTorch models using SciPy optimizers. It bridges the gap between SciPy’s powerful optimization routines and PyTorch’s autograd system, enabling flexible and efficient hybrid optimization workflows.

Key Features:¶

Seamless wrapping of PyTorch-based objective functions for use with SciPy.
Automatic gradient computation using PyTorch’s autograd.
Support for parameter bounds, including custom mask-based bounds.
Caching and reuse of recent function/gradient evaluations.
Integration with SciPy’s minimize.
Optional tracking of optimization history (function values and gradient norms).
Utility functions for flattening/unpacking tensor parameters.
Conversion of PyTorch parameters to SciPy-compatible formats with bounds.
Support for custom constraints and callback functions.

Optimization Constraints in Optical Systems¶

When using optimization procedures to attain parameters of an optical system, it is important to have constraints that ensure that the optical system can be manufactured. The following demonstrates the implementation of different types of constraints in our library, with a specific focus on the positive air spacing and minimum glass thickness constraints.

Constraint optimization problems can often be expressed as a nonlinear program, which is defined as follows (see [GZ18]):

\[\min_{p} \quad m(p)\]

\[\text{subject to} \quad \hat{g}_i(p) \leq 0, \quad i = 1, \ldots, N_1,\]

\[\text{subject to} \quad \hat{h}_j(p) = 0, \quad j = 1, \ldots, N_2,\]

where: - \(p \in \mathbb{R}^n\) is the vector of parameters. - \(m: \mathbb{R}^n \to \mathbb{R}\) is the nonlinear objective (merit) function. - \(\hat{g}_i: \mathbb{R}^n \to \mathbb{R}\) are the inequality constraint functions. - \(\hat{h}_j: \mathbb{R}^n \to \mathbb{R}\) are the equality constraint functions.

For this type of problem, multiple numerical schemes are available in the Python library SciPy. Some optimization schemes also require derivative information for functions that describe constraints. For example, Sequential Least Squares Programming (SLSQP) uses the derivatives of the constraint functions \(\hat{g}_i\) and \(\hat{h}_j\) to find local minima.

By combining the libraries PyTorch and SciPy, we leverage the strengths of two sophisticated and established libraries:

PyTorch: Efficiently calculates the derivatives of the merit function \(m\) and the constraint functions \(\hat{g}_i\) and \(\hat{h}_j\) using automatic differentiation. Additionally, it allows evaluation of these functions and their derivatives on a graphics card, providing significant speedups.
SciPy: Provides well-tested traditional algorithms to find local minima. While PyTorch also has a wide variety of optimization algorithms, its main application is stochastic gradient descent in deep learning, which may not be the best choice for optimizing optical systems.

Types of Constraints¶

In our library, we implemented three ways to define constraints:

Bounds Most numerical schemes in SciPy support bounding box constraints, allowing the definition of minimum and maximum values for each parameter. These bounds can be interpreted as constraints in the form \(\hat{g}_i(p) = p - C_i\) or \(\hat{g}_i(p) = C_i - p\), where \(C_i \in \mathbb{R}\). This is particularly useful for distance transformations, where we can ensure that the distance parameter is never smaller than 0. For example:
```
>>> import diffinytrace as dit
>>> import torch
>>> dist_transform = dit.transforms.Distance(10.)
>>> dist_transform.distance.bounds = torch.tensor([5.0, torch.inf])
```
Here, torch.inf indicates that the distance can be arbitrarily large, with no upper bound.
Constant Variables If a specific parameter should be fixed, PyTorch allows disabling gradient computation for that parameter. For example:
```
>>> import diffinytrace as dit
>>> distance_transform = dit.transforms.Distance(10.)
>>> distance_transform.distance.requires_grad = False
```
Note: While it is easy to set specific parameters as constants, it is not possible to disable gradient computation for individual parameters if the variable contains multiple values. For instance, in the case of a B-spline surface, it is not possible to disable gradient computation for individual B-spline coefficients.
Arbitrary Constraint Functions Our library also supports defining nonlinear inequality constraint functions \(\hat{g}_i\) and equality constraint functions \(\hat{h}_i\). Some local optimization methods require derivative information for these nonlinear constraint functions. To efficiently evaluate these derivatives, we use automatic differentiation. This is achieved by defining the constraint functions \(\hat{g}_i\) with PyTorch and calculating their derivatives with respect to the parameters of the optical system. This approach eliminates the need for finite differences, which could significantly slow down the optimization procedure.

diffinytrace.optimize.make_bounds_from_param(param)[source]¶

Creates default bounds (-∞, ∞) for each element of the input tensor.

This function returns a tensor of shape param.shape + [2], where the last dimension represents the lower and upper bounds for each element in param.

Parameters:: param (torch.Tensor) – A tensor for which bounds should be created.
Returns:: A tensor of shape param.shape + [2] where […, 0] = -inf (lower bounds) and […, 1] = inf (upper bounds), with the same dtype and device as param.
Return type:: torch.Tensor

diffinytrace.optimize.make_parameter_from_input(input, bounds=None, dtype=None, device=None, bounds_attr_name='bounds')[source]¶

Converts input to a torch.nn.Parameter and attaches bounds as an attribute.

Parameters:

input (array-like or torch.Tensor) – Input data.
bounds (torch.Tensor, optional) – Bounds to attach to the parameter.
dtype (torch.dtype, optional) – Desired tensor data type.
device (torch.device, optional) – Device to store the parameter on.
bounds_attr_name (str) – Attribute name used to store bounds.

Returns:

The parameter with bounds attached as an attribute.

Return type:

torch.nn.Parameter

diffinytrace.optimize.pack_tensors(tensor_list: List[Tensor]) → Tensor[source]¶

Flattens and concatenates a list of tensors into a single 1D tensor.

Parameters:: tensor_list (list of torch.Tensor or torch.Tensor) – Input tensor(s).
Returns:: A 1D tensor.
Return type:: torch.Tensor

diffinytrace.optimize.unpack_tensors(packed_tensor: Tensor, shapes: List[Tuple[int]]) → List[Tensor][source]¶

Unpacks a 1D tensor into a list of tensors with specified shapes.

Parameters:

packed_tensor (torch.Tensor) – The flat tensor to unpack.
shapes (list of tuple) – Target shapes for unpacked tensors.

Returns:

Unpacked tensors with original shapes.

Return type:

list of torch.Tensor

diffinytrace.optimize.apply_vec_to_params(vec: ndarray, params: list[Parameter], device=None, dtype=None)[source]¶

Updates PyTorch parameters with values from a flattened NumPy vector.

This function is used in optimization workflows to update parameter values during SciPy optimization. It takes a flat vector of parameter values and distributes them back to the original parameter tensors, preserving their original shapes.

Parameters:

vec (np.ndarray) – A 1D NumPy array containing new parameter values. The length must match the total number of elements across all parameters.
params (list[torch.nn.Parameter]) – List of PyTorch parameters to update. Each parameter will be reshaped from the corresponding portion of vec.
device (torch.device, optional) – Target device for the parameters. If None, uses the device of the first parameter. Defaults to None.
dtype (torch.dtype, optional) – Target data type for the parameters. If None, uses the dtype of the first parameter. Defaults to None.

Raises:

RuntimeError – If vec is not a NumPy array.

Example

>>> import torch
>>> import numpy as np
>>> import diffinytrace as dit
>>>
>>> # Create some parameters
>>> params = [
...     torch.nn.Parameter(torch.ones((2,2)))*0.25,
...     torch.nn.Parameter(torch.ones(3))
... ]
>>> # Flatten parameters to create a vector
>>> vec = dit.optimize.pack_tensors(params).detach().cpu().numpy()
>>> print(f"Vector length: {len(vec)}")  # Should be 2*2 + 3 = 7
>>> # Modify the vector
>>>
>>> print(params)
>>>
>>> vec_new = vec * 2.0
>>> # Update parameters with new values
>>> dit.optimize.apply_vec_to_params(vec_new, params)
>>>
>>> # Parameters are now updated with doubled values
>>> print(params)

Note

This function modifies parameters in-place using param.data = …
The function uses torch.no_grad() to avoid building computation graphs
Parameter shapes are preserved during the update process
Commonly used with pack_tensors() and unpack_tensors() for optimization

diffinytrace.optimize.set_full_if_nan(input: ndarray, fill_value: float) → ndarray[source]¶

Replaces NaNs in input with a specified fill value.

Parameters:

input (np.ndarray) – A NumPy array or scalar.
fill_value (float) – Value to use in place of NaNs.

Returns:

Modified input with no NaNs.

Return type:

np.ndarray or float

class diffinytrace.optimize.ParameterFunHelper(orginal_fun, params, nan_fallback=inf)[source]¶

Bases: object

Helper class for evaluating PyTorch functions and gradients in SciPy optimization.

This class bridges PyTorch’s automatic differentiation with SciPy’s optimization routines by providing function and gradient evaluations in NumPy format. It includes caching to avoid redundant computations and handles NaN values gracefully during optimization.

Parameters:

original_fun (Callable) – PyTorch function to be optimized. Should return a scalar tensor.
params (List[torch.nn.Parameter]) – List of PyTorch parameters to optimize over.
nan_fallback (float, optional) – Value to return if NaN is detected in function or gradient evaluation. Defaults to float(“inf”).

original_fun¶

The objective function being optimized.

Type:: Callable

params¶

Parameters for optimization.

Type:: List[torch.nn.Parameter]

nan_fallback¶

Fallback value for NaN handling.

Type:: float

last_x_fun_numpy¶

Cache of last input for function evaluation.

Type:: np.ndarray

last_fun_val_numpy¶

Cache of last function value in NumPy format.

Type:: float

last_fun_val_torch¶

Cache of last function value as PyTorch tensor.

Type:: torch.Tensor

last_x_grad_numpy¶

Cache of last input for gradient evaluation.

Type:: np.ndarray

last_grad_val_numpy¶

Cache of last gradient in NumPy format.

Type:: np.ndarray

Example

>>> import torch
>>> import diffinytrace as dit
>>> import numpy as np
>>>
>>> # Define parameters and objective function
>>> params = [torch.nn.Parameter(torch.randn(5))]
>>> def objective():
...     return torch.sum(params[0]**2)
>>>
>>> # Create helper for SciPy optimization
>>> helper = dit.optimize.ParameterFunHelper(objective, params)
>>>
>>> # Use with SciPy
>>> x0 = np.ones((5,))*3.
>>> fun_val = helper.fun(x0)        # Evaluate function 5*3^2 = 45
>>> grad_val = helper.jac(x0)       # Evaluate gradient 2*3 = 6
>>> fun_val, grad_val = helper.fun_jac(x0)  # Evaluate both
>>>
>>> print(fun_val, grad_val)  # (45.0, array([6., 6., 6., 6., 6.]))

Note

Function and gradient evaluations are cached to avoid redundant computations when SciPy requests the same point multiple times.
All NaN values in function outputs or gradients are replaced with nan_fallback.
Parameters are automatically updated with new values during evaluation.

fun(x)[source]¶

Evaluates the objective function at a given input.

Parameters:: x (np.ndarray) – Flat input array.
Returns:: Function value with NaNs replaced if needed.
Return type:: float

jac(x)[source]¶

Computes the gradient of the objective function at input x.

Parameters:: x (np.ndarray) – Flat input array.
Returns:: Gradient with NaNs replaced if needed.
Return type:: np.ndarray

fun_jac(x)[source]¶

Evaluates both function value and gradient at once.

Parameters:: x (np.ndarray) – Flat input array.
Returns:: Function value and gradient.
Return type:: Tuple[float, np.ndarray]

diffinytrace.optimize.create_fun_and_gradient(merit_fun, params, nan_fallback, device, dtype)[source]¶

Wraps a PyTorch merit function and returns a callable that evaluates both the function and its gradient in NumPy format.

Parameters:

merit_fun (Callable) – PyTorch function to optimize.
params (list) – List of torch.nn.Parameter objects.
nan_fallback (float) – Value to use if NaNs are encountered.
device (torch.device) – Target device.
dtype (torch.dtype) – Target dtype.

Returns:

Function that returns (value, gradient) as NumPy arrays.

Return type:

Callable

diffinytrace.optimize.remove_bounds(params, bounds_attr_name) → None[source]¶

Removes the bounds attribute from parameters if present.

Parameters:

params (list) – List of torch.nn.Parameter objects.
bounds_attr_name (str) – Attribute name of bounds to remove.

diffinytrace.optimize.get_bounds(params, bounds_attr_name='bounds')[source]¶

Extracts and concatenates bounds for all parameters.

Parameters:

params (list) – List of torch.nn.Parameter objects.
bounds_attr_name (str) – Name of attribute storing bounds.

Returns:

Array of shape (N, 2) with all bounds.

Return type:

np.ndarray

diffinytrace.optimize.get_scipy_constraint(constraint, params, nan_fallback)[source]¶

Converts a constraint into SciPy-compatible format.

Parameters:

constraint (Constraint) – A custom constraint object.
params (list) – List of parameters for the optimization.
nan_fallback (float) – Fallback value for NaNs.

Returns:

A dictionary compatible with SciPy constraints.

Return type:

dict

diffinytrace.optimize.create_callback(callback_fun, params, device, dtype)[source]¶

Wraps a PyTorch callback function for use in SciPy.

Parameters:

callback_fun (Callable) – A function taking no arguments.
params (list) – List of parameters to update before calling.
device (torch.device) – Device of the parameters.
dtype (torch.dtype) – Data type of the parameters.

Returns:

A callback function for SciPy optimizers.

Return type:

Callable

diffinytrace.optimize.minimize(fun, params, constraints: List = [], method=None, tol: float = 1e-09, callback: Callable = <function <lambda>>, options: dict | None = None, nan_fallback: float = inf, bounds_attr_name: str = 'bounds', save_history: bool = False, call_before_minimize: bool = False) → dict[source]¶

Minimizes a function using SciPy’s minimize, supporting bounds and constraints.

Parameters:

fun (Callable) – Objective function.
params (list) – Parameters to optimize.
constraints (list) – List of constraints.
method (str) – SciPy optimization method (e.g., ‘L-BFGS-B’).
tol (float) – Tolerance for convergence.
callback (Callable) – Optional callback function.
options (dict) – Optimizer options.
nan_fallback (float) – Value to use if function returns NaN.
bounds_attr_name (str) – Name of bounds attribute.
save_history (bool) – If True, saves function values and gradient norms.
call_before_minimize (bool) – Whether to evaluate once before optimization.

Returns:

Dictionary containing optimization results (and optionally history).

Return type:

dict

diffinytrace.optimize.copy_bounds_to_attr_name(params, bounds_attr_name_new, bounds_attr_name_old='bounds', replace_existing_once=True)[source]¶

Copies bounds from one attribute name to another.

Parameters:

params (list) – List of parameters.
bounds_attr_name_new (str) – New attribute name.
bounds_attr_name_old (str) – Existing attribute name.
replace_existing_once (bool) – Whether to skip copying if new attribute exists.

diffinytrace.optimize.set_bounds_from_params_mask(params, mask: list | Tensor, bounds_attr_name_new, bounds_attr_name_old='bounds')[source]¶

Sets bounds for parameters based on a mask. Parameters with mask=False get fixed bounds (equal lower and upper bounds).

Parameters:

params (list) – List of parameters.
mask (list or torch.Tensor) – Mask specifying which elements are free.
bounds_attr_name_new (str) – Attribute name to store new bounds.
bounds_attr_name_old (str) – Attribute name to read old bounds from.

Physical Object¶

class diffinytrace.physical_object.PhysicalObject[source]¶

Bases: Module

Abstract base class for physical objects in the optical system. This class can be used to define surface distance constraints and is also used for plotting.

get_transformation_matrix()[source]¶

Returns the transformation matrix of the object.

Returns:: The transformation matrix.
Return type:: torch.Tensor

to_global_dir(direction)[source]¶

Converts a direction vector from local to global coordinates.

Parameters:: direction (torch.Tensor) – Direction vector in local coordinates.
Returns:: Direction vector in global coordinates.
Return type:: torch.Tensor

to_local_dir(direction)[source]¶

Converts a direction vector from global to local coordinates.

Parameters:: direction (torch.Tensor) – Direction vector in global coordinates.
Returns:: Direction vector in local coordinates.
Return type:: torch.Tensor

to_global_pos(position)[source]¶

Converts a position from local to global coordinates.

Parameters:: position (torch.Tensor) – Position in local coordinates.
Returns:: Position in global coordinates.
Return type:: torch.Tensor

to_local_pos(position)[source]¶

Converts a position from global to local coordinates.

Parameters:: position (torch.Tensor) – Position in global coordinates.
Returns:: Position in local coordinates.
Return type:: torch.Tensor

get_transform()[source]¶

Returns the transformation object associated with this physical object.

Raises:: NotImplementedError – If not implemented in subclass.
Returns:: Transformation object.
Return type:: object

class diffinytrace.physical_object.PhysicalSurface[source]¶

Bases: PhysicalObject

Abstract base class for physical surfaces in the optical system. This class can be used to define surface distance constraints and is also used for plotting.

get_constraint_funs_leq_zero()[source]¶

Returns constraint functions for the surface that must be less than or equal to zero.

Raises:: NotImplementedError – If not implemented in subclass.
Returns:: List of constraint functions.
Return type:: list[Callable]

parametric_sample(num_points: int, method: str = 'sobol') → tuple[Tensor, Tensor][source]¶

Samples points on the surface in parameter space.

Parameters:

num_points (int) – Number of points to sample.
method (str, optional) – Sampling method. Defaults to “sobol”.

Raises:

NotImplementedError – If not implemented in subclass.

Returns:

Sampled parameter positions and corresponding surface positions.

Return type:

tuple[torch.Tensor, torch.Tensor]

parametric_surface(parametric_pos: Tensor) → Tensor[source]¶

Maps parameter space positions to surface positions.

Parameters:: parametric_pos (torch.Tensor) – Positions in parameter space.
Raises:: NotImplementedError – If not implemented in subclass.
Returns:: Surface positions.
Return type:: torch.Tensor

Refractive Index¶

class diffinytrace.refractive_index.RefractiveIndex(func, bounds)[source]¶

Bases: Module, PlotableWavelength

This class is used to calculate the refractive index of a material.

\[\mathbf{D'} = \mathbf{N} \sqrt{1 - (1 - \cos^2 \psi_i) \eta^2} + \eta (\mathbf{D} - \mathbf{N} \cos \psi_i),\]

where \(\cos \psi_i = \mathbf{D} \cdot \mathbf{N}\) and \(\eta = n / n'\) is the ratio of the refractive indices of the two materials.

We have implemented the refractive indices as a class that is initialized by a refractive index function and the start and end wavelengths for which the function is valid. This makes it very convenient to use with the RefractiveIndex.info database of optical constants (see [Pol24]), since this database often provides Python functions for wavelength-dependent refractive indices.

Example

Below is an example of how to set up an optical material in our ray tracing library:

>>> import diffinytrace as dit
>>> BaSF = dit.RefractiveIndex(
>>>     lambda x: (1 + 1.65554268 / (1 - 0.0104485644 / x**2) +
>>>                   0.17131977 / (1 - 0.0499394756 / x**2) +
>>>                   1.33664448 / (1 - 118.961472 / x**2))**0.5,
>>>     [0.365, 2.5]
>>> )
>>> dit.plotting.wavelength.plot(
>>>     BaSF, title="Refractive index of BaSF (Barium dense flint)"
>>> )

Parameters:

func (callable) – A function that takes a wavelength in μm and returns the refractive index.
bounds (tuple) – A tuple containing the minimum and maximum wavelength in μm.

forward(wl)[source]¶

Calculates the refractive index for given wavelengths. :param wl: Wavelength in μm. :type wl: torch.Tensor or float

Returns:: Refractive index at the given wavelengths.
Return type:: torch.Tensor

Render¶

diffinytrace.render.smoothed_irradiance(optical_system: SequentialOpticalSystem, sequence: List, source: LightSource, detector: Detector, smoother, num_rays: int, device=device(type='cpu'), method_ray_tracing: str = 'sobol_pow2') → Tensor[source]¶

Calculate the smoothed irradiance on the detector using ray tracing.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system to trace rays through.
sequence – The sequence of elements in the optical system.
source (LightSource) – The light source used for ray tracing.
detector (Detector) – The detector where the rays are traced to.
smoother – The smoother object used for smoothing the irradiance.
num_rays (int, optional) – The number of rays to trace. Default is 100000.
device (torch.device, optional) – The device to perform computations on. Default is the default device.
method_ray_tracing (str, optional) – The method used for ray tracing. Default is “sobol”.

Returns:

The smoothed irradiance on the detector.

Return type:

torch.Tensor

diffinytrace.render.binned_irradiance(optical_system: SequentialOpticalSystem, sequence: List, source: LightSource, detector: Detector, grid, num_rays: int, device=device(type='cpu'), method_ray_tracing: str = 'sobol_pow2') → Tensor[source]¶

Calculate the binned irradiance on the detector using ray tracing.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system to trace rays through.
sequence – The sequence of elements in the optical system.
source (LightSource) – The light source used for ray tracing.
detector (Detector) – The detector where the rays are traced to.
grid – The grid used for binning the irradiance.
num_rays (int, optional) – The number of rays to trace. Default is 100000.
device (torch.device, optional) – The device to perform computations on. Default is the default device.
method_ray_tracing (str, optional) – The method used for ray tracing. Default is “sobol”.

Returns:

The binned irradiance on the detector.

Return type:

torch.Tensor

Source¶

class diffinytrace.source.RaySource(transform)[source]¶

Bases: PhysicalObject

The RaySource is a the base class off all Objects emitting rays. It has a sample function which samples points. These points are used to initialize the rays with the forward function.

sample(num_points, method)[source]¶

Samples points in the parameter space of the ray source. :param num_points: The number of points to sample. :type num_points: int :param method: The sampling method to use. Can be “sobol”, “monte_carlo”, etc. :type method: str

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

get_plot_points_2D(resolution)[source]¶

get_plot_points_3D(resolution)[source]¶

get_volume()[source]¶

get_transform()[source]¶

Returns the transformation object associated with this physical object.

Raises:: NotImplementedError – If not implemented in subclass.
Returns:: Transformation object.
Return type:: object

class diffinytrace.source.LightSource(transform, integrator, flux_func=None, total_power=1.0, num_points_normalize=700000, method_normalize='sobol')[source]¶

Bases: RaySource, Plotable

In our model of an optical system, the light source is composed of the following components:

1. Light Source Domain The light source domain \(\Omega \subseteq \mathbb{R}^N\) represents the spatial domain of the light source.

2. Global Ray Initialization Function The global ray initialization function \(f: \mathbb{R}^N \times \mathbb{R}^{4 \times 4} \to \mathbb{R}^7\) is responsible for initializing the ray’s position \(O\), direction \(D\), and wavelength \(\lambda\) using a point \(x \in \Omega\) from the light source domain \(\Omega\). It consists of:

An affine transformation \(M \in \mathbb{R}^{4 \times 4}\).
A local ray initialization function \(\hat{f}: \mathbb{R}^N \to \mathbb{R}^7\).

The global ray initialization function is defined by:

\[\begin{split}f(x, M) := \begin{bmatrix} M \begin{bmatrix}\hat{O} \\ 1\end{bmatrix} \\ M \begin{bmatrix}\hat{D} \\ 0\end{bmatrix} \\ \lambda \end{bmatrix},\end{split}\]

where:

\[\begin{split}\begin{bmatrix}\hat{O} \\ \hat{D} \\ \lambda \end{bmatrix} = \hat{f}(x).\end{split}\]

3. Flux Function The flux function \(Q: \mathbb{R}^N \to \mathbb{R}\) assigns a flux weight to each point \(x \in \Omega\).

The ray initialization function bridges the gap between the light source domain and the actual ray tracer. This abstraction is also very helpful for implementation, as it allows for modular design. Following this principle, one does not sample rays in our library, but points and weights according to an integration rule. The light source classes also implement a forward function which is equivalent to the ray initialization function.

Example

>>> import diffinytrace as dit
>>> wave_len = 1.0
>>> radius = 8.0
>>> num_rays = 100
>>> transform = dit.transforms.Identity()
>>> light_source = dit.source.CollimatedMonochromatic(transform, radius, wave_len)
>>> x, weights = light_source.sample(num_rays, "sobol")
>>> O, D, wls, _, meta_data = light_source(x, n_func_enviroment=dit.materials["AIR"])

In the following, we introduce two types of light source models along with their corresponding light source domains and local ray initialization functions.

1. Collimated Monochromatic Light Source For this light source model, rays originate from a planar aperture with a uniform direction and a single wavelength. An example for light sources which can be modeled as collimated monochromatic light sources are lasers (see [Dic18]).

The light source domain \(\Omega_{CO} \subseteq \mathbb{R}^2\) can take various shapes.
The local ray initialization function is, given a single wavelength \(\lambda\), always in the form of:

\[\begin{split}\hat{f}(x, \lambda)_{CO} = \begin{bmatrix} x_1 \\ x_2 \\ 0 \\ 0 \\ 0 \\ 1 \\ \lambda \end{bmatrix}.\end{split}\]
The flux function must align with the light source domain \(\Omega_{CO}\).

Examples of light source domains for collimated monochromatic light sources include:

Round Domains

Defined by:

\[\Omega^{round(\hat{r})}_{CO} = \left\{x : ||x||_2 \leq \hat{r}\right\} \subseteq \mathbb{R}^2,\]

where \(\hat{r}\) is the radius.

Example:

>>> import diffinytrace as dit
>>> wl = 0.5
>>> r = 8.0
>>> transform = dit.transforms.Identity()
>>> light_source = dit.source.CollimatedMonochromatic(transform, r, wl)

Square Domains
Defined by:

\[\Omega^{square(\hat{d})}_{CO} = \left\{x : ||x||_\infty \leq \hat{d}\right\} \subseteq \mathbb{R}^2,\]

where \(\hat{d}\) is the half diameter.

2. Monochromatic Light Sources with Etendue Real-world physical light sources are characterized by a finite area-angle product, corresponding to a finite etendue (see [RRHT17]). In this model, rays originate from a planar surface with a specific wavelength, but the directions differ.

The light source domain \(\Omega_{ET} \subseteq \mathbb{R}^4\) has four dimensions.
The flux function maps a 4-dimensional vector to a single value.
The local ray initialization function maps a 4-dimensional vector to a 7-dimensional vector.

Example: Pillbox Sunlight Model The sun has an etendue, meaning both an angular and spatial extent. The sunshape describes how solar radiation is distributed across the solar disk, represented by the normalized radiance profile \(\hat{L}(\theta)\):

\[\begin{split}\hat{L}_{\text{pillbox}}(\theta) = \begin{cases} L_1 & \text{if } 0 \leq \theta \leq \theta_{disc}, \\ 0 & \text{if } \theta > \theta_{disc}. \end{cases}\end{split}\]

Here, \(\theta\) is the angular displacement of a ray compared to the general direction of the sun. Typically, \(\theta_{disc} = 4.65 \, \text{mrad}\) is used (see [WPA+20]).

The light source domain for the pillbox sunlight model is defined as:

\[\Omega^{pillbox(\hat{d})}_{ET} := \left\{x : ||[x_1, x_2]^T||_\infty \leq \hat{d}, ||[x_3, x_4]^T||_\infty \leq 1\right\} \subseteq \mathbb{R}^4,\]

where \(\hat{d}\) is the half diameter.

The local ray initialization function is:

\[\begin{split}\hat{f}(x)_{ET}^{pillbox} := \begin{bmatrix} x_1 \\ x_2 \\ 0 \\ \sin(\theta(x_3))\cos(\phi(x_4)) \\ \sin(\theta(x_3))\sin(\phi(x_4)) \\ \cos(\theta(x_3)) \\ \lambda \end{bmatrix},\end{split}\]

where:

\[\theta(x_3) := -\theta_{disc} + 2\theta_{disc}x_3, \quad \phi(x_4) := 2\pi x_4.\]

The flux function is:

\[Q_{ET}(x)_{pillbox} := \frac{1 \, \text{W}}{|\Omega^{pillbox(\hat{d})}_{ET}|}.\]

sample(num_points, method)[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

get_flux(x)[source]¶

get_volume()[source]¶

class diffinytrace.source.PlaneSource(transform, aperture_radius, integrator, is_square=False, flux_func=None, total_power=1.0, num_points_normalize=700000, method_normalize='sobol')[source]¶

Bases: LightSource

get_plot_points_2D(resolution)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution)[source]¶

Returns a list of 3D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 3D plot points.
Return type:: list

class diffinytrace.source.PlaneSource1D(transform, aperture_radius, integrator, flux_func, total_power=1.0, num_points_normalize=700000, method_normalize='sobol')[source]¶

Bases: LightSource

get_plot_points_2D(resolution)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution)[source]¶

Returns a list of 3D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 3D plot points.
Return type:: list

diffinytrace.source.make_cone_directions(num_rays, unif1, unif2, theta_max_rad)[source]¶

Sample directions uniformly within a cone of angular radius theta_max_rad centered on the z-axis.

Parameters: - num_rays (int): Number of direction vectors to sample within the cone. - unif1 (torch.Tensor): Tensor of uniform samples for cos(theta) sampling. - unif2 (torch.Tensor): Tensor of uniform samples for the azimuthal angle sampling. - theta_max_rad (float): Angular radius of the cone in radians.

Returns: - directions (torch.Tensor): Tensor of shape (num_rays, 3), with each row a direction vector.

class diffinytrace.source.VisibleSunlightSimpleMonochromatic(transform, aperture_radius, wl=0.5, is_square=True, total_power=1.0, theta_max_rad=0.0046500000000000005)[source]¶

Bases: PlaneSource

A class representing a visible sunlight source with a monochromatic spectrum. It also has a cone of 4.65 mrad.

sample(num_points, method='monte_carlo')[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class diffinytrace.source.VisibleSunlightSimple(transform, aperture_radius, is_square=True, total_power=1.0)[source]¶

Bases: PlaneSource

A class representing a visible sunlight source with a spectrum acording to the visble sunlight. It also has a an etendue with cone of 4.65 mrad.

sample(num_points, method='monte_carlo')[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class diffinytrace.source.CollimatedMonochromatic(transform, aperture_radius, wl, is_square=False, flux_func=None, total_power=1.0)[source]¶

Bases: PlaneSource

A class representing a collimated monochromatic light source. This class is a subclass of PlaneSource and is used to generate rays with a specific wavelength and a collimated beam profile.

sample(num_points, method='monte_carlo') → Tensor[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class diffinytrace.source.CollimatedGaussianBeam(transform, aperture_radius, wl, gaussian_constant, total_power=1.0)[source]¶: Bases: CollimatedMonochromatic

class diffinytrace.source.CollimatedMonochromatic1D(transform, aperture_radius, wl, flux_func=None, total_power=1.0)[source]¶

Bases: PlaneSource1D

A class representing a collimated monochromatic light source. This class is a subclass of PlaneSource1D and is used to generate rays with a specific wavelength and a collimated beam profile.

sample(num_points, method='monte_carlo')[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class diffinytrace.source.CollimatedMonochromatic1DRotSym(transform, aperture_radius, wl, flux_func=None, total_power=1.0)[source]¶

Bases: PlaneSource1D

A class representing a collimated monochromatic light source. This class is a subclass of PlaneSource1D and is used to generate rays with a specific wavelength and a collimated beam profile.

sample(num_points, method='monte_carlo')[source]¶

Returns:: A tuple containing the sampled points and their weights.
Return type:: tuple

forward(x, n_func_enviroment)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Spectrum¶

class diffinytrace.spectrum.Spectrum(func: Callable[[Tensor], Tensor], bounds: Tuple[float, float])[source]¶

Bases: Module, PlotableWavelength

A class to represent a spectrum as a function of wavelength.

forward(wl: Tensor) → Tensor[source]¶

Calculate the spectrum for given wavelengths.

Parameters:: wl (torch.Tensor or float) – Wavelength in μm.
Returns:: Spectrum value at the given wavelengths.
Return type:: torch.Tensor

class diffinytrace.spectrum.VisibleSunlight_am15g[source]¶

Bases: Spectrum

A class to represent the AM 1.5 G spectrum. This class uses the pvlib library to calculate the spectrum.

Surface¶

class diffinytrace.surface.Surface[source]¶

Bases: SemiFunctionalModule

While we all have an intuitive idea of what curves and surfaces are, we need a mathematically accurate definition from which we can proceed to illustrate how different types of algorithms are implemented. In the following, we will introduce three common ways of describing curves and surfaces.

1. Parametric Equations Parametric curves are functions that map a single variable \(\theta\) (the parameter) to a vector in \(\mathbb{R}^2\). Thus, such curves are referred to as parametrized or parametrically defined curves (see [GVJ+09]). The variable \(\theta\) is an element of the parametric domain of the parametric curve (see [CHB09]). For example, a circle can be described with the parametric domain \([0, 2\pi]\) and the function \(f: [0, 2\pi] \mapsto \mathbb{R}^2\),

\[\begin{split}f(\theta) = \begin{bmatrix} \cos \theta \\ \sin \theta \end{bmatrix}.\end{split}\]

Similarly, parametric surfaces can be described by a function that maps from a two-dimensional parametric domain to \(\mathbb{R}^3\) (see [GVJ+09]).

2. Explicit Equations Curves and surfaces can also be expressed using explicit equations. When describing a curve with explicit equations, an explicit function \(f: \mathbb{R} \to \mathbb{R}\) of the form \(y = f(x)\) assigns a unique value of \(y\) to each \(x \in \mathbb{R}\). The values of \(y\) can then be seen as a description of the curve. Unfortunately, it is not possible to describe all curves and surfaces with this method. For example, considering the unit circle, only one semicircle can be represented at a time using explicit equations such as:

\[y = \sqrt{1 - x^2} \quad \text{or} \quad y = -\sqrt{1 - x^2}.\]

Similarly, three-dimensional surfaces can be described explicitly using functions of the form \(y = f(x_1, x_2)\), which assign a unique \(y\)-value to each pair of \((x_1, x_2)\)-coordinates (see [GVJ+09]).

3. Implicit Equations A planar curve is defined implicitly, or in Cartesian coordinates, when it is described as the set of solutions to an equation involving two variables, typically expressed as \(f(y_1, y_2) = 0\). For example, the equation

\[y_1^2 + y_2^2 - 1 = 0\]

represents an implicit unit circle in \(\mathbb{R}^2\). Similarly, an implicit surface can be expressed with an equation in the form of (see [GVJ+09]):

\[f(y_1, y_2, y_3) = 0.\]

Optical Surfaces In our ray tracer, we use a less general description of the surfaces. We will call surfaces relevant for ray tracing optical surfaces. Every optical surface is composed of an explicit surface \(\hat{S}: \mathbb{R}^2 \mapsto \mathbb{R}\) and a transformation matrix \(M \in \mathbb{R}^{4 \times 4}\). In the following, we will state the implicit surface description for the ray tracer itself and the parametric surface description for plots and constraint optimization.

1. Implicit Surface Description Here, surfaces are described implicitly by the equation \(s(\hat{y}) = 0\). The function \(s\) is composed of the explicit description \(\hat{S}(\hat{x}_1, \hat{x}_2)\) and an affine transformation matrix \(M\) as follows:

\[\begin{split}\begin{bmatrix} \hat{x} \\ 1 \end{bmatrix} = M^{-1} \begin{bmatrix} \hat{y} \\ 1 \end{bmatrix}^T\end{split}\]

\[s(\hat{y}) = \hat{S}(\hat{x}_1(\hat{y}), \hat{x}_2(\hat{y})) - \hat{x}_3(\hat{y})\]

This description allows us to calculate ray-surface intersections efficiently. Typically, we do not state \(M^{-1}\) explicitly in the implementation but simply apply the transformation itself directly.

2. Parametric Surface Description In this approach, surfaces are defined by parameterizing coordinates. For optical surfaces, the surface is described again as a composition of the explicit description \(\hat{S}\) and a transformation matrix \(M\) as follows:

\[\begin{split}\begin{bmatrix} S(\hat{x}_1, \hat{x}_2) \\ 1 \end{bmatrix} = M \begin{bmatrix} \hat{x}_1 \\ \hat{x}_2 \\ \hat{S}(\hat{x}_1, \hat{x}_2) \\ 1 \end{bmatrix}\end{split}\]

In our library, the parametric domains are defined by the lenses or target surfaces (detectors). For example, in the case of a round lens, the parametric domain would be the disc determined by the aperture radius. This surface description is typically used for plotting but is also useful in the context of constraint optimization.

Examples

>>> import diffinytrace as dit
>>> aperture_radius = 30.
>>> lens_thickness = 8.
>>> material = dit.materials["NBK7"]
>>> transform = dit.transforms.Identity()
>>> asphere = dit.Aspheric(1./40., 0.0, [-0.00001])
>>> plane = dit.Plane()
>>> lens = dit.Lens(transform, lens_thickness,
>>>          asphere, plane,
>>>          material, aperture_radius)
>>> dit.plotting.system2D.plot(lens)

static functional(O, *params_list)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

explicit(local_pos)[source]¶

class diffinytrace.surface.Plane[source]¶

Bases: Surface

A class to represent a plane surface in 3D space. The plane is defined by the equation z = 0, and the functional method returns the z-coordinate of the input points.

static functional(O)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

explicit(local_pos)[source]¶

class diffinytrace.surface.Aspheric(curvature, conic_coeff=None, aspheric_param=None)[source]¶

Bases: Surface

This is the aspheric surface class, implementation follows: https://en.wikipedia.org/wiki/Aspheric_lens.

The surface is parameterized as an implicit function \(f(x, y, z) = 0\). For simplicity, we assume the surface function \(f(x, y, z)\) can be decomposed as:

\[f(x, y, z) = g(x, y) + h(z),\]

where \(g(x, y)\) and \(h(z)\) are explicit functions:

\[r^2 = x^2 + y^2\]

\[g(x, y) = \frac{c \cdot r^2}{1 + \sqrt{1 - (1 + k) \cdot \frac{r^2}{R^2}}} + a_0 \cdot r^4 + a_1 \cdot r^6 + \cdots\]

\[h(z) = -z\]

Parameters:

c (float) – Surface curvature, or one over the radius of curvature.
k (float) – Conic coefficient.
ai (list or None) – Aspheric parameters, could be a vector. When None, the surface is spherical.

static g(x: Tensor, y: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶

static h(z: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶

static functional(O: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args() → List[Tensor][source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

class diffinytrace.surface.Bspline(aperture_radius: float, orders: List[int], ns: List[int])[source]¶

Bases: Surface

A class to represent a B-spline surface in 3D space. The surface is defined by the B-spline basis functions and control points. The functional method returns the z-coordinate of the input points.

get_CAD_coeff(affine_transform: Transform) → ndarray[source]¶

Get the CAD coefficients from the affine transform.

Parameters:: affine_transform (torch.Tensor) – Affine transformation matrix.
Returns:: Control points of the B-spline surface.
Return type:: numpy.ndarray

get_CAD_face(affine_transform)[source]¶

Get the CAD face from the affine transform.

Parameters:: affine_transform (torch.Tensor) – Affine transformation matrix.
Returns:: CAD face object.

refine()[source]¶: Refine the B-spline surface by increasing the number of control points. The number of control points is increased by 1 in each direction.

functional(O, coeff, aperture_radius)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args() → List[Tensor][source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

Convert local position to global position using the B-spline surface functional.

Parameters:: local_pos (torch.Tensor) – Local position in 2D space.
Returns:: Global position in 3D space.
Return type:: torch.Tensor

class diffinytrace.surface.Legendre(aperture_radius: float, degree: int)[source]¶

Bases: Surface

A class to represent a Legendre surface in 3D space. Its kinda work in progress.

refine()[source]¶

functional(O: Tensor, coeffs: Tensor, aperture_radius: float) → Tensor[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

Target Grid¶

This module implements grid-based spatial aggregation for ray optics.

Classes:

Grid: Represents a 2D grid for spatial aggregation and statistics.
GridSquare: Square variant of Grid for symmetric apertures.

Functions:

(none at top level)

Example

>>> grid = Grid([0, 1], [0, 1], 10, 10)
>>> area = grid.get_area()

class diffinytrace.target_grid.Grid(y_range, x_range, y_grid_size, x_grid_size)[source]¶

Bases: object

Represents a 2D grid over a rectangular area with aggregation and indexing utilities.

Parameters:

y_range (tuple[float, float]) – The range in y-direction, as (y_min, y_max).
x_range (tuple[float, float]) – The range in x-direction, as (x_min, x_max).
y_grid_size (int) – Number of grid cells in y-direction.
x_grid_size (int) – Number of grid cells in x-direction.

get_area()[source]¶

Computes the total area of the grid.

Returns:: Total area of the grid.
Return type:: float

\[A = (x_{max} - x_{min}) \cdot (y_{max} - y_{min})\]

get_pixel_area()[source]¶

Returns the area of a single pixel/grid cell.

Returns:: Area of a single grid cell.
Return type:: float

\[A_{pixel} = \Delta x \cdot \Delta y\]

get_yi_xi(local_points, round_to_bounds=True)[source]¶

Converts 2D local coordinates to integer grid indices.

Parameters:

local_points (torch.Tensor) – Tensor of shape (N, 2) representing 2D points.
round_to_bounds (bool) – If True, clamps indices to stay within grid bounds. If False, returns a mask indicating valid indices.

Returns:

Tuple of tensors (yi, xi) of shape (N,).

Return type:

Tuple[torch.Tensor, torch.Tensor]

get_k(local_points, round_to_bounds=True)[source]¶

Maps local coordinates to flattened grid indices.

Parameters:

local_points (torch.Tensor) – Tensor of shape (N, 2).
round_to_bounds (bool) – Whether to clamp indices to grid bounds.

Returns:

If round_to_bounds is True: Tensor of shape (N,).
Otherwise: Tuple (indices, validity_mask).

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

map_matrix_to_ray(local_points, old_matrix)[source]¶

Maps a matrix defined on the grid to the given local points.

Parameters:

local_points (torch.Tensor) – Points of shape (N, 2).
old_matrix (torch.Tensor) – Matrix of shape (H, W, …).

Returns:

Resampled matrix values of shape (N, …).

Return type:

torch.Tensor

sum(local_points: Tensor, values: Tensor, old_matrix=None, round_to_bounds: bool = False)[source]¶

Sums values over the grid based on point locations.

Parameters:

local_points (torch.Tensor) – Points of shape (N, 2).
values (torch.Tensor) – Values of shape (N,) or (N, D).
old_matrix (torch.Tensor or None) – Previous result for accumulation.
round_to_bounds (bool) – Clamp indices to bounds if True.

Returns:

Aggregated result of shape (H, W).

Return type:

torch.Tensor

prod(local_points, values, old_matrix=None, round_to_bounds=False)[source]¶

Multiplies values over the grid based on point locations. :param local_points: Points of shape (N, 2). :type local_points: torch.Tensor :param values: Values of shape (N,) or (N, D). :type values: torch.Tensor :param old_matrix: Previous result for accumulation. :type old_matrix: torch.Tensor or None :param round_to_bounds: Clamp indices to bounds if True. :type round_to_bounds: bool

Returns:: Aggregated result of shape (H, W).
Return type:: torch.Tensor

mean(local_points, values, old_matrix=None, round_to_bounds=False)[source]¶

Computes the mean of values over the grid based on point locations. :param local_points: Points of shape (N, 2). :type local_points: torch.Tensor :param values: Values of shape (N,) or (N, D). :type values: torch.Tensor :param old_matrix: Previous result for accumulation. :type old_matrix: torch.Tensor or None :param round_to_bounds: Clamp indices to bounds if True. :type round_to_bounds: bool

Returns:: Aggregated result of shape (H, W).
Return type:: torch.Tensor

min(local_points, values, old_matrix=None, return_args=False)[source]¶

Finds the minimum value for each grid cell based on local points. :param local_points: Points of shape (N, 2). :type local_points: torch.Tensor :param values: Values of shape (N,) or (N, D). :type values: torch.Tensor :param old_matrix: Previous result for accumulation. :type old_matrix: torch.Tensor or None :param return_args: If True, also return indices. :type return_args: bool

Returns:: Minimum values, optionally with indices.
Return type:: torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

max(local_points, values, old_matrix=None, return_args=False)[source]¶

Finds the maximum value for each grid cell based on local points. :param local_points: Points of shape (N, 2). :type local_points: torch.Tensor :param values: Values of shape (N,) or (N, D). :type values: torch.Tensor :param old_matrix: Previous result for accumulation. :type old_matrix: torch.Tensor or None :param return_args: If True, also return indices. :type return_args: bool

Returns:: Maximum values, optionally with indices.
Return type:: torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

get_y_middle()[source]¶

get_x_middle()[source]¶

nearest(local_points, return_args=False)[source]¶

Finds the nearest pixel for each local point using L2 distance.

Parameters:

local_points (torch.Tensor) – Tensor of shape (N, 2).
return_args (bool) – If True, also return indices.

Returns:

Minimum squared distances, optionally with indices.

Return type:

torch.Tensor or Tuple[torch.Tensor, torch.Tensor]

get_pixel_centers()[source]¶

Returns the 2D center coordinates of each grid cell.

Returns:: Tensor of shape (H, W, 2) with pixel center coordinates.
Return type:: torch.Tensor

get_nearest_ray(local_points)[source]¶

Finds the index of the nearest ray for each grid cell using sklearn.neighbors.NearestNeighbors.

Parameters:: local_points (torch.Tensor) – Tensor of shape (N, 2) representing sampled rays.
Returns:: Tensor of shape (H, W) with ray indices.
Return type:: torch.Tensor

class diffinytrace.target_grid.GridSquare(aperture_radius, grid_size)[source]¶

Bases: Grid

Convenience class for square grids centered at the origin.

Parameters:

aperture_radius (float) – Half-width of the square domain.
grid_size (int) – Number of grid points in each direction.

Transforms¶

class diffinytrace.transforms.Transform[source]¶

Bases: SemiFunctionalModule

Base class for coordinate transformations.

This class provides interfaces to transform directions and positions between local and global coordinate systems using homogeneous coordinates.

get_functional_param_args()[source]¶: Return parameters required for the transformation.

functional(O, *params)[source]¶: Apply transformation in functional style.

get_transformation_matrix()[source]¶: Return the 4x4 transformation matrix.

to_global_dir(direction)[source]¶: Transform direction to global space.

to_local_dir(direction)[source]¶: Transform direction to local space.

to_global_pos(position)[source]¶: Transform position to global space.

to_local_pos(position)[source]¶: Transform position to local space.

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

static functional(O, *params) → Tensor[source]¶

Apply transformation in functional style. This is global to local.

Parameters:

O (torch.Tensor) – Input tensor to be transformed.
*params – Parameters for the transformation.

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

get_transform()[source]¶: Returns itself.

to_global_dir(direction: Tensor) → Tensor[source]¶

Transform direction to global space. :param direction: Direction vector in local space. :type direction: torch.Tensor

Returns:: Direction vector in global space.
Return type:: torch.Tensor

to_local_dir(direction: Tensor) → Tensor[source]¶

Transform direction to local space. :param direction: Direction vector in global space. :type direction: torch.Tensor

Returns:: Direction vector in local space.
Return type:: torch.Tensor

to_global_pos(position: Tensor) → Tensor[source]¶

Transform position to global space. :param position: Position vector in local space. :type position: torch.Tensor

Returns:: Position vector in global space.
Return type:: torch.Tensor

to_local_pos(position: Tensor) → Tensor[source]¶

Transform position to local space. :param position: Position vector in global space. :type position: torch.Tensor

Returns:: Position vector in local space.
Return type:: torch.Tensor

class diffinytrace.transforms.Identity[source]¶

Bases: Transform

Identity transformation that returns input positions unchanged.

Example

>>> import diffinytrace as dit
>>> transf1 = dit.transforms.Identity()

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

static functional(O: Tensor) → Tensor[source]¶

Apply transformation in functional style. This is global to local.

Parameters:

O (torch.Tensor) – Input tensor to be transformed.
*params – Parameters for the transformation.

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

class diffinytrace.transforms.Compose(transform_list)[source]¶

Bases: Transform

Compose multiple transforms in sequence.

Parameters:: transform_list (list[Transform]) – List of transformations to apply in order.

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

class diffinytrace.transforms.Offset(pos, parent_transform=Identity())[source]¶

Bases: Transform

Translation transform using an offset vector.

The offset transformation shifts a position by a specified vector ( vec{w} = (w_x, w_y, w_z) ). The transformation matrix ( M ) for an offset transformation is:

\[\begin{split}M^{offset}(w_x, w_y, w_z) = \begin{bmatrix} 1 & 0 & 0 & w_x \\ 0 & 1 & 0 & w_y \\ 0 & 0 & 1 & w_z \\ 0 & 0 & 0 & 1 \end{bmatrix}\end{split}\]

Example

>>> import diffinytrace as dit
>>> transf1 = dit.transforms.Identity()
>>> transf2 = dit.transforms.Offset([1.0, 2.0, 3.0], parent_transform=transf1)

Parameters:

pos (Tensor or list or float) – The offset position as a 3D vector.
parent_transform (Transform, optional) – Optional parent transformation.

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

functional(O: Tensor, pos: Tensor, *parent_param_args) → Tensor[source]¶

Apply transformation in functional style. This is global to local.

Parameters:

O (torch.Tensor) – Input tensor to be transformed.
*params – Parameters for the transformation.

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

class diffinytrace.transforms.Distance(distance, axis=2, parent_transform=Identity())[source]¶

Bases: Transform

Applies a translation along a specific axis by a given distance.

The distance transformation applies a translation by a specific distance along a given axis (e.g., ( x )-, ( y )-, or ( z )-axis). The transformation matrix ( M ) for a distance transformation along the ( z )-axis is given by:

\[\begin{split}M^{dist}_z(d) = \begin{bmatrix} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ 0 & 0 & 1 & d \\ 0 & 0 & 0 & 1 \end{bmatrix},\end{split}\]

where ( d ) represents the distance of translation along the ( z )-axis.

Parameters:

distance (float or Tensor) – Distance to translate.
axis (int) – Axis along which translation is applied (0=X, 1=Y, 2=Z).
parent_transform (Transform) – Optional parent transformation.

Example

>>> import diffinytrace as dit
>>> transf1 = dit.transforms.Identity()
>>> transf2 = dit.transforms.Distance(10.0,axis=2,parent_transform=transf1)

Notes

For the local to global transformation it applies the following transformation:

\[\mathbf{x}_\text{local} = \mathbf{x}_\text{parent} + d \cdot \mathbf{e}_i\]

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

functional(O: Tensor, distance: Tensor, unit_vec: Tensor, *parent_param_args) → Tensor[source]¶

Apply transformation in functional style. This is global to local.

Parameters:

O (torch.Tensor) – Input tensor to be transformed.
*params – Parameters for the transformation.

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

class diffinytrace.transforms.Rotation(angle: float, axis: int, parent_transform=Identity())[source]¶

Bases: Transform

Applies a 3D rotation around a principal axis.

The rotational transformation rotates a point or direction around a specific axis (e.g., ( x )-, ( y )-, and ( z )-axis). For example, the rotation matrix around the ( z )-axis is:

\[\begin{split}M^{rot}_z(\theta_z) = \begin{bmatrix} \cos \theta_z & -\sin \theta_z & 0 & 0 \\ \sin \theta_z & \cos \theta_z & 0 & 0 \\ 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 1 \end{bmatrix}\end{split}\]

Parameters:

angle (float or Tensor) – Rotation angle in degrees.
axis (int) – Axis index (0=X, 1=Y, 2=Z).
parent_transform (Transform, optional) – Optional parent transformation.

Example

>>> import diffinytrace as dit
>>> transf1 = dit.transforms.Identity()
>>> transf2 = dit.transforms.Distance(10.0,axis=2,parent_transform=transf1)
>>> transf3 = dit.transforms.Rotation(45.,axis=0,parent_transform=transf2)

get_functional_param_args()[source]¶

Return parameters required for the transformation which constructs the surfaces through the functional.

Returns:: List of parameters required for the functional which constructs the surfaces.
Return type:: list

functional(O: Tensor, angle: Tensor, *parent_param_args) → Tensor[source]¶

Apply transformation in functional style. This is global to local.

Parameters:

O (torch.Tensor) – Input tensor to be transformed.
*params – Parameters for the transformation.

get_transformation_matrix(device=None, dtype=None) → Tensor[source]¶

Return the 4x4 transformation matrix.

Parameters:

device (torch.device, optional) – Device for the matrix.
dtype (torch.dtype, optional) – Data type for the matrix.

Returns:

4x4 transformation matrix.

Return type:

torch.Tensor

diffinytrace.transforms.rotation_matrix_x(angle: Tensor) → Tensor[source]¶

Construct a 3x3 rotation matrix around the X-axis.

Parameters:: angle (Tensor) – Angle in degrees.
Returns:: 3x3 rotation matrix.
Return type:: Tensor

diffinytrace.transforms.rotation_matrix_y(angle: Tensor) → Tensor[source]¶

Construct a 3x3 rotation matrix around the Y-axis.

Parameters:: angle (Tensor) – Angle in degrees.
Returns:: 3x3 rotation matrix.
Return type:: Tensor

diffinytrace.transforms.rotation_matrix_z(angle: Tensor) → Tensor[source]¶

Construct a 3x3 rotation matrix around the Z-axis.

Parameters:: angle (Tensor) – Angle in degrees.
Returns:: 3x3 rotation matrix.
Return type:: Tensor

Module Contents¶

This module provides a collection of functions and classes for optical system design and analysis. It includes modules for ray tracing, surface definitions, optimization, and more.

diffinytrace.cat_semi_functionals(functional_modules: List[SemiFunctionalModule]) → Callable[source]¶

Recursively chains a list of `SemiFunctionalModule`s into a single composite function.

Each module’s functional() method is applied in sequence using the respective slice of the parameter list.

Parameters:: functional_modules (list[SemiFunctionalModule]) – List of functional modules.
Returns:: A function f(O, *params) that applies all modules in sequence.
Return type:: Callable

diffinytrace.get_functional_param_args(semi_functional_module_list: List[SemiFunctionalModule]) → List[source]¶

Collects all functional parameters from a list of semi-functional modules.

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: Flattened list of all parameters.
Return type:: list[torch.nn.Parameter]

diffinytrace.construct_surface_and_normal_func(semi_functional_module_list: List[SemiFunctionalModule]) → Callable[source]¶

Constructs a function to evaluate both the surface value and its gradient (normal direction) with respect to the ray origin O.

The surface is defined by composing the provided semi-functional modules.

Returns a callable:

\[(O, p_1, ..., p_n) \mapsto ( s(O), \frac{\partial s}{\partial O} )\]

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: A function s_dsd(O, *params, only_s=False) returning surface value s and optionally gradient ds/dO.
Return type:: Callable

diffinytrace.construct_surface_and_normal_func_with_params(semi_functional_module_list: List[SemiFunctionalModule]) → Tuple[Callable, List][source]¶

Constructs both the surface function and a list of its functional parameters.

Useful for optimization workflows that require parameter tracking.

Parameters:: semi_functional_module_list (list[SemiFunctionalModule]) – List of modules.
Returns:: Callable: A function computing surface and its gradient. list[torch.nn.Parameter]: The list of parameters for the surface.
Return type:: tuple

class diffinytrace.CustomAutogradRule_t(*args, **kwargs)[source]¶

Bases: Function

Custom PyTorch autograd rule for ray-surface intersection.

Computes a differentiable intersection length t such that:

\[s(O + t D) = 0\]

where O is the ray origin, D is the direction, and s is the surface function.

This rule enables backpropagation through t with respect to O, D, and surface parameters.

static forward(ctx, O: Tensor, D: Tensor, surface_and_normal_func: Callable, t_detached: Tensor, *param_args) → Tensor[source]¶

Stores inputs for backward pass and returns precomputed t.

Parameters:

O (torch.Tensor) – Ray origin of shape (N, 3).
D (torch.Tensor) – Ray direction of shape (N, 3).
surface_and_normal_func (Callable) – Surface function returning (s, ds/dR).
t_detached (torch.Tensor) – Estimated intersection length (detached).
*param_args – Surface parameters.

Returns:

Intersection length t.

Return type:

torch.Tensor

static backward(ctx, grad_outputs: Tensor) → Tuple[source]¶

Computes gradients of intersection length t with respect to: - ray origin O - ray direction D - surface parameters

Parameters:: grad_outputs (torch.Tensor) – Gradient of the loss w.r.t. output t.
Returns:: Gradients with respect to inputs (O, D, None, None, *param_args).
Return type:: tuple

diffinytrace.get_ray_intersection_length(O: Tensor, D: Tensor, surface_and_normal_func: Callable, param_args: List, t_init: Tensor | None = None) → Tensor[source]¶

Solves for the intersection length t such that:

\[s(O + t D) = 0\]

using a Newton-style iteration method with damping.

This function finds the length t where a ray intersects a parametric surface, given by a composed function with normal information.

Parameters:

O (torch.Tensor) – Ray origins of shape (N, 3).
D (torch.Tensor) – Ray directions of shape (N, 3).
surface_and_normal_func (Callable) – A function returning (s, ds/dR).
param_args (list) – List of surface parameters.
t_init (torch.Tensor, optional) – Initial guess for t. If None, starts from zero.

Returns:

Estimated intersection lengths t with autograd support.

Return type:

torch.Tensor

Raises:

Warning is printed (not exception) if convergence fails within max_iter. –

class diffinytrace.Plane[source]¶

Bases: Surface

A class to represent a plane surface in 3D space. The plane is defined by the equation z = 0, and the functional method returns the z-coordinate of the input points.

static functional(O)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

explicit(local_pos)[source]¶

class diffinytrace.Aspheric(curvature, conic_coeff=None, aspheric_param=None)[source]¶

Bases: Surface

This is the aspheric surface class, implementation follows: https://en.wikipedia.org/wiki/Aspheric_lens.

The surface is parameterized as an implicit function \(f(x, y, z) = 0\). For simplicity, we assume the surface function \(f(x, y, z)\) can be decomposed as:

\[f(x, y, z) = g(x, y) + h(z),\]

where \(g(x, y)\) and \(h(z)\) are explicit functions:

\[r^2 = x^2 + y^2\]

\[g(x, y) = \frac{c \cdot r^2}{1 + \sqrt{1 - (1 + k) \cdot \frac{r^2}{R^2}}} + a_0 \cdot r^4 + a_1 \cdot r^6 + \cdots\]

\[h(z) = -z\]

Parameters:

c (float) – Surface curvature, or one over the radius of curvature.
k (float) – Conic coefficient.
ai (list or None) – Aspheric parameters, could be a vector. When None, the surface is spherical.

static g(x: Tensor, y: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶

static h(z: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶

static functional(O: Tensor, curvature: Tensor, conic_coeff: Tensor, aspheric_param: Tensor | None) → Tensor[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args() → List[Tensor][source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

class diffinytrace.Bspline(aperture_radius: float, orders: List[int], ns: List[int])[source]¶

Bases: Surface

A class to represent a B-spline surface in 3D space. The surface is defined by the B-spline basis functions and control points. The functional method returns the z-coordinate of the input points.

get_CAD_coeff(affine_transform: Transform) → ndarray[source]¶

Get the CAD coefficients from the affine transform.

Parameters:: affine_transform (torch.Tensor) – Affine transformation matrix.
Returns:: Control points of the B-spline surface.
Return type:: numpy.ndarray

get_CAD_face(affine_transform)[source]¶

Get the CAD face from the affine transform.

Parameters:: affine_transform (torch.Tensor) – Affine transformation matrix.
Returns:: CAD face object.

refine()[source]¶: Refine the B-spline surface by increasing the number of control points. The number of control points is increased by 1 in each direction.

functional(O, coeff, aperture_radius)[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args() → List[Tensor][source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

Convert local position to global position using the B-spline surface functional.

Parameters:: local_pos (torch.Tensor) – Local position in 2D space.
Returns:: Global position in 3D space.
Return type:: torch.Tensor

class diffinytrace.Legendre(aperture_radius: float, degree: int)[source]¶

Bases: Surface

A class to represent a Legendre surface in 3D space. Its kinda work in progress.

refine()[source]¶

functional(O: Tensor, coeffs: Tensor, aperture_radius: float) → Tensor[source]¶: This method provides the implicit surface description. It is a static method. Diffinytrace constructs a function s(R, p) on the fly to describe the surface, allowing better control over derivative calculations.

get_functional_param_args()[source]¶

explicit(local_pos: Tensor) → Tensor[source]¶

diffinytrace.bspline_n_after_refinement(n, k)[source]¶

class diffinytrace.OpticalSystem(modules_dict: Dict)[source]¶

Bases: Module, Plotable

Base class for optical systems composed of multiple optical modules.

This class serves as a container for modules such as lenses, mirrors, and detectors. It supports visualization and modular organization.

modules_dict¶

Dictionary of named optical modules.

Type:: nn.ModuleDict

forward()[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

get_plotable_childs()[source]¶: Returns a list of all plotable child objects of this object. Each child is represented as a list containing the child object and its name.

get_plot_points_2D(resolution: int)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution: int)[source]¶

Returns a list of 3D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 3D plot points.
Return type:: list

class diffinytrace.SequentialOpticalSystem(modules_dict: Dict, n_func_enviroment=RefractiveIndex())[source]¶

Bases: OpticalSystem

Optical system that processes rays in a defined sequence.

Useful for simulating light propagation through a sequence of elements, e.g., source → lens → detector.

n_func_enviroment¶

Function returning refractive index of the surrounding medium.

Type:: Callable

forward(x, mapping_sequence: List[str])[source]¶

Propagates rays through the defined sequence of modules.

Parameters:

x (Any) – Input rays or sampling data.
mapping_sequence (list[str]) – Ordered list of module names defining propagation sequence.

Returns:

Output after final module in the sequence.

Return type:

Any

class diffinytrace.OpticalElement(fill_color='white', outline_color='black', is_volume=False)[source]¶

Bases: PhysicalObject, Plotable

Abstract base class for optical elements like lenses, mirrors, and detectors.

Provides interface for geometric transformation and ray propagation.

forward(O2: Tensor, D2: Tensor, wl: Tensor, n_func_enviroment, meta_data)[source]¶

Propagates rays through the optical element.

Parameters:

O2 (torch.Tensor) – Ray origins.
D2 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment (Callable) – Function returning environmental refractive index.
meta_data (dict) – Dictionary with path length and validity information.

Raises:

NotImplementedError – Must be overridden by subclasses.

get_transform()[source]¶

Returns the transformation associated with the surface.

Returns:: The local-to-global transformation object.
Return type:: Transform

class diffinytrace.OpticalSurface(transform: Transform, surface, aperture_radius: float, is_square: bool = False, fill_color: str = 'white', outline_color: str = 'black')[source]¶

Bases: OpticalElement, PhysicalSurface

Represents a surface in 3D space with a defined aperture and transformation.

Supports both square and circular apertures, and provides methods for parametric sampling, CAD conversion, ray intersection, and plotting.

surface¶

Object with a method explicit(parametric_pos) returning z-values.

Type:: object

aperture_radius¶

Radius of the circular or square aperture.

Type:: float

is_square¶

Whether the aperture is square-shaped.

Type:: bool

transform¶

Local-to-global transformation.

Type:: Transform

integrator¶

Integration object (Disc or Cube) for parametric sampling.

Type:: Integrator

get_constraint_funs_leq_zero()[source]¶

Returns constraint functions used for integration and optimization over the surface.

Returns:: List of functions f(param_pos) <= 0 indicating valid parametric regions.
Return type:: list[Callable]
Raises:: RuntimeError – If is_square is True (not yet implemented).

get_plot_points_2D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 2D slices through the surface (z-y plane) for plotting.

Parameters:: resolution (int) – Number of sample points along the y-axis.
Returns:: List of (z, y) coordinate tuples.
Return type:: List[Tuple[torch.Tensor]]

get_plot_points_3D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 3D grid of surface points for visualization.

Parameters:: resolution (int) – Grid resolution in x and y.
Returns:: List of (x, y, z) meshgrids as torch tensors.
Return type:: List[Tuple[torch.Tensor]]

get_CAD_points(resolution: int) → List[Tuple[Tensor]][source]¶

Generates a 3D surface point grid for CAD conversion.

Parameters:: resolution (int) – Sampling resolution.
Returns:: (x, y, z) coordinate grids for CAD modeling.
Return type:: Tuple[torch.Tensor]

get_CAD_face(resolution: int, tol: float = 0.001, smoothing=None, minDeg: int = 1, maxDeg: int = 3)[source]¶

Converts the surface into a CAD face using B-spline approximation.

Parameters:

resolution (int) – Sampling resolution.
tol (float, optional) – Approximation tolerance. Defaults to 0.001.
smoothing (Optional[int]) – Smoothing value for fitting.
minDeg (int) – Minimum degree of the spline.
maxDeg (int) – Maximum degree of the spline.

Returns:

CAD face object.

Return type:

cadquery.Face

parametric_sample(num_points: int, method: str = 'sobol') → tuple[Tensor, Tensor][source]¶

Samples parametric positions on the aperture using the integrator.

Parameters:

num_points (int) – Number of sample points.
method (str) – Sampling method. Options: “sobol”, “monte_carlo”, “midpoint”, etc.

Returns:

Sampled positions and integration weights.

Return type:

Tuple[torch.Tensor, torch.Tensor]

parametric_surface(parametric_pos: Tensor) → Tensor[source]¶

Maps 2D parametric coordinates to 3D global coordinates using the surface height and transform.

Parameters:: parametric_pos (torch.Tensor) – 2D parametric positions of shape (N, 2).
Returns:: 3D positions of shape (N, 3) in global space.
Return type:: torch.Tensor
Raises:: RuntimeError – If input does not have shape […, 2].

get_surface_and_normal_func_with_params()[source]¶

Constructs a callable for surface position and normal computation with parameter tracking.

Returns:: Callable computes (position, normal), and the list contains parameters to be optimized.
Return type:: Tuple[Callable, List]

get_ray_intersect_length(O, D) → Tensor[source]¶

Computes intersection length along ray until hitting the surface.

Parameters:

O (torch.Tensor) – Ray origins of shape (N, 3).
D (torch.Tensor) – Ray directions of shape (N, 3).

Returns:

Intersection distances t such that O + t*D lies on the surface.

Return type:

torch.Tensor

get_new_is_valid(O, valid) → Tensor[source]¶

Updates a boolean mask indicating which rays are still valid after hitting the aperture.

Parameters:

O (torch.Tensor) – Ray intersection points.
valid (torch.Tensor) – Previous boolean validity mask.

Returns:

Updated validity mask.

Return type:

torch.Tensor

get_transform() → Transform[source]¶

Returns the transformation associated with the surface.

Returns:: The local-to-global transformation object.
Return type:: Transform

class diffinytrace.LensSurfaceTransmissionEnter(transform: Transform, surface, aperture_radius: float, n_func, is_square: bool = False)[source]¶

Bases: OpticalSurface

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the lens entry surface.

Parameters:

O1 (torch.Tensor) – Ray origins.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.LensSurfaceTransmissionLeave(transform: Transform, surface, aperture_radius: float, n_func, is_square: bool = False)[source]¶

Bases: OpticalSurface

forward(O2, D2, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the lens exit surface.

Parameters:

O2 (torch.Tensor) – Ray origins.
D2 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.Lens(transform: Transform, lens_thickness: float, surface1, surface2, n_func, aperture_radius: float, is_square=False)[source]¶

Bases: OpticalElement

Represents a transmissive lens consisting of two refractive surfaces.

Example

Below is an example of initializing a square lens:

>>> import diffinytrace as dit
>>> aperture_half = 30.
>>> lens_thickness = 8.
>>> material = dit.materials["NBK7"]
>>> transform = dit.transforms.Identity()
>>> bspline = dit.Bspline(aperture_half, [3, 3], [8, 8])
>>> plane = dit.Plane()
>>> lens = dit.Lens(transform, lens_thickness,
>>>          bspline, plane,
>>>          material, aperture_half, is_square=True)

n_func¶

Function mapping wavelength to refractive index of the lens material.

Type:: Callable

_transform1¶

Transform for the first surface.

Type:: Transform

_transform2¶

Transform for the second surface.

Type:: Transform

lens_thickness¶

Learnable thickness of the lens.

Type:: torch.nn.Parameter

surface1¶

Entry surface.

Type:: LensSurfaceTransmissionEnter

surface2¶

Exit surface.

Type:: LensSurfaceTransmissionLeave

lens_surface_side¶

Side surface (for 3D rendering).

Type:: LensSurfaceSide

aperture_radius¶

Radius (or half-width) of aperture.

Type:: float

is_square¶

Whether the aperture is square.

Type:: bool

get_plot_points_2D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 2D slices through the lens for plotting.

Parameters:: resolution (int) – Number of sample points.
Returns:: List of (z, y) coordinate tuples.
Return type:: List[Tuple[torch.Tensor]]

get_plot_points_3D(resolution: int) → List[Tuple[Tensor]][source]¶

Returns 3D grid of lens surface points for visualization.

Parameters:: resolution (int) – Grid resolution.
Returns:: List of (x, y, z) meshgrids.
Return type:: List[Tuple[torch.Tensor]]

get_plotly_color_scale() → List[source]¶

Returns color scale for plotly visualization.

Returns:: Color scale values.
Return type:: List

get_plotable_childs() → List[source]¶

Returns plotable child elements.

Returns:: List of child elements.
Return type:: List

forward(O1: Tensor, D1: Tensor, wl: Tensor, n_func_enviroment, meta_data)[source]¶

Simulates light passing through the lens.

Parameters:

O1 (torch.Tensor) – Ray origin positions.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment (Callable) – Function returning external medium refractive index.
meta_data (dict) – Ray metadata (PL, OPL, paths, valid).

Returns:

Updated ray origins, directions, etc.

Return type:

Tuple[torch.Tensor]

get_transform()[source]¶

Returns the transformation of the lens exit surface.

Returns:: The transformation object.
Return type:: Transform

class diffinytrace.Mirror(transform, surface, aperture_radius, is_square=False)[source]¶

Bases: OpticalSurface

Reflective optical element that reflects rays according to the law of reflection.

Visualization is colored in a warm gold tone.

Inherits:: OpticalSurface: Full support for surface transformation and intersection.

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Propagates rays through the mirror surface.

Parameters:

O1 (torch.Tensor) – Ray origins.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning environmental refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, wavelengths, environment function, and metadata.

Return type:

Tuple

class diffinytrace.Detector(transform, surface, aperture_radius, is_square=True)[source]¶

Bases: OpticalSurface

Represents a terminal optical element that collects ray data.

Example

Below is an example of how to initialize a detector:

>>> import diffinytrace as dit
>>> aperture_half = 30.
>>> transform = dit.transforms.Identity()
>>> plane = dit.Plane()
>>> detector = dit.Detector(transform, plane,
>>>                         aperture_half, is_square=False)

forward(O1, D1, wl, n_func_enviroment, meta_data)[source]¶

Captures the final ray interaction without altering its direction.

Parameters:

O1 (torch.Tensor) – Ray origin.
D1 (torch.Tensor) – Ray direction.
wl (torch.Tensor) – Wavelength.
n_func_enviroment (Callable) – Function for surrounding medium.
meta_data (dict) – Ray tracing metadata.

Returns:

Final ray data.

Return type:

Tuple[torch.Tensor]

diffinytrace.trace_to_detector(optical_system: SequentialOpticalSystem, sequence: List, source, detector: Detector, num_rays: int, device=device(type='cpu'), method_ray_tracing: str = 'sobol_pow2')[source]¶

Traces rays through a system to a detector and returns the impact coordinates.

Parameters:

optical_system (SequentialOpticalSystem) – Ray-tracing pipeline.
sequence (list[str]) – Ordered names of system modules.
source – Source object with .sample() method.
detector (Detector) – Final surface to collect rays.
num_rays (int) – Number of rays to simulate.
device – Torch device (CPU/GPU).
method_ray_tracing (str) – Sampling method for source rays.

Returns:

(input samples, weights, detector plane hits, wavelengths)

Return type:

Tuple[torch.Tensor]

diffinytrace.get_unused_params_mask(optical_system: SequentialOpticalSystem, sequence: List[str], source, params, num_rays: int = 100000, method_ray_tracing='sobol') → List[BoolTensor][source]¶

Returns a boolean mask identifying which parameters are unused in the ray tracing process.

Parameters:

optical_system (SequentialOpticalSystem) – Full system.
sequence (list) – Ordered module names.
source – Ray source.
params (list[torch.nn.Parameter]) – Parameter list.
num_rays (int) – Number of rays to test.
method_ray_tracing (str) – Sampling method.

Returns:

Masks of the same shape as each parameter.

Return type:

list[torch.BoolTensor]

diffinytrace.set_used_params_bounds_to_constant(optical_system, sequence, source, params, bounds_attr_name_new, bounds_attr_name_old='bounds', num_rays=100000, method_ray_tracing='sobol')[source]¶

Locks unused parameters by copying their current value as bounds, making them constant.

Parameters:

bounds_attr_name_new (str) – Name of the new bounds attribute to write.
bounds_attr_name_old (str) – Name of the original bounds attribute.

diffinytrace.set_unused_params_to_zero(optical_system: SequentialOpticalSystem, sequence, source, params, num_rays=200000, method_ray_tracing='sobol')[source]¶

Sets unused parameters (those with zero gradient across ray paths) to zero.

Parameters:

optical_system (SequentialOpticalSystem) – Full system.
sequence (list) – Ordered module names.
source – Ray source.
params (list[torch.nn.Parameter] or torch.nn.Parameter) – Parameters to clean.
num_rays (int) – Ray sample count.
method_ray_tracing (str) – Sampling method.

diffinytrace.set_unused_bspline_coeff_to_nearest(optical_system, sequence: list[str], source, bspline_surface, num_rays=100000, method_ray_tracing='sobol')[source]¶

Fills only the unused B-spline coefficients with the nearest used value.

This is useful for having geometry that is simple to manifacture while not tempering with the overall performance.

Parameters:

optical_system (SequentialOpticalSystem) – The optical system used for tracing.
sequence (list[str]) – Ordered list of module names for ray propagation.
source – Ray source with a .sample() method.
bspline_surface – Surface object with a .coeff tensor.
num_rays (int, optional) – Number of rays used to detect unused coefficients. Default is 100000.
method_ray_tracing (str, optional) – Sampling method (e.g., “sobol”). Default is “sobol”.

Raises:

RuntimeError – If all coefficients are unused — likely due to insufficient ray coverage.

class diffinytrace.FresnelVirtualLens(transform, lens_thickness, surface1, surface2, n_func, aperture_radius, surface1_derivative_x=None, surface1_derivative_y=None, surface2_derivative_x=None, surface2_derivative_y=None, is_square=False)[source]¶

Bases: OpticalElement

get_plot_points_2D(resolution: int)[source]¶

Returns a list of 2D plot points for the object.

Parameters:: resolution (int) – The resolution for the plot points.
Returns:: A list of 2D plot points.
Return type:: list

get_plot_points_3D(resolution)[source]¶

Returns 3D grid of Fresnel lens surface points for visualization.

Parameters:: resolution (int) – Grid resolution.
Returns:: List of (x, y, z) meshgrids.
Return type:: List[Tuple[torch.Tensor]]

get_plotly_color_scale()[source]¶

Returns color scale for plotly visualization.

Returns:: Color scale values.
Return type:: List

get_plotable_childs() → List[source]¶

Returns plotable child elements.

Returns:: List of child elements.
Return type:: List

forward(O1: Tensor, D1: Tensor, wl: Tensor, n_func_enviroment, meta_data) → Tensor[source]¶

Simulates light passing through the Fresnel lens.

Parameters:

O1 (torch.Tensor) – Ray origin positions.
D1 (torch.Tensor) – Ray directions.
wl (torch.Tensor) – Wavelengths.
n_func_enviroment – Function returning external medium refractive index.
meta_data (dict) – Ray metadata.

Returns:

Updated ray origins, directions, etc.

Return type:

Tuple

get_transform() → Transform[source]¶

Returns the transformation of the Fresnel lens exit surface.

Returns:: The transformation object.
Return type:: Transform

diffinytrace.set_tolerance(new_tolerance)[source]¶

Set the tolerance for ray intersection calculations.

Parameters:: new_tolerance (float) – The new tolerance value (must be > 0).
Raises:: ValueError – If new_tolerance is not greater than 0.

diffinytrace.get_tolerance()[source]¶

Get the current tolerance for ray intersection calculations.

Returns:: The current tolerance value.
Return type:: float

diffinytrace.set_max_iterations(new_max_iterations)[source]¶

Set the maximum number of iterations for the ray intersection solver.

Parameters:: new_max_iterations (int) – The new maximum number of iterations (must be > 0).
Raises:: ValueError – If new_max_iterations is not greater than 0.

diffinytrace.get_max_iterations()[source]¶

Get the current maximum number of iterations for the ray intersection solver.

Returns:: The current maximum number of iterations.
Return type:: int

diffinytrace.restore_default_settings()[source]¶

Reset to the default configuration settings for the ray tracer.

This will reset all configuration parameters to their default values.

diffinytrace.get_damping_factor()[source]¶

Get the current damping factor for the Newton method.

Returns:: The current damping factor.
Return type:: float

diffinytrace.set_damping_factor(new_damping_factor)[source]¶

Set the damping factor for the Newton method used in ray intersections.

Parameters:: new_damping_factor (float) – The new damping factor (0 < new_damping_factor <= 1).
Raises:: ValueError – If new_damping_factor is not between 0 and 1 (exclusive of 0, inclusive of 1).

diffinytrace.get_show_iteration_count()[source]¶

Check if the number of iterations should be shown.

Returns:: True if the number of iterations should be shown, False otherwise.
Return type:: bool

diffinytrace.set_show_iteration_count(flag)[source]¶

Set the option to show the number of iterations for each intersection.

Parameters:: flag (bool) – True to show the number of iterations, False otherwise.

diffinytrace.minimize(fun, params, constraints: List = [], method=None, tol: float = 1e-09, callback: Callable = <function <lambda>>, options: dict | None = None, nan_fallback: float = inf, bounds_attr_name: str = 'bounds', save_history: bool = False, call_before_minimize: bool = False) → dict[source]¶

Minimizes a function using SciPy’s minimize, supporting bounds and constraints.

Parameters:

fun (Callable) – Objective function.
params (list) – Parameters to optimize.
constraints (list) – List of constraints.
method (str) – SciPy optimization method (e.g., ‘L-BFGS-B’).
tol (float) – Tolerance for convergence.
callback (Callable) – Optional callback function.
options (dict) – Optimizer options.
nan_fallback (float) – Value to use if function returns NaN.
bounds_attr_name (str) – Name of bounds attribute.
save_history (bool) – If True, saves function values and gradient norms.
call_before_minimize (bool) – Whether to evaluate once before optimization.

Returns:

Dictionary containing optimization results (and optionally history).

Return type:

dict

diffinytrace.make_parameter_from_input(input, bounds=None, dtype=None, device=None, bounds_attr_name='bounds')[source]¶

Converts input to a torch.nn.Parameter and attaches bounds as an attribute.

Parameters:

input (array-like or torch.Tensor) – Input data.
bounds (torch.Tensor, optional) – Bounds to attach to the parameter.
dtype (torch.dtype, optional) – Desired tensor data type.
device (torch.device, optional) – Device to store the parameter on.
bounds_attr_name (str) – Attribute name used to store bounds.

Returns:

The parameter with bounds attached as an attribute.

Return type:

torch.nn.Parameter

class diffinytrace.RefractiveIndex(func, bounds)[source]¶

Bases: Module, PlotableWavelength

This class is used to calculate the refractive index of a material.

\[\mathbf{D'} = \mathbf{N} \sqrt{1 - (1 - \cos^2 \psi_i) \eta^2} + \eta (\mathbf{D} - \mathbf{N} \cos \psi_i),\]

where \(\cos \psi_i = \mathbf{D} \cdot \mathbf{N}\) and \(\eta = n / n'\) is the ratio of the refractive indices of the two materials.

Example

Below is an example of how to set up an optical material in our ray tracing library:

>>> import diffinytrace as dit
>>> BaSF = dit.RefractiveIndex(
>>>     lambda x: (1 + 1.65554268 / (1 - 0.0104485644 / x**2) +
>>>                   0.17131977 / (1 - 0.0499394756 / x**2) +
>>>                   1.33664448 / (1 - 118.961472 / x**2))**0.5,
>>>     [0.365, 2.5]
>>> )
>>> dit.plotting.wavelength.plot(
>>>     BaSF, title="Refractive index of BaSF (Barium dense flint)"
>>> )

Parameters:

func (callable) – A function that takes a wavelength in μm and returns the refractive index.
bounds (tuple) – A tuple containing the minimum and maximum wavelength in μm.

forward(wl)[source]¶

Calculates the refractive index for given wavelengths. :param wl: Wavelength in μm. :type wl: torch.Tensor or float

Returns:: Refractive index at the given wavelengths.
Return type:: torch.Tensor

diffinytrace.grad(outputs: Tensor | Sequence[Tensor], inputs: Tensor | Sequence[Tensor] | GradientEdge | Sequence[GradientEdge], grad_outputs: Tensor | Sequence[Tensor] | None = None, retain_graph: bool | None = True, create_graph: bool = True, only_inputs: bool = True, is_grads_batched: bool = False, materialize_grads: bool = False, remove_no_grad_outputs: bool = True)[source]¶

Computes the gradients of the outputs with respect to the inputs.

Parameters:

outputs (torch.Tensor or tuple of torch.Tensor) – The output tensors.
inputs (torch.Tensor or tuple of torch.Tensor) – The input tensors.
grad_outputs (torch.Tensor or tuple of torch.Tensor, optional) – The gradients of the outputs.
retain_graph (bool, optional) – Whether to retain the graph after computing gradients.
create_graph (bool, optional) – Whether to create the graph for higher-order gradients.
only_inputs (bool, optional) – Whether to only compute gradients for the inputs.
is_grads_batched (bool, optional) – Whether the gradients are batched.
materialize_grads (bool, optional) – Whether to materialize the gradients.
remove_no_grad_outputs (bool, optional) – Whether to remove outputs that do not require gradients.

Returns:

A list of gradients for each input tensor.

Return type:

list