"""Robust Ray Aiming Module
This module implements a chief-ray calibrated robust ray aiming algorithm.
For each field, a cheap chief-ray solve plus four cardinal edge probes
(``(+-1, 0)``, ``(0, +-1)`` on the stop) are fit to a 2x2 affine launch
model (see ``pupil_map.py``). That model seeds every requested pupil point,
which is then driven to its exact target by the same Newton/Broyden polish
used by ``IterativeRayAimer``. Calibration is warm-started from the
previous fit (or the nearest already-solved field) rather than a paraxial
guess extrapolated across the whole field range, which is what allows this
to converge cold at extreme field angles without the recursive homotopy
subdivision the previous implementation relied on.
See ``optiland/jupyter/SPEC_ray_aiming_20260703.md`` for the full design.
Kramer Harrison, 2026
"""
from __future__ import annotations
import contextlib
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import optiland.backend as be
from optiland.paraxial_path import paraxial_seed_scope
from optiland.rays.ray_aiming.base import BaseRayAimer
from optiland.rays.ray_aiming.initialization import get_stop_radius_strategy
from optiland.rays.ray_aiming.iterative import IterativeRayAimer
from optiland.rays.ray_aiming.parameterization import (
LaunchParameterization,
SolveReport,
)
from optiland.rays.ray_aiming.pupil_map import PupilMap, PupilMapCache, to_float
from optiland.rays.ray_aiming.registry import register_aimer
if TYPE_CHECKING:
from optiland.optic import Optic
# Cardinal edge probes on the stop, in (Px, Py) order: east, west, north, south.
_EDGE_PROBES = ((1.0, 0.0), (-1.0, 0.0), (0.0, 1.0), (0.0, -1.0))
[docs]
@dataclass(frozen=True)
class RobustFieldReport:
"""Per-field outcome of one robust aiming call.
Attributes:
Hx, Hy: Normalized field coordinates of the group.
wavelength: Wavelength of the group in micrometers.
final_polish: :class:`SolveReport` of the group's final Newton
polish (or of the failed chief solve when calibration itself
failed).
chief_seed_strategy: How the chief anchor was obtained --
``"initial_guess"`` (caller-supplied guess solved directly),
``"cached_map"`` (fresh cached map reused, no calibration),
``"warm_map"`` (chief solved from a stale/nearest cached map),
``"direct_paraxial"`` (chief solved from the fresh paraxial
seed), ``"marching"`` (field-marching fallback), ``"scan"``
(transverse scan fallback) or ``"failed"``.
used_cached_map: Whether a fresh cached pupil map was reused
without recalibration.
edge_probe_fallbacks: Number of cardinal edge probes that failed to
converge and fell back to the chief launch in the affine fit.
calibration_used: Whether a fresh chief-plus-probes calibration ran
for this field.
fallback_used: Whether any substitute strategy was used for this
field: a failed caller-supplied guess, chief marching or scan,
edge-probe fallbacks, or a Jacobian conditioning fallback
inside the final polish.
"""
Hx: float
Hy: float
wavelength: float
final_polish: SolveReport
chief_seed_strategy: str
used_cached_map: bool
edge_probe_fallbacks: int
calibration_used: bool
fallback_used: bool
[docs]
@dataclass(frozen=True)
class RobustSolveReport:
"""Aggregate outcome of one :meth:`RobustRayAimer.aim_rays` call.
Robust aiming may return NaN for individual vignetted/unreachable
rays; ``converged`` is therefore defined as
``num_converged == num_rays`` and the exact counts are retained.
Attributes:
field_reports: One :class:`RobustFieldReport` per processed field
group, in processing order.
num_rays: Total number of requested rays.
num_converged: Number of rays that met the solver tolerance.
converged: ``num_converged == num_rays``.
seed_residual: Worst (max) seed residual across the field polishes.
final_residual: Worst (max) final residual across the field
polishes.
final_polish_iterations: Largest Newton iteration count among the
field polishes.
fallback_used: Whether any field used a substitute strategy (see
:attr:`RobustFieldReport.fallback_used`).
"""
field_reports: tuple
num_rays: int
num_converged: int
converged: bool
seed_residual: float
final_residual: float
final_polish_iterations: int
fallback_used: bool
@dataclass(frozen=True)
class _CalibrationRecord:
"""Bookkeeping of one fresh per-field calibration."""
chief_strategy: str
chief_report: SolveReport | None
edge_probe_fallbacks: int
def _scan_candidate_offsets(
g_xi: float, g_eta: float, Hx: float, Hy: float, n: int
) -> tuple:
"""Transverse ``(xi, eta)`` offsets of the chief-scan candidates.
The offsets parameterize a line *through the fresh paraxial seed*: the
candidate at index ``n // 2`` is exactly ``(0, 0)`` (the seed itself),
and every other candidate displaces the seed within the entry frame's
transverse plane only. The sweep direction comes from the local field
coordinates ``(Hx, Hy)`` normalized -- the meridional ``eta`` axis for
a zero/degenerate field -- and the sweep half-width scales with the
seed's own transverse offset from the first-surface anchor
(``g_xi, g_eta``), which is invariant under a rigid translation of the
system.
Args:
g_xi, g_eta: Transverse components of (seed - first-surface vertex)
in the entry frame.
Hx, Hy: Normalized field coordinates of the chief solve.
n: Number of candidates (odd keeps the exact seed in the sweep).
Returns:
tuple: ``(xi_offsets, eta_offsets)`` backend arrays of length ``n``.
"""
h_norm = (Hx * Hx + Hy * Hy) ** 0.5
if h_norm < 1e-12:
dir_xi, dir_eta = 0.0, 1.0
else:
dir_xi, dir_eta = Hx / h_norm, Hy / h_norm
seed_offset = (g_xi * g_xi + g_eta * g_eta) ** 0.5
scale = max(50.0, 20.0 * seed_offset)
r = be.linspace(-scale, scale, n)
# Force the center candidate to exactly zero so it reproduces the seed
# bit-for-bit (linspace's midpoint is only zero to round-off).
center = be.arange_indices(n) == n // 2
r = be.where(center, be.zeros_like(r), r)
return dir_xi * r, dir_eta * r
def _offset_from_seed(
param: LaunchParameterization, seed: tuple, launch: tuple
) -> tuple[float, float]:
"""Free offset ``(xi, eta)`` of a launch state from a paraxial seed.
For infinite conjugates this is the transverse displacement of the
launch point; for finite ones, the tangent-plane rotation of the
direction at the seed's object point. For a chief ray it measures how
far the real chief deviates from the paraxial one (its pupil
aberration).
Args:
param: The shared launch parameterization.
seed: Paraxial seed state ``(x, y, z, L, M, N)`` (backend arrays).
launch: Launch state ``(x, y, z, L, M, N)`` as plain floats.
Returns:
tuple: ``(xi, eta)`` as plain floats.
"""
bound = param.bind(*seed)
xi, eta = bound.project(*(be.array([v]) for v in launch))
return to_float(xi), to_float(eta)
def _launch_from_seed(
param: LaunchParameterization, seed: tuple, offset: tuple[float, float]
) -> tuple:
"""Launch state at a paraxial seed displaced by a free offset.
Inverse of :func:`_offset_from_seed` for the same seed; applied to a
different field's seed it transports the offset to that field.
"""
bound = param.bind(*seed)
return bound.launch(be.array([offset[0]]), be.array([offset[1]]))
def _carry_to_seed(
param: LaunchParameterization,
seed: tuple,
launch: tuple,
offset: tuple[float, float] | None,
) -> tuple:
"""Warm-start a fresh paraxial seed from another field's converged chief.
Only the free DOF carries over; the fixed DOF comes from ``seed``. What
is carried depends on what the free DOF is:
- infinite conjugate: the launch point, projected onto the seed's
transverse plane. The chief then crosses the launch plane in front of
the system where the previous chief did. For wide-angle designs,
whose real entrance pupil moves forward with field, that is a closer
seed at the extreme fields than carrying the offset from the seed.
- finite conjugate: ``offset``, the previous chief's rotation away from
*its own* paraxial seed direction. The absolute free DOF here is the
launch direction itself, and launching the previous field's direction
from this field's object point misses the system at wide fields.
Args:
param: The shared launch parameterization.
seed: This field's paraxial chief seed (backend arrays).
launch: The converged chief launch state, as plain floats.
offset: That chief's offset from its own seed (see
:func:`_offset_from_seed`), or ``None`` when unknown, in which
case the launch is projected as for an infinite conjugate.
Returns:
tuple: Launch state ``(x, y, z, L, M, N)`` (backend arrays).
"""
if param.is_infinite or offset is None:
offset = _offset_from_seed(param, seed, launch)
return _launch_from_seed(param, seed, offset)
@contextlib.contextmanager
def _cached_paraxial_constants(optic: Optic):
"""Temporarily memoize ``Paraxial.EPD``/``EPL`` on this optic.
Both are system-wide constants (independent of field/pupil), but each
call re-traces the whole system. The chief-ray field-marching fallback
(:meth:`RobustRayAimer._march_chief`) calls ``ParaxialRayAimer.aim_rays``
-- which calls these -- once per marching attempt, and a cold extreme
field can need dozens of attempts; without caching, that cost dominates
total aiming time. Scoped and reversible: the original bound methods are
restored on exit, so this never leaks stale values past one calibration.
"""
para = optic.paraxial
orig_epd = para.EPD
orig_epl = para.EPL
cache: dict[str, Any] = {}
def cached_epd():
if "epd" not in cache:
cache["epd"] = orig_epd()
return cache["epd"]
def cached_epl(path=None):
# ``path`` only avoids a rebuild inside the wrapped call; the cached
# value is the same either way (geometry is fixed within this scope).
if "epl" not in cache:
cache["epl"] = orig_epl(path=path)
return cache["epl"]
para.EPD = cached_epd
para.EPL = cached_epl
try:
yield
finally:
para.EPD = orig_epd
para.EPL = orig_epl
@contextlib.contextmanager
def _relaxed_tolerance(iterative, tol: float):
"""Temporarily loosen ``iterative.tol`` for cheap intermediate solves.
Used only by the chief-ray marching fallback: an intermediate marching
step just needs to be "good enough" to warm-start the next one, since
the final per-ray polish (elsewhere, always at full tolerance) is what
actually guarantees exactness (D1). Tighter-than-needed intermediate
tolerance costs several extra Newton iterations per step for no
accuracy benefit that survives to the final result.
"""
orig_tol = iterative.tol
iterative.tol = max(tol, orig_tol)
try:
yield
finally:
iterative.tol = orig_tol
[docs]
@register_aimer("robust")
class RobustRayAimer(BaseRayAimer):
"""Chief-ray calibrated robust ray aiming algorithm.
Designed to handle challenging optical systems (wide-angle, fisheye)
where a cold paraxial seed for the iterative solver is too far from the
real solution to converge directly. Per field, a chief-ray calibration
(§4.2 of the spec) produces a cheap affine seed model; every requested
ray is then polished to exactness (§4.3) via the reused
``IterativeRayAimer`` Newton/Broyden core. Individual ray failures
(vignetting, TIR) are reported as NaN rather than aborting the batch.
Attributes:
optic (Optic): The optical system instance.
max_iter (int): Maximum number of iterations for the internal solver.
tol (float): Numerical tolerance for convergence.
scale_fields (bool): Retained for constructor-signature stability;
the calibration-based algorithm does not use homotopy
field-scaling, so this is a no-op.
"""
def __init__(
self,
optic: Optic,
max_iter: int = 20,
tol: float = 1e-8,
scale_fields: bool = True,
**kwargs: Any,
) -> None:
"""Initialize the RobustRayAimer.
Args:
optic (Optic): The optical system to aim rays for.
max_iter (int, optional): Maximum number of iterations. Defaults to 20.
tol (float, optional): Error tolerance for convergence. Defaults to 1e-8.
scale_fields (bool, optional): No-op, retained for backward
compatibility. Defaults to True.
**kwargs: Additional keyword arguments passed to BaseRayAimer.
"""
super().__init__(optic, **kwargs)
self.scale_fields = scale_fields
self.max_iter = max_iter
self.tol = tol
self._iterative = IterativeRayAimer(optic, max_iter=max_iter, tol=tol)
self._paraxial = self._iterative._paraxial_aimer
self._cache = PupilMapCache()
#: Aggregate report of the most recent :meth:`aim_rays` call. Set
#: before raising on a total field failure, so the failure can be
#: inspected.
self.last_report: RobustSolveReport | None = None
# Report of the last chief-solve attempt, kept for the failure path.
self._last_chief_failure_report: SolveReport | None = None
[docs]
def aim_rays(
self,
fields: tuple,
wavelengths: Any,
pupil_coords: tuple,
initial_guess: tuple | None = None,
) -> tuple:
"""Calculate ray starting coordinates using chief-ray calibration.
Args:
fields (tuple): Field coordinates ``(Hx, Hy)``.
wavelengths (Any): Wavelengths in microns.
pupil_coords (tuple): Normalized pupil coordinates ``(Px, Py)``.
initial_guess (tuple | None, optional): Optional starting guess.
If provided, the method first attempts to solve directly
using the iterative solver with this guess; only on failure
does it fall back to the full calibrated solve below.
Returns:
tuple: Solved ray parameters ``(x, y, z, L, M, N)``.
Raises:
ValueError: If every ray for a field fails to converge (a
misconfiguration, not ordinary partial vignetting).
"""
# Scalar paraxial values are only seeds here -- every returned ray
# is Newton-polished against real traces -- so out-of-domain
# geometries warn instead of raising inside this scope.
with paraxial_seed_scope():
return self._aim_rays_scoped(
fields, wavelengths, pupil_coords, initial_guess
)
def _aim_rays_scoped(
self,
fields: tuple,
wavelengths: Any,
pupil_coords: tuple,
initial_guess: tuple | None = None,
) -> tuple:
"""Body of :meth:`aim_rays`, run inside the paraxial seed scope."""
guess_failed = False
if initial_guess is not None:
try:
result = self._iterative.aim_rays(
fields, wavelengths, pupil_coords, initial_guess=initial_guess
)
except ValueError:
# Fall through to the calibrated solve below; using
# calibration after a failed requested guess is a fallback
# and is reported as such.
guess_failed = True
else:
self.last_report = self._report_for_initial_guess(fields, wavelengths)
return result
Px, Py = pupil_coords
Px = be.as_array_1d(Px)
Py = be.as_array_1d(Py)
n = len(Px)
Hx, Hy = fields
Hx = be.as_array_1d(Hx)
Hy = be.as_array_1d(Hy)
if len(Hx) == 1 and n > 1:
Hx = Hx * be.ones(n)
if len(Hy) == 1 and n > 1:
Hy = Hy * be.ones(n)
if hasattr(wavelengths, "__len__"):
wl_arr = be.as_array_1d(wavelengths)
if len(wl_arr) == 1 and n > 1:
wl_arr = wl_arr * be.ones(n)
else:
wl_arr = be.ones(n) * float(wavelengths)
Hx_list = be.to_numpy(Hx).reshape(-1).tolist()
Hy_list = be.to_numpy(Hy).reshape(-1).tolist()
wl_list = be.to_numpy(wl_arr).reshape(-1).tolist()
stop_idx = self.optic.surfaces.stop_index
is_inf = getattr(self.optic.object_surface, "is_infinite", False)
r_stop = get_stop_radius_strategy(self.optic, "robust").calculate_stop_radius()
# One launch parameterization per aiming call: all solves below share
# the same entry frame, so the local transverse basis is built once
# and passed through rather than re-derived per solve.
param = LaunchParameterization.for_optic(self.optic, bool(is_inf))
self._cache.sync(self.optic)
# Group rays by field (D3: reuse the same pupil map across pupil
# distributions for the same field), then process fields ordered
# by radial magnitude so later (larger) fields can warm-start from
# already-solved smaller ones (D8 field marching).
groups: dict[tuple[float, float, float], list[int]] = {}
for i in range(n):
key = (Hx_list[i], Hy_list[i], wl_list[i])
groups.setdefault(key, []).append(i)
ordered_keys = sorted(groups, key=lambda k: k[0] ** 2 + k[1] ** 2)
order_parts: list[list[int]] = []
x_parts: list[Any] = []
y_parts: list[Any] = []
z_parts: list[Any] = []
L_parts: list[Any] = []
M_parts: list[Any] = []
N_parts: list[Any] = []
field_reports: list[RobustFieldReport] = []
num_converged_total = 0
def _finalize_report() -> None:
"""Assemble the aggregate report from the field reports so far."""
polishes = [fr.final_polish for fr in field_reports]
self.last_report = RobustSolveReport(
field_reports=tuple(field_reports),
num_rays=n,
num_converged=num_converged_total,
converged=num_converged_total == n,
seed_residual=max(
(p.seed_residual for p in polishes), default=float("inf")
),
final_residual=max(
(p.final_residual for p in polishes), default=float("inf")
),
final_polish_iterations=max(
(p.iterations for p in polishes), default=0
),
fallback_used=guess_failed
or any(fr.fallback_used for fr in field_reports),
)
for key in ordered_keys:
idx = groups[key]
Hxk, Hyk, wlk = key
used_cached_map = False
calibration: _CalibrationRecord | None = None
pmap = self._cache.get_fresh(Hxk, Hyk, wlk)
if pmap is not None:
used_cached_map = True
else:
seed_map = self._cache.get_stale(Hxk, Hyk, wlk)
if seed_map is None:
seed_map = self._cache.nearest(Hxk, Hyk)
try:
with _cached_paraxial_constants(self.optic):
pmap, calibration = self._calibrate_field(
Hxk,
Hyk,
wlk,
stop_idx,
is_inf,
r_stop,
seed_map,
param,
)
except ValueError:
# Total chief failure for this field: record what was
# attempted, publish the partial report, then re-raise
# so the failure can be inspected via last_report.
failed_polish = self._last_chief_failure_report
if failed_polish is None:
failed_polish = SolveReport(
seed_residual=float("inf"),
final_residual=float("inf"),
converged=False,
iterations=0,
num_rays=len(idx),
num_converged=0,
)
field_reports.append(
RobustFieldReport(
Hx=Hxk,
Hy=Hyk,
wavelength=wlk,
final_polish=failed_polish,
chief_seed_strategy="failed",
used_cached_map=False,
edge_probe_fallbacks=0,
calibration_used=True,
fallback_used=True,
)
)
_finalize_report()
raise
self._cache.put(Hxk, Hyk, wlk, pmap)
Px_g = Px[idx]
Py_g = Py[idx]
wl_g = wl_arr[idx]
x0, y0, z0, L0, M0, N0 = pmap.seed(Px_g, Py_g)
tx = Px_g * r_stop
ty = Py_g * r_stop
x, y, z, L, M, N, converged, _, polish = self._iterative._solve_core(
x0, y0, z0, L0, M0, N0, wl_g, stop_idx, is_inf, tx, ty, param=param
)
group_converged = int(be.to_numpy(converged).reshape(-1).sum())
num_converged_total += group_converged
if calibration is None:
strategy = "cached_map"
probe_fallbacks = 0
else:
strategy = calibration.chief_strategy
probe_fallbacks = calibration.edge_probe_fallbacks
field_reports.append(
RobustFieldReport(
Hx=Hxk,
Hy=Hyk,
wavelength=wlk,
final_polish=polish,
chief_seed_strategy=strategy,
used_cached_map=used_cached_map,
edge_probe_fallbacks=probe_fallbacks,
calibration_used=calibration is not None,
fallback_used=guess_failed
or strategy in ("marching", "scan")
or probe_fallbacks > 0
or polish.fallback_used,
)
)
if not be.any(converged):
# Publish the report -- including this field's failed
# polish -- before raising, so the failure is inspectable.
_finalize_report()
raise ValueError(
"RobustRayAimer: every ray failed to converge for field "
f"(Hx={Hxk}, Hy={Hyk}); check the system configuration."
)
# Renormalize direction cosines (G3 invariant).
norm = be.sqrt(L**2 + M**2 + N**2)
L = L / norm
M = M / norm
N = N / norm
# Graceful per-ray failure (D6): non-converged rays -> NaN.
x = be.where(converged, x, be.nan)
y = be.where(converged, y, be.nan)
z = be.where(converged, z, be.nan)
L = be.where(converged, L, be.nan)
M = be.where(converged, M, be.nan)
N = be.where(converged, N, be.nan)
order_parts.append(idx)
x_parts.append(x)
y_parts.append(y)
z_parts.append(z)
L_parts.append(L)
M_parts.append(M)
N_parts.append(N)
order = [i for part in order_parts for i in part]
inv_perm = [0] * n
for pos, orig in enumerate(order):
inv_perm[orig] = pos
x_out = be.concatenate(x_parts)[inv_perm]
y_out = be.concatenate(y_parts)[inv_perm]
z_out = be.concatenate(z_parts)[inv_perm]
L_out = be.concatenate(L_parts)[inv_perm]
M_out = be.concatenate(M_parts)[inv_perm]
N_out = be.concatenate(N_parts)[inv_perm]
_finalize_report()
return x_out, y_out, z_out, L_out, M_out, N_out
def _report_for_initial_guess(
self, fields: tuple, wavelengths: Any
) -> RobustSolveReport:
"""Aggregate report for a call solved directly from a caller guess."""
polish = self._iterative.last_report
if polish is None: # pragma: no cover - aim_rays always sets it
polish = SolveReport(
seed_residual=float("nan"),
final_residual=float("nan"),
converged=True,
iterations=0,
num_rays=0,
num_converged=0,
)
Hx, Hy = fields
hx0 = float(be.to_numpy(be.as_array_1d(Hx)).reshape(-1)[0])
hy0 = float(be.to_numpy(be.as_array_1d(Hy)).reshape(-1)[0])
if hasattr(wavelengths, "__len__"):
wl0 = float(be.to_numpy(be.as_array_1d(wavelengths)).reshape(-1)[0])
else:
wl0 = float(wavelengths)
field_report = RobustFieldReport(
Hx=hx0,
Hy=hy0,
wavelength=wl0,
final_polish=polish,
chief_seed_strategy="initial_guess",
used_cached_map=False,
edge_probe_fallbacks=0,
calibration_used=False,
fallback_used=polish.fallback_used,
)
return RobustSolveReport(
field_reports=(field_report,),
num_rays=polish.num_rays,
num_converged=polish.num_converged,
converged=polish.num_converged == polish.num_rays,
seed_residual=polish.seed_residual,
final_residual=polish.final_residual,
final_polish_iterations=polish.iterations,
fallback_used=polish.fallback_used,
)
def _calibrate_field(
self,
Hx: float,
Hy: float,
wl: float,
stop_idx: int,
is_inf: bool,
r_stop: float,
seed_map: PupilMap | None,
param: LaunchParameterization | None = None,
) -> tuple[PupilMap, _CalibrationRecord]:
"""Chief solve + 4 edge probes -> affine :class:`PupilMap` (§4.2).
Returns:
tuple: The fitted map and the :class:`_CalibrationRecord`
describing how the chief anchor was obtained and how many edge
probes fell back to the chief launch.
"""
if param is None:
param = LaunchParameterization.for_optic(self.optic, bool(is_inf))
seed = self._chief_seed(Hx, Hy, be.array([wl]))
chief, strategy, chief_report = self._solve_chief(
Hx, Hy, wl, stop_idx, is_inf, seed_map, param, seed=seed
)
solved = self._solve_probes(wl, stop_idx, is_inf, r_stop, chief, param)
probes = [probe for probe, _ in solved]
probe_fallbacks = sum(int(fell_back) for _, fell_back in solved)
record = _CalibrationRecord(
chief_strategy=strategy,
chief_report=chief_report,
edge_probe_fallbacks=probe_fallbacks,
)
chief_offset = None
if not param.is_infinite:
chief_offset = _offset_from_seed(param, seed, chief)
return self._fit_affine(chief, probes, param, chief_offset), record
def _chief_seed(self, Hx: float, Hy: float, wl_a: Any) -> tuple:
"""Fresh paraxial chief-ray seed (pupil center) for field (Hx, Hy)."""
return self._paraxial.aim_rays(
(be.array([Hx]), be.array([Hy])),
wl_a,
(be.array([0.0]), be.array([0.0])),
)
def _solve_chief(
self,
Hx: float,
Hy: float,
wl: float,
stop_idx: int,
is_inf: bool,
seed_map: PupilMap | None,
param: LaunchParameterization | None = None,
seed: tuple | None = None,
) -> tuple[tuple, str, SolveReport | None]:
"""Solve the chief ray (stop target (0, 0)) for this field.
Seed order: warm-started map for this field or the nearest
already-solved field, then a direct paraxial guess. If both fail --
the paraxial seed can be too far from the real solution at extreme
field angles to converge in one Newton solve, or may not trace at
all -- fall back to marching the chief ray outward in field angle
from the axis (:meth:`_march_chief`), which is what makes a *cold*
extreme-field solve (e.g. WideAngle170FOV) converge without
recursive subdivision, and finally to a scan of the free launch DOF
through the paraxial seed (:meth:`_scan_chief`).
The fixed launch components (direction for infinite conjugates,
object position for finite ones) always come fresh from *this*
field's paraxial trace, never from ``seed_map`` -- they encode the
field angle itself, so reusing another field's fixed components
would silently solve the wrong (e.g. on-axis) problem even though
Newton still converges. Only the free transverse 2-DOF is
warm-started (see :func:`_carry_to_seed`): the stored chief's launch
point for an infinite conjugate, and for a finite one its rotation
away from its own paraxial seed (``PupilMap.chief_offset``), since
there the free DOF is the launch direction and its absolute value
would carry the other field along.
``seed`` is this field's paraxial chief seed when the caller already
has it (see :meth:`_chief_seed`); it is computed here otherwise.
Returns:
tuple: ``(launch, strategy, report)`` -- the solved chief
launch state, which strategy produced it (``"warm_map"``,
``"direct_paraxial"``, ``"marching"`` or ``"scan"``), and the
:class:`SolveReport` of the producing solve (``None`` when the
producing step exposes no single report).
Raises:
ValueError: If every strategy fails. The report of the last
attempted solve is stored on
``self._last_chief_failure_report`` before raising.
"""
if param is None:
param = LaunchParameterization.for_optic(self.optic, bool(is_inf))
wl_a = be.array([wl])
tx = be.array([0.0])
ty = be.array([0.0])
self._last_chief_failure_report = None
if seed is None:
seed = self._chief_seed(Hx, Hy, wl_a)
px0, py0, pz0, pL0, pM0, pN0 = seed
if seed_map is not None:
x0, y0, z0, L0, M0, N0 = _carry_to_seed(
param, seed, seed_map.base, seed_map.chief_offset
)
x, y, z, L, M, N, converged, _, report = self._iterative._solve_core(
x0, y0, z0, L0, M0, N0, wl_a, stop_idx, is_inf, tx, ty, param=param
)
self._last_chief_failure_report = report
if be.any(converged):
launch = (
to_float(x),
to_float(y),
to_float(z),
to_float(L),
to_float(M),
to_float(N),
)
return launch, "warm_map", report
x, y, z, L, M, N, converged, _, report = self._iterative._solve_core(
px0, py0, pz0, pL0, pM0, pN0, wl_a, stop_idx, is_inf, tx, ty, param=param
)
self._last_chief_failure_report = report
if be.any(converged):
launch = (
to_float(x),
to_float(y),
to_float(z),
to_float(L),
to_float(M),
to_float(N),
)
return launch, "direct_paraxial", report
marched = self._march_chief(Hx, Hy, wl_a, stop_idx, is_inf, tx, ty, param)
if marched is not None:
launch, report = marched
if report is not None:
self._last_chief_failure_report = report
return launch, "marching", report
scanned = self._scan_chief(
px0, py0, pz0, pL0, pM0, pN0, wl_a, stop_idx, tx, ty, param, Hx, Hy
)
if scanned is not None:
launch, report = scanned
self._last_chief_failure_report = report
return launch, "scan", report
raise ValueError(
f"RobustRayAimer: chief ray failed to converge for field "
f"(Hx={Hx}, Hy={Hy}) after marching from the axis and scanning "
f"around the paraxial seed; check the system configuration."
)
def _scan_chief(
self,
px0: Any,
py0: Any,
pz0: Any,
pL0: Any,
pM0: Any,
pN0: Any,
wl_a: Any,
stop_idx: int,
tx: Any,
ty: Any,
param: LaunchParameterization,
Hx: float = 0.0,
Hy: float = 0.0,
n: int = 2001,
) -> tuple[tuple, SolveReport] | None:
"""Last-resort chief-ray seed search for when neither the paraxial
guess nor field marching converges (extreme field angles, or a
paraxial seed that does not trace at all).
Sweeps candidate rays along a transverse line *through the fresh
paraxial seed* ``(px0 ... pN0)`` and returns the converged candidate
nearest the seed. For an infinite conjugate the candidates are
parallel rays whose launch points are displaced transverse to the
entry direction. For a finite conjugate they fan out from the fixed
object point: each candidate crosses the transverse plane through
the first-surface vertex at a point displaced along the same line,
so both conjugates sweep the same length window across the front of
the system. The ``r = 0`` candidate reproduces the seed exactly.
The sweep direction comes from the field coordinates ``(Hx, Hy)``
(the meridional axis for a degenerate field) and the sweep width
from the seed's transverse offset relative to the first-surface
anchor -- both invariant under a rigid translation of the system
(see :func:`_scan_candidate_offsets`).
"""
# Seed offset relative to the first physical surface's vertex,
# expressed on the entry frame's transverse basis. Anchoring at the
# vertex (not the global origin) keeps the sweep width invariant
# under rigid translations.
path = self.optic.surfaces.build_paraxial_path()
anchor = tuple(to_float(c) for c in path.vertices_gcs[1])
u, v = param.u, param.v
origin = (to_float(px0), to_float(py0), to_float(pz0))
seed_dir = (to_float(pL0), to_float(pM0), to_float(pN0))
if param.is_infinite:
crossing = origin
else:
# Where the seed ray crosses the transverse plane through the
# anchor. A seed that never reaches that plane leaves nothing
# to sweep.
d0 = tuple(to_float(c) for c in path.entry_direction)
along = sum(seed_dir[i] * d0[i] for i in range(3))
reach = sum((anchor[i] - origin[i]) * d0[i] for i in range(3))
if along <= 1e-12 or reach <= 0.0:
return None
s = reach / along
crossing = tuple(origin[i] + s * seed_dir[i] for i in range(3))
g_rel = tuple(crossing[i] - anchor[i] for i in range(3))
g_xi = g_rel[0] * u[0] + g_rel[1] * u[1] + g_rel[2] * u[2]
g_eta = g_rel[0] * v[0] + g_rel[1] * v[1] + g_rel[2] * v[2]
xi_off, eta_off = _scan_candidate_offsets(g_xi, g_eta, Hx, Hy, n)
ones = be.ones(n)
sx = ones * origin[0]
sy = ones * origin[1]
sz = ones * origin[2]
L0 = ones * seed_dir[0]
M0 = ones * seed_dir[1]
N0 = ones * seed_dir[2]
wl_b = ones * to_float(wl_a)
tx_b = ones * to_float(tx)
ty_b = ones * to_float(ty)
if param.is_infinite:
# Candidates sit on the transverse line through the seed: binding
# the parameterization at repeated copies of the seed makes
# launch(0, 0) the seed itself and keeps the xi/eta displacements
# in the transverse plane for any entry direction.
bound = param.bind(sx, sy, sz, L0, M0, N0)
x0, y0, z0, L0, M0, N0 = bound.launch(xi_off, eta_off)
else:
# Aim from the object point through each displaced crossing;
# the center candidate keeps the seed direction bit-for-bit.
qx = crossing[0] + xi_off * u[0] + eta_off * v[0] - origin[0]
qy = crossing[1] + xi_off * u[1] + eta_off * v[1] - origin[1]
qz = crossing[2] + xi_off * u[2] + eta_off * v[2] - origin[2]
norm = be.sqrt(qx**2 + qy**2 + qz**2)
center = be.arange_indices(n) == n // 2
x0, y0, z0 = sx, sy, sz
L0 = be.where(center, L0, qx / norm)
M0 = be.where(center, M0, qy / norm)
N0 = be.where(center, N0, qz / norm)
# A candidate whose seed ray does not reach the stop can never
# converge -- the Newton core holds a seed that does not trace where
# it is -- but it would still be traced on every iteration. At wide
# fields most of the sweep is such rays, so one screening trace
# drops them before the polish; which candidates converge does not
# change.
rays = self._iterative._trace_subset(
x0, y0, z0, L0, M0, N0, wl_b, stop_idx, param.is_infinite
)
lx, ly = self._iterative._get_local_stop_coords(rays, stop_idx)
live = be.logical_and(be.isfinite(lx), be.isfinite(ly))
if not be.any(live):
return None
keep = be.arange_indices(n)[live]
x, y, z, L, M, N, converged, _, report = self._iterative._solve_core(
x0[keep],
y0[keep],
z0[keep],
L0[keep],
M0[keep],
N0[keep],
wl_b[keep],
stop_idx,
param.is_infinite,
tx_b[keep],
ty_b[keep],
param=param,
)
if not be.any(converged):
return None
# Among the converged candidates, take the one whose starting
# offset from the fresh seed is smallest: a wide sweep can also
# converge onto physically extreme solution branches far from the
# seed, and an extreme "chief" would poison the probe/affine
# calibration built on top of it. Nearest-to-seed stays on the
# seed's own branch.
conv_np = be.to_numpy(converged).reshape(-1)
keep_np = be.to_numpy(keep).reshape(-1)
offsets_np = (
be.to_numpy(xi_off).reshape(-1) ** 2 + be.to_numpy(eta_off).reshape(-1) ** 2
)[keep_np]
candidates = conv_np.nonzero()[0]
idx = int(candidates[offsets_np[candidates].argmin()])
launch = (
to_float(x[idx : idx + 1]),
to_float(y[idx : idx + 1]),
to_float(z[idx : idx + 1]),
to_float(L[idx : idx + 1]),
to_float(M[idx : idx + 1]),
to_float(N[idx : idx + 1]),
)
return launch, report
def _march_chief(
self,
Hx: float,
Hy: float,
wl_a: Any,
stop_idx: int,
is_inf: bool,
tx: Any,
ty: Any,
param: LaunchParameterization,
max_attempts: int = 150,
min_dt: float = 1e-4,
) -> tuple[tuple, SolveReport | None] | None:
"""March the chief ray from the axis out to (Hx, Hy) in field angle.
A step-halving walk -- each step a single-ray 2-DOF solve
warm-started from the *last successfully converged* launch, never
from a failed one -- replaces the old recursive homotopy as the
cold-start robustness mechanism (D8). It is bounded (a fixed attempt
budget, no recursion) and physically monotonic: only the free launch
DOF carries over between steps (see :func:`_carry_to_seed`), while
the fixed DOF (z, and direction for infinite conjugates / object
position for finite ones) is refreshed from the paraxial trace at
each step's actual field angle. A failed step is retried once from
the fresh paraxial seed before its stride is halved.
For a finite conjugate the walk starts at half the field: the axis
chief's offset from its seed is zero, so a full stride would repeat
exactly the direct paraxial solve the caller has already tried.
A step size is never grown back up after a success: this system's
maximum reliable step tends to shrink (never grow) as the field
angle increases, so re-attempting a larger step every time just
wastes evaluations that repeatedly fail the same way.
Returns ``None`` if the walk cannot reach ``t=1`` (the actual
target field) within the attempt budget -- the caller must treat
this as a hard failure, not silently accept whatever intermediate
field angle happened to converge. Returning a wrong-but-converged
intermediate result here is exactly the failure mode this method
exists to prevent (see SPEC_ray_aiming_20260703.md D8): the fixed
launch DOF encodes the field angle itself, so a caller that used a
partial march's result as the final chief ray would be aiming at
the wrong field entirely, not just aiming imprecisely.
On success, returns ``(launch, report)`` where ``report`` is the
:class:`SolveReport` of the final full-tolerance solve (``None``
when the attempt budget ran out exactly at ``t = 1`` before that
solve could run).
"""
t = 0.0
# t=0 (the axis) is trivial and always converges: L=M=0, N=+-1. The
# axis chief is its own paraxial seed, so its offset is zero.
launch = tuple(to_float(v) for v in self._chief_seed(0.0, 0.0, wl_a))
offset = (0.0, 0.0)
dt = 1.0 if param.is_infinite else 0.5
relaxed_tol = max(self._iterative.tol, 1e-4)
for _attempt in range(max_attempts):
if t >= 1.0:
# Intermediate steps used a relaxed tolerance as a cheap
# warm-start; do one final full-tolerance solve so the
# chief anchor itself is exact, not just "close enough".
x0 = be.array([launch[0]])
y0 = be.array([launch[1]])
z0 = be.array([launch[2]])
L0 = be.array([launch[3]])
M0 = be.array([launch[4]])
N0 = be.array([launch[5]])
x, y, z, L, M, N, converged, _, report = self._iterative._solve_core(
x0,
y0,
z0,
L0,
M0,
N0,
wl_a,
stop_idx,
is_inf,
tx,
ty,
param=param,
)
if be.any(converged):
return (
to_float(x),
to_float(y),
to_float(z),
to_float(L),
to_float(M),
to_float(N),
), report
return launch, report
t_next = min(t + dt, 1.0)
seed = self._chief_seed(Hx * t_next, Hy * t_next, wl_a)
px0, py0, pz0, pL0, pM0, pN0 = seed
# Carry only the free DOF of the last converged launch onto this
# step's fresh paraxial seed (the fixed DOF -- direction for
# infinite conjugates, object point for finite ones -- encodes
# the field angle and must stay fresh).
x0, y0, z0, L0, M0, N0 = _carry_to_seed(param, seed, launch, offset)
with _relaxed_tolerance(self._iterative, relaxed_tol):
x, y, z, L, M, N, converged, _, _r = self._iterative._solve_core(
x0,
y0,
z0,
L0,
M0,
N0,
wl_a,
stop_idx,
is_inf,
tx,
ty,
param=param,
)
if not be.any(converged) and (param.is_infinite or offset != (0.0, 0.0)):
# The launch warm-started from the previous step can
# occasionally be a worse seed than a fresh paraxial guess at
# the new angle (e.g. right where marching first takes a
# large stride); retry once from the fresh guess before
# giving up and shrinking the step. A finite-conjugate carry
# with a zero offset already is the fresh guess.
with _relaxed_tolerance(self._iterative, relaxed_tol):
x, y, z, L, M, N, converged, _, _r = self._iterative._solve_core(
px0,
py0,
pz0,
pL0,
pM0,
pN0,
wl_a,
stop_idx,
is_inf,
tx,
ty,
param=param,
)
if be.any(converged):
launch = (
to_float(x),
to_float(y),
to_float(z),
to_float(L),
to_float(M),
to_float(N),
)
if not param.is_infinite:
offset = _offset_from_seed(param, seed, launch)
t = t_next
# Do not grow dt back up -- see docstring.
else:
# Retry the SAME target angle at half the step, warm-started
# from the last known-good launch -- never advance t on a
# failed step.
dt /= 2.0
if dt < min_dt:
return None
return (launch, None) if t >= 1.0 else None
def _solve_probes(
self,
wl: float,
stop_idx: int,
is_inf: bool,
r_stop: float,
chief: tuple[float, float, float, float, float, float],
param: LaunchParameterization,
) -> list[tuple[tuple, bool]]:
"""Solve the four cardinal edge probes in one batch.
Every probe starts from the chief launch and the solver core treats
each ray independently (its own Jacobian, line search and
convergence), so this gives the same probes as :meth:`_solve_probe`
called once per probe, with a quarter of the traces.
Returns:
list: One ``(launch, fell_back)`` pair per entry of
``_EDGE_PROBES``, as returned by :meth:`_solve_probe`.
"""
n = len(_EDGE_PROBES)
x0, y0, z0, L0, M0, N0 = (be.array([v] * n) for v in chief)
wl_a = be.array([wl] * n)
tx = be.array([px * r_stop for px, _ in _EDGE_PROBES])
ty = be.array([py * r_stop for _, py in _EDGE_PROBES])
x, y, z, L, M, N, converged, _, _r = self._iterative._solve_core(
x0, y0, z0, L0, M0, N0, wl_a, stop_idx, is_inf, tx, ty, param=param
)
converged = be.to_numpy(converged).reshape(-1)
state = [be.to_numpy(v).reshape(-1) for v in (x, y, z, L, M, N)]
solved = []
for i in range(n):
if not converged[i]:
# See _solve_probe: an unreachable probe falls back to the
# chief launch.
solved.append((chief, True))
else:
solved.append((tuple(float(v[i]) for v in state), False))
return solved
def _solve_probe(
self,
wl: float,
stop_idx: int,
is_inf: bool,
Px_e: float,
Py_e: float,
r_stop: float,
chief: tuple[float, float, float, float, float, float],
param: LaunchParameterization,
) -> tuple[tuple, bool]:
"""Solve one cardinal edge probe, seeded from the chief launch.
Returns:
tuple: ``(launch, fell_back)`` -- the probe launch state and
whether it fell back to the chief launch because the probe
failed to converge.
"""
x0, y0, z0, L0, M0, N0 = (be.array([v]) for v in chief)
wl_a = be.array([wl])
tx = be.array([Px_e * r_stop])
ty = be.array([Py_e * r_stop])
x, y, z, L, M, N, converged, _, _r = self._iterative._solve_core(
x0, y0, z0, L0, M0, N0, wl_a, stop_idx, is_inf, tx, ty, param=param
)
if not be.any(converged):
# An unreachable edge probe still yields a usable (if less
# accurate) linear seed by falling back to the chief launch,
# rather than propagating NaN into the affine fit.
return chief, True
return (
to_float(x),
to_float(y),
to_float(z),
to_float(L),
to_float(M),
to_float(N),
), False
def _fit_affine(
self,
chief: tuple[float, float, float, float, float, float],
probes: list[tuple[float, float, float, float, float, float]],
param: LaunchParameterization,
chief_offset: tuple[float, float] | None = None,
) -> PupilMap:
"""Fit the 2x2 affine launch model from the chief ray + 4 probes.
The probes' launch states are projected into the chief-bound local
transverse parameterization, so the fitted offsets are (xi, eta)
coordinates -- valid for any entry direction, and stored as plain
floats (detached by design). ``chief_offset`` (a finite-conjugate
chief's offset from its own paraxial seed) is stored on the map for
warm starts.
"""
p_east, p_west, p_north, p_south = probes
bound = param.bind(*(be.array([v]) for v in chief))
def free(state: tuple) -> tuple[float, float]:
xi, eta = bound.project(*(be.array([v]) for v in state))
return to_float(xi), to_float(eta)
e1, e2 = free(p_east)
w1, w2 = free(p_west)
n1, n2 = free(p_north)
s1, s2 = free(p_south)
A = (
((e1 - w1) / 2.0, (n1 - s1) / 2.0),
((e2 - w2) / 2.0, (n2 - s2) / 2.0),
)
return PupilMap(base=chief, A=A, param=param, chief_offset=chief_offset)
