Adding functionality to represent non linear spectral dispersion

changelog/713.feature.rst

@@ -0,0 +1 @@
+Added dkist.wcs.models.generate_grating_spectral_transform() for building FITS -GRA/-GRI spectral transforms from header-derived grating parameters.

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
@@ -31,10 +32,115 @@
     "VaryingCelestialTransform2D",
     "VaryingCelestialTransform3D",
     "generate_celestial_transform",
+    "generate_grating_spectral_transform",
     "varying_celestial_transform_from_tables",
 ]
 
 
+def _grating_equation_constants(
+        reference_wavelength,
+        grating_density,
+        spectral_order,
+        incident_angle,
+        refractive_index,
+        refractive_index_derivative,
+        out_of_plane_angle,
+        camera_angle,
+):
+    alpha = incident_angle.to(u.rad)
+    epsilon = out_of_plane_angle.to(u.rad)
+    theta = camera_angle.to(u.rad)
+
+    grism_parameter = (grating_density * spectral_order) / np.cos(epsilon)
+    adjusted_incident_angle_sine = (
+        refractive_index - (refractive_index_derivative * reference_wavelength).decompose()
+    ) * np.sin(alpha)
+    wavelength_parameter = grism_parameter - refractive_index_derivative * np.sin(alpha)
+    reference_refracted_angle = np.arcsin(
+        (grism_parameter * reference_wavelength).decompose() - refractive_index * np.sin(alpha)
+    )
+    angle_slope = wavelength_parameter / (
+        np.cos(reference_refracted_angle) * np.cos(theta) ** 2
+    )
+
+    return (
+        adjusted_incident_angle_sine,
+        wavelength_parameter,
+        reference_refracted_angle,
+        theta,
+        angle_slope,
+    )
+
+
+def _refracted_angle_sine_model(
+        reference_pixel, dispersion, reference_refracted_angle, theta, angle_slope
+):
+    @custom_model
+    def refracted_angle_sine(pixel):
+        wavelength_offset = ((pixel - reference_pixel) * u.pix) * dispersion
+        output_angle = (
+            np.arctan(-np.tan(theta) + wavelength_offset * angle_slope)
+            + reference_refracted_angle
+            + theta
+        )
+        return np.sin(output_angle)
+
+    return refracted_angle_sine()
+
+
+def generate_grating_spectral_transform(
+        reference_pixel,
+        reference_wavelength,
+        dispersion,
+        grating_density,
+        spectral_order,
+        incident_angle,
+        refractive_index=1 * u.one,
+        refractive_index_derivative=0 / u.m,
+        out_of_plane_angle=0 * u.deg,
+        camera_angle=0 * u.deg,
+) -> CompoundModel:
+    """
+    Generate a one-dimensional FITS ``-GRA``/``-GRI`` spectral transform.
+
+    This function handles the FITS grating/grism intermediate-angle terms and
+    delegates the final wavelength calculation to
+    ``gwcs.spectroscopy.WavelengthFromGratingEquation``.
+    """
+    (
+        adjusted_incident_angle_sine,
+        wavelength_parameter,
+        reference_refracted_angle,
+        theta,
+        angle_slope,
+    ) = _grating_equation_constants(
+        reference_wavelength=reference_wavelength,
+        grating_density=grating_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,
+    )
+
+    refracted_angle_sine = _refracted_angle_sine_model(
+        reference_pixel=reference_pixel,
+        dispersion=dispersion,
+        reference_refracted_angle=reference_refracted_angle,
+        theta=theta,
+        angle_slope=angle_slope,
+    )
+    alpha_in = m.Const1D(amplitude=adjusted_incident_angle_sine)
+    wavelength_from_grating = WavelengthFromGratingEquation(
+        groove_density=wavelength_parameter / spectral_order,
+        spectral_order=spectral_order,
+        name="Spectral",
+    )
+
+    return m.Mapping((0, 0)) | (alpha_in & refracted_angle_sine) | wavelength_from_grating
+
+
 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,
+                              update_celestial_transform_parameters,
                               varying_celestial_transform_from_tables)
 
 
@@ -52,6 +54,36 @@ def test_generate_celestial_unitless():
     assert u.allclose(shift1.offset, 0)
 
 
+def test_generate_grating_spectral_transform():
+    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,
+    }
+    transform = generate_grating_spectral_transform(
+        reference_pixel=header["CRPIX1"] - 1,
+        reference_wavelength=header["CRVAL1"] * u.nm,
+        dispersion=header["CDELT1"] * u.nm / u.pix,
+        grating_density=header["PV1_0"] / u.m,
+        spectral_order=header["PV1_1"] * u.one,
+        incident_angle=header["PV1_2"] * u.deg,
+    )
+
+    pixels = np.array([0, 100, 217, 300, 511], dtype=float)
+    expected = WCS(header).all_pix2world(pixels, 0)[0] * u.m
+    result = transform(pixels)
+
+    assert isinstance(transform, CompoundModel)
+    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,