Adding functionality to represent non linear spectral dispersion

changelog/713.feature.rst

@@ -0,0 +1 @@
+Updated `dkist.wcs.models.generate_grating_spectral_transform()` to build FITS ``-GRA``/``-GRI`` spectral transforms by composing a constant incident-angle term and pixel-dependent refracted-angle model with `gwcs.spectroscopy.WavelengthFromGratingEquation`.

dkist/wcs/models.py

@@ -12,8 +12,9 @@
 
 import astropy.modeling.models as m
 import astropy.units as u
-from astropy.modeling import CompoundModel, Model, Parameter, separable
+from astropy.modeling import CompoundModel, Model, Parameter, custom_model, separable
 from astropy.utils.decorators import deprecated_renamed_argument
+from gwcs.spectroscopy import WavelengthFromGratingEquation
 
 from dkist.utils.decorators import deprecated
 from dkist.utils.exceptions import DKISTDeprecationWarning
@@ -30,11 +31,126 @@
     "VaryingCelestialTransform",
     "VaryingCelestialTransform2D",
     "VaryingCelestialTransform3D",
+    "build_grating_spectral_transform",
     "generate_celestial_transform",
+    "generate_grating_spectral_transform",
+    "refracted_angle_sine_model",
     "varying_celestial_transform_from_tables",
 ]
 
 
