Added support for using the object surface as the field stop of a SequentialSystem

optika/systems/_sequential.py

@@ -178,6 +178,26 @@ def pupil_stop(self) -> optika.surfaces.AbstractSurface:
         """
         return self.surfaces_all[self.index_pupil_stop]
 
+    @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,
@@ -288,7 +308,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 +320,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)
 
@@ -347,7 +421,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)
@@ -419,6 +499,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)
 

optika/systems/_sequential_test.py

@@ -268,3 +268,165 @@ def test_spot_diagram(self, a: optika.systems.AbstractSequentialSystem):
 )
 class TestSequentialSystem(AbstractTestAbstractSequentialSystem):
     pass
+
+
+def test__anchor_surface():
+    first = optika.surfaces.Surface(name="first")
+    last = optika.surfaces.Surface(name="last")
+    mirror = optika.surfaces.Surface(
+        name="mirror",
+        material=optika.materials.Mirror(),
+    )
+    curved = optika.surfaces.Surface(
+        name="curved",
+        sag=optika.sags.SphericalSag(radius=-100 * u.mm),
+    )
+    grating = optika.surfaces.Surface(
+        name="grating",
+        rulings=optika.rulings.Rulings(spacing=1 * u.um, diffraction_order=1),
+    )
+    flat = optika.surfaces.Surface(name="flat")
+
+    anchor = optika.systems.SequentialSystem._anchor_surface
+    assert anchor([first, flat, mirror, last]) is mirror
+    assert anchor([first, curved, last]) is curved
+    assert anchor([first, grating, last]) is grating
+    assert anchor([first, flat, last]) is last
+
+
+# small enough that the image of the field fits on the sensor
+_radius_field_newtonian = 0.05 * u.deg
+
+_system_newtonian = optika.systems.SequentialSystem(
+    object=optika.surfaces.Surface(
+        name="source",
+        aperture=optika.apertures.CircularAperture(
+            radius=np.sin(_radius_field_newtonian),
+        ),
+        is_field_stop=True,
+    ),
+    surfaces=[
+        optika.surfaces.Surface(
+            name="primary",
+            sag=optika.sags.SphericalSag(radius=-2000 * u.mm),
+            material=optika.materials.Mirror(),
+            aperture=optika.apertures.CircularAperture(radius=50 * u.mm),
+            transformation=na.transformations.Cartesian3dTranslation(
+                z=500 * u.mm,
+            ),
+        ),
+        optika.surfaces.Surface(
+            name="aperture",
+            aperture=optika.apertures.CircularAperture(radius=10 * u.mm),
+            transformation=na.transformations.Cartesian3dTranslation(
+                z=250 * u.mm,
+            ),
+            is_pupil_stop=True,
+        ),
+    ],
+    sensor=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(128, 128),
+        transformation=na.transformations.Cartesian3dTranslation(
+            z=-500 * u.mm,
+        ),
+    ),
+    grid_input=_grid_input,
+)
+
+
+@pytest.mark.parametrize(argnames="a", argvalues=[_system_newtonian])
+class TestSequentialSystemNewtonian(
+    AbstractTestAbstractSequentialSystem,
+):
+    """
+    A Newtonian-style telescope where the pupil stop is downstream of the
+    primary mirror, so that the initial guess of the stop root-finding
+    problem must be aimed at the center of the primary instead of directly
+    at its own target on the pupil stop.
+    """
+
+    def test_field_max_matches_source_aperture(
+        self,
+        a: optika.systems.AbstractSequentialSystem,
+    ):
+        result = a.field_max
+        assert np.abs(result.x - _radius_field_newtonian) < 1e-6 * u.deg
+        assert np.abs(result.y - _radius_field_newtonian) < 1e-6 * u.deg
+
+
+_radius_field_grazing = 0.25 * u.deg
+
+_system_grazing = optika.systems.SequentialSystem(
+    object=optika.surfaces.Surface(
+        name="source",
+        aperture=optika.apertures.CircularAperture(
+            radius=np.sin(_radius_field_grazing),
+        ),
+        is_field_stop=True,
+    ),
+    surfaces=[
+        optika.surfaces.Surface(
+            name="paraboloid",
+            sag=optika.sags.ParabolicSag(focal_length=-2000 * u.mm),
+            material=optika.materials.Mirror(),
+            aperture=optika.apertures.CircularAperture(radius=260 * u.mm),
+            transformation=na.transformations.Cartesian3dTranslation(
+                z=2500 * u.mm,
+            ),
+            is_pupil_stop=True,
+        ),
+        optika.surfaces.Surface(
+            name="grating",
+            rulings=optika.rulings.Rulings(
+                spacing=10 * u.um,
+                diffraction_order=1,
+            ),
+            aperture=optika.apertures.RectangularAperture(
+                half_width=60 * u.mm,
+            ),
+            transformation=na.transformations.Cartesian3dTranslation(
+                z=1000 * u.mm,
+            ),
+        ),
+    ],
+    sensor=optika.sensors.ImagingSensor(
+        name="sensor",
+        width_pixel=15 * u.um,
+        axis_pixel=na.Cartesian2dVectorArray("detector_x", "detector_y"),
+        # short exposure so that the Poisson lam stays representable for the
+        # large collecting area of the grazing primary
+        timedelta_exposure=1 * u.us,
+        num_pixel=na.Cartesian2dVectorArray(2048, 1024),
+        # offset by the deflection of the first diffraction order,
+        # (z_grating - z_sensor) * wavelength / spacing
+        transformation=na.transformations.Cartesian3dTranslation(
+            x=26 * u.mm,
+            z=480 * u.mm,
+        ),
+    ),
+    grid_input=_grid_input,
+)
+
+
+@pytest.mark.parametrize(argnames="a", argvalues=[_system_grazing])
+class TestSequentialSystemGrazingSpectrograph(
+    AbstractTestAbstractSequentialSystem,
+):
+    """
+    A grazing-incidence spectrograph with a transmission grating, where the
+    object surface (with an angular aperture) is the field stop. This guards
+    against regressions in the object-as-field-stop code path of the stop
+    root-finding problem.
+    """
+
+    def test_field_max_matches_source_aperture(
+        self,
+        a: optika.systems.AbstractSequentialSystem,
+    ):
+        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