Repository: spacetelescope/jwst

Issue:
All frames of the psfalign QuadModel are identical
_Issue [JP-3727](https://jira.stsci.edu/browse/JP-3727) was created on JIRA by [Ned Molter](https://jira.stsci.edu/secure/ViewProfile.jspa?name=emolter):_

[This Pull Request](https://github.com/spacetelescope/jwst/pull/8643) implemented a change that simplified the pipeline's coronagraphic data processing such that the alignment of the PSF to the science data is only done on the first science integration, then assumed to be the same for all subsequent integrations.  Prior to this change, the shifted PSF was stored in a 4-dimensional data structure (QuadModel of shape n_sci_integrations, n_psf_integrations, n_rows, n_cols), which made sense when the shift was different on each science integration.  Now, though, this model is identical over the zeroth axis and therefore redundant with itself.  Furthermore, this model uses a lot of memory (see [JP-3724](https://jira.stsci.edu/browse/JP-3724)), and is likely a (the?) cause of the out-of-memory failures reported on multiple datasets in [JWSTDMS-921](https://jira.stsci.edu/browse/JWSTDMS-921) (see comments on that ticket for additional datasets).  The psfalign product should be stored as a CubeModel instead.

Discussion:

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index cb5c7ea9d..fbec6e44c 100644
--- 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
index 54965aa76..e9b5efb48 100644
--- 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
index 38eb4c068..12fc18527 100644
--- a/jwst/coron/imageregistration.py
+++ b/jwst/coron/imageregistration.py
@@ -3,7 +3,7 @@ import numpy as np
 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
index ae6aab4d8..e99efc8cb 100644
--- 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])


Test patch:
diff --git a/jwst/coron/tests/test_coron.py b/jwst/coron/tests/test_coron.py
index c33b56fa6..b113e9bb0 100644
--- 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