Add optika.systems.LinearSystem, a linear forward-model approximation to SequentialSystem

docs/index.rst

@@ -33,6 +33,8 @@ Features
 * Sequential raytrace modeling of an optical system
 * Stratified random sampling of input rays for faster convergence
 * Image simulation of a given scene using an optical system
+* Fast linear forward model approximating a raytraced system, for imaging many
+  scenes without raytracing each one
 * Spherical, conical, and toroidal surface sag profiles
 * Circular, rectangular, and polygonal apertures
 * Support for mirrors and arbitrary multilayer coatings

optika/distortion/_distortion.py

@@ -24,6 +24,7 @@
 @dataclasses.dataclass(eq=False, repr=False)
 class AbstractDistortionModel(
     optika.mixins.Printable,
+    optika.mixins.Shaped,
 ):
     """
     An interface describing an arbitrary distortion model,
@@ -189,6 +190,15 @@ class SimpleDistortionModel(
     """The reference wavelength and the sensor position that the field center
     maps to at that wavelength."""
 
+    @property
+    def shape(self) -> dict[str, int]:
+        return na.broadcast_shapes(
+            optika.shape(self.plate_scale),
+            optika.shape(self.dispersion),
+            optika.shape(self.angle),
+            optika.shape(self.reference),
+        )
+
     @functools.cached_property
     def matrix(self) -> na.SpectralPositionalMatrixArray:
         cos = np.cos(self.angle)
@@ -344,6 +354,14 @@ class PolynomialDistortionModel(
     where: bool | na.AbstractScalar = True
     """A boolean mask selecting which calibration points to use for fitting."""
 
+    @property
+    def shape(self) -> dict[str, int]:
+        shape = na.broadcast_shapes(
+            optika.shape(self.coordinates_scene),
+            optika.shape(self.coordinates_sensor),
+        )
+        return {ax: n for ax, n in shape.items() if ax not in self._axis_scene}
+
     @property
     def _axis_scene(self) -> tuple[str, ...]:
         """The logical axes over which the calibration points are distributed."""
@@ -358,6 +376,7 @@ def fit(self) -> na.PolynomialFitFunctionArray:
             outputs=self.coordinates_sensor,
             center=scene.mean(self._axis_scene),
             degree=self.degree,
+            axis_polynomial=self._axis_scene,
             where_polynomial=self.where,
         )
 
@@ -374,6 +393,7 @@ def fit_inverse(self) -> na.PolynomialFitFunctionArray:
             outputs=scene.position,
             center=inputs.mean(self._axis_scene),
             degree=self.degree,
+            axis_polynomial=self._axis_scene,
             where_polynomial=self.where,
         )
 

optika/distortion/_distortion_test.py

@@ -22,6 +22,7 @@ def _scene() -> na.SpectralPositionalVectorArray:
 
 class AbstractTestAbstractDistortionModel(
     test_mixins.AbstractTestPrintable,
+    test_mixins.AbstractTestShaped,
 ):
     def test_distort(self, a: optika.distortion.AbstractDistortionModel):
         coordinates = _scene()
@@ -151,3 +152,43 @@ def test_plot_residual(
         assert isinstance(ax, na.ScalarArray)
         assert a.axis_wavelength in na.shape(ax)
         plt.close(fig)
+
+
+def test_polynomial_distortion_model_channel():
+    """
+    Calibration points that vary along an axis orthogonal to the scene axes
+    (e.g. the channel axis of a multi-channel instrument) must be fit with an
+    independent polynomial per channel, not one polynomial averaged over all
+    the channels.
+    """
+    scene = _scene()
+
+    scale = na.ScalarArray([10, 12, 8] * u.mm / u.deg, axes="channel")
+    angle = na.ScalarArray([0, 10, -15] * u.deg, axes="channel")
+
+    cos, sin = np.cos(angle), np.sin(angle)
+    sensor = na.Cartesian2dVectorArray(
+        x=scale * (cos * scene.position.x - sin * scene.position.y),
+        y=scale * (sin * scene.position.x + cos * scene.position.y),
+    )
+
+    a = optika.distortion.PolynomialDistortionModel(
+        coordinates_scene=scene,
+        coordinates_sensor=sensor,
+        axis_wavelength="wavelength",
+        axis_field=("field_x", "field_y"),
+        degree=1,
+    )
+
+    distorted = a.distort(scene).position
+    assert "channel" in distorted.shape
+    assert np.all((distorted - sensor).length < 1e-9 * u.mm)
+
+    undistorted = a.undistort(
+        na.SpectralPositionalVectorArray(
+            wavelength=scene.wavelength,
+            position=sensor,
+        )
+    ).position
+    assert "channel" in undistorted.shape
+    assert np.all((undistorted - scene.position).length < 1e-9 * u.deg)

optika/radiometry/_effective_area.py

@@ -12,6 +12,7 @@
 @dataclasses.dataclass(eq=False, repr=False)
 class AbstractEffectiveAreaModel(
     optika.mixins.Printable,
+    optika.mixins.Shaped,
 ):
     """
     An interface describing the effective area of an optical system as a
