Fixed the stop root-finding problem for transmissive stop surfaces

optika/systems/_sequential.py

@@ -178,6 +178,28 @@ 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],
@@ -203,7 +225,8 @@ 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,
@@ -215,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)
 
@@ -390,6 +412,18 @@ def zfunc(xy: na.AbstractCartesian2dVectorArray):
             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
@@ -408,7 +442,8 @@ def zfunc(xy: na.AbstractCartesian2dVectorArray):
             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,
@@ -477,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,

optika/systems/_sequential_test.py

@@ -294,6 +294,25 @@ def test__anchor_surface():
     assert anchor([first, flat, last]) is last
 
 
+def test__stop_is_involutory():
+    involutory = optika.systems.SequentialSystem._stop_is_involutory
+    assert involutory(
+        optika.surfaces.Surface(material=optika.materials.Mirror()),
+    )
+    assert not involutory(
+        optika.surfaces.Surface(
+            rulings=optika.rulings.Rulings(spacing=1 * u.um, diffraction_order=1),
+        ),
+    )
+    assert not involutory(
+        optika.surfaces.Surface(material=optika.materials.MultilayerFilm()),
+    )
+    assert involutory(
+        optika.surfaces.Surface(material=optika.materials.Vacuum()),
+    )
+    assert involutory(optika.surfaces.Surface())
+
+
 # small enough that the image of the field fits on the sensor
 _radius_field_newtonian = 0.05 * u.deg
 
@@ -430,3 +449,69 @@ def test_field_max_matches_source_aperture(
         result = a.field_max
         assert np.abs(result.x - _radius_field_grazing) < 1e-6 * u.deg
         assert np.abs(result.y - _radius_field_grazing) < 1e-6 * u.deg
+
+
+_focal_length_fzp = 100 * u.mm
+_wavelength_fzp = 500 * u.nm
+
+_sensor_fzp = optika.sensors.ImagingSensor(
+    name="sensor",
+    width_pixel=15 * u.um,
+    axis_pixel=na.Cartesian2dVectorArray("detector_x", "detector_y"),
+    timedelta_exposure=1 * u.s,
+    num_pixel=na.Cartesian2dVectorArray(64, 64),
+    transformation=na.transformations.Cartesian3dTranslation(
+        z=_focal_length_fzp,
+    ),
+    is_field_stop=True,
+)
+
+_system_fzp = optika.systems.SequentialSystem(
+    object=optika.surfaces.Surface(
+        name="source",
+        aperture=optika.apertures.CircularAperture(
+            radius=np.sin(0.3 * u.deg),
+        ),
+    ),
+    surfaces=[
+        optika.surfaces.Surface(
+            name="fzp",
+            rulings=optika.rulings.Rulings(
+                spacing=optika.rulings.HolographicRulingSpacing(
+                    x1=na.Cartesian3dVectorArray(0, 0, -1) * u.AU,
+                    x2=na.Cartesian3dVectorArray(0, 0, 1) * _focal_length_fzp,
+                    wavelength=_wavelength_fzp,
+                ),
+                diffraction_order=1,
+            ),
+            aperture=optika.apertures.CircularAperture(radius=5 * u.mm),
+            is_pupil_stop=True,
+        ),
+    ],
+    sensor=_sensor_fzp,
+    grid_input=_grid_input,
+)
+
+
+@pytest.mark.parametrize(argnames="a", argvalues=[_system_fzp])
+class TestSequentialSystemFZP(
+    AbstractTestAbstractSequentialSystem,
+):
+    """
+    A telescope whose only optical element is a transmissive Fresnel zone
+    plate acting as the pupil stop. This guards against regressions in the
+    handling of non-involutory (transmissive) stop surfaces by the stop
+    root-finding problem, which previously re-applied the diffraction of the
+    stop surface and computed wildly incorrect field extents.
+    """
+
+    def test_field_max_matches_plate_scale(
+        self,
+        a: optika.systems.AbstractSequentialSystem,
+    ):
+        sensor = a.sensor
+        half_width = sensor.width_pixel * sensor.num_pixel / 2
+        field_expected = np.arctan2(half_width, _focal_length_fzp)
+        result = a.field_max
+        assert np.abs(result.x - field_expected.x) < 1e-3 * u.deg
+        assert np.abs(result.y - field_expected.y) < 1e-3 * u.deg