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

CHANGES.rst

@@ -40,6 +40,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)
 ===================
@@ -73,7 +81,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
 -----
 
@@ -125,7 +133,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,171 @@
 from os import path
+from copy import deepcopy
 
+import pytest
 import numpy as np
 import asdf
+from astropy.modeling.rotations import RotationSequence3D
 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.tweakreg.utils import adjust_wcs, transfer_wcs_correction
+from jwst.datamodels import ImageModel
 
 
 data_path = path.split(path.abspath(data.__file__))[0]
 
 
+def _wcsinfo_from_wcs_transform(wcs):
+    frms = wcs.available_frames
+    if 'v2v3' not in frms or 'world' not in frms or frms[-1] != 'world':
+        raise ValueError(
+            "Unsupported WCS structure."
+        )
+
+    # Initially get v2_ref, v3_ref, and roll_ref from
+    # the v2v3 to world transform. Also get ra_ref, dec_ref
+    t = wcs.get_transform(frms[-2], 'world')
+    for m in t:
+        if isinstance(m, RotationSequence3D) and m.parameters.size == 5:
+            v2_ref, nv3_ref, roll_ref, dec_ref, nra_ref = m.angles.value
+            break
+    else:
+        raise ValueError(
+            "Unsupported WCS structure."
+        )
+
+    # overwrite v2_ref, v3_ref, and roll_ref with
+    # values from the tangent plane when available:
+    if 'v2v3corr' in frms:
+        # get v2_ref, v3_ref, and roll_ref from
+        # the v2v3 to v2v3corr transform:
+        frm1 = 'v2v3vacorr' if 'v2v3vacorr' in frms else 'v2v3'
+        tpcorr = wcs.get_transform(frm1, 'v2v3corr')
+        v2_ref, nv3_ref, roll_ref = tpcorr['det_to_optic_axis'].angles.value
+
+    wcsinfo = {
+        'v2_ref' : 3600 * v2_ref,
+        'v3_ref' : -3600 * nv3_ref,
+        'roll_ref' : roll_ref,
+        'ra_ref' : -nra_ref,
+        'dec_ref' : dec_ref
+    }
+
+    return wcsinfo
+
+
+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 +210,70 @@ def test_adjust_wcs():
         rtol=1.0e-5,
         atol=0.0
     )
+
+
+def test_transfer_wcs_correction_model(nircam_rate):
+    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
+
+    # transfer correction to m1:
+    transfer_wcs_correction(m1, m2)
+
+    # test:
+    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=1e-11
+    )
+
+
+def test_transfer_wcs_correction_file(nircam_rate, tmp_path):
+    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
+
+    # write to files:
+    mf1 = str(tmp_path / 'model1.fits')
+    mf2 = str(tmp_path / 'model2.fits')
+
+    m1.write(mf1)
+    m2.write(mf2)
+
+    # transfer correction to mf1 file:
+    transfer_wcs_correction(mf1, mf2)
+
+    # test:
+    m1c = ImageModel(mf1)
+
+    x, y = _get_test_pts(m1, 5)
+
+    assert np.allclose(
+        np.linalg.norm(
+            np.subtract(m2.meta.wcs(x, y), m1c.meta.wcs(x, y)),
+            axis=0
+        ),
+        0,
+        rtol=0,
+        atol=1e-11
+    )