diffinytrace 2.1__py3-none-any.whl

diffinytrace/physical_object.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "PhysicalObject",
+    "PhysicalSurface"
+]
+
+import torch
+import torch.nn as nn
+
+
+class PhysicalObject(nn.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.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def get_transformation_matrix(self):
+        """
+        Returns the transformation matrix of the object.
+
+        Returns:
+            torch.Tensor: The transformation matrix.
+        """
+        return self.get_transform().get_transformation_matrix()
+
+    def to_global_dir(self,direction):
+        """
+        Converts a direction vector from local to global coordinates.
+
+        Args:
+            direction (torch.Tensor): Direction vector in local coordinates.
+
+        Returns:
+            torch.Tensor: Direction vector in global coordinates.
+        """
+        return self.get_transform().to_global_dir(direction)
+
+    def to_local_dir(self,direction):
+        """
+        Converts a direction vector from global to local coordinates.
+
+        Args:
+            direction (torch.Tensor): Direction vector in global coordinates.
+
+        Returns:
+            torch.Tensor: Direction vector in local coordinates.
+        """
+        return self.get_transform().to_local_dir(direction)
+
+    def to_global_pos(self,position):
+        """
+        Converts a position from local to global coordinates.
+
+        Args:
+            position (torch.Tensor): Position in local coordinates.
+
+        Returns:
+            torch.Tensor: Position in global coordinates.
+        """
+        return self.get_transform().to_global_pos(position)
+
+    def to_local_pos(self,position):
+        """
+        Converts a position from global to local coordinates.
+
+        Args:
+            position (torch.Tensor): Position in global coordinates.
+
+        Returns:
+            torch.Tensor: Position in local coordinates.
+        """
+        return self.get_transform().to_local_pos(position)
+
+    def get_transform(self):
+        """
+        Returns the transformation object associated with this physical object.
+
+        Raises:
+            NotImplementedError: If not implemented in subclass.
+
+        Returns:
+            object: Transformation object.
+        """
+        raise NotImplementedError("PhysicalObject: get_transform not implemented")
+
+class PhysicalSurface(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.
+    """
+    def __init__(self):
+        super().__init__()
+    
+    def get_constraint_funs_leq_zero(self):
+        """
+        Returns constraint functions for the surface that must be less than or equal to zero.
+
+        Raises:
+            NotImplementedError: If not implemented in subclass.
+
+        Returns:
+            list[Callable]: List of constraint functions.
+        """
+        raise NotImplementedError("PhysicalSurface: get_constraint_funs_geq_zero is not implemented")
+    
+    """
+    def get_corners_in_parameter_space(self):
+        raise NotImplementedError("get_corners: is not implemented")
+    
+    def get_edge_funcs_in_parameter_space(self):
+        raise NotImplementedError("get_edge_funcs_in_parameter_space: is not implemented")
+    """
+    
+    def parametric_sample(self, num_points: int, method: str = "sobol") -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Samples points on the surface in parameter space.
+
+        Args:
+            num_points (int): Number of points to sample.
+            method (str, optional): Sampling method. Defaults to "sobol".
+
+        Raises:
+            NotImplementedError: If not implemented in subclass.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Sampled parameter positions and corresponding surface positions.
+        """
+        raise NotImplementedError("PhysicalSurface: sample() not implemented")
+
+    def parametric_surface(self, parametric_pos: torch.Tensor) -> torch.Tensor:
+        """
+        Maps parameter space positions to surface positions.
+
+        Args:
+            parametric_pos (torch.Tensor): Positions in parameter space.
+
+        Raises:
+            NotImplementedError: If not implemented in subclass.
+
+        Returns:
+            torch.Tensor: Surface positions.
+        """
+        raise NotImplementedError("PhysicalSurface: parametric_surface is not implemented")
+

diffinytrace/plotting/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "quantity2D",
+    "system2D",
+    "system3D",
+    "Plotable",
+    "wavelength"
+]
+
+from . import quantity2D
+from . import system2D
+from . import system3D
+from .core import Plotable
+from . import wavelength

diffinytrace/plotting/core.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "Plotable"
+]
+
+from matplotlib.colors import to_rgb
+from typing import List,Tuple,Optional,Union
+
+class Plotable:
+    """
+    Base class for objects that can be visualized in the optical system.
+
+    This class provides a common interface for objects that support 2D/3D plotting,
+    color scale configuration, and hierarchical visualization. Subclasses should implement
+    methods for generating plot points and color scales as needed.
+
+    Attributes:
+        fill_color (str): Color used to fill the object in plots.
+        outline_color (str): Color used for the object's outline in plots.
+        is_volume (bool): If True, the object is treated as a volumetric entity.
+    """
+    def __init__(self, fill_color:str = "white", outline_color:str = "black", is_volume:bool = False):
+        """
+        Initialize the plotable object with fill and outline colors.
+        
+        Args:
+            fill_color (str): The color used to fill the object.
+            outline_color (str): The color used for the outline of the object.
+            is_volume (bool): If True, the object is treated as a volume.
+        
+        """
+        self.fill_color = fill_color
+        self.outline_color = outline_color
+        self.is_volume = is_volume
+
+    def get_plotly_color_scale(self)->List[List[Union[float,str]]]:
+        """
+        Returns a color scale for Plotly, based on the fill and outline colors.
+        """
+        
+        fill_color_rgb = to_rgb(self.fill_color)
+        outline_color_rgb = to_rgb(self.outline_color)
+
+        fill_color_text = f'rgb({int(fill_color_rgb[0] * 255)}, {int(fill_color_rgb[1] * 255)}, {int(fill_color_rgb[2] * 255)})'
+        outline_color_text = f'rgb({int(outline_color_rgb[0] * 255)}, {int(outline_color_rgb[1] * 255)}, {int(outline_color_rgb[2] * 255)})'
+
+        colorscale = [[0,fill_color_text],
+                      [1,outline_color_text]]
+        return [colorscale]
+    
+    def get_plotable_childs(self)->List:
+        """
+        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.
+        """
+        out = []
+        for attr_name in dir(self):
+            attr = getattr(self, attr_name)
+            if isinstance(attr, Plotable):
+                out.append([attr,attr_name])
+        return out
+    
+    def get_plot_points_2D(self, resolution:int)->List:
+        """
+        Returns a list of 2D plot points for the object.
+        
+        Args:
+            resolution (int): The resolution for the plot points.
+        
+        
+        Returns:
+            list: A list of 2D plot points.
+        """
+        
+        print("get_plot_points_2D not implemented")
+        return []
+    
+    def get_plot_points_3D(self, resolution:int)->List:
+        """
+        Returns a list of 3D plot points for the object.
+        
+        Args:
+            resolution (int): The resolution for the plot points.
+        
+        Returns:
+            list: A list of 3D plot points.
+        """
+        print("get_plot_points_3D not implemented")
+        return []
+

