refactor. support new header blender

jwst/resample/gwcs_drizzle.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from drizzle import util
 from drizzle import cdrizzle
 from . import resample_utils
 
@@ -336,9 +335,8 @@ def dodrizzle(insci, input_wcs, inwht, output_wcs, outsci, outwht, outcon,
     output input image.
 
     """
-
     # Insure that the fillval parameter gets properly interpreted for use with tdriz
-    if util.is_blank(str(fillval)):
+    if resample_utils.is_blank(str(fillval)):
         fillval = 'NAN'
     else:
         fillval = str(fillval)

jwst/resample/resample.py

@@ -9,9 +9,8 @@
 from astropy.io import fits
 
 from stdatamodels.jwst import datamodels
-from stcal.resample import LibModelAccess, Resample, OutputTooLargeError
+from stcal.resample import LibModelAccessBase, Resample, OutputTooLargeError
 
-from drizzle.resample import Drizzle
 from stdatamodels.jwst.datamodels.dqflags import pixel
 from stdatamodels.properties import ObjectNode
 
@@ -20,6 +19,7 @@
 
 from . import gwcs_drizzle
 from jwst.model_blender.blender import ModelBlender
+from jwst.assign_wcs import util
 from jwst.resample import resample_utils
 
 log = logging.getLogger(__name__)
@@ -29,12 +29,13 @@
 __all__ = [
     "ResampleData",
     "OutputTooLargeError",
-    "JWSTLibModelAccess",
-    "ResampleJWST",
+    "LibModelAccess",
+    "ResampleImage",
+    "is_imaging_wcs",
 ]
 
 
-class JWSTLibModelAccess(LibModelAccess):
+class LibModelAccess(LibModelAccessBase):
     attributes_path = {
         "data": "data",
         "dq": "dq",
@@ -64,6 +65,10 @@ class JWSTLibModelAccess(LibModelAccess):
         "weight_type": "meta.resample.weight_type",
         "pointings": "meta.resample.pointings",
         "n_coadds": "meta.resample.n_coadds",
+
+        # spectroscopy-specific:
+        "instrument_name": "meta.instrument.name",
+        "exposure_type": "meta.exposure.type",
     }
 
     def __new__(cls, *args, **kwargs):
@@ -135,12 +140,11 @@ def set_active_group(self, group_id=None):
         self._active_group = group_id
 
 
-class ResampleJWST(Resample):
+class ResampleImage(Resample):
     dq_flag_name_map = pixel
 
     def __init__(self, input_models, *args, blendheaders=True,
-                 output_model=None, output_file_name=None, in_memory=True,
-                 **kwargs):
+                 output_model=None, **kwargs):
         if output_model is None:
             self.resampled_model = datamodels.ImageModel()
             self._update_output_meta_with_first_model = True
@@ -153,11 +157,20 @@ def __init__(self, input_models, *args, blendheaders=True,
                 enable_ctx=kwargs.get("enable_ctx", True),
                 enable_var=kwargs.get("enable_var", True),
             )
-            output_model = JWSTLibModelAccess.get_model_attributes(
+            output_model = LibModelAccess.get_model_attributes(
                 output_model,
                 attributes=attributes,
             )
 
+        if blendheaders:
+            self.blender = ModelBlender(
+                blend_ignore_attrs=[
+                    'meta.photometry.pixelarea_steradians',
+                    'meta.photometry.pixelarea_arcsecsq',
+                    'meta.filename',
+                ]
+            )
+
         super().__init__(
             input_models,
             *args,
@@ -175,11 +188,15 @@ def add_model(self, model_info, image_model):
             self._update_output_meta_with_first_model = False
 
         # blend headers if needed:
-        # TODO: add blenheaders operation here
+        if self._blendheaders:
+            self.blender.accumulate(image_model)
 
     def update_output_model_data(self):
         # update data and meta for the output model:
         # * arrays:
+        if self._blendheaders:
+            self.blender.finalize_model(self.resampled_model)
+
         self.resampled_model.data = self.output_model["data"]
         self.resampled_model.wht = self.output_model["wht"]
 
@@ -197,8 +214,13 @@ def update_output_model_data(self):
         self.resampled_model.meta.cal_step.resample = 'COMPLETE'
         self.resampled_model.meta.resample.pixel_scale_ratio = self._pixel_scale_ratio
         self.resampled_model.meta.resample.pixfrac = self.pixfrac
-        _update_fits_wcsinfo(self.resampled_model)
-        util.update_s_region_imaging(self.resampled_model)
+
+        if is_imaging_wcs(self.resampled_model.meta.wcs):
+            # only for an imaging WCS:
+            self.update_fits_wcsinfo(self.resampled_model)
+            util.update_s_region_imaging(self.resampled_model)
+        else:
+            util.update_s_region_spectral(self.resampled_model)
 
         self.resampled_model.meta.asn.pool_name = self._input_models.asn.get(
             "pool_name",
@@ -231,49 +253,55 @@ def run(self):
         self.update_output_model_data()
         return self.resampled_model
 
+    @staticmethod
+    def update_fits_wcsinfo(model):
+        """ Update FITS WCS keywords of the resampled image. """
+        # Delete any SIP-related keywords first
+        pattern = r"^(cd[12]_[12]|[ab]p?_\d_\d|[ab]p?_order)$"
+        regex = re.compile(pattern)
+
+        keys = list(model.meta.wcsinfo.instance.keys())
+        for key in keys:
+            if regex.match(key):
+                del model.meta.wcsinfo.instance[key]
+
+        # Write new PC-matrix-based WCS based on GWCS model
+        transform = model.meta.wcs.forward_transform
+        model.meta.wcsinfo.crpix1 = -transform[0].offset.value + 1
+        model.meta.wcsinfo.crpix2 = -transform[1].offset.value + 1
+        model.meta.wcsinfo.cdelt1 = transform[3].factor.value
+        model.meta.wcsinfo.cdelt2 = transform[4].factor.value
+        model.meta.wcsinfo.ra_ref = transform[6].lon.value
+        model.meta.wcsinfo.dec_ref = transform[6].lat.value
+        model.meta.wcsinfo.crval1 = model.meta.wcsinfo.ra_ref
+        model.meta.wcsinfo.crval2 = model.meta.wcsinfo.dec_ref
+        model.meta.wcsinfo.pc1_1 = transform[2].matrix.value[0][0]
+        model.meta.wcsinfo.pc1_2 = transform[2].matrix.value[0][1]
+        model.meta.wcsinfo.pc2_1 = transform[2].matrix.value[1][0]
+        model.meta.wcsinfo.pc2_2 = transform[2].matrix.value[1][1]
+        model.meta.wcsinfo.ctype1 = "RA---TAN"
+        model.meta.wcsinfo.ctype2 = "DEC--TAN"
+
+        # Remove no longer relevant WCS keywords
+        rm_keys = [
+            'v2_ref',
+            'v3_ref',
+            'ra_ref',
+            'dec_ref',
+            'roll_ref',
+            'v3yangle',
+            'vparity',
+        ]
+        for key in rm_keys:
+            if key in model.meta.wcsinfo.instance:
+                del model.meta.wcsinfo.instance[key]
 
-def _update_fits_wcsinfo(model):
-    """ Update FITS WCS keywords of the resampled image. """
-    # Delete any SIP-related keywords first
-    pattern = r"^(cd[12]_[12]|[ab]p?_\d_\d|[ab]p?_order)$"
-    regex = re.compile(pattern)
-
-    keys = list(model.meta.wcsinfo.instance.keys())
-    for key in keys:
-        if regex.match(key):
-            del model.meta.wcsinfo.instance[key]
-
-    # Write new PC-matrix-based WCS based on GWCS model
-    transform = model.meta.wcs.forward_transform
-    model.meta.wcsinfo.crpix1 = -transform[0].offset.value + 1
-    model.meta.wcsinfo.crpix2 = -transform[1].offset.value + 1
-    model.meta.wcsinfo.cdelt1 = transform[3].factor.value
-    model.meta.wcsinfo.cdelt2 = transform[4].factor.value
-    model.meta.wcsinfo.ra_ref = transform[6].lon.value
-    model.meta.wcsinfo.dec_ref = transform[6].lat.value
-    model.meta.wcsinfo.crval1 = model.meta.wcsinfo.ra_ref
-    model.meta.wcsinfo.crval2 = model.meta.wcsinfo.dec_ref
-    model.meta.wcsinfo.pc1_1 = transform[2].matrix.value[0][0]
-    model.meta.wcsinfo.pc1_2 = transform[2].matrix.value[0][1]
-    model.meta.wcsinfo.pc2_1 = transform[2].matrix.value[1][0]
-    model.meta.wcsinfo.pc2_2 = transform[2].matrix.value[1][1]
-    model.meta.wcsinfo.ctype1 = "RA---TAN"
-    model.meta.wcsinfo.ctype2 = "DEC--TAN"
-
-    # Remove no longer relevant WCS keywords
-    rm_keys = [
-        'v2_ref',
-        'v3_ref',
-        'ra_ref',
-        'dec_ref',
-        'roll_ref',
-        'v3yangle',
-        'vparity',
-    ]
-    for key in rm_keys:
-        if key in model.meta.wcsinfo.instance:
-            del model.meta.wcsinfo.instance[key]
 
+def is_imaging_wcs(wcs):
+    imaging = all(
+        ax == 'SPATIAL' for ax in wcs.output_frame.axes_type
+    )
+    return imaging
 
 ####################################################
 #  Code below was left for spectral data for now   #
@@ -1005,7 +1033,7 @@ def drizzle_arrays(insci, inwht, input_wcs, output_wcs, outsci, outwht,
         """
 
         # Insure that the fillval parameter gets properly interpreted for use with tdriz
-        if util.is_blank(str(fillval)):
+        if resample_utils.is_blank(str(fillval)):
             fillval = 'NAN'
         else:
             fillval = str(fillval)

jwst/resample/resample_spec.py

@@ -17,18 +17,25 @@
 
 from jwst.assign_wcs.util import compute_scale, wrap_ra
 from jwst.resample import resample_utils
-from jwst.resample.resample import ResampleData
+from jwst.resample.resample import (
+    ResampleImage,
+    LibModelAccess,
+)
+from stcal.alignment.util import (
+    compute_scale,
+    wcs_bbox_from_shape,
+)
 
 
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 
 _S2C = SphericalToCartesian()
 
-__all__ = ["ResampleSpecData"]
+__all__ = ["ResampleSpec"]
 
 
-class ResampleSpecData(ResampleData):
+class ResampleSpec(ResampleImage):
     """
     This is the controlling routine for the resampling process for spectral data.
 
@@ -45,9 +52,12 @@ class ResampleSpecData(ResampleData):
          a record of metadata from all input models.
     """
 
-    def __init__(self, input_models, output=None, single=False, blendheaders=False,
-                 pixfrac=1.0, kernel="square", fillval=0, wht_type="ivm",
-                 good_bits=0, pscale_ratio=1.0, pscale=None, **kwargs):
+    def __init__(self, input_models, *args, output_model=None, **kwargs):
+    # def __init__(self, input_models, pixfrac=1.0, kernel="square",
+    #              fillval=0.0, wht_type="ivm", good_bits=0,
+    #              output_wcs=None, wcs_pars=None, output_model=None,
+    #              accumulate=False, enable_ctx=True, enable_var=True,
+    #              allowed_memory=None):
         """
         Parameters
         ----------
@@ -60,6 +70,30 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False,
         kwargs : dict
             Other parameters
         """
+        if output_model is None:
+            self.resampled_model = datamodels.SlitModel()
+            self._update_output_meta_with_first_model = True
+        else:
+            self.resampled_model = output_model
+            self._update_output_meta_with_first_model = False
+            # convert output_model to dictionary:
+            attributes = ResampleImage.output_model_attributes(
+                accumulate=False,
+                enable_ctx=kwargs.get("enable_ctx", True),
+                enable_var=kwargs.get("enable_var", True),
+            )
+            output_model = LibModelAccess.get_model_attributes(
+                output_model,
+                attributes=attributes,
+            )
+
+        super().__init__(
+            input_models,
+            *args,
+            output_model=output_model,
+            **kwargs
+        )
+
         self.output_dir = None
         self.output_filename = output
         if output is not None and '.fits' not in str(output):
@@ -68,8 +102,6 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False,
         self.intermediate_suffix = 'outlier_s2d'
 
         self.pscale_ratio = pscale_ratio
-        self.single = single
-        self.blendheaders = blendheaders
         self.pixfrac = pixfrac
         self.kernel = kernel
         self.fillval = fillval
@@ -441,7 +473,7 @@ def build_nirspec_output_wcs(self, input_models, refmodel=None):
 
         # Compute bounding box and output array shape.
         self.data_size = (ny, n_lam)
-        bounding_box = resample_utils.wcs_bbox_from_shape(self.data_size)
+        bounding_box = wcs_bbox_from_shape(self.data_size)
         output_wcs.bounding_box = bounding_box
         output_wcs.array_shape = self.data_size
 
@@ -728,7 +760,7 @@ def build_interpolated_output_wcs(self, input_models):
         # turn the size into a numpy shape in (y, x) order
         output_wcs.array_shape = output_array_size[::-1]
         output_wcs.pixel_shape = output_array_size
-        bounding_box = resample_utils.wcs_bbox_from_shape(output_array_size[::-1])
+        bounding_box = wcs_bbox_from_shape(output_array_size[::-1])
         output_wcs.bounding_box = bounding_box
 
         return output_wcs
@@ -837,7 +869,7 @@ def build_nirspec_lamp_output_wcs(self, input_models):
         # turn the size into a numpy shape in (y, x) order
         output_wcs.array_shape = output_array_size[::-1]
         output_wcs.pixel_shape = output_array_size
-        bounding_box = resample_utils.wcs_bbox_from_shape(output_array_size[::-1])
+        bounding_box = wcs_bbox_from_shape(output_array_size[::-1])
         output_wcs.bounding_box = bounding_box
 
         return output_wcs