diff --git a/optika/apertures/_apertures.py b/optika/apertures/_apertures.py index 71f36074..033551a2 100644 --- a/optika/apertures/_apertures.py +++ b/optika/apertures/_apertures.py @@ -135,6 +135,7 @@ def plot( transformation: None | na.transformations.AbstractTransformation = None, components: None | tuple[str, ...] = None, sag: None | optika.sags.AbstractSag = None, + unit: None | u.UnitBase = None, **kwargs, ) -> None | na.ScalarArray[npt.NDArray[None | matplotlib.lines.Line2D]]: if ax is None: @@ -149,6 +150,9 @@ def plot( if sag is not None: wire.z = sag(wire) + if unit is not None: + wire = wire.to(unit) + kwargs_plot = self.kwargs_plot if kwargs_plot is None: kwargs_plot = dict() diff --git a/optika/apertures/_apertures_test.py b/optika/apertures/_apertures_test.py index bdf01db7..dbb2d690 100644 --- a/optika/apertures/_apertures_test.py +++ b/optika/apertures/_apertures_test.py @@ -2,6 +2,7 @@ import pytest import numpy as np import matplotlib.axes +import matplotlib.pyplot as plt import astropy.units as u import named_arrays as na import optika @@ -487,3 +488,52 @@ class TestIsoscelesTrapezoidalAperture( AbstractTestAbstractIsoscelesTrapezoidalAperture, ): pass + + +def _extent_plotted( + aperture: optika.apertures.AbstractAperture, + unit: None | u.UnitBase, +) -> float: + """ + How far from the origin the aperture is drawn. + + The bare magnitude handed to matplotlib, which is what sets the scale it is + drawn at on an axes which does not understand units. + """ + fig, ax = plt.subplots() + try: + aperture.plot(ax=ax, components=("x", "y"), unit=unit) + x = np.concatenate( + [ + np.asarray(getattr(line.get_xdata(), "value", line.get_xdata())) + for line in ax.lines + ] + ) + return float(np.max(np.abs(x))) + finally: + plt.close(fig) + + +def test_plot_unit(): + """ + Asking for a unit draws every aperture to one scale. + + Matplotlib is told about units by + :func:`astropy.visualization.quantity_support`, but that reconciles them + only on a 2D axes. On a 3D axes an aperture measured in microns is drawn a + thousand times larger than the same aperture measured in millimeters, and + the instrument around it collapses to a speck. + """ + millimeters = optika.apertures.RectangularAperture(1 * u.mm) + micrometers = optika.apertures.RectangularAperture(1000 * u.um) + + # the same aperture, described in two different units + assert _extent_plotted(millimeters, None) == pytest.approx(1) + assert _extent_plotted(micrometers, None) == pytest.approx(1000) + + # which are drawn to one scale when a unit is asked for + assert _extent_plotted(millimeters, u.mm) == pytest.approx(1) + assert _extent_plotted(micrometers, u.mm) == pytest.approx(1) + + # and converted, not merely relabelled + assert _extent_plotted(millimeters, u.um) == pytest.approx(1000) diff --git a/optika/surfaces.py b/optika/surfaces.py index ae25f1d5..704f2d29 100644 --- a/optika/surfaces.py +++ b/optika/surfaces.py @@ -202,6 +202,7 @@ def plot( ax: None | matplotlib.axes.Axes | na.ScalarArray[npt.NDArray] = None, transformation: None | na.transformations.AbstractTransformation = None, components: None | tuple[str, ...] = None, + unit: None | u.UnitBase = None, **kwargs, ) -> dict[str, na.AbstractScalar]: sag = self.sag @@ -228,6 +229,7 @@ def plot( transformation=transformation, components=components, sag=sag, + unit=unit, **kwargs, ) @@ -237,6 +239,7 @@ def plot( transformation=transformation, components=components, sag=sag, + unit=unit, **kwargs, ) diff --git a/optika/systems/_sequential.py b/optika/systems/_sequential.py index a1e617a6..e901f7a5 100644 --- a/optika/systems/_sequential.py +++ b/optika/systems/_sequential.py @@ -1696,6 +1696,7 @@ def plot( plot_rays: bool = True, plot_rays_vignetted: bool = False, kwargs_rays: None | dict[str, Any] = None, + unit: None | u.UnitBase = None, **kwargs, ) -> na.AbstractScalar | dict[str, na.AbstractScalar]: """ @@ -1715,6 +1716,19 @@ def plot( Boolean flag indicating whether to plot the vignetted rays. kwargs_rays Any additional keyword arguments to use when plotting the rays. + unit + The unit to express every plotted length in. + + A system is free to describe its parts in whichever units suit + them, and a detector measured in microns alongside optics measured + in millimeters is quite normal. Matplotlib is told about units by + :func:`astropy.visualization.quantity_support`, which reconciles + them on a 2D axes but not on a 3D one, where a surface given in + microns is drawn a thousand times too large. + + Giving a unit here converts everything to it first, so the system + is drawn to one scale whatever the axes. Leaving it as + :obj:`None`, the default, passes the lengths through as they are. kwargs Any additional keyword arguments to use when plotting the surfaces. """ @@ -1741,6 +1755,7 @@ def plot( ax=ax, transformation=transformation, components=components, + unit=unit, **kwargs, ) ) @@ -1755,8 +1770,12 @@ def plot( if not plot_rays_vignetted: where = raytrace.outputs.unvignetted[{self.axis_surface: ~0}] + position = raytrace.outputs.position + if unit is not None: + position = position.to(unit) + result["rays"] = na.plt.plot( - raytrace.outputs.position, + position, ax=ax, axis=self.axis_surface, where=where, diff --git a/optika/systems/_sequential_test.py b/optika/systems/_sequential_test.py index 306278df..fea93b16 100644 --- a/optika/systems/_sequential_test.py +++ b/optika/systems/_sequential_test.py @@ -716,3 +716,30 @@ 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 + + +def test_plot_unit(): + """ + The whole system is drawn in the unit asked for, rays included. + + :func:`astropy.visualization.quantity_support` reconciles units on a 2D + axes but not on a 3D one, where a part described in microns is drawn a + thousand times larger than the millimeters around it. + """ + + def extent(unit: u.UnitBase) -> float: + """The largest magnitude handed to matplotlib, surfaces and rays alike.""" + fig, ax = plt.subplots() + try: + _system_newtonian.plot(ax=ax, components=("z", "x"), unit=unit) + x = np.concatenate( + [ + np.asarray(getattr(line.get_xdata(), "value", line.get_xdata())) + for line in ax.lines + ] + ) + return float(np.max(np.abs(x))) + finally: + plt.close(fig) + + assert extent(u.um) == pytest.approx(1000 * extent(u.mm))