diffinytrace 2.1__py3-none-any.whl

diffinytrace/basis_functions/legendre.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "precompute_legendre_polynomials",
+    "basis_2D",
+    "get_num_coeff"
+]
+
+import torch
+
+def precompute_legendre_polynomials(x: torch.Tensor, degree: int) -> list[torch.Tensor]:
+    """
+    Precomputes all Legendre polynomials up to a given degree.
+
+    Args:
+    x (torch.Tensor): Input tensor for x-coordinates.
+    degree (int): Maximum degree of the Legendre polynomials.
+    
+    Returns:
+    list of torch.Tensor: List of precomputed Legendre polynomials [P_0(x), P_1(x), ..., P_degree(x)].
+    """
+    P = [torch.ones_like(x)]  # P_0(x) = 1
+    if degree >= 1:
+        P.append(x)  # P_1(x) = x
+    
+    for n in range(2, degree + 1):
+        Pn = ((2 * n - 1) * x * P[-1] - (n - 1) * P[-2]) / n
+        P.append(Pn)
+    
+    return P
+
+def basis_2D(points: torch.Tensor, degree: int) -> torch.Tensor:
+    """
+    Generates 2D Legendre polynomial basis functions up to a given degree using precomputed 1D polynomials.
+
+    Args:
+    degree (int): Maximum degree of the Legendre polynomials.
+    x (torch.Tensor): x-coordinates as a torch tensor.
+    y (torch.Tensor): y-coordinates as a torch tensor.
+
+    Returns:
+    torch.Tensor: Tensor of shape (num_basis_functions, *x.shape) with all 2D basis functions.
+    """
+    x = points[:, 0]
+    y = points[:, 1]
+
+    # Precompute 1D Legendre polynomials for x and y
+    Px = precompute_legendre_polynomials(x, degree)
+    Py = precompute_legendre_polynomials(y, degree)
+
+    basis_functions = {}
+
+    for i in range(degree + 1):
+        for j in range(degree + 1 - i):  # Ensure that i + j <= degree
+            if not i+j in basis_functions.keys():
+                basis_functions[i+j] = []
+            basis_functions[i+j] += [Px[i] * Py[j]]
+    out = []
+    for key in basis_functions.keys():
+        out += basis_functions[key]
+
+    # Stack basis functions along a new dimension
+    return torch.stack(out, dim=1)
+
+def get_num_coeff(degree: int) -> int:
+    """
+    Returns the number of coefficients for a given degree of Legendre polynomials.
+    The number of coefficients is given by the formula (degree + 1) * (degree + 2) / 2.
+    
+    Args:
+        degree (int): Degree of the Legendre polynomial.
+        
+    Returns:
+        int: Number of coefficients.
+    """
+    return (degree + 1) * (degree + 2) // 2

diffinytrace/basis_functions/zernike.py

