Add distortion transform to assign_wcs

CHANGES.rst

@@ -1,6 +1,11 @@
 0.7.2 (unreleased)
 ==================
 
+assign_wcs
+----------
+
+- Add distortion transform to assign_wcs step. [#510]
+
 photom
 ------
 

romancal/assign_wcs/assign_wcs_step.py

@@ -27,7 +27,7 @@ class AssignWcsStep(RomanStep):
     """ Assign a gWCS object to a science image.
     """
 
-    reference_file_types = []
+    reference_file_types = ['distortion']
 
     def process(self, input):
         reference_file_names = {}
@@ -41,38 +41,87 @@ def process(self, input):
         return result
 
 
-def load_wcs(dmodel, reference_files):
+def load_wcs(input_model, reference_files=None):
     """ Create a gWCS object and store it in ``Model.meta``.
 
     Parameters
     ----------
-    dmodel : `~roman_datamodels.datamodels.WfiImage`
+    input_model : `~roman_datamodels.datamodels.WfiImage`
         The exposure.
     reference_files : dict
         A dict {reftype: reference_file_name} containing all
         reference files that apply to this exposure.
 
     Returns
     -------
-    dmodel : `~roman_datamodels.ImageModel`
+    output_model : `~roman_datamodels.ImageModel`
         The input image file with attached gWCS object.
         The data is not modified.
     """
-    if reference_files:
+    output_model = input_model.copy()
+
+    if reference_files is not None:
         for ref_type, ref_file in reference_files.items():
             if ref_file not in ["N/A", ""]:
                 reference_files[ref_type] = ref_file
             else:
                 reference_files[ref_type] = None
+    else:
+        reference_files = {}
 
+    # Frames
     detector = cf.Frame2D(name='detector', axes_order=(0, 1), unit=(u.pix, u.pix))
+    v2v3 = cf.Frame2D(name='v2v3', axes_order=(0, 1), axes_names=('v2', 'v3'), unit=(u.arcsec, u.arcsec))
     world = cf.CelestialFrame(reference_frame=coord.ICRS(), name='world')
-    tel2sky = pointing.v23tosky(dmodel)
-    pipeline = [Step(detector, tel2sky),
+
+    # Transforms between frames
+    distortion = wfi_distortion(output_model, reference_files)
+    tel2sky = pointing.v23tosky(output_model)
+
+    pipeline = [Step(detector, distortion),
+                Step(v2v3, tel2sky),
                 Step(world, None)]
     wcs = WCS(pipeline)
     if wcs.bounding_box is None:
-        wcs.bounding_box = wcs_bbox_from_shape(dmodel.data.shape)
-    dmodel.meta['wcs'] = wcs
-    dmodel.meta.cal_step['assign_wcs'] = 'COMPLETE'
-    return dmodel
+        wcs.bounding_box = wcs_bbox_from_shape(output_model.data.shape)
+
+    output_model.meta['wcs'] = wcs
+    output_model.meta.cal_step['assign_wcs'] = 'COMPLETE'
+
+    return output_model
+
+def wfi_distortion(model, reference_files):
+    """
+    Create the "detector" to "v2v3" transform for WFI
+
+    Parameters
+    ----------
+    model : `~roman_datamodels.datamodels.WfiImage`
+        The data model for processing
+    reference_files : dict
+        A dict {reftype: reference_file_name} containing all
+        reference files that apply to this exposure.
+
+    Returns
+    -------
+    The transform model
+    """
+
+    dist = rdm.DistortionRefModel(reference_files['distortion'])
+    transform = dist.coordinate_distortion_transform
+
+    try:
+        bbox = transform.bounding_box
+    except NotImplementedError:
+        # Check if the transform in the reference file has a ``bounding_box``.
+        # If not set a ``bounding_box`` equal to the size of the image after
+        # assembling all distortion corrections.
+        bbox = None
+    dist.close()
+
+    if bbox is None:
+        transform.bounding_box = wcs_bbox_from_shape(model.data.shape)
+    else:
+        transform.bounding_box = bbox
+
+    return transform

romancal/assign_wcs/tests/test_wcs.py

@@ -1,28 +1,64 @@
 import numpy as np
+import os
+import pytest
+
 from numpy.testing import assert_allclose
 
 from gwcs.wcstools import grid_from_bounding_box
 
-from romancal.assign_wcs.assign_wcs_step import load_wcs
+from romancal.assign_wcs.assign_wcs_step import AssignWcsStep, load_wcs
 from roman_datamodels import datamodels as rdm
 from roman_datamodels.testing import utils as testutil
 
+from romancal.assign_wcs.utils import wcs_bbox_from_shape
+
 
 def create_image():
     l2 = testutil.mk_level2_image()
+
     l2.meta.wcsinfo.v2_ref = -503
     l2.meta.wcsinfo.v3_ref = -318
     l2.meta.wcsinfo.ra_ref = 156
     l2.meta.wcsinfo.dec_ref = 54.2
     l2.meta.wcsinfo.vparity = -1
     l2.meta.wcsinfo.roll_ref = 0.15
+
     l2im = rdm.ImageModel(l2)
     return l2im
 
 
-def test_wcs():
+def create_distortion():
+    distortions = [testutil.mk_distortion()]
+
+    model = create_image()
+    dist = testutil.mk_distortion()
+    dist.coordinate_distortion_transform.bounding_box = wcs_bbox_from_shape(model.data.shape)
+    distortions.append(dist)
+
+    return distortions
+
+
+def create_step():
+    def load_wcs_step(image, file_name):
+        return load_wcs(image, {'distortion': file_name})
+
+    def assign_wcs_step(image, file_name):
+        if os.environ.get("CI") == "true":
+            pytest.skip("Roman CRDS servers are not currently available outside the internal network")
+        return AssignWcsStep.call(image, override_distortion=file_name)
+
+    return [load_wcs_step, assign_wcs_step]
+
+
+@pytest.mark.parametrize("distortion", create_distortion())
+@pytest.mark.parametrize("step", create_step())
+def test_wcs(tmpdir, distortion, step):
+    file_name = str(tmpdir / 'distortion.asdf')
+    dist = rdm.DistortionRefModel(distortion)
+    dist.save(file_name)
+
     l2im = create_image()
-    l2_wcs = load_wcs(l2im, {})
+    l2_wcs = step(l2im, file_name)
 
     assert l2_wcs.meta.wcs is not None
     assert l2_wcs.meta.cal_step.assign_wcs == 'COMPLETE'