JP-3784 Fix Pixel replace and cube_build output file names

changes/8979.pixel_replace.rst

@@ -0,0 +1 @@
+fix a bug in pixel_replace and cube_build output filenames

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':
 
@@ -178,9 +179,10 @@ def define_cubename(self):
                 # that the remaining suffixes created below form the entire
                 # list of optical elements in the final output name.
                 suffix = self.output_name_base[self.output_name_base.rfind('_') + 1:]
+
                 if suffix in ['clear']:
                     self.output_name_base = self.output_name_base[:self.output_name_base.rfind('_')]
-
+                    
                 # Now compose the appropriate list of optical element suffix names
                 # based on MRS channel and sub-channel
                 channels = []

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

@@ -5,6 +5,7 @@
 from jwst.stpipe import record_step_status
 from jwst import datamodels
 from .pixel_replace import PixelReplacement
+from jwst.stpipe.utilities import invariant_filename
 
 __all__ = ["PixelReplaceStep"]
 
@@ -72,19 +73,27 @@
                 'n_adjacent_cols': self.n_adjacent_cols,
             }
 
-            # ___________________
-            # calewbb_spec3 case
-            # ___________________
+            # ___________________________________
+            # calewbb_spec3 case / ModelContainer
+            # __________________________________
             if isinstance(input_model, datamodels.ModelContainer):
-                output_model = input_model
                 # Setup output path naming if associations are involved.
+                # if input is ModelContainer do not copy the input_model
+                # because assocation information is not copied.
+                # Instead update input_modelin place. 
                 asn_id = None
                 try:
                     asn_id = input_model.asn_table["asn_id"]
                 except (AttributeError, KeyError):
                     pass
                 if asn_id is None:
                     asn_id = self.search_attr('asn_id')
+                if asn_id is None: # it is still None. It does not exist. A ModelContainer was passed in with
+                                   # no asn information
+                    # to create the correct output name set the following.
+                    self.output_use_model = True
+                    self.save_model = invariant_filename(self.save_model)
+
                 if asn_id is not None:
                     _make_output_path = self.search_attr(
                         '_make_output_path', parent_first=True
@@ -93,8 +102,9 @@
                         _make_output_path,
                         asn_id=asn_id
                     )
+
                 # Check models to confirm they are the correct type
-                for i, model in enumerate(output_model):
+                for i, model in enumerate(input_model):
                     run_pixel_replace = True
                     if model.meta.model_type in ['MultiSlitModel', 'SlitModel',
                                                  'ImageModel', 'IFUImageModel', 'CubeModel']:
@@ -110,9 +120,9 @@
                     if run_pixel_replace:
                         replacement = PixelReplacement(model, **pars)
                         replacement.replace()
-                        output_model[i] = replacement.output
-                        record_step_status(output_model[i], 'pixel_replace', success=True)
-                return output_model
+                        model = replacement.output
+                        record_step_status(model, 'pixel_replace', success=True)
+                return input_model
             # ________________________________________
             # calewbb_spec2 case - single input model
             # ________________________________________

jwst/pixel_replace/tests/test_pixel_replace.py

@@ -1,13 +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'):
     if model == 'image':
@@ -99,6 +101,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
@@ -212,11 +215,15 @@ 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'
+
     for ext in ['data', 'err', 'var_poisson', 'var_rnoise', 'var_flat']:
         # non-science edges are uncorrected
         assert np.all(np.isnan(getattr(result, ext)[..., :, 1]))
@@ -238,3 +245,83 @@ def test_pixel_replace_nirspec_ifu(input_model_function, algorithm):
 
     result.close()
     input_model.close()
