Fixed the stop root-finding problem for transmissive stop surfaces

optika/systems/_sequential.py

@@ -178,12 +178,55 @@ def pupil_stop(self) -> optika.surfaces.AbstractSurface:
         """
         return self.surfaces_all[self.index_pupil_stop]
 
+    @staticmethod
+    def _stop_is_involutory(
+        surface: optika.surfaces.AbstractSurface,
+    ) -> bool:
+        """
+        Whether applying the given stop surface to its own forward output
+        exactly undoes its effect on the rays. This is true for mirrors
+        (reflection is an involution) and for surfaces that don't modify the
+        rays at all (e.g. a vacuum object surface), but false for surfaces
+        with rulings or refraction, whose effect would be applied a second
+        time instead of undone.
+        """
+        material = surface.material
+        if material is not None and material.is_mirror:
+            return True
+        if surface.rulings is not None:
+            return False
+        if material is not None:
+            if not isinstance(material, optika.materials.Vacuum):
+                return False
+        return True
+
+    @staticmethod
+    def _anchor_surface(
+        subsystem: list[optika.surfaces.AbstractSurface],
+    ) -> optika.surfaces.AbstractSurface:
+        """
+        The first surface after the start of the given subsystem with optical
+        power (a mirror, a curved sag, or rulings), used to aim the initial
+        guess of the stop root-finding problem.
+        """
+        for surface in subsystem[1:]:
+            material = surface.material
+            if material is not None and material.is_mirror:
+                return surface
+            sag = surface.sag
+            if sag is not None and not isinstance(sag, optika.sags.NoSag):
+                return surface
+            if surface.rulings is not None:
+                return surface
+        return subsystem[~0]
+
     @classmethod
     def _ray_error(
         cls,
         a: na.Cartesian2dVectorArray,
         rays: optika.rays.RayVectorArray,
-        subsystem: list[optika.surfaces.AbstractSurface],
+        propagators: list[optika.surfaces.AbstractSurface],
+        transformation_last: None | na.transformations.AbstractTransformation,
         grid_last: na.Cartesian2dVectorArray,
         component_variable: str,
         component_target: str,
@@ -195,11 +238,10 @@ def _ray_error(
         rays_component_variable.z = zfunc(a)
 
         rays = optika.propagators.propagate_rays(
-            propagators=subsystem[1:],
+            propagators=propagators,
             rays=rays,
         )
 
-        transformation_last = subsystem[~0].transformation
         if transformation_last is not None:
             rays = transformation_last.inverse(rays)
 
@@ -288,7 +330,7 @@ def _calc_rayfunction_stops_only(
                 result.outputs.direction = na.Cartesian3dVectorArray(0, 0, 1)
                 component_variable = "direction"
 
-                def zfunc(xy: na.AbstractCartesian3dVectorArray):
+                def zfunc(xy: na.AbstractCartesian2dVectorArray):
                     return np.sqrt(1 - np.square(xy.length))
 
             elif na.unit(grid_first).is_equivalent(u.dimensionless_unscaled):
@@ -300,12 +342,66 @@ def zfunc(xy: na.AbstractCartesian3dVectorArray):
                 result.outputs.position = na.Cartesian3dVectorArray() * u.mm
                 component_variable = "position"
 
-                def zfunc(xy: na.AbstractCartesian3dVectorArray):
-                    return surface_first.sag(xy)
+                def zfunc(xy: na.AbstractCartesian2dVectorArray):
+                    position = na.Cartesian3dVectorArray(
+                        x=xy.x,
+                        y=xy.y,
+                        z=0 * na.unit_normalized(xy.x),
+                    )
+                    return surface_first.sag(position)
 
             else:
                 raise ValueError(f"unrecognized input grid unit, {na.unit(grid_first)}")
 
+            # Seed the free ray component by aiming each ray at its own
+            # target point on the last stop surface when no surface with
+            # optical power lies between the two stops (the seed is then
+            # nearly exact), and otherwise at the center of the first powered
+            # surface, so that the root-finding starts inside its basin of
+            # convergence even for surfaces far from the axis of the first
+            # stop (e.g. an off-axis fold or feed mirror).
+            anchor = self._anchor_surface(subsystem)
+            if anchor is surface_last and na.unit(grid_last).is_equivalent(u.mm):
+                aim = na.Cartesian3dVectorArray(
+                    x=grid_last.x,
+                    y=grid_last.y,
+                    z=0 * na.unit_normalized(grid_last.x),
+                )
+                aim.z = surface_last.sag(aim)
+                if surface_last.transformation is not None:
+                    aim = surface_last.transformation(aim)
+            else:
+                aim = na.Cartesian3dVectorArray() * u.mm
+                if anchor.transformation is not None:
+                    aim = anchor.transformation(aim)
+            if surface_first.transformation is not None:
+                aim = surface_first.transformation.inverse(aim)
+
+            if component_variable == "direction":
+                d = aim - result.outputs.position
+                d = d / d.length
+                # the sign of the seed direction is irrelevant to the
+                # root-finding problem (surface intercepts may have negative
+                # distance), but the z-component must be positive to be
+                # consistent with `zfunc`
+                flip = np.sign(d.z)
+                where = d.z != 0
+                result.outputs.direction = na.Cartesian3dVectorArray(
+                    x=np.where(where, flip * d.x, 0),
+                    y=np.where(where, flip * d.y, 0),
+                    z=np.where(where, flip * d.z, 1),
+                )
+            else:
+                d = result.outputs.direction
+                t = aim.z / d.z
+                position_seed = na.Cartesian3dVectorArray(
+                    x=aim.x - d.x * t,
+                    y=aim.y - d.y * t,
+                    z=0 * u.mm,
+                )
+                position_seed.z = surface_first.sag(position_seed)
+                result.outputs.position = position_seed
+
             if surface_first.transformation is not None:
                 result.outputs = surface_first.transformation(result.outputs)
 
@@ -316,6 +412,18 @@ def zfunc(xy: na.AbstractCartesian3dVectorArray):
             else:
                 raise ValueError(f"unrecognized output grid unit, {na.unit(grid_last)}")
 
+            # A mirror stop reverses the rays, so the solved variable can be
+            # the post-reflection ray and the stop surface itself can be
+            # excluded from the root-finding propagators. A transmissive stop
+            # that modifies the rays (e.g. a transmission grating or Fresnel
+            # zone plate) does not reverse them, so its own diffraction or
+            # refraction must be included in the solve, and the solved
+            # variable is the pre-stop ray.
+            if self._stop_is_involutory(surface_first):
+                propagators = subsystem[1:]
+            else:
+                propagators = subsystem
+
             # The residual of the root-finding problem has the same units as
             # the target grid, so the convergence tolerance must scale with
             # the size of the target aperture to be achievable in floating
@@ -334,7 +442,8 @@ def zfunc(xy: na.AbstractCartesian3dVectorArray):
             function = functools.partial(
                 self._ray_error,
                 rays=result.outputs,
-                subsystem=subsystem,
+                propagators=propagators,
+                transformation_last=surface_last.transformation,
                 grid_last=grid_last,
                 component_variable=component_variable,
                 component_target=component_target,
@@ -347,7 +456,13 @@ def zfunc(xy: na.AbstractCartesian3dVectorArray):
             # measured in physical units, yielding a Jacobian made of noise.
             # Use a perturbation proportional to the scale of the problem
             # instead.
-            dx = 1e-6
+            if component_variable == "direction":
+                dx = 1e-6
+            else:
+                dx = 1e-6 * np.maximum(
+                    scale,
+                    1 * na.unit_normalized(scale),
+                )
 
             def jacobian(x, _function=function, _dx=dx):
                 return na.jacobian(function=_function, x=x, dx=_dx)
@@ -397,6 +512,26 @@ def _calc_rayfunction_stops(
 
         subsystem = surfaces[index_stop::-1]
 
+        # This reversed trace works for a mirror stop because re-applying the
+        # stop surface reflects the rays back toward the object, and
+        # reflection is an involution (it exactly undoes itself on the second
+        # application). A transmissive stop (e.g. a transmission grating or a
+        # Fresnel zone plate) is not an involution: re-applying it would
+        # diffract/refract the rays a second time instead of undoing the
+        # first pass. (A true time-reversed trace through transmissive
+        # rulings would require negating the diffraction order, while
+        # reflective rulings are time-reversal symmetric with the same
+        # order.) So for a non-mirror stop, keep its geometry but strip its
+        # optical effect; `_calc_rayfunction_stops_only` solves for the
+        # pre-stop rays in this case.
+        surface_stop = subsystem[0]
+        if not self._stop_is_involutory(surface_stop):
+            subsystem[0] = dataclasses.replace(
+                surface_stop,
+                material=None,
+                rulings=None,
+            )
+
         rays_stop = self._calc_rayfunction_stops_only(
             wavelength_input=wavelength_input,
             axis_pupil_stop=axis_pupil_stop,
@@ -419,6 +554,17 @@ def _calc_rayfunction_stops(
         where = rays.direction @ obj.sag.normal(rays.position) > 0
         result.outputs.direction[where] = -result.outputs.direction[where]
 
+        # If the first stop is the object surface, the solved variable is the
+        # position and the direction retains only the field-stop axis, so
+        # broadcast both components against each other to guarantee that
+        # reductions over both stop axes are well-defined downstream.
+        shape = na.shape_broadcasted(
+            result.outputs.position,
+            result.outputs.direction,
+        )
+        result.outputs.position = result.outputs.position.broadcast_to(shape)
+        result.outputs.direction = result.outputs.direction.broadcast_to(shape)
+
         if self.transformation is not None:
             result.outputs = self.transformation(result.outputs)