JP-3628: Revise coron3 alignment step

CHANGES.rst

@@ -1,6 +1,10 @@
 1.15.2 (unreleased)
 ===================
 
+align_refs
+----------
+
+- Compute alignment shifts from the first integration of the science exposure only. [#8643]
 ami_average
 -----------
 

docs/jwst/align_refs/description.rst

@@ -8,14 +8,19 @@ The ``align_refs`` step is one of the coronagraphic-specific steps in the ``coro
 sub-package that is part of Stage 3 :ref:`calwebb_coron3 <calwebb_coron3>` processing.
 It computes offsets between science target
 images and reference PSF images, and shifts the PSF images into
-alignment. This is performed on a per-integration basis for both the science target
-data and the reference PSF data. Each integration contained in the stacked PSF data
+alignment. The alignment shifts are computed from the first integration and applied to all the
+subsequent ones for both the science target data and the reference PSF data.
+Each integration contained in the stacked PSF data
 (the result of the :ref:`stack_refs <stack_refs_step>`) step is
-aligned to each integration within a given science target exposure.
+aligned to the first integration within a given science target exposure.
 This results in a new product for each science target exposure that contains a stack
-of individual PSF images that have been aligned to each integration in the science
+of individual PSF images that have been aligned to the first integration in the science
 target exposure.
 
+Note that aligning to the first science integration is sufficient because flight data
+shows that there are minimal drifts during an observation in line-of-sight pointing, or in PSF
+properties.
+
 Shifts between each PSF and target image are computed using the
 ``scipy.optimize.leastsq`` function. A 2D mask, supplied via a PSFMASK reference file, 
 is used to indicate pixels to ignore when performing the minimization in the

jwst/coron/imageregistration.py

@@ -127,7 +127,7 @@
     return offset
 
 
-def align_array(reference, target, mask=None):
+def align_array(reference, target, mask=None, return_aligned=True):
     """
     Computes shifts between target image (or image "slices") and the reference
     image and re-aligns input images to the target.
@@ -160,21 +160,26 @@
 
     if len(target.shape) == 2:
         shifts = align_fourierLSQ(reference, target, mask=mask)
-        aligned = fourier_imshift(target, -shifts)
+        if return_aligned:
+            aligned = fourier_imshift(target, -shifts)
 
     elif len(target.shape) == 3:
         nslices = target.shape[0]
         shifts = np.empty((nslices, 3), dtype=float)
-        aligned = np.empty_like(target)
+        if return_aligned:
+            aligned = np.empty_like(target)
 
         for m in range(nslices):
             sh = align_fourierLSQ(reference, target[m], mask=mask)
             shifts[m, :] = sh
-            aligned[m, :, :] = fourier_imshift(target[m], -sh)
+            if return_aligned:
+                aligned[m, :, :] = fourier_imshift(target[m], -sh)
 
     else:
         raise ValueError("Input target image must be either a 2D or 3D array.")
 
+    if not return_aligned:
+        return shifts
     return aligned, shifts
 
 
@@ -216,17 +221,20 @@
     output_model = QuadModel(quad_shape)
     output_model.update(target)
 
+    # Compute the shifts of the PSF ("target") images relative to
+    # the science ("reference") image in the first integration
+    shifts = align_array(
+        reference.data[0].astype(np.float64),
+        target.data.astype(np.float64),
+        mask=mask.data, return_aligned=False)
+
     # Loop over all integrations of the science exposure
     for k in range(nrefslices):
 
-        # Compute the shifts of the PSF ("target") images relative to
-        # the science ("reference") image in this integration, and apply
-        # the shifts to the PSF images
-        d, shifts = align_array(
-            reference.data[k].astype(np.float64),
+        # Apply the shifts to the PSF images
+        output_model.data[k] = fourier_imshift(
             target.data.astype(np.float64),
-            mask.data)
-        output_model.data[k] = d
+            -shifts)
 
         # Apply the same shifts to the PSF error arrays, if they exist
         if target.err is not None:
@@ -237,5 +245,4 @@
         # TODO: in the future we need to add shifts and other info (such as
         # slice ID from the reference image to which target was aligned)
         # to output cube metadata (or property).
-
     return output_model

jwst/coron/tests/test_coron.py

@@ -183,8 +183,8 @@ def test_align_models():
     truth_results_sub = np.array(
         [
             [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
-            [[10.053974, 11.1953335, 11.993213], [10.36224, 10.805556, 10.274276]],
-            [[9.988604, 11.33026, 11.968155], [10.024722, 10.971058, 10.108071]],
+            [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
+            [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
         ]
     )