+
+
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize('input_model_function',
+                         [nirspec_ifu])
+@pytest.mark.parametrize('algorithm', ['fit_profile', 'mingrad'])
+def test_pixel_replace_nirspec_ifu_container_names(tmp_cwd, tmp_path, input_model_function, algorithm):
+    """
+    Test pixel replacement for NIRSpec IFU using a container
+
+    Larger data and more WCS operations required for testing make
+    this test take more than a minute, so marking this test 'slow'.
+
+    The test is otherwise the same as for other modes.
+    """
+    output_dir = tmp_path / 'output'
+    output_dir.mkdir(exist_ok=True)
+    output_dir = str(output_dir)
+
+    input_model, bad_idx = input_model_function()
+    input_model.meta.filename = 'jwst_nirspec_1_cal.fits'
+    input_model2, bad_idx2 = input_model_function()
+    input_model2.meta.filename = 'jwst_nirspec_2_cal.fits'
+    cfiles = [input_model, input_model2]
+    container = ModelContainer(cfiles) 
+
+    expected_name = []
+    expected_name.append('jwst_nirspec_1_pixelreplacestep.fits')
+    expected_name.append('jwst_nirspec_2_pixelreplacestep.fits')
+
+    return_files = []
+    # for this simple case, the results from either algorithm should
+    # be the same
+    result = PixelReplaceStep.call(container, skip=False, algorithm=algorithm,
+                                   save_results=True)
+    result_files = glob(os.path.join(tmp_cwd, '*pixelreplacestep.fits'))
+    for file in result_files:
+        basename = os.path.basename(file)
+        return_files.append(basename)
+
+    assert expected_name[0] == return_files[0]
+    assert expected_name[1] == return_files[1]
+
+    result.close()
+    input_model.close()
+
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize('input_model_function',
+                         [nirspec_ifu])
+@pytest.mark.parametrize('algorithm', ['fit_profile', 'mingrad'])
+def test_pixel_replace_nirspec_ifu_name(tmp_cwd, tmp_path, input_model_function, algorithm):
+    """
+    Test pixel replacement for NIRSpec IFU using a single file
+
+    Larger data and more WCS operations required for testing make
+    this test take more than a minute, so marking this test 'slow'.
+
+    The test is otherwise the same as for other modes.
+    """
+    output_dir = tmp_path / 'output'
+    output_dir.mkdir(exist_ok=True)
+    output_dir = str(output_dir)
+
+    input_model, bad_idx = input_model_function()
+    input_model.meta.filename = 'jwst_nirspec_cal.fits'
+    expected_name = 'jwst_nirspec_pixelreplacestep.fits'
+
+    # for this simple case, the results from either algorithm should
+    # be the same
+    result = PixelReplaceStep.call(input_model, skip=False, algorithm=algorithm,save_results=True)
+
+    assert expected_name == result.meta.filename
+
+    result.close()
+    input_model.close()
+

jwst/stpipe/tests/test_utilities.py

@@ -5,6 +5,10 @@
 from jwst.stpipe.utilities import all_steps, NOT_SET
 import jwst.pipeline
 import jwst.step
+import pytest
+from stdatamodels.jwst import datamodels
+from jwst.stpipe.utilities import invariant_filename
+
 from jwst import datamodels as dm
 
 
@@ -54,3 +58,32 @@ def test_record_query_step_status():
     # test query not set
     model3 = dm.MultiSpecModel()
     assert query_step_status(model3, 'test_step') == NOT_SET
+
+# make up a datamodel for testing filename
+@pytest.fixture(scope='function')
+def miri_ifushort():
+    """ Generate input model IFU image """
+
+
+    mirifushort_short = {
+        'detector': 'MIRIFUSHORT',
+        'channel': '12',
+        'band': 'SHORT',
+        'name': 'MIRI'
+
+    }
+    input_model = datamodels.IFUImageModel()
+    input_model.meta.exposure.type = 'MIR_MRS'
+    input_model.meta.instrument._instance.update(mirifushort_short)
+    input_model.meta.filename = 'test1.fits'
+    return input_model
+
+
+def change_name_func(model):
+    model.meta.filename = "changed"
+    return model
+
+def test_invariant_filename(miri_ifushort):
+    invariant_save_func = invariant_filename(change_name_func)
+    output_model = invariant_save_func(miri_ifushort)
+    assert output_model.meta.filename == "test1.fits"

jwst/stpipe/utilities.py

@@ -30,6 +30,7 @@
 Utilities
 """
 import importlib.util
+from functools import wraps
 from importlib import import_module
 import inspect
 import logging
@@ -211,3 +212,23 @@
         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