spacetelescope · melanieclarke · Sep 5, 2024 · Sep 4, 2024 · Sep 4, 2024 · Sep 4, 2024
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -6,6 +6,9 @@ align_refs
 
 - Compute alignment shifts from the first integration of the science exposure only. [#8643]
    
+- Fixed a bug where the aligned PSF was being saved as a 4-D array with the first axis being
+  identical; now it is saved as a 3-D array. [#8747]
+
 ami_average
 -----------
 
@@ -51,6 +54,12 @@ general
 
 - Increase minimum required stpipe. [#8713]
 
+klip
+----
+
+- Allowed klip to ingest a single shifted 3-D PSF model instead of a 4-D structure
+  containing one shifted PSF per science integration. [#8747]
+
 master_background
 -----------------
 

diff --git a/docs/jwst/pipeline/calwebb_coron3.rst b/docs/jwst/pipeline/calwebb_coron3.rst
@@ -32,7 +32,7 @@ The high-level processing provided by these steps is:
 
 1) CR-flag all PSF and science target exposures
 2) Accumulate all reference PSF images into a single product
-3) Align every PSF image to every science target image
+3) Align every PSF image to the first science integration image
 4) Compute an optimal PSF fit and subtract it from every science target image
 5) Combine the PSF-subtracted and CR-flagged images into a single resampled image
 
@@ -115,17 +115,21 @@ steps. The stacked PSF data get written to disk in the form of a "_psfstack" pro
 The output file name is source-based, using the product name specified in the
 ASN file, e.g. "jw86073-a3001_t001_nircam_f140m-maskbar_psfstack.fits."
 
-4D aligned PSF images
+3D aligned PSF images
 ^^^^^^^^^^^^^^^^^^^^^
 
-:Data model: `~jwst.datamodels.QuadModel`
+:Data model: `~jwst.datamodels.CubeModel`
 :File suffix: _psfalign
 
-For each science target exposure, all of the reference PSF images in the
-"_psfstack" product are aligned to each science target integration and saved to
-a 4D "_psfalign" product by the :ref:`align_refs <align_refs_step>` step. The output file
+All of the reference PSF images in the
+"_psfstack" product are aligned to the first science target integration and saved to
+a 3D "_psfalign" product by the :ref:`align_refs <align_refs_step>` step. The output file
 name is exposure-based, with the addition of the associated candidate ID, e.g.
-"jw8607342001_02102_00001_nrcb3_a3001_psfalign.fits."
+"jw8607342001_02102_00001_nrcb3_a3001_psfalign.fits." In-flight data show that
+neither line-of-sight pointing drift nor PSF property changes are significant over the
+timescales relevant to JWST coronagraphic observations; therefore it's sufficient to align
+the PSF images to the first science target integration rather than separately for each science
+target integration as was done in the pre-flight pipeline.
 
 3D PSF-subtracted images
 ^^^^^^^^^^^^^^^^^^^^^^^^

diff --git a/jwst/coron/imageregistration.py b/jwst/coron/imageregistration.py
@@ -3,7 +3,7 @@
 from scipy import optimize
 from scipy.ndimage import fourier_shift
 
-from stdatamodels.jwst.datamodels import QuadModel
+from stdatamodels.jwst.datamodels import CubeModel
 
 import logging
 log = logging.getLogger(__name__)
