rcal-697 Update exposure pipeline to add tweakreg step (spacetelescop…

…e#960)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CHANGES.rst

@@ -14,6 +14,8 @@ dark
 general
 -------
 
+- Update pipeline code to run through tweakreg with single files and associations [#960]
+
 - Update regression tests with new data and update ramp fitting tests to use ols_cas22 [#911]
 
 - Fix bug with ``ModelContainer.get_crds_parameters`` being a property not a method [#846]
@@ -42,7 +44,6 @@ ramp_fitting
 
 - Make uneven ramp fitting the default [#877]
 
-
 - Update Ramp fitting code to support the ``stcal`` changes to the ramp fitting
   interface which were necessary to support jump detection on uneven ramps [#933]
 

romancal/pipeline/exposure_pipeline.py

@@ -9,9 +9,11 @@
 
 # step imports
 from romancal.assign_wcs import AssignWcsStep
-from romancal.associations.exceptions import AssociationNotValidError
-from romancal.associations.load_as_asn import LoadAsLevel2Asn
+
+# from romancal.associations.exceptions import AssociationNotValidError
+# from romancal.associations.load_as_asn import LoadAsLevel2Asn
 from romancal.dark_current import DarkCurrentStep
+from romancal.datamodels import ModelContainer
 from romancal.dq_init import dq_init_step
 from romancal.flatfield import FlatFieldStep
 from romancal.jump import jump_step
@@ -23,6 +25,7 @@
 from romancal.refpix import RefPixStep
 from romancal.saturation import SaturationStep
 from romancal.source_detection import SourceDetectionStep
+from romancal.tweakreg import TweakRegStep
 
 from ..stpipe import RomanPipeline
 
@@ -62,6 +65,7 @@ class ExposurePipeline(RomanPipeline):
         "flatfield": FlatFieldStep,
         "photom": PhotomStep,
         "source_detection": SourceDetectionStep,
+        "tweakreg": TweakRegStep,
     }
 
     # start the actual processing
@@ -85,22 +89,21 @@ def process(self, input):
                 return
 
         if file_type == "asn":
-            try:
-                asn = LoadAsLevel2Asn.load(input, basename=self.output_file)
-            except AssociationNotValidError:
-                log.debug("Error opening file:")
-                return
+            asn = ModelContainer.read_asn(input)
 
         # Build a list of observations to process
         expos_file = []
+        n_members = 0
         if file_type == "asdf":
             expos_file = [input]
         elif file_type == "asn":
             for product in asn["products"]:
+                n_members = len(product["members"])
                 for member in product["members"]:
                     expos_file.append(member["expname"])
 
         results = []
+        tweakreg_input = ModelContainer()
         for in_file in expos_file:
             if isinstance(in_file, str):
                 input_filename = basename(in_file)
@@ -117,7 +120,6 @@ def process(self, input):
             if input_filename:
                 result.meta.filename = input_filename
             result = self.saturation(result)
-            # pdb.set_trace()
 
             # Test for fully saturated data
             if is_fully_saturated(result):
@@ -149,29 +151,45 @@ def process(self, input):
             result = self.dark_current(result)
             result = self.jump(result)
             result = self.rampfit(result)
-
             result = self.assign_wcs(result)
-            if result.meta.exposure.type == "WFI_IMAGE":
-                result = self.flatfield(result)
-            else:
-                log.info("Flat Field step is being SKIPPED")
-                result.meta.cal_step.flat_field = "SKIPPED"
 
             if result.meta.exposure.type == "WFI_IMAGE":
+                result = self.flatfield(result)
                 result = self.photom(result)
                 result = self.source_detection(result)
+                if file_type == "asn":
+                    tweakreg_input.append(result)
+                    log.info(
+                        f"Number of models to tweakreg:   {len(tweakreg_input._models), n_members}"
+                    )
             else:
-                log.info("Photom and source detection steps are being SKIPPED")
+                log.info("Flat Field step is being SKIPPED")
+                log.info("Photom step is being SKIPPED")
+                log.info("Source Detection step is being SKIPPED")
+                log.info("Tweakreg step is being SKIPPED")
+                result.meta.cal_step.flat_field = "SKIPPED"
                 result.meta.cal_step.photom = "SKIPPED"
                 result.meta.cal_step.source_detection = "SKIPPED"
+                result.meta.cal_step.tweakreg = "SKIPPED"
 
             # setup output_file for saving
             self.setup_output(result)
-            log.info("Roman exposure calibration pipeline ending...")
 
             self.output_use_model = True
             results.append(result)
 
+        # Now that all the exposures are collated, run tweakreg
+        # Note: this does not cover the case where the asn mixes imaging and spectral
+        #          observations. This should not occur on-prem
+        if result.meta.exposure.type == "WFI_IMAGE":
+            if file_type == "asdf":
+                mc_result = self.tweakreg([result])
+                result = mc_result._models.pop()
+            if file_type == "asn":
+                result = self.tweakreg(tweakreg_input)
+
+            log.info("Roman exposure calibration pipeline ending...")
+
         return results
 
     def setup_output(self, input):

romancal/regtest/test_wfi_pipeline.py

@@ -258,9 +258,9 @@ def test_level2_image_processing_pipeline(rtdata, ignore_asdf_paths):
 
 @pytest.mark.bigdata
 @pytest.mark.soctests
-@metrics_logger("DMS90", "DMS91", "DMS9")
+@metrics_logger("DMS278", "DMS90", "DMS91", "DMS9")
 def test_level2_grism_processing_pipeline(rtdata, ignore_asdf_paths):
-    """Tests for flat field grism processing requirements DMS90 & DMS 91"""
+    """Tests for flat field grism processing requirements DMS90, DMS91 & DMS 278"""
     input_data = "r0000201001001001002_01101_0001_WFI01_uncal.asdf"
     rtdata.get_data(f"WFI/grism/{input_data}")
     rtdata.input = input_data
@@ -356,6 +356,33 @@ def test_level2_grism_processing_pipeline(rtdata, ignore_asdf_paths):
     )
     assert model.meta.cal_step.saturation == "COMPLETE"
 
+    pipeline.log.info(
+        "DMS278 MSG: Testing WFI Level 2 Data Generation - Spectroscopy in Level 2 spectroscopic"
+        " output......."
+        + passfail(model.meta.cal_step.dq_init == "COMPLETE")
+        + passfail(model.meta.cal_step.saturation == "COMPLETE")
+        + passfail(model.meta.cal_step.refpix == "COMPLETE")
+        + passfail(model.meta.cal_step.linearity == "COMPLETE")
+        + passfail(model.meta.cal_step.jump == "COMPLETE")
+        + passfail(model.meta.cal_step.ramp_fit == "COMPLETE")
+        + passfail(model.meta.cal_step.assign_wcs == "COMPLETE")
+        + passfail(model.meta.cal_step.flat_field == "SKIPPED")
+        + passfail(model.meta.cal_step.photom == "SKIPPED")
+        + passfail(model.meta.cal_step.source_detection == "SKIPPED")
+        + passfail(model.meta.cal_step.tweakreg == "SKIPPED")
+    )
+    assert model.meta.cal_step.dq_init == "COMPLETE"
+    assert model.meta.cal_step.saturation == "COMPLETE"
+    assert model.meta.cal_step.refpix == "COMPLETE"
+    assert model.meta.cal_step.linearity == "COMPLETE"
+    assert model.meta.cal_step.jump == "COMPLETE"
+    assert model.meta.cal_step.ramp_fit == "COMPLETE"
+    assert model.meta.cal_step.assign_wcs == "COMPLETE"
+    assert model.meta.cal_step.flat_field == "SKIPPED"
+    assert model.meta.cal_step.photom == "SKIPPED"
+    assert model.meta.cal_step.source_detection == "SKIPPED"
+    assert model.meta.cal_step.tweakreg == "SKIPPED"
+
     # DMS91 data quality tests
     pipeline.log.info(
         "DMS91 MSG: Testing existence of data quality array (dq) in Level 2"