diffinytrace/plotting/quantity2D.py

@@ -0,0 +1,188 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "plot"
+]
+
+import matplotlib.pyplot as plt
+from torch import linspace,meshgrid,zeros,no_grad,is_tensor
+from copy import deepcopy
+import numpy as np
+from typing import Callable,Tuple,Optional,Union
+
+def plot(val: Union[Callable, np.ndarray], 
+         title: str = "",
+         x_range: Optional[Tuple[float, float]] = None,
+         y_range: Optional[Tuple[float, float]] = None,
+         cmap: str = "jet",
+         subtitle: str = "",
+         title_fontsize: int = 14,
+         suptitle_fontsize: int = 12,
+         interpolation: str = "none",
+         xlabel: str = "x [mm]",
+         ylabel: str = "y [mm]",
+         colorbar: bool = True,
+         norm = None,
+         show: bool = True,
+         vmin: float | None = None,
+         vmax: float | None = None,
+         resolution: int = 501,
+         **kwargs):
+    """
+    Plot a 2D quantity using matplotlib.
+    This function handles both callable quantities and numpy arrays.
+    If a callable is provided, it should accept a 2D array of coordinates and return a 2D array of values.
+    The function will create a 2D plot with the specified parameters.
+    If a numpy array is provided, it will be plotted directly.
+    
+    Args:
+        val (callable or np.ndarray): The quantity to plot. If callable, it should accept a 2D array of coordinates.
+        title (str): Title of the plot.
+        x_range (tuple): Range of x-axis.
+        y_range (tuple): Range of y-axis.
+        cmap (str): Colormap to use.
+        subtitle (str): Subtitle of the plot.
+        title_fontsize (int): Font size of the title.
+        suptitle_fontsize (int): Font size of the subtitle.
+        interpolation (str): Interpolation method.
+        xlabel (str): Label for x-axis.
+        ylabel (str): Label for y-axis.
+        colorbar (bool): Whether to show colorbar.
+        show (bool): Whether to show the plot.
+        vmin (float): Minimum value for color normalization.
+        vmax (float): Maximum value for color normalization.
+        resolution (int): Resolution of the plot.
+        **kwargs: Additional arguments for imshow.
+    
+    Returns:
+        None
+    """
+    
+    if is_tensor(val):
+        val = val.detach().cpu().numpy()#.T
+    
+    if not (x_range is None):
+        if isinstance(x_range,float):
+            x_range = [-x_range,x_range]
+        if y_range is None:
+            y_range = x_range
+    if y_range is None:
+        y_range = [0,1.]
+    
+    if x_range is None:
+        x_range = [0,1.]
+
+    if not isinstance(val,np.ndarray):
+        _y = linspace(*x_range,resolution)
+        _x = linspace(*y_range,resolution)
+        mesh = meshgrid(_y,_x)
+        y = mesh[0].reshape(-1)
+        x = mesh[1].reshape(-1)
+        O = zeros((x.shape[0],2))        
+        O[:,0] = x
+        O[:,1] = y
+        val = val(O).reshape(resolution,resolution)
+        if is_tensor(val):
+            val = val.detach().cpu().numpy()
+    
+    val = val[::-1]
+    
+    plt.cla()   # Clear axis
+    plt.clf()   # Clear figure
+    fig, ax = plt.subplots()
+    mappable = ax.imshow(val,cmap=cmap,interpolation=interpolation,extent=list(x_range)+list(y_range),norm=norm,vmin=vmin,vmax=vmax,**kwargs)
+    
+    
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel(xlabel)
+    if subtitle != "":
+        fig.suptitle(title,fontsize=title_fontsize)
+        ax.set_title(subtitle, fontsize=suptitle_fontsize)
+    else:
+        ax.set_title(title,fontsize=title_fontsize)
+
+    if colorbar:
+        plt.colorbar(mappable, ax=ax)
+
+    if show:
+        plt.show()
+
+
+def intensity(val,title="",x_range=None,y_range=None,cmap="jet",interpolation="none",xlabel="x [mm]",ylabel="y [mm]",norm=None,show=True,vmin: float | None = None,vmax: float | None = None,**kwargs):
+    """
+    Plot a 2D intensity distribution using matplotlib.
+
+    Args:
+        val (callable, np.ndarray, or torch.Tensor): The intensity data to plot. If callable, it should accept a 2D array of coordinates and return a 2D array of values.
+        title (str, optional): Title of the plot.
+        x_range (tuple or None, optional): Range of x-axis. If None, defaults to [0, 1].
+        y_range (tuple or None, optional): Range of y-axis. If None, defaults to [0, 1].
+        cmap (str, optional): Colormap to use. Default is "jet".
+        interpolation (str, optional): Interpolation method for imshow. Default is "none".
+        xlabel (str, optional): Label for x-axis. Default is "x [mm]".
+        ylabel (str, optional): Label for y-axis. Default is "y [mm]".
+        norm (matplotlib.colors.Normalize or None, optional): Normalization for color mapping.
+        show (bool, optional): Whether to display the plot. Default is True.
+        vmin (float or None, optional): Minimum value for color normalization.
+        vmax (float or None, optional): Maximum value for color normalization.
+        **kwargs: Additional keyword arguments passed to matplotlib's imshow.
+
+    Returns:
+        None
+    """
+    plot(val,f"{title} [$W/mm^2$]",x_range,y_range,cmap=cmap,interpolation=interpolation,xlabel=xlabel,ylabel=ylabel,norm=norm,show=show,vmin=vmin,vmax=vmax,**kwargs)
+
+def height(val,title="",x_range=None,y_range=None,cmap="cool",interpolation="none",xlabel="x [mm]",ylabel="y [mm]",norm=None,show=True,vmin: float | None = None,vmax: float | None = None,**kwargs):
+    """
+    Plot a 2D height distribution using matplotlib.
+
+    Args:
+        val (callable, np.ndarray, or torch.Tensor): The height data to plot. If callable, it should accept a 2D array of coordinates and return a 2D array of values.
+        title (str, optional): Title of the plot.
+        x_range (tuple or None, optional): Range of x-axis. If None, defaults to [0, 1].
+        y_range (tuple or None, optional): Range of y-axis. If None, defaults to [0, 1].
+        cmap (str, optional): Colormap to use. Default is "cool".
+        interpolation (str, optional): Interpolation method for imshow. Default is "none".
+        xlabel (str, optional): Label for x-axis. Default is "x [mm]".
+        ylabel (str, optional): Label for y-axis. Default is "y [mm]".
+        norm (matplotlib.colors.Normalize or None, optional): Normalization for color mapping.
+        show (bool, optional): Whether to display the plot. Default is True.
+        vmin (float or None, optional): Minimum value for color normalization.
+        vmax (float or None, optional): Maximum value for color normalization.
+        **kwargs: Additional keyword arguments passed to matplotlib's imshow.
+
+    Returns:
+        None
+    """
+    plot(val,f"Height z [$mm$] of {title}",x_range,y_range,cmap=cmap,interpolation=interpolation,xlabel=xlabel,ylabel=ylabel,norm=norm,show=show,vmin=vmin,vmax=vmax,**kwargs)
+
+"""
+TODO: these functions need testing again before integration
+
+def surface(surface,name,aperture_radius,resolution=256,is_square=True,norm=None,show=True,**kwargs):
+    surface = deepcopy(surface)
+    surface = surface.cpu()
+    x_range = (-aperture_radius,aperture_radius)
+    y_range = (-aperture_radius,aperture_radius)
+    _x = linspace(-aperture_radius,aperture_radius,resolution)
+    _y = linspace(-aperture_radius,aperture_radius,resolution)
+    mesh = meshgrid(_x,_y)
+    x = mesh[0].reshape(-1)
+    y = mesh[1].reshape(-1)
+    O = zeros((x.shape[0],3))        
+    
+    O[:,0] = x
+    O[:,1] = y
+    z = None
+    
+    with no_grad():
+        z = surface.functional(O,*surface.get_functional_param_args())
+    
+    if not is_square:
+        z[O[:,[0,1]].norm(dim=-1)>aperture_radius] = float("nan")
+
+    z = z.detach().reshape(resolution,resolution)
+    height(z,name,x_range,y_range,norm=norm,show=show,**kwargs)
+"""
+    