@@ -0,0 +1,235 @@
+"""
+Zernike polynomial basis functions for optical wavefront representation.
+
+This module provides functions to compute Zernike polynomials, which are commonly used
+in optics for describing wavefront aberrations over a circular aperture. The polynomials
+are orthogonal over the unit disk and are indexed by radial order (n) and azimuthal 
+frequency (m).
+
+.. figure:: _static/zernike_plot1.png
+   :alt: Zernike polynomials visualization
+   :width: 60%
+   :align: center
+
+   Visualization of Zernike polynomials organized by radial order (rows) and azimuthal frequency (columns).
+
+
+Example:
+    Basic usage for computing and visualizing Zernike polynomials:
+
+    >>> import torch
+    >>> import numpy as np
+    >>> import matplotlib.pyplot as plt
+    >>> import diffinytrace.basis_functions.zernike as zernike
+    >>> 
+    >>> # Create unit circle grid
+    >>> grid_size = 256
+    >>> x = torch.linspace(-1, 1, grid_size)
+    >>> y = torch.linspace(-1, 1, grid_size)
+    >>> X, Y = torch.meshgrid(x, y, indexing='ij')
+    >>> 
+    >>> # Create mask for unit circle
+    >>> mask = (X**2 + Y**2) <= 1.0
+    >>> x_points = X[mask]
+    >>> y_points = Y[mask]
+    >>> points = torch.stack([x_points, y_points], dim=1)
+    >>> 
+    >>> # Evaluate Zernike polynomials
+    >>> max_n = 6  # Maximum radial degree
+    >>> basis_values = zernike.basis_function(max_n, points)
+    >>> 
+    >>> # Group basis functions by radial degree
+    >>> basis_by_degree = {}
+    >>> for basis_idx in range(basis_values.shape[1]):
+    ...     radial_order = zernike.get_radial_order(basis_idx)
+    ...     if radial_order not in basis_by_degree:
+    ...         basis_by_degree[radial_order] = []
+    ...     basis_by_degree[radial_order].append(basis_idx)
+    >>> 
+    >>> # Visualize the polynomials
+    >>> max_cols = max(len(indices) for indices in basis_by_degree.values())
+    >>> num_rows = len(basis_by_degree)
+    >>> fig, axes = plt.subplots(num_rows, max_cols, figsize=(3*max_cols, 3*num_rows))
+    >>> 
+    >>> for row_idx, (radial_order, basis_indices) in enumerate(sorted(basis_by_degree.items())):
+    ...     for col_idx, basis_idx in enumerate(basis_indices):
+    ...         # Create 2D array with NaN outside unit circle
+    ...         tmp = torch.full((grid_size, grid_size), float('nan'))
+    ...         tmp[mask] = basis_values[:, basis_idx]
+    ...         
+    ...         # Plot
+    ...         ax = axes[row_idx, col_idx]
+    ...         im = ax.imshow(tmp.numpy(), extent=[-1, 1, -1, 1], 
+    ...                       origin='lower', cmap='jet', vmin=-1, vmax=1)
+    ...         azimuthal = zernike.get_azimuthal_frequency(basis_idx)
+    ...         ax.set_title(f"$Z^{{{azimuthal}}}_{{{radial_order}}}$", fontsize=25)
+    ...         ax.set_xticks([])
+    ...         ax.set_yticks([])
+    ...         ax.set_aspect('equal')
+    >>> 
+    >>> plt.tight_layout()
+    >>> plt.show()
+
+Notes:
+    - Zernike polynomials are only defined for points within the unit circle (r ≤ 1)
+    - Radial order n determines the number of radial variations
+    - Azimuthal frequency m determines the angular variations and symmetry
+"""
+
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "basis_function",
+    "get_num_basis",
+    "get_radial_order",
+    "get_azimuthal_frequency"
+]
+
+import torch
+import math
+
+def __zernike_calc(n:int, m:int, r_powers: list) -> torch.Tensor:
+    radial_sum = torch.zeros_like(r_powers[0])
+    m = abs(m)
+    
+    for k in range((n - m) // 2 + 1):
+        coef = math.factorial(n - k) / (
+            math.factorial(k) * math.factorial((n + m) // 2 - k) * math.factorial((n - m) // 2 - k))
+        if k%2==1:
+            coef = -coef
+        power_idx = n - 2 * k - m
+        
+        if power_idx < 0:
+            raise RuntimeError("Potential zero division!")
+            #tmp = r_powers[abs(power_idx)]
+            #radial_sum += coef / tmp
+        if power_idx % 2 == 1:
+            raise RuntimeError("tried to acces odd power idx!")
+            
+        radial_sum += coef*r_powers[abs(power_idx)]
+    
+    return radial_sum
+
+def basis_2D(points: torch.Tensor, max_radial_order: int) -> torch.Tensor:
+    """
+    Compute Zernike polynomials for a given set of points.
+    
+    Args:
+        max_radial_order (int): Maximum radial order.
+        points (torch.Tensor): Tensor of shape (N, 2) containing the x and y coordinates of the points.
+    
+    Returns:
+        torch.Tensor: Tensor of shape (N, num_coeffs) containing the Zernike polynomial values.
+    """
+    x = points[:, 0]
+    y = points[:, 1]
+    
+    #r = torch.sqrt(x**2 + y**2)
+    r2 = x**2 + y**2
+    # Precompute powers of r from r^0 to r^max_radial_order
+    r_powers = [] #[r ** i for i in range(max_radial_order+1)]
+    for i in range(max_radial_order+1):
+        if i%2 == 0:
+            r_powers += [r2**(i/2.0)]
+        else:
+            r_powers += [None]
+        
+    # List to store Zernike polynomial results
+    zernike_polynomials = []
+    
+    # Loop over radial and azimuthal degrees
+    for n in range(max_radial_order + 1):
+        for m in range(-n, n + 1, 2):  # m must have the same parity as n
+            #r_m = r_powers[abs(m)]  # Precompute r^m for both cos and sin components
+            
+            if m >= 0:
+                #TODO Remove weird complex number stuff!
+                multiplier = torch.real((y + 1j * x)**abs(m))#TODO multiply after zerinke_calc!
+                #this is to slow!!
+                zernike_polynomials.append(multiplier*__zernike_calc(n, m, r_powers))
+            else:
+                multiplier = torch.imag((y + 1j * x)**abs(m))
+                zernike_polynomials.append(multiplier*__zernike_calc(n, m, r_powers))
+    
+    # Stack all Zernike polynomials into a single tensor
+    zernike_basis = torch.stack(zernike_polynomials, dim=1)
+    
+    return zernike_basis
+
+def get_num_basis(max_radial_order:int) -> int:
+    """
+    Calculate the number of basis functions for Zernike polynomials up to a given radial order.
+    The number of basis functions is given by the formula (n + 1) * (n + 2) / 2.
+
+    Args:
+        max_radial_order (int): Maximum radial order.
+
+    Returns:
+        int: Number of coefficients.
+    """
+
+    n = max_radial_order+1
+    return int(n*(n+1) / 2)
+
+
+def get_radial_order(basis_idx:int) -> int:
+    """
+    Calculate the radial degree from the basis function index.
+    
+    Args:
+        basis_idx (int): Index of the basis function.
+    
+    Returns:
+        int: Radial degree.
+    """
+    basis_idx_runner = basis_idx
+
+    num_azimuthal_frequencies = 1
+    row_idx = 0
+    while True:
+        if basis_idx_runner < num_azimuthal_frequencies:
+            return row_idx
+
+        basis_idx_runner = basis_idx_runner-num_azimuthal_frequencies
+        
+        num_azimuthal_frequencies += 1
+        row_idx += 1
+
+def get_azimuthal_frequency(basis_idx:int) -> int:
+    """
+    Calculate the azimuthal frequency from the basis function index.
+
+    Args:
+        basis_idx (int): Index of the basis function.
+    
+    Returns:
+        int: Azimuthal frequency (m value).
+    """
+    # First get the radial degree n
+    row_idx = get_radial_order(basis_idx)
+    
+    
+    num_azimuthal_frequencies = 1
+    basis_idx_runner = basis_idx
+    for k in range(row_idx):
+        basis_idx_runner = basis_idx_runner-num_azimuthal_frequencies
+        num_azimuthal_frequencies += 1
+    
+    x_idx_start = None
+    if num_azimuthal_frequencies % 2 == 0:
+        half = num_azimuthal_frequencies // 2
+        tmp = (half-1)*2+1
+        x_idx_start = - tmp
+    else:
+        half = num_azimuthal_frequencies // 2
+        tmp = (half)*2
+        x_idx_start = - tmp
+        #x_idx_start = - (num_azimuthal_frequencies // 2)
+    # For radial degree n, we have m values: -n, -n+2, ..., -2, 0, 2, ..., n-2, n
+    # The azimuthal frequency m follows the pattern:
+    # pos_in_row = 0 -> m = -n
+    # pos_in_row = 1 -> m = -n + 2
+    # ...
+    # pos_in_row = n -> m = n
+    return x_idx_start + basis_idx_runner*2

diffinytrace/config.py

@@ -0,0 +1,140 @@
+"""
+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()
+"""
+
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "set_show_iteration_count",
+    "get_show_iteration_count",
+    "set_tolerance",
+    "get_tolerance",
+    "set_max_iterations",
+    "get_max_iterations",
+    "set_damping_factor",
+    "get_damping_factor",
+    "restore_default_settings"
+]
+
+# Default settings for ray intersection parameters
+tolerance:float = 1e-6  # Default tolerance for ray intersection
+max_iterations:int = 100  # Default maximum number of iterations for the solver
+damping_factor:float = 1.0  # Default damping factor for the Newton method (1.0 means no damping)
+show_iteration_count:bool = False  # Default setting to not show the number of iterations
+
+def set_show_iteration_count(flag):
+    """
+    Set the option to show the number of iterations for each intersection.
+
+    Args:
+        flag (bool): True to show the number of iterations, False otherwise.
+    """
+    global show_iteration_count
+    show_iteration_count = flag
+
+def get_show_iteration_count():
+    """
+    Check if the number of iterations should be shown.
+
+    Returns:
+        bool: True if the number of iterations should be shown, False otherwise.
+    """
+    return show_iteration_count
+
+def set_tolerance(new_tolerance):
+    """
+    Set the tolerance for ray intersection calculations.
+
+    Args:
+        new_tolerance (float): The new tolerance value (must be > 0).
+
+    Raises:
+        ValueError: If `new_tolerance` is not greater than 0.
+    """
+    if new_tolerance <= 0:
+        raise ValueError("Tolerance must be greater than 0.")
+    global tolerance
+    tolerance = new_tolerance
+
+def get_tolerance():
+    """
+    Get the current tolerance for ray intersection calculations.
+
+    Returns:
+        float: The current tolerance value.
+    """
+    return tolerance
+
+def set_max_iterations(new_max_iterations):
+    """
+    Set the maximum number of iterations for the ray intersection solver.
+
+    Args:
+        new_max_iterations (int): The new maximum number of iterations (must be > 0).
+
+    Raises:
+        ValueError: If `new_max_iterations` is not greater than 0.
+    """
+    if new_max_iterations <= 0:
+        raise ValueError("Maximum iterations must be greater than 0.")
+    global max_iterations
+    max_iterations = new_max_iterations
+
+def get_max_iterations():
+    """
+    Get the current maximum number of iterations for the ray intersection solver.
+
+    Returns:
+        int: The current maximum number of iterations.
+    """
+    global max_iterations
+    return max_iterations
+
+def set_damping_factor(new_damping_factor):
+    """
+    Set the damping factor for the Newton method used in ray intersections.
+
+    Args:
+        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).
+    """
+    global damping_factor
+    if 0 < new_damping_factor <= 1:
+        damping_factor = new_damping_factor
+    else:
+        raise ValueError("Damping factor must be between 0 and 1.")
+
+def get_damping_factor():
+    """
+    Get the current damping factor for the Newton method.
+
+    Returns:
+        float: The current damping factor.
+    """
+    return damping_factor
+
+def restore_default_settings():
+    """
+    Reset to the default configuration settings for the ray tracer.
+
+    This will reset all configuration parameters to their default values.
+    """
+    global tolerance, max_iterations, damping_factor, show_iteration_count
+    tolerance = 1e-6
+    max_iterations = 100
+    damping_factor = 1.0
+    show_iteration_count = False  # Reset to default (not showing)

diffinytrace/constraints.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = ["Constraint", "EqualZero", "GEQZero", "LEQZero"]
+
+from .physical_object import PhysicalSurface
+import torch
+from . integrators import Cube
+from .optimize import minimize
+from .utils.autograd import grad
+#from .element import OpticalSurface
+
+class Constraint():
+    """
+    Base class for optimization constraints.
+
+    Attributes:
+        fun (Callable): Function defining the constraint.
+        type (str): Type of constraint ('eq' or 'ineq').
+    """
+    def __init__(self,fun,type):
+        self.fun = fun
+        self.type = type
+    
+class EqualZero(Constraint):
+    """
+    Equality constraint enforcing `fun() == 0`.
+
+    Args:
+        fun (Callable): The constraint function.
+    """
+    def __init__(self,fun):
+        super().__init__(fun,'eq')
+
+class GEQZero(Constraint):
+    """
+    Inequality constraint enforcing `fun() >= 0`.
+
+    Args:
+        fun (Callable): The constraint function.
+    """
+    def __init__(self,fun):
+        super().__init__(fun,'ineq')
+
+class LEQZero(Constraint):
+    """
+    Inequality constraint enforcing `fun() <= 0`.
+
+    Args:
+        fun (Callable): The constraint function.
+    """
+    def __init__(self,fun):
+        super().__init__(lambda: -fun(),'ineq')
+