JP-3784: Fix off-nominal pixel_replace and cube_build output names (#9019)

Author: melanieclarke

changes/9019.pixel_replace.rst

@@ -0,0 +1 @@
+Base output filenames on input filenames when the step is called outside the pipeline, in order to create sensible names when the input is a list of models.

jwst/cube_build/data_types.py

@@ -98,7 +98,6 @@ def __init__(self, input, single, output_file, output_dir):
         # Suffixes will be added to this name later, to designate the
         # channel+subchannel (MIRI MRS) or grating+filter (NRS IFU) the output cube covers.
 
-
         if output_file is not None:
             basename, ext = os.path.splitext(os.path.basename(output_file))
             self.output_name = basename

jwst/cube_build/ifu_cube.py

@@ -53,6 +53,7 @@ def __init__(self,
 
         self.input_models = input_models  # needed when building single mode IFU cubes
         self.output_name_base = output_name_base
+        
         self.num_files = None
 
         self.instrument = instrument
@@ -168,7 +169,7 @@ def define_cubename(self):
         """ Define the base output name
         """
         if self.pipeline == 2:
-            newname = self.output_name_base + self.suffix + '.fits'
+            newname = self.output_name_base + '_' + self.suffix + '.fits'
         else:
             if self.instrument == 'MIRI':
 

jwst/cube_build/tests/test_cube_build_step.py

@@ -3,6 +3,7 @@
 """
 
 import numpy as np
+import os
 import pytest
 from astropy.io import fits
 
@@ -13,6 +14,7 @@
 from jwst.cube_build import CubeBuildStep
 from jwst.cube_build.file_table import ErrorNoAssignWCS
 from jwst.cube_build.cube_build import ErrorNoChannels
+from jwst.datamodels import ModelContainer
 
 
 @pytest.fixture(scope='module')
@@ -158,7 +160,7 @@ def nirspec_data():
     image.meta.instrument.name = 'NIRSPEC'
     image.meta.instrument.detector = 'NRS1'
     image.meta.exposure.type = 'NRS_IFU'
-    image.meta.filename = 'test_nirspec.fits'
+    image.meta.filename = 'test_nirspec_cal.fits'
     image.meta.observation.date = '2023-10-06'
     image.meta.observation.time = '00:00:00.000'
     # below values taken from regtest using file
@@ -182,12 +184,42 @@ def nirspec_data():
 @pytest.mark.parametrize("as_filename", [True, False])
 def test_call_cube_build_nirspec(tmp_cwd, nirspec_data, tmp_path, as_filename):
     if as_filename:
-        fn = tmp_path / 'test_nirspec.fits'
+        fn = tmp_path / 'test_nirspec_cal.fits'
         nirspec_data.save(fn)
         step_input = fn
     else:
         step_input = nirspec_data
     step = CubeBuildStep()
     step.channel = '1'
     step.coord_system = 'internal_cal'
-    step.run(step_input)
+    step.save_results = True
+    result = step.run(step_input)
+
+    assert isinstance(result, ModelContainer)
+    assert len(result) == 1
+    model = result[0]
+    assert model.meta.cal_step.cube_build == 'COMPLETE'
+    assert model.meta.filename == 'test_nirspec_g395h-f290lp_internal_s3d.fits'
+    assert os.path.isfile(model.meta.filename)
+
+
+@pytest.mark.parametrize("as_filename", [True, False])
+def test_call_cube_build_nirspec_multi(tmp_cwd, nirspec_data, tmp_path, as_filename):
+    if as_filename:
+        fn = tmp_path / 'test_nirspec_cal.fits'
+        nirspec_data.save(fn)
+        step_input = fn
+    else:
+        step_input = nirspec_data
+    step = CubeBuildStep()
+    step.channel = '1'
+    step.coord_system = 'internal_cal'
+    step.save_results = True
+    step.output_type = 'multi'
+    result = step.run(step_input)
+
+    assert isinstance(result, ModelContainer)
+    assert len(result) == 1
+    model = result[0]
+    assert model.meta.cal_step.cube_build == 'COMPLETE'
+    assert model.meta.filename == 'test_nirspec_s3d.fits'

jwst/pipeline/calwebb_spec3.py

@@ -1,13 +1,12 @@
 #!/usr/bin/env python
 from collections import defaultdict
-from functools import wraps
 import os.path as op
 
 from stdatamodels.jwst import datamodels
 
 from jwst.datamodels import SourceModelContainer
 from jwst.stpipe import query_step_status
-
+from jwst.stpipe.utilities import invariant_filename
 from ..associations.lib.rules_level3_base import format_product
 from ..exp_to_source import multislit_to_container
 from ..master_background.master_background_step import split_container
@@ -350,25 +349,3 @@ def _create_nrsmos_source_id(self, source_models):
             srcid = f's{str(source_id):>09s}'
 
         return srcid
-
-# #########
-# Utilities
-# #########
-def invariant_filename(save_model_func):
-    """Restore meta.filename after save_model"""
-
-    @wraps(save_model_func)
-    def save_model(model, **kwargs):
-        try:
-            filename = model.meta.filename
-        except AttributeError:
-            filename = None
-
-        result = save_model_func(model, **kwargs)
-
-        if filename:
-            model.meta.filename = filename
-
-        return result
-
-    return save_model

jwst/pixel_replace/pixel_replace.py

@@ -29,8 +29,6 @@ class PixelReplacement:
     VERTICAL = 2
     LOG_SLICE = ['column', 'row']
 
-    default_suffix = 'pixrep'
-
     def __init__(self, input_model, **pars):
         """
         Initialize the class with input data model.

jwst/pixel_replace/pixel_replace_step.py

@@ -1,10 +1,8 @@
-#! /usr/bin/env python
 from functools import partial
 
-from ..stpipe import Step
-from jwst.stpipe import record_step_status
 from jwst import datamodels
-from .pixel_replace import PixelReplacement
+from jwst.pixel_replace.pixel_replace import PixelReplacement
+from jwst.stpipe import record_step_status, Step
 
 __all__ = ["PixelReplaceStep"]
 
@@ -33,6 +31,7 @@ class PixelReplaceStep(Step):
         algorithm = option("fit_profile", "mingrad", "N/A", default="fit_profile")
         n_adjacent_cols = integer(default=3)    # Number of adjacent columns to use in creation of profile
         skip = boolean(default=True) # Step must be turned on by parameter reference or user
+        output_use_model = boolean(default=True) # Use input filenames in the output models
     """
 
     def process(self, input):
@@ -54,7 +53,7 @@ def process(self, input):
 
             if isinstance(input_model, (datamodels.MultiSlitModel,
                                         datamodels.SlitModel,
-                                        datamodels.ImageModel, 
+                                        datamodels.ImageModel,
                                         datamodels.IFUImageModel,
                                         datamodels.CubeModel)):
                 self.log.debug(f'Input is a {input_model.meta.model_type}.')
@@ -72,12 +71,12 @@ def process(self, input):
                 'n_adjacent_cols': self.n_adjacent_cols,
             }
 
-            # ___________________
-            # calewbb_spec3 case
-            # ___________________
+            # calwebb_spec3 case / ModelContainer
             if isinstance(input_model, datamodels.ModelContainer):
                 output_model = input_model
-                # Setup output path naming if associations are involved.
+
+                # Set up output path name to include the ASN ID
+                # if associations are involved
                 asn_id = None
                 try:
                     asn_id = input_model.asn_table["asn_id"]
@@ -93,6 +92,7 @@ def process(self, input):
                         _make_output_path,
                         asn_id=asn_id
                     )
+
                 # Check models to confirm they are the correct type
                 for i, model in enumerate(output_model):
                     run_pixel_replace = True
@@ -112,10 +112,10 @@ def process(self, input):
                         replacement.replace()
                         output_model[i] = replacement.output
                         record_step_status(output_model[i], 'pixel_replace', success=True)
+
                 return output_model
-            # ________________________________________
-            # calewbb_spec2 case - single input model
-            # ________________________________________
+
+            # calwebb_spec2 case / single input model
             else:
                 # Make copy of input to prevent overwriting
                 result = input_model.copy()

jwst/pixel_replace/tests/test_pixel_replace.py

@@ -1,12 +1,15 @@
+import os
 import numpy as np
 import pytest
 
 from stdatamodels.jwst import datamodels
+from jwst.datamodels import ModelContainer
 from stdatamodels.jwst.datamodels.dqflags import pixel as flags
 
 from jwst.assign_wcs import AssignWcsStep
 from jwst.assign_wcs.tests.test_nirspec import create_nirspec_ifu_file
 from jwst.pixel_replace.pixel_replace_step import PixelReplaceStep
+from glob import glob
 
 
 def cal_data(shape, bad_idx, dispaxis=1, model='slit'):
@@ -99,6 +102,7 @@ def nirspec_ifu():
     model.var_poisson = test_data.var_poisson
     model.var_rnoise = test_data.var_rnoise
     model.var_flat = test_data.var_flat
+
     test_data.close()
 
     return model, bad_idx
@@ -199,10 +203,9 @@ def test_pixel_replace_multislit(input_model_function, algorithm):
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize('input_model_function',
-                         [nirspec_ifu])
+@pytest.mark.parametrize('input_model_function', [nirspec_ifu])
 @pytest.mark.parametrize('algorithm', ['fit_profile', 'mingrad'])
-def test_pixel_replace_nirspec_ifu(input_model_function, algorithm):
+def test_pixel_replace_nirspec_ifu(tmp_cwd, input_model_function, algorithm):
     """
     Test pixel replacement for NIRSpec IFU.
 
@@ -212,10 +215,16 @@ def test_pixel_replace_nirspec_ifu(input_model_function, algorithm):
     The test is otherwise the same as for other modes.
     """
     input_model, bad_idx = input_model_function()
+    input_model.meta.filename = 'jwst_nirspec_cal.fits'
 
     # for this simple case, the results from either algorithm should
     # be the same
-    result = PixelReplaceStep.call(input_model, skip=False, algorithm=algorithm)
+    result = PixelReplaceStep.call(input_model, skip=False,
+                                   algorithm=algorithm, save_results=True)
+
+    assert result.meta.filename == 'jwst_nirspec_pixelreplacestep.fits'
+    assert result.meta.cal_step.pixel_replace == 'COMPLETE'
+    assert os.path.isfile(result.meta.filename)
 
     for ext in ['data', 'err', 'var_poisson', 'var_rnoise', 'var_flat']:
         # non-science edges are uncorrected
@@ -238,3 +247,33 @@ def test_pixel_replace_nirspec_ifu(input_model_function, algorithm):
 
     result.close()
     input_model.close()
+
+
+@pytest.mark.parametrize('input_model_function', [nirspec_fs_slitmodel])
+def test_pixel_replace_container_names(tmp_cwd, input_model_function):
+    """Test pixel replace output names for input container."""
+    input_model, _ = input_model_function()
+    input_model.meta.filename = 'jwst_nirspec_1_cal.fits'
+    input_model2, _ = input_model_function()
+    input_model2.meta.filename = 'jwst_nirspec_2_cal.fits'
+    cfiles = [input_model, input_model2]
+    container = ModelContainer(cfiles) 
+
+    expected_name = ['jwst_nirspec_1_pixelreplacestep.fits',
+                     'jwst_nirspec_2_pixelreplacestep.fits']
+
+    result = PixelReplaceStep.call(container, skip=False, save_results=True)
+    for i, model in enumerate(result):
+        assert model.meta.filename == expected_name[i]
+        assert model.meta.cal_step.pixel_replace == 'COMPLETE'
+
+    result_files = glob(os.path.join(tmp_cwd, '*pixelreplacestep.fits'))
+    for i, file in enumerate(sorted(result_files)):
+        basename = os.path.basename(file)
+        assert expected_name[i] == basename
+        with datamodels.open(file) as model:
+            assert model.meta.cal_step.pixel_replace == 'COMPLETE'
+            assert model.meta.filename == expected_name[i]
+
+    result.close()
+    input_model.close()

jwst/stpipe/tests/test_utilities.py

@@ -5,6 +5,9 @@
 from jwst.stpipe.utilities import all_steps, NOT_SET
 import jwst.pipeline
 import jwst.step
+from stdatamodels.jwst import datamodels
+from jwst.stpipe.utilities import invariant_filename
+
 from jwst import datamodels as dm
 
 
@@ -54,3 +57,31 @@ def test_record_query_step_status():
     # test query not set
     model3 = dm.MultiSpecModel()
     assert query_step_status(model3, 'test_step') == NOT_SET
+
+
+def change_name_func(model):
+    model.meta.filename = "changed"
+    model.meta.cal_step.pixel_replace = "COMPLETE"
+    return model
+
+
+def test_invariant_filename():
+    # Make sure the change_name_func changes the name and has side effects
+    # (here, setting a status variable, but normally, actually saving the file)
+    input_model = datamodels.IFUImageModel()
+    input_model.meta.filename = 'test1.fits'
+    change_name_func(input_model)
+    assert input_model.meta.filename == 'changed'
+    assert input_model.meta.cal_step.pixel_replace == 'COMPLETE'
+
+    # When the function is wrapped with invariant_filename,
+    # the filename is not changed, but the side effect still happens
+    input_model = datamodels.IFUImageModel()
+    input_model.meta.filename = 'test2.fits'
+    invariant_save_func = invariant_filename(change_name_func)
+    output_model = invariant_save_func(input_model)
+    assert output_model.meta.filename == "test2.fits"
+    assert output_model.meta.cal_step.pixel_replace == 'COMPLETE'
+
+    # The output model is not a copy - the name is reset in place
+    assert output_model is input_model

jwst/stpipe/utilities.py

@@ -30,12 +30,14 @@
 Utilities
 """
 import importlib.util
-from importlib import import_module
 import inspect
 import logging
 import os
 import re
 from collections.abc import Sequence
+from functools import wraps
+from importlib import import_module
+
 from jwst import datamodels
 
 # Configure logging
@@ -211,3 +213,23 @@ def query_step_status(datamodel, cal_step):
         return getattr(datamodel[0].meta.cal_step, cal_step, NOT_SET)
     else:
         return getattr(datamodel.meta.cal_step, cal_step, NOT_SET)
+
+
+def invariant_filename(save_model_func):
+    """Restore meta.filename after save_model"""
+
+    @wraps(save_model_func)
+    def save_model(model, **kwargs):
+        try:
+            filename = model.meta.filename
+        except AttributeError:
+            filename = None
+
+        result = save_model_func(model, **kwargs)
+
+        if filename:
+            model.meta.filename = filename
+
+        return result
+
+    return save_model