diffinytrace/plotting/system2D.py

@@ -0,0 +1,220 @@
+# Copyright (c) 2025 Martin Pflaum
+# This file is part of the diffinytrace project, licensed under the MIT License.
+
+__all__ = [
+    "annotate_position_simple",
+    "annotate_position",
+    "annotated_arrow",
+    "layout",
+    "ray_paths",
+    "_plot_surface",
+    "_plot_surface_recursively",
+    "plot"
+]
+
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+import matplotlib.patches as patches
+import matplotlib.colors as mcolors
+from copy import deepcopy
+
+def annotate_position_simple(nz,ny,name):
+    """
+    Annotate the position of a point in 2D space using its coordinates.
+
+    Args:
+        nz (torch.Tensor): z-coordinates of the point.
+        ny (torch.Tensor): y-coordinates of the point.
+        name (str): Text label to annotate at the position.
+
+    Returns:
+        None
+    """
+    zdiff = (torch.max(nz)-torch.min(nz))
+    ydiff = (torch.max(ny)-torch.min(ny))
+    offset = max(ydiff*0.05,zdiff*0.025)
+    argmax = torch.argmax(ny)
+    zpos = nz[argmax]#torch.min(nz)+zdiff*0.5
+    ypos = ny[argmax]
+    fontsize = 10
+    #-len(name)*fontsize/4.0
+    plt.annotate(name,xy=(zpos, ypos),fontsize=fontsize,xytext=(0.0,fontsize*0.5), textcoords='offset points')
+
+
+def annotate_position(position,offset,name,color="black",**kwargs):
+    """
+    Annotate a point in 2D space with an arrow and label.
+
+    Args:
+        position (tuple): (z, y) coordinates of the point.
+        offset (tuple): Offset for the annotation text.
+        name (str): Text label to annotate.
+        color (str): Color of the annotation and arrow.
+
+    Returns:
+        None
+    """
+    plt.annotate(name,color=color,xy=position,xytext=offset, textcoords='offset points',arrowprops=dict(arrowstyle="->",color=color,linewidth=1.5, mutation_scale=10), **kwargs)
+
+
+def annotated_arrow(start,end,offset,name,arrowstyle,color="black",**kwargs):
+    """
+    Draw and annotate an arrow between two points in 2D space.
+
+    Args:
+        start (tuple): Start position (z, y) of the arrow.
+        end (tuple): End position (z, y) of the arrow.
+        offset (tuple): Offset for the annotation text.
+        name (str): Text label to annotate.
+        arrowstyle (str): Matplotlib arrow style string.
+        color (str): Color of the arrow and annotation.
+
+    Returns:
+        None
+    """
+    
+    arrow_patch = patches.FancyArrowPatch(start, end, arrowstyle=arrowstyle,linewidth=1.5, mutation_scale=10,color=color)
+    plt.gca().add_patch(arrow_patch)
+    middle = (start[0]+ (end[0]-start[0])*0.5,start[1]+ (end[1]-start[1])*0.5)
+
+    plt.annotate(name,xy=middle,xytext=offset, textcoords='offset points',color=color,**kwargs)
+
+def layout():
+    """
+    Set up the layout for the 2D plot, including margins, aspect ratio, and axis labels.
+
+    Returns:
+        None
+    """
+    #plt.grid(True)
+    plt.margins(x=0.1,y=0.1)
+    plt.gca().set_aspect('equal')
+    plt.ylabel("y [mm]")
+    plt.xlabel("z [mm]")
+
+def ray_paths(rays,ray_color="#85549c",ray_linewidth=1.25):
+    """
+    Plot ray paths projected onto the y-z plane.
+
+    Args:
+        rays (list[torch.Tensor]): List of ray paths to plot.
+        ray_color (str): Color of the rays.
+        ray_linewidth (float): Line width of the rays.
+
+    Returns:
+        None
+    """
+    ray_color = mcolors.to_hex(ray_color)
+    print("WARNING: ray_paths will project the ray position onto the y-z plane!")
+    pathsA = rays
+    if torch.is_tensor(rays[0]):
+        pathsA = np.array([elem.numpy() for elem in rays])
+    pathsA = np.array(pathsA)
+
+    for iray in range(pathsA.shape[1]):
+        plt.plot(pathsA[:,iray,2],pathsA[:,iray,1],color=ray_color,linewidth=ray_linewidth)
+
+
+
+def _plot_surface(surface,name,resolution,annotate,fill_color,outline_color,linewidth):
+    """
+    Plot a 2D surface and optionally annotate it.
+
+    Args:
+        surface: Object with get_plot_points_2D method.
+        name (str): Name for annotation.
+        resolution (int): Resolution for the surface plot.
+        annotate (bool): Whether to annotate the surface.
+        fill_color (str): Fill color for the surface.
+        outline_color (str): Outline color for the surface.
+        linewidth (float): Line width for the surface.
+
+    Returns:
+        None
+    """
+    surface_list = surface.get_plot_points_2D(resolution)
+    if len(surface_list)==0:
+        return
+    if fill_color is None:
+        fill_color = surface.fill_color
+    if outline_color is None:
+        outline_color = surface.outline_color
+
+    zs,ys = torch.cat([z for z,y in surface_list]),torch.cat([y for z,y in surface_list])
+    if annotate:
+        annotate_position_simple(zs,ys,name)
+    if surface.is_volume:
+        ax = plt.gca()
+        ax.fill(zs, ys, facecolor=fill_color, edgecolor=outline_color, linewidth=linewidth)
+    else:
+        for z,y in surface_list:
+            plt.plot(z,y,color=outline_color,label="",linewidth=linewidth)
+    
+
+def _plot_surface_recursively(current_elem,name,resolution=200,annotate=False,fill_color=None,outline_color=None,linewidth=None):
+    """
+    Recursively plot a surface and its plotable children in 2D.
+
+    Args:
+        current_elem: The current plotable element.
+        name (str): Name for annotation.
+        resolution (int): Resolution for the surface plot.
+        annotate (bool): Whether to annotate the surface.
+        fill_color (str): Fill color for the surface.
+        outline_color (str): Outline color for the surface.
+        linewidth (float): Line width for the surface.
+
+    Returns:
+        None
+    """
+    _plot_surface(current_elem,name,resolution,annotate,fill_color,outline_color,linewidth)
+    for elem,elem_name in current_elem.get_plotable_childs():
+        _plot_surface_recursively(elem,elem_name,resolution,annotate,fill_color,outline_color,linewidth)
+        
+
+def plot(element=None,rays=None,resolution=200,annotate=False,ray_color="#85549c",ray_linewidth=1.25,fill_color=None,outline_color=None,linewidth=None,show=True):
+    """
+    Plot a 2D surface and optionally ray paths.
+
+    Args:
+        element: The element to plot (must implement Plotable interface).
+        rays (list[torch.Tensor]): List of ray paths to plot.
+        resolution (int): Resolution for the surface plot.
+        annotate (bool): Whether to annotate the surface.
+        ray_color (str): Color of the rays.
+        ray_linewidth (float): Line width of the rays.
+        fill_color (str): Fill color for the surface.
+        outline_color (str): Outline color for the surface.
+        linewidth (float): Line width for the surface.
+        show (bool): Whether to display the plot immediately.
+
+    Returns:
+        None
+    """
+    
+    layout()
+    
+    if isinstance(element,(list,tuple)):
+        for subelem in element:
+            subelem = deepcopy(subelem)
+            subelem = subelem.to("cpu")
+            _plot_surface_recursively(subelem,"",resolution,annotate,fill_color,outline_color,linewidth)
+
+    elif not element is None:    
+        element = deepcopy(element)
+        element = element.to("cpu")
+        _plot_surface_recursively(element,"",resolution,annotate,fill_color,outline_color,linewidth)
+    
+
+    if not rays is None:
+        if isinstance(rays,dict):
+            rays = rays["ray_paths"]
+
+        if torch.is_tensor(rays[0]):
+            rays = [elem.cpu() for elem in rays]
+        
+        ray_paths(rays,ray_color=ray_color,ray_linewidth=ray_linewidth)
+    if show:
+        plt.show()
+