Let a plot be drawn in a unit of the caller's choosing

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()

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)

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,
             )
 

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,

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))