"""Optic Updater Module
This module contains the OpticModifier class, which is responsible for updating the
optical system properties, such as the surface radii of curvature, thicknesses,
materials, conic constants, polarization, etc.
Kramer Harrison, 2024
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import optiland.backend as be
from optiland._suggest import options_hint
from optiland.apodization import BaseApodization
from optiland.geometries import StandardGeometry
from optiland.materials import IdealMaterial
if TYPE_CHECKING:
from optiland.materials.base import BaseMaterial
from optiland.rays import PolarizationState
[docs]
class OpticUpdater:
"""Class to update or modify an optical system
This class is responsible for updating the optical system properties, such as
the surface radii of curvature, thicknesses, materials, conic constants,
polarization, etc.
Args:
optic (Optic): The optical system to be modified.
"""
def __init__(self, optic):
self.optic = optic
[docs]
def set_radius(self, value, surface_number):
"""Set the radius of curvature of a surface.
Args:
value (float): The new radius of curvature.
surface_number (int): The index of the surface to modify.
"""
surface = self.optic.surfaces[surface_number]
try:
surface.geometry.set_radius(value)
except AttributeError:
# Plane geometry does not support set_radius; replace with StandardGeometry
cs = surface.geometry.cs
surface.geometry = StandardGeometry(cs, radius=value, conic=0)
[docs]
def set_conic(self, value, surface_number):
"""Set the conic constant of a surface.
Args:
value (float): The new conic constant.
surface_number (int): The index of the surface to modify.
"""
surface = self.optic.surfaces[surface_number]
surface.geometry.k = value
[docs]
def set_thickness(self, value, surface_number):
"""Set the thickness of a surface.
Args:
value (float): The new thickness value to set for the space
following the specified surface.
surface_number (int): The index of the surface *before* the
thickness to be modified.
"""
from optiland.paraxial_path import require_global_z_geometry
# Thickness updates rebuild downstream cs.z from cumulative
# thicknesses, which is only meaningful while the beam path runs
# along global +z. Guard before touching any geometry.
require_global_z_geometry(self.optic.surfaces, "set_thickness")
if surface_number == 0:
# First surface thickness sets the object distance.
# We treat this specially to avoid issues with infinite values.
self.optic.surfaces[0].thickness = value
self.optic.surfaces[0].geometry.cs.z = be.array(-value)
# No need to shift other surfaces as they are relative to S1 at z=0
return
# Source of truth for downstream positions is each surface's `.thickness`
# attribute. Update the requested one first, then rebuild every cs.z
# downstream from those thickness values (rather than incrementally
# mutating cs.z). This keeps every current-iteration thickness tensor
# in the autograd graph — the prior incremental form needed `detach()`
# to drop stale graphs from earlier iterations, but that also severed
# in-iteration grad paths when multiple thickness variables were
# updated sequentially (only the last one kept its gradient). See #569.
surfaces = self.optic.surfaces
n = len(surfaces)
if surface_number < n:
surfaces[surface_number].thickness = value
# Surface 1 anchored at z=0; cs.z[k] = sum of upstream thicknesses.
if n >= 2:
z = be.array(0.0)
surfaces[1].geometry.cs.z = be.array(z)
for k in range(2, n):
t_prev = surfaces[k - 1].thickness
z = z + (t_prev if hasattr(t_prev, "detach") else be.array(t_prev))
surfaces[k].geometry.cs.z = be.array(z)
[docs]
def set_index(self, value: float, surface_number: int) -> None:
"""Set the index of refraction of a surface.
Args:
value (float): The new refractive index value.
surface_number (int): The index of the surface whose *post-material*
(material after the surface) will be updated. This also updates
the *pre-material* of the subsequent surface.
"""
new_material = IdealMaterial(n=value, k=0)
self.set_material(new_material, surface_number)
[docs]
def set_material(self, material: BaseMaterial, surface_number: int) -> None:
"""Set the material of a surface.
Args:
value (BaseMaterial): The new material.
surface_number (int): The material of the surface whose *post-material*
(material after the surface) will be updated. This also updates
the *pre-material* of the subsequent surface.
"""
surface = self.optic.surfaces[surface_number]
surface.material_post = material
[docs]
def set_norm_radius(self, value, surface_number, is_fixed=True):
"""Set the normalization radius on a surface.
Args:
value (float): The new value for the normalization radius.
surface_number (int): The index of the surface to modify.
is_fixed (bool, optional): Whether to lock the normalization radius
from automatic paraxial updates. Defaults to True.
"""
surface = self.optic.surfaces[surface_number]
if hasattr(surface.geometry, "norm_radius"):
surface.geometry.norm_radius = value
surface.geometry.normalization_mode = "manual" if is_fixed else "auto"
else:
# This error is useful for debugging if used on an incorrect surface type
raise AttributeError(
f"Surface {surface_number} with geometry type "
f"'{surface.geometry.__class__.__name__}' has no "
"'norm_radius' attribute."
)
[docs]
def set_asphere_coeff(self, value, surface_number, aspher_coeff_idx):
"""Set the asphere coefficient on a surface
Args:
value (float): The new value for the aspheric coefficient.
surface_number (int): The index of the surface to modify.
aspher_coeff_idx (int): The index of the aspheric coefficient
within the surface's coefficient list to set.
"""
surface = self.optic.surfaces[surface_number]
surface.geometry.coefficients[aspher_coeff_idx] = value
[docs]
def set_polarization(self, polarization: PolarizationState | str):
"""Set the polarization state of the optic.
Args:
polarization (PolarizationState or str): The polarization
state to set. Can be a `PolarizationState` object or the string
'ignore'.
"""
if isinstance(polarization, str) and polarization != "ignore":
raise ValueError(
f"Invalid polarization state, got {polarization!r}. Pass a "
"PolarizationState instance, or the string 'ignore' to "
"disable polarization ray tracing.",
)
self.optic.polarization = polarization
[docs]
def scale_system(self, scale_factor):
"""Scales the optical system by a given scale factor.
Args:
scale_factor (float): The factor by which to scale all relevant
system dimensions (radii, thicknesses, EPD, physical apertures).
Raises:
UnsupportedParaxialGeometryError: If the beam path is folded off
global +z (or entered along another direction), before any
value is read or any geometry is touched. Scaling rebuilds
positions from cumulative thicknesses along global z, which
would move folded surfaces off their physical legs; guarding
only inside ``set_thickness`` would leave the system
partially scaled.
"""
from optiland.paraxial_path import require_global_z_geometry
# Preflight-atomic: reject before reading thicknesses (which walk
# the unfolded path) and before the first geometry.scale() call, so
# a rejected system is left exactly as it was.
require_global_z_geometry(self.optic.surfaces, "scale_system")
num_surfaces = self.optic.surfaces.num_surfaces
thicknesses = [
self.optic.surfaces.get_thickness(surf_idx)[0]
for surf_idx in range(num_surfaces - 1)
]
# Scale radii, geometries, and thicknesses
for surf_idx in range(num_surfaces):
surface = self.optic.surfaces[surf_idx]
surface.geometry.scale(scale_factor)
if surf_idx != num_surfaces - 1 and not be.isinf(thicknesses[surf_idx]):
self.set_thickness(thicknesses[surf_idx] * scale_factor, surf_idx)
# Scale aperture if the aperture type supports scaling
if self.optic.aperture and self.optic.aperture.is_scalable:
self.optic.aperture = self.optic.aperture.scale(scale_factor)
# Scale physical apertures
for surface in self.optic.surfaces:
if surface.aperture is not None:
surface.aperture.scale(scale_factor)
[docs]
def update_paraxial(self):
"""Update the semi-aperture of all surfaces based on paraxial marginal
and chief ray heights. Also updates normalization radii for relevant
surface types.
"""
ya, _ = self.optic.paraxial.marginal_ray()
yb, _ = self.optic.paraxial.chief_ray()
ya = be.abs(be.ravel(ya))
yb = be.abs(be.ravel(yb))
for k, surface in enumerate(self.optic.surfaces):
r_max = ya[k] + yb[k]
if surface.aperture is not None:
extent_max = be.max(be.abs(be.array(surface.aperture.extent)))
if be.isfinite(be.array(extent_max)):
r_max = be.max(be.array([r_max, extent_max]))
surface.set_semi_aperture(r_max=r_max)
self.update_normalization(surface)
[docs]
def update_normalization(self, surface) -> None:
"""Update the normalization radius/factors of a given non-spherical surface.
Args:
surface (Surface): The surface whose normalization factors are to be
updated.
"""
# Skip updating normalization when the normalization radius is a variable
if getattr(surface, "is_norm_radius_variable", False):
return
# Delegate updating normalization directly to the geometry
surface.geometry.update_normalization(surface.semi_aperture)
[docs]
def update(self) -> None:
"""Update the optical system by applying all defined pickups and solves.
If certain surface types requiring paraxial updates are present,
`update_paraxial()` is also called.
"""
self.optic.pickups.apply()
self.optic.solves.apply()
if any(
surface.surface_type
in ["chebyshev", "zernike", "forbes_qbfs", "forbes_q2d"]
for surface in self.optic.surfaces
):
self.update_paraxial()
[docs]
def image_solve(self):
"""Adjusts the position of the image surface (last surface) such that
the paraxial marginal ray crosses the optical axis at this new location.
This effectively sets the paraxial focus.
Raises:
UnsupportedParaxialGeometryError: If the system's unfolded axial
coordinate is not global z (folded or off-axis-entered
systems), before any surface is mutated -- the focus offset
is an unfolded axial distance and writing it into ``cs.z``
would move the image plane off its physical leg.
"""
from optiland.paraxial_path import require_global_z_geometry
# Guard before any computation touches geometry: no partial mutation.
require_global_z_geometry(self.optic.surfaces, "image_solve")
ya, ua = self.optic.paraxial.marginal_ray()
offset = float(ya[-1, 0] / ua[-1, 0])
surfaces = self.optic.surfaces
old_z = float(surfaces[-1].geometry.cs.z)
surfaces[-1].geometry.cs.z = be.array(old_z - offset)
surfaces[-2].thickness = float(surfaces[-1].geometry.cs.z) - float(
surfaces[-2].geometry.cs.z
)
[docs]
def flip(self):
"""Flips the optical system, reversing the order of surfaces (excluding
object and image planes), their geometries, and materials. Pickups and
solves referencing surface indices are updated accordingly.
The new first optical surface (originally the last) is placed at z=0.0.
"""
if self.optic.surfaces.num_surfaces < 3:
raise ValueError(
f"Cannot flip a system with "
f"{self.optic.surfaces.num_surfaces} surface(s): at least 3 "
"are required (object, one optical surface, image). Add the "
"missing surfaces with lens.add_surface(...) first."
)
# 1. Call SurfaceGroup.flip()
self.optic.surfaces.flip()
# 2. Define remapping function for indices
num_surfaces = self.optic.surfaces.num_surfaces
def remap_index_func(old_idx): # pragma: no cover
if old_idx == 0 or old_idx == num_surfaces - 1: # Object or Image surface
return old_idx
if 1 <= old_idx <= num_surfaces - 2:
return num_surfaces - 1 - old_idx
return old_idx # Should not happen if indices are valid
# 3. Handle Pickups
if self.optic.pickups and len(self.optic.pickups.pickups) > 0:
self.optic.pickups.remap_surface_indices(remap_index_func)
# 4. Handle Solves
if (
hasattr(self.optic, "solves")
and self.optic.solves
and hasattr(self.optic.solves, "solves")
and len(self.optic.solves.solves) > 0
):
self.optic.solves.remap_surface_indices(remap_index_func)
# 5. Update Optic instance
self.update()
[docs]
def set_apodization(
self, apodization: BaseApodization | str | dict = None, **kwargs
):
"""Sets the apodization for the optical system.
This method supports setting the apodization in multiple ways:
1. By providing an instance of a `BaseApodization` subclass.
2. By providing a string identifier (e.g., "GaussianApodization")
and keyword arguments for its parameters.
3. By providing a dictionary that can be passed to `from_dict`.
4. By passing `None` to remove any existing apodization.
Args:
apodization (BaseApodization | str | dict, optional): The
apodization to apply. Defaults to None.
**kwargs: Additional keyword arguments used to initialize the
apodization class when `apodization` is a string.
Raises:
TypeError: If the provided `apodization` is not a supported type.
ValueError: If the string identifier is not found in the registry.
"""
if apodization is None:
self.optic.apodization = None
elif isinstance(apodization, BaseApodization):
self.optic.apodization = apodization
elif isinstance(apodization, str):
if apodization in BaseApodization._registry:
apodization_class = BaseApodization._registry[apodization]
self.optic.apodization = apodization_class(**kwargs)
else:
raise ValueError(
f"Unknown apodization type, got {apodization!r}."
f"{options_hint(apodization, BaseApodization._registry)}"
)
elif isinstance(apodization, dict):
self.optic.apodization = BaseApodization.from_dict(apodization)
else:
raise TypeError(
"apodization must be a string, a dict, a BaseApodization "
"instance, or None."
)