@@ -90,6 +91,14 @@ class InterpolatedEffectiveAreaModel(
     axis_wavelength: str = dataclasses.MISSING
     """The logical axis corresponding to changing wavelength."""
 
+    @property
+    def shape(self) -> dict[str, int]:
+        shape = na.broadcast_shapes(
+            optika.shape(self.wavelength),
+            optika.shape(self.area),
+        )
+        return {ax: n for ax, n in shape.items() if ax != self.axis_wavelength}
+
     def __call__(
         self,
         wavelength: na.AbstractScalar,

optika/radiometry/_effective_area_test.py

@@ -16,6 +16,7 @@ def _area() -> na.AbstractScalar:
 
 class AbstractTestAbstractEffectiveAreaModel(
     test_mixins.AbstractTestPrintable,
+    test_mixins.AbstractTestShaped,
 ):
     def test__call__(self, a: optika.radiometry.AbstractEffectiveAreaModel):
         wavelength = na.linspace(200, 900, axis="wavelength", num=5) * u.AA

optika/radiometry/_vignetting.py

@@ -20,6 +20,7 @@
 @dataclasses.dataclass(eq=False, repr=False)
 class AbstractVignettingModel(
     optika.mixins.Printable,
+    optika.mixins.Shaped,
 ):
     """
     An interface describing an arbitrary vignetting model, which maps scene
@@ -161,6 +162,14 @@ class PolynomialVignettingModel(
     where: bool | na.AbstractScalar = True
     """A boolean mask selecting which calibration points to use for fitting."""
 
+    @property
+    def shape(self) -> dict[str, int]:
+        shape = na.broadcast_shapes(
+            optika.shape(self.coordinates_scene),
+            optika.shape(self.illumination),
+        )
+        return {ax: n for ax, n in shape.items() if ax not in self._axis_scene}
+
     @property
     def _axis_scene(self) -> tuple[str, ...]:
         """The logical axes over which the calibration points are distributed."""
@@ -175,6 +184,7 @@ def fit(self) -> na.PolynomialFitFunctionArray:
             outputs=self.illumination,
             center=scene.mean(self._axis_scene),
             degree=self.degree,
+            axis_polynomial=self._axis_scene,
             where_polynomial=self.where,
         )
 

optika/radiometry/_vignetting_test.py

@@ -26,6 +26,7 @@ def _illumination() -> na.AbstractScalar:
 
 class AbstractTestAbstractVignettingModel(
     test_mixins.AbstractTestPrintable,
+    test_mixins.AbstractTestShaped,
 ):
     def test__call__(self, a: optika.radiometry.AbstractVignettingModel):
         scene = _scene()
@@ -125,3 +126,28 @@ def test_plot_residual(
         assert isinstance(ax, na.ScalarArray)
         assert a.axis_wavelength in na.shape(ax)
         plt.close(fig)
+
+
+def test_polynomial_vignetting_model_channel():
+    """
+    Calibration points that vary along an axis orthogonal to the scene axes
+    (e.g. the channel axis of a multi-channel instrument) must be fit with an
+    independent polynomial per channel, not one polynomial averaged over all
+    the channels.
+    """
+    scene = _scene()
+
+    coefficient = na.ScalarArray([0.1, 0.2, 0.05] / u.deg**2, axes="channel")
+    illumination = 1 - coefficient * scene.position.length**2
+
+    a = optika.radiometry.PolynomialVignettingModel(
+        coordinates_scene=scene,
+        illumination=illumination,
+        axis_wavelength="wavelength",
+        axis_field=("field_x", "field_y"),
+        degree=2,
+    )
+
+    result = a(scene)
+    assert "channel" in result.shape
+    assert np.all(np.abs(result - illumination) < 1e-9)

optika/rays/_ray_vectors.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from typing import TypeVar, Generic
+from collections.abc import Sequence
 import abc
 import dataclasses
 import numpy as np
@@ -98,6 +99,13 @@ def type_explicit(self) -> type[na.AbstractExplicitArray]:
     def type_matrix(self) -> type[na.AbstractMatrixArray]:
         raise NotImplementedError
 
+    def volume_cell(self, axis: None | str | Sequence[str]) -> na.AbstractScalar:
+        """
+        A ray bundle is a scattered collection of rays rather than a
+        logically-rectangular grid, so the per-voxel volume is undefined.
+        """
+        raise NotImplementedError
+
     @property
     def explicit(self) -> RayVectorArray:
         return super().explicit