diff --git a/optika/systems/_sequential.py b/optika/systems/_sequential.py index f75fa9cb..8e2eed3f 100644 --- a/optika/systems/_sequential.py +++ b/optika/systems/_sequential.py @@ -340,6 +340,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, @@ -450,7 +470,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): @@ -462,12 +482,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) @@ -509,7 +583,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) @@ -581,6 +661,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) diff --git a/optika/systems/_sequential_test.py b/optika/systems/_sequential_test.py index 98dda69d..665f8b56 100644 --- a/optika/systems/_sequential_test.py +++ b/optika/systems/_sequential_test.py @@ -455,3 +455,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