@@ -212,13 +212,10 @@ def align_models(reference, target, mask):
 
     """
 
-    # Get the number of integrations in the science exposure
-    nrefslices = reference.data.shape[0]
-
-    # Create output QuadModel of required dimensions
-    quad_shape = (nrefslices,
-                  target.shape[0], target.shape[1], target.shape[2])
-    output_model = QuadModel(quad_shape)
+    # Create output CubeModel of required dimensions. Since all science integrations
+    # are assumed to have the same shift, the output is just a shifted copy of the
+    # 3-D PSF cube
+    output_model = CubeModel(target.shape)
     output_model.update(target)
 
     # Compute the shifts of the PSF ("target") images relative to
@@ -228,21 +225,18 @@ def align_models(reference, target, mask):
         target.data.astype(np.float64),
         mask=mask.data, return_aligned=False)
 
-    # Loop over all integrations of the science exposure
-    for k in range(nrefslices):
+    # Apply the shifts to the PSF images
+    output_model.data = fourier_imshift(
+        target.data.astype(np.float64),
+        -shifts)
 
-        # Apply the shifts to the PSF images
-        output_model.data[k] = fourier_imshift(
-            target.data.astype(np.float64),
+    # Apply the same shifts to the PSF error arrays, if they exist
+    if target.err is not None:
+        output_model.err = fourier_imshift(
+            target.err.astype(np.float64),
             -shifts)
 
-        # Apply the same shifts to the PSF error arrays, if they exist
-        if target.err is not None:
-            output_model.err[k] = fourier_imshift(
-                target.err.astype(np.float64),
-                -shifts)
-
-        # 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).
+    # 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
diff --git a/jwst/coron/klip.py b/jwst/coron/klip.py
@@ -22,14 +22,10 @@ def klip(target_model, refs_model, truncate):
         a single exposure are stacked along the first (NINTS) axis of
         the data arrays.
 
-    refs_model : QuadModel (NTARG x NINTS x NROWS x NCOLS)
-        The input 4D stack of reference images. The first (NTARG) axis
-        corresponds to the index of each target integration. The second
-        (NINTS) axis is the stack of aligned PSF integrations for that
-        target image. The length of the target_model first (NINTS) axis
-        should be equal to the length of the refs_model first (NTARG)
-        axis (i.e. one stack of aligned PSF images for each target
-        integration).
+    refs_model : CubeModel (NINTS_PSF x NROWS x NCOLS)
+        The input 3D stack of reference images. The first
+        (NINTS_PSF) axis is the stack of aligned PSF integrations for that
+        target image.
 
     truncate : int
         Indicates how many rows to keep in the Karhunen-Loeve transform.
@@ -48,7 +44,7 @@ def klip(target_model, refs_model, truncate):
         target = target.reshape(-1)
 
         # Load the reference psf arrays and flatten them from 3-D to 2-D
-        refs = refs_model.data[i].astype(np.float64)
+        refs = refs_model.data.astype(np.float64)
         rshape = refs.shape
         nrefs = rshape[0]
         refs = refs.reshape(nrefs, rshape[1] * rshape[2])

diff --git a/jwst/coron/tests/test_coron.py b/jwst/coron/tests/test_coron.py
@@ -178,14 +178,10 @@ def test_align_models():
     ref_mod = datamodels.CubeModel(data=ref)
 
     am_results = imageregistration.align_models(ref_mod, targ_mod, mask_mod)
-    results_sub = am_results.data[:3, :2, 2, :3]
+    results_sub = am_results.data[:2, 2, :3]
 
     truth_results_sub = np.array(
-        [
-            [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
-            [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
             [[10.0, 11.7, 12.0], [10.036278, 11.138131, 10.180669]],
-        ]
     )
 
     npt.assert_allclose(results_sub, truth_results_sub, atol=1e-6)
@@ -313,7 +309,6 @@ def test_klip():
     target_model = datamodels.CubeModel(data=target_model_data)
 
     refs_model_data = np.array(
-        [
             [
                 [
                     [0.8174741, 0.74938107, 0.73527235, 1.3193785],
@@ -331,28 +326,10 @@ def test_klip():
                     [1.419994, 1.1546139, 0.961317, 0.95088667],
                 ],
             ],
-            [
-                [
-                    [0.8174741, 0.74938107, 0.73527235, 1.3193785],
-                    [1.0032778, 0.8247719, 0.78944355, 0.99227476],
-                    [1.4609907, 1.1605016, 0.9564753, 0.9186427],
-                ],
-                [
-                    [0.86789674, 0.7998908, 0.8557136, 1.2926395],
-                    [0.97756547, 0.7788742, 0.75892323, 0.9819151],
-                    [1.495664, 1.1455023, 1.002115, 0.92159164],
-                ],
-                [
-                    [0.84426856, 0.7719569, 0.8088021, 1.2781427],
-                    [0.98734635, 0.8125992, 0.77424014, 0.9934157],
-                    [1.419994, 1.1546139, 0.961317, 0.95088667],
-                ],
-            ],
-        ],
         dtype=np.float32,
     )
 
-    refs_model = datamodels.QuadModel(data=refs_model_data)
+    refs_model = datamodels.CubeModel(data=refs_model_data)
 
     # Call the KLIP routine
     truncate = 50