JP-3227: Add analog of HST's tweakback tool

CHANGES.rst

@@ -42,6 +42,14 @@ ramp_fitting
   of groups in a segment are computed when applying optimal weighting to line
   fit segments. [#7560, spacetelescope/stcal#163]
 
+tweakreg
+--------
+
+- Added a tool ``transfer_wcs_correction`` to ``jwst.tweakreg.utils`` that
+  allows transferring alignment corrections from one file/data model to
+  another. It is an analog of the ``tweakback`` task in the
+  ``drizzlepac``. [#7573]
+
 
 1.10.2 (2023-04-14)
 ===================
@@ -75,7 +83,7 @@ jump
   in the STCAL jump code but not in JWST. This allows the parameter to be changed.
   Also, scaled two input parameters that are listed as radius to be a factor of two
   higher to match the opencv code that uses diameter. [#7545]
-  
+
 other
 -----
 
@@ -127,7 +135,7 @@ cube_build
 - Windows: MSVC: Allocate ``wave_slice_dq`` array using ``mem_alloc_dq()`` [#7491]
 
 - Memory improvements, do not allow MIRI and 'internal_cal', allow user to set suffix. [#7521]
-  
+
 datamodels
 ----------
 

jwst/tweakreg/tests/test_utils.py

@@ -1,20 +1,135 @@
 from os import path
+from copy import deepcopy
 
+import pytest
 import numpy as np
 import asdf
 from astropy.coordinates import SkyCoord
 from astropy import units as u
+from tweakwcs.correctors import JWSTWCSCorrector
+from tweakwcs.linearfit import build_fit_matrix
 
 from jwst.tweakreg.tests import data
-from jwst.tweakreg.utils import adjust_wcs
+from jwst.datamodels import ImageModel
+from jwst.tweakreg.utils import (
+    adjust_wcs,
+    transfer_wcs_correction,
+    _wcsinfo_from_wcs_transform
+)
 
 
 data_path = path.split(path.abspath(data.__file__))[0]
 
 
+def _get_test_pts(model, npts=5):
+    assert npts >= 5
+    ysize, xsize = model.data.shape
+    x = np.empty(npts, dtype=float)
+    y = np.empty(npts, dtype=float)
+    x[:5] = [-0.499, -0.499, xsize - 0.501, xsize - 0.501, (xsize - 1.0) / 2.0]
+    y[:5] = [-0.499, ysize - 0.501, -0.499, ysize - 0.501, (ysize - 1.0) / 2.0]
+    if npts > 5:
+        npts -= 5
+        np.random.seed(0)
+        x[5:] = xsize * np.random.random(npts) - 0.5
+        y[5:] = ysize * np.random.random(npts) - 0.5
+    return x, y
+
+
+@pytest.fixture
+def nircam_rate():
+    wcs_file = path.join(data_path, 'nrcb1-wcs.asdf')
+    with asdf.open(wcs_file) as af:
+        wcs = deepcopy(af['wcs'])
+
+    xsize = 204
+    ysize = 204
+    shape = (ysize, xsize)
+    im = ImageModel(shape)
+    im.var_rnoise += 0
+
+    wcsinfo = _wcsinfo_from_wcs_transform(wcs)
+    im.meta.wcsinfo = {
+        'ctype1': 'RA---TAN',
+        'ctype2': 'DEC--TAN',
+        'v2_ref': wcsinfo['v2_ref'],
+        'v3_ref': wcsinfo['v3_ref'],
+        'roll_ref': wcsinfo['roll_ref'],
+        'ra_ref': wcsinfo['ra_ref'],
+        'dec_ref': wcsinfo['dec_ref'],
+        'v3yangle': -0.07385127,
+        'vparity': -1,
+        'wcsaxes': 2
+    }
+
+    im.meta.wcs = wcs
+
+    im.meta.instrument = {
+        'channel': 'LONG',
+        'detector': 'NRCALONG',
+        'filter': 'F444W',
+        'lamp_mode': 'NONE',
+        'module': 'A',
+        'name': 'NIRCAM',
+        'pupil': 'CLEAR'
+    }
+    im.meta.subarray = {
+        'fastaxis': -1,
+        'name': 'FULL',
+        'slowaxis': 2,
+        'xsize': xsize,
+        'xstart': 1,
+        'ysize': ysize,
+        'ystart': 1
+    }
+
+    im.meta.observation = {
+        'activity_id': '01',
+        'date': '2021-10-25',
+        'exposure_number': '00001',
+        'obs_id': 'V42424001001P0000000001101',
+        'observation_label': 'nircam_ptsrc_only',
+        'observation_number': '001',
+        'program_number': '42424',
+        'sequence_id': '1',
+        'time': '16:58:27.258',
+        'visit_group': '01',
+        'visit_id': '42424001001',
+        'visit_number': '001'
+    }
+
+    im.meta.exposure = {
+        'duration': 161.05155,
+        'end_time': 59512.70899968495,
+        'exposure_time': 150.31478,
+        'frame_time': 10.73677,
+        'group_time': 21.47354,
+        'groupgap': 1,
+        'integration_time': 150.31478,
+        'mid_time': 59512.70812980775,
+        'nframes': 1,
+        'ngroups': 7,
+        'nints': 1,
+        'nresets_at_start': 1,
+        'nresets_between_ints': 1,
+        'readpatt': 'BRIGHT1',
+        'sample_time': 10,
+        'start_time': 59512.70725993055,
+        'type': 'NRC_IMAGE'
+    }
+
+    im.meta.photometry = {
+        'pixelarea_steradians': 1e-13,
+        'pixelarea_arcsecsq': 4e-3,
+    }
+
+    return im
+
+
 def test_adjust_wcs():
     wcs_file = path.join(data_path, 'nrcb1-wcs.asdf')
-    w0 = asdf.open(wcs_file)['wcs']
+    with asdf.open(wcs_file) as af:
+        w0 = deepcopy(af['wcs'])
 
     crval20, crval10 = w0.pipeline[-2].transform.angles_3.value.tolist()[-2:]
     crval10 = -crval10
@@ -59,3 +174,54 @@ def test_adjust_wcs():
         rtol=1.0e-5,
         atol=0.0
     )
+
+
+@pytest.mark.parametrize("input_type", ["wcs", "model", "file"])
+def test_transfer_wcs_correction(nircam_rate, tmp_path, input_type):
+    m1 = nircam_rate.copy()
+    m2 = nircam_rate.copy()
+
+    # apply some correction to m2 (from_image)
+    wcsinfo = _wcsinfo_from_wcs_transform(m2.meta.wcs)
+    corr = JWSTWCSCorrector(m2.meta.wcs, wcsinfo)
+    mat = build_fit_matrix(23, 1.0001456)
+    shift = corr.tanp_center_pixel_scale * np.array([3.5, 2.0])
+    corr.set_correction(matrix=mat, shift=shift)
+    m2.meta.wcs = corr.wcs
+
+    if input_type == "file":
+        # write to files:
+        model_file1 = str(tmp_path / 'model1.fits')
+        model_file2 = str(tmp_path / 'model2.fits')
+
+        m1.write(model_file1)
+        m2.write(model_file2)
+        m1 = model_file1
+        m2 = model_file2
+        from_model = model_file2
+
+    elif input_type == "wcs":
+        from_model = m2.meta.wcs
+
+    else:
+        from_model = m2
+
+    # transfer correction to m1:
+    transfer_wcs_correction(m1, from_model)
+
+    # test:
+    if input_type == "file":
+        m1 = ImageModel(m1)
+        m2 = ImageModel(m2)
+
+    x, y = _get_test_pts(m1, 5)
+
+    assert np.allclose(
+        np.linalg.norm(
+            np.subtract(m2.meta.wcs(x, y), m1.meta.wcs(x, y)),
+            axis=0
+        ),
+        0,
+        rtol=0,
+        atol=5e-11
+    )