diff --git a/optika/systems/_sequential.py b/optika/systems/_sequential.py index a0dbbbe4..269f6d41 100644 --- a/optika/systems/_sequential.py +++ b/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) diff --git a/optika/systems/_sequential_test.py b/optika/systems/_sequential_test.py index 8b129765..ea5607c8 100644 --- a/optika/systems/_sequential_test.py +++ b/optika/systems/_sequential_test.py @@ -268,3 +268,250 @@ 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 + + +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 + +_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 + + +_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