From 13d7c94efec4594b529eaadf88712f99a53ca157 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 31 May 2026 12:48:48 -0600 Subject: [PATCH 1/5] Add `SimpleDistortionModel`, an analytic linear distortion. Introduce `AbstractLinearDistortionModel`, an affine distortion (`distort = matrix @ (c - center) + intercept`) whose `undistort` is the exact inverse via `matrix.inverse`, and a concrete `SimpleDistortionModel` parameterized by `plate_scale`, `dispersion`, `angle`, `wavelength_ref`, and `position_ref`. It builds the rotation + dispersion + plate-scale `SpectralPositionalMatrixArray`, capturing the distortion of an idealized spectrograph (cf. the ctis `IdealInstrument` distortion). Co-Authored-By: Claude Opus 4.8 --- optika/distortion/__init__.py | 4 + optika/distortion/_distortion.py | 162 ++++++++++++++++++++++++++ optika/distortion/_distortion_test.py | 32 +++++ 3 files changed, 198 insertions(+) diff --git a/optika/distortion/__init__.py b/optika/distortion/__init__.py index 01f41c1..9f6d39a 100644 --- a/optika/distortion/__init__.py +++ b/optika/distortion/__init__.py @@ -2,12 +2,16 @@ from ._distortion import ( AbstractDistortionModel, + AbstractLinearDistortionModel, + SimpleDistortionModel, AbstractInterpolatedDistortionModel, PolynomialDistortionModel, ) __all__ = [ "AbstractDistortionModel", + "AbstractLinearDistortionModel", + "SimpleDistortionModel", "AbstractInterpolatedDistortionModel", "PolynomialDistortionModel", ] diff --git a/optika/distortion/_distortion.py b/optika/distortion/_distortion.py index b67534e..984eedb 100644 --- a/optika/distortion/_distortion.py +++ b/optika/distortion/_distortion.py @@ -1,17 +1,21 @@ import abc import dataclasses import functools +import numpy as np import matplotlib.axes import matplotlib.cm import matplotlib.colors import matplotlib.figure import matplotlib.pyplot as plt +import astropy.units as u import astropy.visualization import named_arrays as na import optika __all__ = [ "AbstractDistortionModel", + "AbstractLinearDistortionModel", + "SimpleDistortionModel", "AbstractInterpolatedDistortionModel", "PolynomialDistortionModel", ] @@ -61,6 +65,164 @@ def undistort( """ +@dataclasses.dataclass(eq=False, repr=False) +class AbstractLinearDistortionModel( + AbstractDistortionModel, +): + r""" + A distortion model which is an affine transformation of the scene + coordinates, + + .. math:: + + \text{distort}(\vec{c}) = \mathbf{M} \, (\vec{c} - \vec{c}_0) + \vec{b}, + + where :math:`\mathbf{M}` is :attr:`matrix`, :math:`\vec{c}_0` is + :attr:`center`, and :math:`\vec{b}` is :attr:`intercept`. + Since the transformation is linear, :meth:`undistort` is its *exact* + inverse (unlike a polynomial fit). + """ + + @property + @abc.abstractmethod + def matrix(self) -> na.AbstractSpectralPositionalMatrixArray: + """The linear part of the affine transformation.""" + + @property + @abc.abstractmethod + def center(self) -> na.AbstractSpectralPositionalVectorArray: + """The reference point subtracted from the coordinates before + applying :attr:`matrix`.""" + + @property + @abc.abstractmethod + def intercept(self) -> na.AbstractSpectralPositionalVectorArray: + """The constant offset added after applying :attr:`matrix`.""" + + def distort( + self, + coordinates: na.AbstractSpectralPositionalVectorArray, + ) -> na.SpectralPositionalVectorArray: + return self.matrix @ (coordinates - self.center) + self.intercept + + def undistort( + self, + coordinates: na.AbstractSpectralPositionalVectorArray, + ) -> na.SpectralPositionalVectorArray: + return self.matrix.inverse @ (coordinates - self.intercept) + self.center + + +@dataclasses.dataclass(eq=False, repr=False) +class SimpleDistortionModel( + AbstractLinearDistortionModel, +): + r""" + A simple analytic distortion model consisting of a rotation of the field, + an isotropic plate scale, and a linear spectral dispersion along the + rotated :math:`x` axis. + + This captures the distortion of an idealized spectrograph: the field at + :attr:`wavelength_ref` maps to :attr:`position_ref` on the sensor, and + other wavelengths are displaced along the dispersion direction. + + Examples + -------- + + Distort a grid of scene coordinates and confirm the inverse recovers them. + + .. jupyter-execute:: + + import astropy.units as u + import named_arrays as na + import optika + + model = optika.distortion.SimpleDistortionModel( + plate_scale=1 * u.arcsec / u.pix, + dispersion=0.1 * u.nm / u.pix, + angle=15 * u.deg, + wavelength_ref=550 * u.nm, + position_ref=na.Cartesian2dVectorArray(0, 0) * u.pix, + ) + + scene = na.SpectralPositionalVectorArray( + wavelength=na.linspace(500, 600, axis="wavelength", num=3) * u.nm, + position=na.Cartesian2dVectorLinearSpace( + start=-10 * u.arcsec, + stop=+10 * u.arcsec, + axis=na.Cartesian2dVectorArray("field_x", "field_y"), + num=5, + ), + ) + + sensor = model.distort(scene) + scene_roundtrip = model.undistort(sensor) + """ + + plate_scale: u.Quantity | na.AbstractScalar = dataclasses.MISSING + """The spatial plate scale, in units such as :math:`\\text{arcsec} / \\text{pix}`.""" + + dispersion: u.Quantity | na.AbstractScalar = dataclasses.MISSING + """The magnitude of the spectral dispersion, in units such as :math:`\\text{nm} / \\text{pix}`.""" + + angle: u.Quantity | na.AbstractScalar = dataclasses.MISSING + """The angle of the dispersion direction with respect to the scene.""" + + wavelength_ref: u.Quantity | na.AbstractScalar = dataclasses.MISSING + """The reference wavelength which maps the field center to :attr:`position_ref`.""" + + position_ref: u.Quantity | na.AbstractScalar | na.AbstractCartesian2dVectorArray = ( + dataclasses.MISSING + ) + """The sensor position of the field center at :attr:`wavelength_ref`.""" + + @functools.cached_property + def matrix(self) -> na.SpectralPositionalMatrixArray: + cos = np.cos(self.angle) + sin = np.sin(self.angle) + plate_scale = self.plate_scale + dispersion = self.dispersion + unit_wavelength = na.unit(self.wavelength_ref) + return na.SpectralPositionalMatrixArray( + wavelength=na.SpectralPositionalVectorArray( + wavelength=1, + position=na.Cartesian2dVectorArray( + x=0 * unit_wavelength / u.arcsec, + y=0 * unit_wavelength / u.arcsec, + ), + ), + position=na.Cartesian2dMatrixArray( + x=na.SpectralPositionalVectorArray( + wavelength=1 / dispersion, + position=na.Cartesian2dVectorArray( + x=cos / plate_scale, + y=-sin / plate_scale, + ), + ), + y=na.SpectralPositionalVectorArray( + wavelength=0 / dispersion, + position=na.Cartesian2dVectorArray( + x=sin / plate_scale, + y=cos / plate_scale, + ), + ), + ), + ) + + @property + def center(self) -> na.SpectralPositionalVectorArray: + return na.SpectralPositionalVectorArray( + wavelength=self.wavelength_ref, + position=na.Cartesian2dVectorArray(0, 0) * u.arcsec, + ) + + @property + def intercept(self) -> na.SpectralPositionalVectorArray: + return na.SpectralPositionalVectorArray( + wavelength=self.wavelength_ref, + position=self.position_ref, + ) + + @dataclasses.dataclass(eq=False, repr=False) class AbstractInterpolatedDistortionModel( AbstractDistortionModel, diff --git a/optika/distortion/_distortion_test.py b/optika/distortion/_distortion_test.py index f2a77d0..fae3134 100644 --- a/optika/distortion/_distortion_test.py +++ b/optika/distortion/_distortion_test.py @@ -44,6 +44,38 @@ def test_roundtrip(self, a: optika.distortion.AbstractDistortionModel): assert np.all(error < 1e-9 * u.deg) +class AbstractTestAbstractLinearDistortionModel( + AbstractTestAbstractDistortionModel, +): + def test_matrix(self, a: optika.distortion.AbstractLinearDistortionModel): + assert isinstance(a.matrix, na.AbstractSpectralPositionalMatrixArray) + + def test_center(self, a: optika.distortion.AbstractLinearDistortionModel): + assert isinstance(a.center, na.AbstractSpectralPositionalVectorArray) + + def test_intercept(self, a: optika.distortion.AbstractLinearDistortionModel): + assert isinstance(a.intercept, na.AbstractSpectralPositionalVectorArray) + + +@pytest.mark.parametrize( + argnames="a", + argvalues=[ + optika.distortion.SimpleDistortionModel( + plate_scale=1 * u.arcsec / u.pix, + dispersion=0.1 * u.nm / u.pix, + angle=angle, + wavelength_ref=550 * u.nm, + position_ref=na.Cartesian2dVectorArray(0, 0) * u.pix, + ) + for angle in [0 * u.deg, 15 * u.deg] + ], +) +class TestSimpleDistortionModel( + AbstractTestAbstractLinearDistortionModel, +): + pass + + class AbstractTestAbstractInterpolatedDistortionModel( AbstractTestAbstractDistortionModel, ): From d85ea95ae8407836527921953886d3e32f62dadb Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 31 May 2026 13:06:33 -0600 Subject: [PATCH 2/5] Merge wavelength_ref/position_ref into a single `reference` coordinate. `SimpleDistortionModel` now takes one `reference: SpectralPositionalVectorArray` (the reference wavelength + the sensor position the field center maps to) instead of separate `wavelength_ref`/`position_ref` fields; `intercept` is then just `reference`. The class example now distorts a grid and scatter-plots the result on the sensor, colored by wavelength. Co-Authored-By: Claude Opus 4.8 --- optika/distortion/_distortion.py | 48 +++++++++++++++------------ optika/distortion/_distortion_test.py | 6 ++-- 2 files changed, 31 insertions(+), 23 deletions(-) diff --git a/optika/distortion/_distortion.py b/optika/distortion/_distortion.py index 984eedb..f6f5bac 100644 --- a/optika/distortion/_distortion.py +++ b/optika/distortion/_distortion.py @@ -121,17 +121,20 @@ class SimpleDistortionModel( an isotropic plate scale, and a linear spectral dispersion along the rotated :math:`x` axis. - This captures the distortion of an idealized spectrograph: the field at - :attr:`wavelength_ref` maps to :attr:`position_ref` on the sensor, and - other wavelengths are displaced along the dispersion direction. + This captures the distortion of an idealized spectrograph: the field + center at the :attr:`reference` wavelength maps to the :attr:`reference` + position on the sensor, and other wavelengths are displaced along the + dispersion direction. Examples -------- - Distort a grid of scene coordinates and confirm the inverse recovers them. + Distort a grid of scene coordinates and plot the result on the sensor, + colored by wavelength. .. jupyter-execute:: + import matplotlib.pyplot as plt import astropy.units as u import named_arrays as na import optika @@ -140,8 +143,10 @@ class SimpleDistortionModel( plate_scale=1 * u.arcsec / u.pix, dispersion=0.1 * u.nm / u.pix, angle=15 * u.deg, - wavelength_ref=550 * u.nm, - position_ref=na.Cartesian2dVectorArray(0, 0) * u.pix, + reference=na.SpectralPositionalVectorArray( + wavelength=550 * u.nm, + position=na.Cartesian2dVectorArray(0, 0) * u.pix, + ), ) scene = na.SpectralPositionalVectorArray( @@ -155,7 +160,15 @@ class SimpleDistortionModel( ) sensor = model.distort(scene) - scene_roundtrip = model.undistort(sensor) + + fig, ax = plt.subplots(constrained_layout=True) + ax.set_aspect("equal") + na.plt.scatter( + sensor.position.x, + sensor.position.y, + c=sensor.wavelength, + ax=ax, + ); """ plate_scale: u.Quantity | na.AbstractScalar = dataclasses.MISSING @@ -167,13 +180,9 @@ class SimpleDistortionModel( angle: u.Quantity | na.AbstractScalar = dataclasses.MISSING """The angle of the dispersion direction with respect to the scene.""" - wavelength_ref: u.Quantity | na.AbstractScalar = dataclasses.MISSING - """The reference wavelength which maps the field center to :attr:`position_ref`.""" - - position_ref: u.Quantity | na.AbstractScalar | na.AbstractCartesian2dVectorArray = ( - dataclasses.MISSING - ) - """The sensor position of the field center at :attr:`wavelength_ref`.""" + reference: na.AbstractSpectralPositionalVectorArray = dataclasses.MISSING + """The reference wavelength and the sensor position that the field center + maps to at that wavelength.""" @functools.cached_property def matrix(self) -> na.SpectralPositionalMatrixArray: @@ -181,7 +190,7 @@ def matrix(self) -> na.SpectralPositionalMatrixArray: sin = np.sin(self.angle) plate_scale = self.plate_scale dispersion = self.dispersion - unit_wavelength = na.unit(self.wavelength_ref) + unit_wavelength = na.unit(self.reference.wavelength) return na.SpectralPositionalMatrixArray( wavelength=na.SpectralPositionalVectorArray( wavelength=1, @@ -211,16 +220,13 @@ def matrix(self) -> na.SpectralPositionalMatrixArray: @property def center(self) -> na.SpectralPositionalVectorArray: return na.SpectralPositionalVectorArray( - wavelength=self.wavelength_ref, + wavelength=self.reference.wavelength, position=na.Cartesian2dVectorArray(0, 0) * u.arcsec, ) @property - def intercept(self) -> na.SpectralPositionalVectorArray: - return na.SpectralPositionalVectorArray( - wavelength=self.wavelength_ref, - position=self.position_ref, - ) + def intercept(self) -> na.AbstractSpectralPositionalVectorArray: + return self.reference @dataclasses.dataclass(eq=False, repr=False) diff --git a/optika/distortion/_distortion_test.py b/optika/distortion/_distortion_test.py index fae3134..721b12f 100644 --- a/optika/distortion/_distortion_test.py +++ b/optika/distortion/_distortion_test.py @@ -64,8 +64,10 @@ def test_intercept(self, a: optika.distortion.AbstractLinearDistortionModel): plate_scale=1 * u.arcsec / u.pix, dispersion=0.1 * u.nm / u.pix, angle=angle, - wavelength_ref=550 * u.nm, - position_ref=na.Cartesian2dVectorArray(0, 0) * u.pix, + reference=na.SpectralPositionalVectorArray( + wavelength=550 * u.nm, + position=na.Cartesian2dVectorArray(0, 0) * u.pix, + ), ) for angle in [0 * u.deg, 15 * u.deg] ], From 5d90edf3615a32521a588318bd641064c17995ef Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 31 May 2026 13:43:22 -0600 Subject: [PATCH 3/5] Fix SimpleDistortionModel example for the docs build. `na.plt.scatter(c=...)` lets matplotlib autoscale a colormap `Normalize`, which calls `float()` on the color values; a `Quantity` wavelength raises "only dimensionless scalar quantities can be converted to Python scalars" when the figure is rendered (as it is by jupyter-execute, breaking the Read the Docs build). Pass `sensor.wavelength.to_value(u.nm)` so `c` is unitless. Co-Authored-By: Claude Opus 4.8 --- optika/distortion/_distortion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optika/distortion/_distortion.py b/optika/distortion/_distortion.py index f6f5bac..d267aa6 100644 --- a/optika/distortion/_distortion.py +++ b/optika/distortion/_distortion.py @@ -166,7 +166,7 @@ class SimpleDistortionModel( na.plt.scatter( sensor.position.x, sensor.position.y, - c=sensor.wavelength, + c=sensor.wavelength.to_value(u.nm), ax=ax, ); """ From a87e18f8b642dc6dd753b42e127164e7d192381c Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 31 May 2026 14:28:58 -0600 Subject: [PATCH 4/5] Tune SimpleDistortionModel example: less dispersion, axis labels. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Increase the example dispersion to 2 nm/pix so the wavelength channels are less spread out (the previous 0.1 nm/pix dispersed them ~±500 px versus the ~±10 px field), and label the detector x/y axes. Co-Authored-By: Claude Opus 4.8 --- optika/distortion/_distortion.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/optika/distortion/_distortion.py b/optika/distortion/_distortion.py index d267aa6..1d6fbfe 100644 --- a/optika/distortion/_distortion.py +++ b/optika/distortion/_distortion.py @@ -141,7 +141,7 @@ class SimpleDistortionModel( model = optika.distortion.SimpleDistortionModel( plate_scale=1 * u.arcsec / u.pix, - dispersion=0.1 * u.nm / u.pix, + dispersion=2 * u.nm / u.pix, angle=15 * u.deg, reference=na.SpectralPositionalVectorArray( wavelength=550 * u.nm, @@ -168,7 +168,9 @@ class SimpleDistortionModel( sensor.position.y, c=sensor.wavelength.to_value(u.nm), ax=ax, - ); + ) + ax.set_xlabel(f"detector $x$ ({na.unit(sensor.position.x):latex_inline})") + ax.set_ylabel(f"detector $y$ ({na.unit(sensor.position.y):latex_inline})"); """ plate_scale: u.Quantity | na.AbstractScalar = dataclasses.MISSING From c200ac2e99a720a01c3906fca717b09bdfc71987 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sun, 31 May 2026 15:09:54 -0600 Subject: [PATCH 5/5] Label wavelengths with a legend in the SimpleDistortionModel example. Scatter one series per wavelength (selected with a `where` mask) so each gets its own color and a legend entry, instead of an unlabeled color mapping. Co-Authored-By: Claude Opus 4.8 --- optika/distortion/_distortion.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/optika/distortion/_distortion.py b/optika/distortion/_distortion.py index 1d6fbfe..ddd5197 100644 --- a/optika/distortion/_distortion.py +++ b/optika/distortion/_distortion.py @@ -163,14 +163,17 @@ class SimpleDistortionModel( fig, ax = plt.subplots(constrained_layout=True) ax.set_aspect("equal") - na.plt.scatter( - sensor.position.x, - sensor.position.y, - c=sensor.wavelength.to_value(u.nm), - ax=ax, - ) + for wavelength in scene.wavelength.ndarray: + na.plt.scatter( + sensor.position.x, + sensor.position.y, + where=scene.wavelength == wavelength, + label=f"{wavelength}", + ax=ax, + ) ax.set_xlabel(f"detector $x$ ({na.unit(sensor.position.x):latex_inline})") - ax.set_ylabel(f"detector $y$ ({na.unit(sensor.position.y):latex_inline})"); + ax.set_ylabel(f"detector $y$ ({na.unit(sensor.position.y):latex_inline})") + ax.legend(); """ plate_scale: u.Quantity | na.AbstractScalar = dataclasses.MISSING