+@custom_model
+def refracted_angle_sine_model(
+    pixel,
+    reference_pixel=0,
+    reference_refracted_angle=0 * u.rad,
+    dispersion=0 * u.nm / u.pix,
+    grism_parameter_per_wavelength=0 / u.nm,
+    camera_angle=0 * u.deg,
+):
+    """
+    Compute the refracted-angle sine term for FITS ``-GRA``/``-GRI`` spectra.
+
+    This model is defined in ``dkist`` so its import path remains stable for
+    ASDF serialization and deserialization.
+    """
+    wavelength_offset = ((pixel - reference_pixel) * u.pix) * dispersion
+    output_angle = (
+        np.arctan(-np.tan(camera_angle) + wavelength_offset * grism_parameter_per_wavelength)
+        + reference_refracted_angle
+        + camera_angle
+    )
+    return np.sin(output_angle)
+
+
+def build_grating_spectral_transform(
+    reference_pixel: float,
+    reference_wavelength: u.Quantity,
+    dispersion: u.Quantity,
+    groove_density: u.Quantity,
+    spectral_order: u.Quantity,
+    incident_angle: u.Quantity,
+    refractive_index: u.Quantity = 1 * u.one,
+    refractive_index_derivative: u.Quantity = 0 / u.m,
+    out_of_plane_angle: u.Quantity = 0 * u.deg,
+    camera_angle: u.Quantity = 0 * u.deg,
+) -> CompoundModel:
+    """
+    Build a FITS grating spectral transform from header-derived parameters.
+
+    This is the low-level implementation used by the compatibility entry
+    points in `dkist.wcs.models`. It composes a constant incident-angle sine
+    term, a pixel-dependent refracted-angle sine model, and
+    `~gwcs.spectroscopy.WavelengthFromGratingEquation`, following the FITS
+    grating/grism spectral-coordinate formalism described by Greisen et al.
+    (2006):
+    https://scixplorer.org/abs/2006A%26A...446..747G/abstract
+    """
+    model = WavelengthFromGratingEquation(
+        groove_density=groove_density,
+        spectral_order=spectral_order,
+        reference_wavelength=reference_wavelength,
+        refractive_index=refractive_index,
+        refractive_index_derivative=refractive_index_derivative,
+        out_of_plane_angle=out_of_plane_angle,
+    )
+
+    alpha_in = m.Const1D(amplitude=np.sin(incident_angle))
+
+    grism_constant = (groove_density * spectral_order) / np.cos(out_of_plane_angle)
+    reference_refracted_angle = np.arcsin(
+        (grism_constant * reference_wavelength) - refractive_index * np.sin(incident_angle)
+    )
+    grism_parameter_per_wavelength = (
+        grism_constant - refractive_index_derivative * np.sin(incident_angle)
+    ) / (np.cos(reference_refracted_angle) * np.cos(camera_angle) ** 2)
+
+    alpha_out = refracted_angle_sine_model(
+        reference_pixel=reference_pixel,
+        reference_refracted_angle=reference_refracted_angle,
+        dispersion=dispersion,
+        grism_parameter_per_wavelength=grism_parameter_per_wavelength,
+        camera_angle=camera_angle,
+    )
+
+    return m.Mapping((0, 0)) | (alpha_in & alpha_out) | model
+
+
+def generate_grating_spectral_transform(
+        reference_pixel: float,
+        reference_wavelength: u.Quantity,
+        dispersion: u.Quantity,
+        groove_density: u.Quantity,
+        spectral_order: u.Quantity,
+        incident_angle: u.Quantity,
+        refractive_index: u.Quantity = 1 * u.one,
+        refractive_index_derivative: u.Quantity = 0 / u.m,
+        out_of_plane_angle: u.Quantity = 0 * u.deg,
+        camera_angle: u.Quantity = 0 * u.deg,
+) -> CompoundModel:
+    """
+    Build a FITS ``-GRA``/``-GRI`` spectral transform from header-derived parameters.
+
+    Composes a constant incident-angle sine term (`~astropy.modeling.models.Const1D`),
+    a pixel-dependent refracted-angle sine model (`~dkist.wcs.models.refracted_angle_sine_model`),
+    and `~gwcs.spectroscopy.WavelengthFromGratingEquation` following the FITS
+    grating/grism spectral-coordinate formalism described by Greisen et al. (2006):
+    https://scixplorer.org/abs/2006A%26A...446..747G/abstract
+    """
+    return build_grating_spectral_transform(
+        reference_pixel=reference_pixel,
+        reference_wavelength=reference_wavelength,
+        dispersion=dispersion,
+        groove_density=groove_density,
+        spectral_order=spectral_order,
+        incident_angle=incident_angle,
+        refractive_index=refractive_index,
+        refractive_index_derivative=refractive_index_derivative,
+        out_of_plane_angle=out_of_plane_angle,
+        camera_angle=camera_angle,
+    )
+
+
 def generate_celestial_transform(
         crpix: Iterable[float] | u.Quantity,
         cdelt: Iterable[float] | u.Quantity,

dkist/wcs/tests/test_models.py

@@ -9,10 +9,12 @@
 from astropy.coordinates.matrix_utilities import rotation_matrix
 from astropy.modeling import CompoundModel
 from astropy.modeling.models import Tabular1D
+from astropy.wcs import WCS
 
 from dkist.wcs.models import (AsymmetricMapping, Ravel, Unravel, VaryingCelestialTransform,
                               VaryingCelestialTransform2D, VaryingCelestialTransform3D,
-                              generate_celestial_transform, update_celestial_transform_parameters,
+                              generate_celestial_transform, generate_grating_spectral_transform,
+                              refracted_angle_sine_model, update_celestial_transform_parameters,
                               varying_celestial_transform_from_tables)
 
 
@@ -52,6 +54,48 @@ def test_generate_celestial_unitless():
     assert u.allclose(shift1.offset, 0)
 
 
+def test_generate_grating_spectral_transform() -> None:
+    header = {
+        "CTYPE1": "AWAV-GRA",
+        "CUNIT1": "nm",
+        "CRPIX1": 218,
+        "CRVAL1": 854.1738582455826,
+        "CDELT1": 0.0022975580183395555,
+        "PV1_0": 23000.0,
+        "PV1_1": 90,
+        "PV1_2": 65.696,
+        "PV1_3": 1.25,
+        "PV1_4": 1000.0,
+        "PV1_5": 1.5,
+        "PV1_6": 0.8,
+    }
+    transform = generate_grating_spectral_transform(
+        reference_pixel=header["CRPIX1"] - 1,
+        reference_wavelength=header["CRVAL1"] * u.nm,
+        dispersion=header["CDELT1"] * u.nm / u.pix,
+        groove_density=header["PV1_0"] / u.m,
+        spectral_order=header["PV1_1"] * u.one,
+        incident_angle=header["PV1_2"] * u.deg,
+        refractive_index=header["PV1_3"] * u.one,
+        refractive_index_derivative=header["PV1_4"] / u.m,
+        out_of_plane_angle=header["PV1_5"] * u.deg,
+        camera_angle=header["PV1_6"] * u.deg,
+    )
+
+    pixels = np.array([0, 100, 217, 300, 511], dtype=float)
+    expected = WCS(header).spectral.pixel_to_world(pixels)
+    result = transform(pixels)
+
+    assert isinstance(transform, CompoundModel)
+    assert any(
+        isinstance(submodel, refracted_angle_sine_model)
+        for submodel in transform.traverse_postorder()
+    )
+    np.testing.assert_allclose(
+        result.to_value(u.nm), expected.to_value(u.nm), rtol=1e-10, atol=1e-10
+    )
+
+
 def test_update_celestial():
     trsfm  = generate_celestial_transform(
               crpix=[0, 0] * u.pix,