"""Optical System Visualization Module
This module provides tools for visualizing optical systems.
It utilizes Matplotlib to render optical components and ray tracing paths.
The `OpticViewer` class is the primary interface for generating these visualizations,
offering customization for ray properties, field of view, and display parameters.
Kramer Harrison, 2024
re-worked by Manuel Fragata Mendes, june 2025
"""
from __future__ import annotations
import numpy as np
import optiland.backend as be
from optiland.visualization.base import BaseViewer2D
from optiland.visualization.system.interaction import InteractionManager
from optiland.visualization.system.lens import Lens2D
from optiland.visualization.system.rays import Rays2D
from optiland.visualization.system.surface import Surface2D
from optiland.visualization.system.system import OpticalSystem
from optiland.visualization.system.utils import transform
[docs]
class OpticViewer(BaseViewer2D):
"""A class used to visualize optical systems.
Args:
optic: The optical system to be visualized.
Attributes:
optic: The optical system to be visualized.
rays: An instance of Rays2D for ray tracing.
system: An instance of OpticalSystem for system representation.
Methods:
See :meth:`view` for visualizing the optical system with specified
parameters.
"""
def __init__(self, optic):
self.optic = optic
self.rays = Rays2D(optic)
self.system = OpticalSystem(optic, self.rays, projection="2d")
self.legend_artist_map = {}
[docs]
def view(
self,
fields="all",
wavelengths="primary",
num_rays=3,
distribution=None,
show_apertures=True,
hide_vignetted=False,
figsize=None,
xlim=None,
ylim=None,
title=None,
reference=None,
tooltip_format=None,
show_legend=True,
projection="YZ",
ax=None,
theme=None,
show=True,
):
"""Visualizes the optical system.
Args:
fields (str, optional): The fields to be visualized.
Defaults to 'all'.
wavelengths (str, optional): The wavelengths to be visualized.
Defaults to 'primary'.
num_rays (int, optional): The number of rays to be visualized.
Defaults to 3.
distribution (str | None, optional): The distribution of rays.
Defaults to None, which selects a default based on projection.
show_apertures (bool, optional): If True, overlays aperture graphics
on the system view. Defaults to True.
hide_vignetted (bool, optional): If True, rays that vignette at any
surface are not shown. Defaults to False.
figsize (tuple, optional): The size of the figure.
Defaults to None, which uses the theme's default.
xlim (tuple, optional): The x-axis limits. Defaults to None.
ylim (tuple, optional): The y-axis limits. Defaults to None.
reference (str, optional): The reference rays to plot. Options
include "chief" and "marginal". Defaults to None.
projection (str, optional): The projection plane. Must be 'XY',
'XZ', or 'YZ'. Defaults to 'YZ'.
ax (matplotlib.axes.Axes, optional): The axes to plot on.
If None, a new figure and axes are created. Defaults to None.
show (bool, optional): If True (default), calls plt.show(). Set
False for headless use (e.g. saving to file, CI environments)
or when embedding in an existing figure via ``ax``.
"""
if projection not in ["XY", "XZ", "YZ"]:
raise ValueError("Invalid projection type. Must be 'XY', 'XZ', or 'YZ'.")
if distribution is None:
if projection == "XY":
distribution = "hexapolar"
elif projection == "XZ":
distribution = "line_x"
else:
distribution = "line_y"
is_gui_embedding = ax is not None
theme = self._resolve_theme(theme)
fig, ax = self._make_figure(theme, figsize, ax)
interaction_manager = InteractionManager(fig, ax, self.optic, tooltip_format)
ray_artists = self.rays.plot(
ax,
fields=fields,
wavelengths=wavelengths,
num_rays=num_rays,
distribution=distribution,
reference=reference,
theme=theme,
projection=projection,
hide_vignetted=hide_vignetted,
)
for artist, ray_bundle in ray_artists.items():
interaction_manager.register_artist(artist, ray_bundle)
system_artists = self.system.plot(
ax, theme=theme, projection=projection, show_apertures=show_apertures
)
for artist, surface in system_artists.items():
interaction_manager.register_artist(artist, surface)
ax.axis("image")
if xlim is None or ylim is None:
auto_xlim, auto_ylim = self._default_axis_limits(projection)
xlim = xlim or auto_xlim
ylim = ylim or auto_ylim
self._apply_axes_style(ax, projection, theme, title=title, xlim=xlim, ylim=ylim)
if show and not is_gui_embedding:
import matplotlib.pyplot as plt
plt.show()
# Return the figure, axes and interaction_manager
return fig, ax, interaction_manager
def _default_axis_limits(self, projection):
"""Compute default axis limits sized to the lens system rather than
the full ray extent.
For an infinite-conjugate (angle field) system, the object-side ray
segment is drawn from an arbitrary, often very distant, launch
point -- fine for aiming, but if left to matplotlib's autoscale it
dominates the view for wide-FOV systems, squeezing the actual lens
system down to an unreadable sliver. Sizing instead from the real
surfaces (z) and ``r_extent`` -- the same radially-symmetric,
per-surface ray extent ``OpticalSystem`` already uses to size lens
and mirror components (see ``system.py``) -- keeps the transverse
limits centered on the optical axis and guaranteed to cover
whatever radius the lens components are actually drawn at (an
asymmetric min/max of the traced rays' signed coordinates isn't
enough: a field that's only ever traced on one side of the axis
would leave the *other*, still-drawn side of the lens clipped),
while still leaving margin to see rays approaching and entering
the first surface.
Args:
projection (str): The projection plane, 'XY', 'XZ', or 'YZ'.
Returns:
tuple: ``(xlim, ylim)``, each either a ``(min, max)`` tuple or
``None`` if there isn't enough information to size that axis
(falls back to matplotlib's own autoscale).
"""
if projection == "XY":
# A single cross-section at one z -- no object-segment issue.
return None, None
# Axis limits are drawn in the global frame, so use the full 3-D
# vertex chain rather than the unfolded axial coordinate
# ``positions`` reports. A folded path can double back in z (first/
# last surface do not bound the span) and walk off-axis
# transversely (the vertical axis is not centered on 0), so both
# axes are sized from the min/max vertex coordinates. For a
# straight +z system every vertex has x = y = 0 and z increases
# monotonically, which reduces to the previous first-to-last-z,
# symmetric-about-the-axis behavior exactly.
vertices = be.to_numpy(self.optic.surfaces.vertices_gcs)
start_idx = 1 if self.optic.object_surface.is_infinite else 0
if start_idx >= vertices.shape[0] - 1:
return None, None
vertices = vertices[start_idx:]
# Resolve bounds from the components that will actually be drawn. In
# particular, a ray's first blocking hit can be far outside a finite
# physical aperture and must not enlarge the auto-fit region.
surfaces = []
for component in self.system.components:
if isinstance(component, Surface2D):
surfaces.append(component)
elif isinstance(component, Lens2D):
surfaces.extend(component.surfaces)
r_extent = (
np.asarray([float(be.to_numpy(surface.extent)) for surface in surfaces])
if surfaces
else be.to_numpy(self.rays.r_extent)[start_idx:]
)
# A tilted image marker can extend past the first/last vertex along Z.
# Include its transformed outline, without tracing any additional rays.
z_bounds = [vertices[:, 2].ravel()]
for surface in surfaces:
x, y, z = surface._compute_sag(projection)
_, _, z = transform(x, y, z, surface.surf, is_global=False)
z = be.to_numpy(z)
z_bounds.append(z[np.isfinite(z)])
z_bounds = np.concatenate(z_bounds)
z_min, z_max = float(z_bounds.min()), float(z_bounds.max())
z_margin = max(0.15 * (z_max - z_min), 1e-6)
auto_xlim = (z_min - z_margin, z_max + z_margin)
r_extent = r_extent[np.isfinite(r_extent)]
if r_extent.size == 0:
return auto_xlim, None
r_max = float(r_extent.max())
if r_max <= 0:
return auto_xlim, None
# Vertical axis: transverse component of the drawn projection.
t_comp = 0 if projection == "XZ" else 1
t_min = float(vertices[:, t_comp].min())
t_max = float(vertices[:, t_comp].max())
r_margin = max(0.15 * max(r_max, t_max - t_min), 1e-6)
auto_ylim = (t_min - r_max - r_margin, t_max + r_max + r_margin)
return auto_xlim, auto_ylim
