Forward axis_wavelength from image() through measure() to expose()

optika/sensors/_sensors.py

@@ -241,6 +241,7 @@ def measure(
         wavelength: na.AbstractScalar,
         axis: None | str | Sequence[str] = None,
         where: bool | na.AbstractScalar = True,
+        axis_wavelength: None | str = None,
         timedelta: None | u.Quantity | na.AbstractScalar = None,
         noise: bool = True,
     ) -> na.FunctionArray[
@@ -264,6 +265,11 @@ def measure(
             The logical axes along which to collect photons.
         where
             A boolean mask used to indicate which rays should be considered.
+        axis_wavelength
+            The logical axis of `wavelength` corresponding to changing
+            wavelength coordinate, forwarded to :meth:`expose`.
+            If :obj:`None` (the default), `wavelength` must have only one
+            logical axis.
         timedelta
             The exposure time of the measurement.
             If :obj:`None` (the default), the value in :attr:`timedelta_exposure`
@@ -280,6 +286,7 @@ def measure(
         return self.expose(
             image,
             direction=direction,
+            axis_wavelength=axis_wavelength,
             timedelta=timedelta,
             noise=noise,
         )

optika/sensors/_sensors_test.py

@@ -60,6 +60,42 @@ def test_measure(
         assert a.axis_pixel.x in result.outputs.shape
         assert a.axis_pixel.y in result.outputs.shape
 
+    def test_measure_axis_wavelength(
+        self,
+        a: optika.sensors.AbstractImagingSensor,
+    ):
+        """
+        Test measuring rays whose wavelength varies along more than one axis,
+        for example a scene composed of several disjoint spectral lines, each
+        sampled by its own set of wavelength bins.
+        """
+        line = na.ScalarArray([500, 600] * u.nm, axes="line")
+        wavelength = line + na.linspace(-1, 1, axis="wavelength", num=3) * u.nm
+
+        rays = optika.rays.RayVectorArray(
+            intensity=100 * u.photon / u.s,
+            wavelength=line + na.linspace(-0.5, 0.5, axis="wavelength", num=2) * u.nm,
+            position=na.Cartesian3dVectorArray(
+                x=na.random.uniform(-1, 1, shape_random=dict(wavelength=2, t=11)),
+                y=na.random.uniform(-1, 1, shape_random=dict(wavelength=2, t=11)),
+                z=0,
+            )
+            * u.mm,
+            direction=na.Cartesian3dVectorArray(0, 0, 1),
+        )
+
+        result = a.measure(
+            rays,
+            wavelength,
+            axis=("wavelength", "t"),
+            axis_wavelength="wavelength",
+        )
+        assert isinstance(result, na.FunctionArray)
+        assert result.outputs.unit.is_equivalent(u.electron)
+        assert "line" in result.outputs.shape
+        assert a.axis_pixel.x in result.outputs.shape
+        assert a.axis_pixel.y in result.outputs.shape
+
 
 @pytest.mark.parametrize(
     argnames="a",

optika/systems/_sequential.py

@@ -919,6 +919,7 @@ def image(
             rays=rayfunction.outputs,
             wavelength=wavelength,
             axis=(axis_wavelength,) + axis_field + axis_pupil,
+            axis_wavelength=axis_wavelength,
             noise=noise,
         )