From 55f3a344d252aef6ddee54548a28227ba3ca579b Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Tue, 8 Oct 2024 11:42:26 -0400 Subject: [PATCH 01/19] Switch resample to the new drizzle package API --- jwst/resample/gwcs_drizzle.py | 9 ++++++--- jwst/resample/resample.py | 14 +++++++------- jwst/resample/resample_utils.py | 15 +++++---------- pyproject.toml | 2 +- 4 files changed, 19 insertions(+), 21 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 9221cbdf01..63f0a84ad7 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -1,7 +1,6 @@ import numpy as np -from drizzle import util -from drizzle import cdrizzle +from drizzle import cdrizzle, utils from . import resample_utils import logging @@ -338,7 +337,7 @@ def dodrizzle(insci, input_wcs, inwht, output_wcs, outsci, outwht, outcon, """ # Insure that the fillval parameter gets properly interpreted for use with tdriz - if util.is_blank(str(fillval)): + if str(fillval).strip() == '': fillval = 'NAN' else: fillval = str(fillval) @@ -381,9 +380,13 @@ def dodrizzle(insci, input_wcs, inwht, output_wcs, outsci, outwht, outcon, # Compute the mapping between the input and output pixel coordinates # for use in drizzle.cdrizzle.tdriz + + # pixmap = utils.calc_pixmap(input_wcs, output_wcs) pixmap = resample_utils.calc_gwcs_pixmap(input_wcs, output_wcs, insci.shape) + # inwht[np.isnan(pixmap[:,:,0])] = 0. + print(f"Pixmap shape: {pixmap.shape}") log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") log.debug(f"Input Sci shape: {insci.shape}") log.debug(f"Output Sci shape: {outsci.shape}") diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index b3e7e7f485..6671580d6b 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -5,7 +5,7 @@ import numpy as np import psutil -from drizzle import cdrizzle, util +from drizzle import cdrizzle from spherical_geometry.polygon import SphericalPolygon from stdatamodels.jwst import datamodels @@ -251,10 +251,10 @@ def _get_intensity_scale(self, img): else: iscale = 1.0 return iscale - + def resample_group(self, input_models, indices, compute_error=False): """Apply resample_many_to_many for one group - + Parameters ---------- input_models : ModelLibrary @@ -302,7 +302,7 @@ def resample_group(self, input_models, indices, compute_error=False): if isinstance(img, datamodels.SlitModel): # must call this explicitly to populate area extension # although the existence of this extension may not be necessary - img.area = img.area + img.area = img.area iscale = self._get_intensity_scale(img) log.debug(f'Using intensity scale iscale={iscale}') @@ -371,7 +371,7 @@ def resample_many_to_many(self, input_models): """ output_models = [] for group_id, indices in input_models.group_indices.items(): - + output_model = self.resample_group(input_models, indices) if not self.in_memory: @@ -555,7 +555,7 @@ def resample_variance_arrays(self, output_model, input_models): axis=0 ) total_weight_flat_var[mask] += weight[mask] - + del model.meta.iscale del weight input_models.shelve(model, i) @@ -786,7 +786,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 str(fillval).strip() == '': fillval = 'NAN' else: fillval = str(fillval) diff --git a/jwst/resample/resample_utils.py b/jwst/resample/resample_utils.py index c527b450f1..0ff50902de 100644 --- a/jwst/resample/resample_utils.py +++ b/jwst/resample/resample_utils.py @@ -5,6 +5,7 @@ import numpy as np from astropy import units as u import gwcs +from drizzle.utils import calc_pixmap from stdatamodels.dqflags import interpret_bit_flags from stdatamodels.jwst.datamodels.dqflags import pixel @@ -124,17 +125,11 @@ def shape_from_bounding_box(bounding_box): def calc_gwcs_pixmap(in_wcs, out_wcs, shape=None): """ Return a pixel grid map from input frame to output frame. """ - if shape: - bb = wcs_bbox_from_shape(shape) - log.debug("Bounding box from data shape: {}".format(bb)) - else: - bb = in_wcs.bounding_box - log.debug("Bounding box from WCS: {}".format(in_wcs.bounding_box)) - - grid = gwcs.wcstools.grid_from_bounding_box(bb) - pixmap = np.dstack(reproject(in_wcs, out_wcs)(grid[0], grid[1])) + if shape and not np.array_equiv(shape, in_wcs.array_shape): + in_wcs = deepcopy(in_wcs) + in_wcs.array_shape = shape - return pixmap + return calc_pixmap(wcs_from=in_wcs, wcs_to=out_wcs) def reproject(wcs1, wcs2): diff --git a/pyproject.toml b/pyproject.toml index 150bd6e376..584a785efd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "astropy>=5.3", "BayesicFitting>=3.0.1", "crds>=11.17.14", - "drizzle>=1.15.0", + "drizzle @ git+https://github.com/mcara/drizzle.git@redesign-api", "gwcs>=0.21.0,<0.23.0", "numpy>=1.22,<2.0", "opencv-python-headless>=4.6.0.66", From 9f530987fe68f37691cc561a40fae3fae35dfb96 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Tue, 8 Oct 2024 23:32:31 -0400 Subject: [PATCH 02/19] changelog; flake8 --- changes/8866.resample.rst | 1 + jwst/resample/gwcs_drizzle.py | 7 +------ jwst/resample/resample.py | 35 +++++++++++++++++++-------------- jwst/resample/resample_utils.py | 3 +-- 4 files changed, 23 insertions(+), 23 deletions(-) create mode 100644 changes/8866.resample.rst diff --git a/changes/8866.resample.rst b/changes/8866.resample.rst new file mode 100644 index 0000000000..f31024d0ee --- /dev/null +++ b/changes/8866.resample.rst @@ -0,0 +1 @@ +Updated resample code to support the new ``drizzle`` API, see https://github.com/spacetelescope/drizzle/pull/134 for more details. diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 63f0a84ad7..22450fc87e 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -1,6 +1,6 @@ import numpy as np -from drizzle import cdrizzle, utils +from drizzle import cdrizzle from . import resample_utils import logging @@ -380,13 +380,8 @@ def dodrizzle(insci, input_wcs, inwht, output_wcs, outsci, outwht, outcon, # Compute the mapping between the input and output pixel coordinates # for use in drizzle.cdrizzle.tdriz - - # pixmap = utils.calc_pixmap(input_wcs, output_wcs) pixmap = resample_utils.calc_gwcs_pixmap(input_wcs, output_wcs, insci.shape) - # inwht[np.isnan(pixmap[:,:,0])] = 0. - - print(f"Pixmap shape: {pixmap.shape}") log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") log.debug(f"Input Sci shape: {insci.shape}") log.debug(f"Output Sci shape: {outsci.shape}") diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index 6671580d6b..2c5a7c314c 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -192,7 +192,6 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, output_pix_area * np.rad2deg(3600)**2 ) - def do_drizzle(self, input_models): """Pick the correct drizzling mode based on self.single """ @@ -389,7 +388,7 @@ def resample_many_to_many(self, input_models): # build ModelLibrary as an association from the output files # this saves memory if there are multiple groups asn = asn_from_list(output_models, product_name='outlier_i2d') - asn_dict = json.loads(asn.dump()[1]) # serializes the asn and converts to dict + asn_dict = json.loads(asn.dump()[1]) # serializes the asn and converts to dict return ModelLibrary(asn_dict, on_disk=True) # otherwise just build it as a list of in-memory models @@ -430,9 +429,11 @@ def resample_many_to_one(self, input_models): log.debug(f'Using intensity scale iscale={iscale}') img.meta.iscale = iscale - inwht = resample_utils.build_driz_weight(img, - weight_type=self.weight_type, - good_bits=self.good_bits) + inwht = resample_utils.build_driz_weight( + img, + weight_type=self.weight_type, + good_bits=self.good_bits, + ) # apply sky subtraction blevel = img.meta.background.level if not img.meta.background.subtracted and blevel is not None: @@ -479,14 +480,13 @@ def resample_many_to_one(self, input_models): return ModelLibrary([output_model,], on_disk=False) - def resample_variance_arrays(self, output_model, input_models): """Resample variance arrays from input_models to the output_model. Variance images from each input model are resampled individually and - added to a weighted sum. If weight_type is 'ivm', the inverse of the + added to a weighted sum. If ``weight_type`` is 'ivm', the inverse of the resampled read noise variance is used as the weight for all the variance - components. If weight_type is 'exptime', the exposure time is used. + components. If ``weight_type`` is 'exptime', the exposure time is used. The output_model is modified in place. """ @@ -526,8 +526,10 @@ def resample_variance_arrays(self, output_model, input_models): if rn_var is not None: mask = (rn_var >= 0) & np.isfinite(rn_var) & (weight > 0) weighted_rn_var[mask] = np.nansum( - [weighted_rn_var[mask], - rn_var[mask] * weight[mask] * weight[mask]], + [ + weighted_rn_var[mask], + rn_var[mask] * weight[mask] * weight[mask] + ], axis=0 ) total_weight_rn_var[mask] += weight[mask] @@ -539,8 +541,10 @@ def resample_variance_arrays(self, output_model, input_models): if pn_var is not None: mask = (pn_var >= 0) & np.isfinite(pn_var) & (weight > 0) weighted_pn_var[mask] = np.nansum( - [weighted_pn_var[mask], - pn_var[mask] * weight[mask] * weight[mask]], + [ + weighted_pn_var[mask], + pn_var[mask] * weight[mask] * weight[mask] + ], axis=0 ) total_weight_pn_var[mask] += weight[mask] @@ -550,8 +554,10 @@ def resample_variance_arrays(self, output_model, input_models): if flat_var is not None: mask = (flat_var >= 0) & np.isfinite(flat_var) & (weight > 0) weighted_flat_var[mask] = np.nansum( - [weighted_flat_var[mask], - flat_var[mask] * weight[mask] * weight[mask]], + [ + weighted_flat_var[mask], + flat_var[mask] * weight[mask] * weight[mask] + ], axis=0 ) total_weight_flat_var[mask] += weight[mask] @@ -582,7 +588,6 @@ def resample_variance_arrays(self, output_model, input_models): del weighted_rn_var, weighted_pn_var, weighted_flat_var del total_weight_rn_var, total_weight_pn_var, total_weight_flat_var - def _resample_one_variance_array(self, name, input_model, output_model): """Resample one variance image from an input model. diff --git a/jwst/resample/resample_utils.py b/jwst/resample/resample_utils.py index 0ff50902de..3775cb5377 100644 --- a/jwst/resample/resample_utils.py +++ b/jwst/resample/resample_utils.py @@ -4,7 +4,6 @@ import numpy as np from astropy import units as u -import gwcs from drizzle.utils import calc_pixmap from stdatamodels.dqflags import interpret_bit_flags @@ -125,7 +124,7 @@ def shape_from_bounding_box(bounding_box): def calc_gwcs_pixmap(in_wcs, out_wcs, shape=None): """ Return a pixel grid map from input frame to output frame. """ - if shape and not np.array_equiv(shape, in_wcs.array_shape): + if shape is not None and not np.array_equiv(shape, in_wcs.array_shape): in_wcs = deepcopy(in_wcs) in_wcs.array_shape = shape From b658229ece1f68e03d23fb41c47c55f36c930721 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Mon, 14 Oct 2024 21:51:20 -0400 Subject: [PATCH 03/19] revert to original calc pixmap code --- jwst/resample/resample_utils.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/jwst/resample/resample_utils.py b/jwst/resample/resample_utils.py index 3775cb5377..c527b450f1 100644 --- a/jwst/resample/resample_utils.py +++ b/jwst/resample/resample_utils.py @@ -4,7 +4,7 @@ import numpy as np from astropy import units as u -from drizzle.utils import calc_pixmap +import gwcs from stdatamodels.dqflags import interpret_bit_flags from stdatamodels.jwst.datamodels.dqflags import pixel @@ -124,11 +124,17 @@ def shape_from_bounding_box(bounding_box): def calc_gwcs_pixmap(in_wcs, out_wcs, shape=None): """ Return a pixel grid map from input frame to output frame. """ - if shape is not None and not np.array_equiv(shape, in_wcs.array_shape): - in_wcs = deepcopy(in_wcs) - in_wcs.array_shape = shape + if shape: + bb = wcs_bbox_from_shape(shape) + log.debug("Bounding box from data shape: {}".format(bb)) + else: + bb = in_wcs.bounding_box + log.debug("Bounding box from WCS: {}".format(in_wcs.bounding_box)) + + grid = gwcs.wcstools.grid_from_bounding_box(bb) + pixmap = np.dstack(reproject(in_wcs, out_wcs)(grid[0], grid[1])) - return calc_pixmap(wcs_from=in_wcs, wcs_to=out_wcs) + return pixmap def reproject(wcs1, wcs2): From c5e74c4307b02b2ab71fd0f10049c43699ab7d5c Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Tue, 15 Oct 2024 21:26:20 -0400 Subject: [PATCH 04/19] Convert GWCSDrizzle to new API --- jwst/resample/gwcs_drizzle.py | 327 +++++++--------------------------- 1 file changed, 68 insertions(+), 259 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 22450fc87e..7e579ead6e 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -1,6 +1,6 @@ import numpy as np -from drizzle import cdrizzle +from drizzle.resample import Drizzle from . import resample_utils import logging @@ -8,13 +8,14 @@ log.setLevel(logging.DEBUG) -class GWCSDrizzle: +class GWCSDrizzle(Drizzle): """ Combine images using the drizzle algorithm - """ + """ def __init__(self, product, outwcs=None, wt_scl=None, - pixfrac=1.0, kernel="square", fillval="NAN"): + pixfrac=1.0, kernel="square", fillval="NAN", + disable_ctx=False, n_max_images=None): """ Create a new Drizzle output object and set the drizzle parameters. @@ -55,66 +56,68 @@ def __init__(self, product, outwcs=None, wt_scl=None, not overlap it. The default value of NAN sets NaN values. """ + # Initialize the object fields self._product = product - self.outsci = None - self.outwht = None - self.outcon = None - self.uniqid = 0 - if wt_scl is None: - self.wt_scl = "" - else: - self.wt_scl = wt_scl - self.kernel = kernel - self.fillval = fillval self.pixfrac = pixfrac self.sciext = "SCI" self.whtext = "WHT" self.conext = "CON" - out_units = "cps" + self.out_units = "cps" self.outexptime = product.meta.exposure.measurement_time or 0.0 self.outsci = product.data + self.outwht = product.wht + if outwcs: self.outwcs = outwcs else: self.outwcs = product.meta.wcs - self.outwht = product.wht - self.outcon = product.con - - if self.outcon.ndim == 2: - self.outcon = np.reshape(self.outcon, (1, - self.outcon.shape[0], - self.outcon.shape[1])) - elif self.outcon.ndim != 3: - raise ValueError("Drizzle context image has wrong dimensions: \ - {0}".format(product)) - # Check field values if not self.outwcs: raise ValueError("Either an existing file or wcs must be supplied") - if out_units == "counts": - np.divide(self.outsci, self.outexptime, self.outsci) - elif out_units != "cps": - raise ValueError("Illegal value for out_units: %s" % out_units) + if self.outexptime == 0.0: + ctx = None + begin_ctx_id = 0 + else: + ctx = product.con + if ctx.ndim == 2: + begin_ctx_id = 1 + elif ctx.ndim == 3: + begin_ctx_id = ctx.shape[0] + 1 + else: + # TODO: this message seems odd + raise ValueError( + f"Drizzle context image has wrong dimensions: {product}" + ) + + super().__init__( + kernel=kernel, + fillval=fillval, + out_shape=None, + out_img=product.data, + out_wht=product.wht, + out_ctx=ctx, + exptime=self.outexptime, + begin_ctx_id=begin_ctx_id, + max_ctx_id=n_max_images, + disable_ctx=disable_ctx + ) + + product.con = self.out_ctx # Since the context array is dynamic, it must be re-assigned # back to the product's `con` attribute. @property def outcon(self): """Return the context array""" - return self._product.con - - @outcon.setter - def outcon(self, value): - """Set new context array""" - self._product.con = value + return self.out_ctx def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, expin=1.0, in_units="cps", wt_scl=1.0, iscale=1.0): @@ -176,231 +179,37 @@ def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, or "cps" (counts per second.) If the value is counts, before using the input image it is scaled by dividing it by the exposure time. - wt_scl : float, optional - If drizzle was initialized with wt_scl left blank, this value will - set a scaling factor for the pixel weighting. If drizzle was - initialized with wt_scl set to "exptime" or "expsq", the exposure time - will be used to set the weight scaling and the value of this parameter - will be ignored. - iscale : float, optional A scale factor to be applied to pixel intensities of the input image before resampling. """ - if self.wt_scl == "exptime": - wt_scl = expin - elif self.wt_scl == "expsq": - wt_scl = expin * expin - - wt_scl = 1.0 # hard-coded for JWST count-rate data - self.increment_id() - - dodrizzle(insci, inwcs, inwht, self.outwcs, self.outsci, self.outwht, - self.outcon, expin, in_units, wt_scl, uniqid=self.uniqid, - xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, - iscale=iscale, pixfrac=self.pixfrac, kernel=self.kernel, - fillval=self.fillval) - - def increment_id(self): - """ - Increment the id count and add a plane to the context image if needed - - Drizzle tracks which input images contribute to the output image - by setting a bit in the corresponding pixel in the context image. - The uniqid indicates which bit. So it must be incremented each time - a new image is added. Each plane in the context image can hold 32 bits, - so after each 32 images, a new plane is added to the context. - """ - - # Compute what plane of the context image this input would - # correspond to: - planeid = int(self.uniqid / 32) - - # Add a new plane to the context image if planeid overflows - - if self.outcon.shape[0] == planeid: - plane = np.zeros_like(self.outcon[0]) - plane = plane.reshape((1, plane.shape[0], plane.shape[1])) - self.outcon = np.concatenate((self.outcon, plane)) - - # Increment the id - self.uniqid += 1 - - -def dodrizzle(insci, input_wcs, inwht, output_wcs, outsci, outwht, outcon, - expin, in_units, wt_scl, uniqid=1, xmin=0, xmax=0, ymin=0, ymax=0, - iscale=1.0, pixfrac=1.0, kernel='square', fillval="NAN"): - """ - Low level routine for performing 'drizzle' operation on one image. - - Parameters - ---------- - - insci : 2d array - A 2d numpy array containing the input image to be drizzled. - - input_wcs : gwcs.WCS object - The world coordinate system of the input image. - - inwht : 2d array - A 2d numpy array containing the pixel by pixel weighting. - Must have the same dimensions as insci. If none is supplied, - the weighting is set to one. - - output_wcs : gwcs.WCS object - The world coordinate system of the output image. - - outsci : 2d array - A 2d numpy array containing the output image produced by - drizzling. On the first call it should be set to zero. - Subsequent calls it will hold the intermediate results - - outwht : 2d array - A 2d numpy array containing the output counts. On the first - call it should be set to zero. On subsequent calls it will - hold the intermediate results. - - outcon : 2d or 3d array, optional - A 2d or 3d numpy array holding a bitmap of which image was an input - for each output pixel. Should be integer zero on first call. - Subsequent calls hold intermediate results. - - expin : float - The exposure time of the input image, a positive number. The - exposure time is used to scale the image if the units are counts. - - in_units : str - The units of the input image. The units can either be "counts" - or "cps" (counts per second.) - - wt_scl : float - A scaling factor applied to the pixel by pixel weighting. - - uniqid : int, optional - The id number of the input image. Should be one the first time - this function is called and incremented by one on each subsequent - call. - - xmin : int, optional - This and the following three parameters set a bounding rectangle - on the input image. Only pixels on the input image inside this - rectangle will have their flux added to the output image. Xmin - sets the minimum value of the x dimension. The x dimension is the - dimension that varies quickest on the image. All four parameters - are zero based, counting starts at zero. - - xmax : int, optional - Sets the maximum value of the x dimension on the bounding box - of the input image. If ``xmax = 0``, no maximum will - be set in the x dimension (all pixels in a row of the input image - will be resampled). - - ymin : int, optional - Sets the minimum value in the y dimension on the bounding box. The - y dimension varies less rapidly than the x and represents the line - index on the input image. - - ymax : int, optional - Sets the maximum value in the y dimension. If ``ymax = 0``, - no maximum will be set in the y dimension (all pixels in a column - of the input image will be resampled). - - iscale : float, optional - A scale factor to be applied to pixel intensities of the - input image before resampling. - - pixfrac : float, optional - The fraction of a pixel that the pixel flux is confined to. The - default value of 1 has the pixel flux evenly spread across the image. - A value of 0.5 confines it to half a pixel in the linear dimension, - so the flux is confined to a quarter of the pixel area when the square - kernel is used. - - kernel: str, optional - The name of the kernel used to combine the input. The choice of - kernel controls the distribution of flux over the kernel. The kernel - names are: "square", "gaussian", "point", "turbo", "lanczos2", - and "lanczos3". The square kernel is the default. - - fillval: str, optional - The value a pixel is set to in the output if the input image does - not overlap it. The default value of NAN sets NaN values. - - Returns - ------- - A tuple with three values: a version string, the number of pixels - on the input image that do not overlap the output image, and the - number of complete lines on the input image that do not overlap the - output input image. - - """ - - # Insure that the fillval parameter gets properly interpreted for use with tdriz - if str(fillval).strip() == '': - fillval = 'NAN' - else: - fillval = str(fillval) - - if in_units == 'cps': - expscale = 1.0 - else: - expscale = expin - - if insci.dtype > np.float32: - insci = insci.astype(np.float32) - - # Add input weight image if it was not passed in - if inwht is None: - inwht = np.ones_like(insci) - - if xmax is None or xmax == xmin: - xmax = insci.shape[1] - if ymax is None or ymax == ymin: - ymax = insci.shape[0] - - # Compute what plane of the context image this input would - # correspond to: - planeid = int((uniqid - 1) / 32) - - # Check if the context image has this many planes - if outcon.ndim == 3: - nplanes = outcon.shape[0] - elif outcon.ndim == 2: - nplanes = 1 - else: - nplanes = 0 - - if nplanes <= planeid: - raise IndexError("Not enough planes in drizzle context image") - - # Alias context image to the requested plane if 3d - if outcon.ndim == 3: - outcon = outcon[planeid] - - # Compute the mapping between the input and output pixel coordinates - # for use in drizzle.cdrizzle.tdriz - pixmap = resample_utils.calc_gwcs_pixmap(input_wcs, output_wcs, insci.shape) - - log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") - log.debug(f"Input Sci shape: {insci.shape}") - log.debug(f"Output Sci shape: {outsci.shape}") - - # Call 'drizzle' to perform image combination - log.info(f"Drizzling {insci.shape} --> {outsci.shape}") - - _vers, nmiss, nskip = cdrizzle.tdriz( - insci.astype(np.float32), inwht.astype(np.float32), pixmap, - outsci, outwht, outcon, - uniqid=uniqid, - xmin=xmin, xmax=xmax, - ymin=ymin, ymax=ymax, - scale=iscale, - pixfrac=pixfrac, - kernel=kernel, - in_units=in_units, - expscale=expscale, - wtscale=wt_scl, - fillstr=fillval - ) - return _vers, nmiss, nskip + # Compute the mapping between the input and output pixel coordinates + # for use in drizzle.cdrizzle.tdriz + pixmap = resample_utils.calc_gwcs_pixmap( + inwcs, + self.outwcs, + insci.shape + ) + + log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") + log.debug(f"Input Sci shape: {insci.shape}") + log.debug(f"Output Sci shape: {self.outsci.shape}") + + # Call 'drizzle' to perform image combination + log.info(f"Drizzling {insci.shape} --> {self.outsci.shape}") + + super().add_image( + data=insci, + exptime=expin, + pixmap=pixmap, + scale=iscale, + weight_map=inwht, + wht_scale=wt_scl, # hard-coded for JWST count-rate data + pixfrac=self.pixfrac, + in_units=in_units, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax, + ) From 6dcde77e2b0134a8c06ed22d1e20393cc57243f5 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Wed, 16 Oct 2024 01:36:55 -0400 Subject: [PATCH 05/19] fix incorrect start context id --- jwst/resample/gwcs_drizzle.py | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 7e579ead6e..dcd7528660 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -62,12 +62,6 @@ def __init__(self, product, outwcs=None, wt_scl=None, self.pixfrac = pixfrac - self.sciext = "SCI" - self.whtext = "WHT" - self.conext = "CON" - - self.out_units = "cps" - self.outexptime = product.meta.exposure.measurement_time or 0.0 self.outsci = product.data @@ -82,16 +76,12 @@ def __init__(self, product, outwcs=None, wt_scl=None, if not self.outwcs: raise ValueError("Either an existing file or wcs must be supplied") - if self.outexptime == 0.0: - ctx = None - begin_ctx_id = 0 - else: - ctx = product.con - if ctx.ndim == 2: - begin_ctx_id = 1 - elif ctx.ndim == 3: - begin_ctx_id = ctx.shape[0] + 1 - else: + ctx = product.con + begin_ctx_id = 0 + if ctx is not None: + if ctx.size == 0: + ctx = None + elif ctx.ndim not in [2, 3]: # TODO: this message seems odd raise ValueError( f"Drizzle context image has wrong dimensions: {product}" From 3b91efee37e81c3945e900d3029b3717c94cf7de Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Wed, 16 Oct 2024 20:24:37 -0400 Subject: [PATCH 06/19] Replace crdizzle.tdriz calls with new Drizzle class --- jwst/resample/gwcs_drizzle.py | 32 ++-- jwst/resample/resample.py | 312 ++++++++++++-------------------- jwst/resample/resample_utils.py | 7 + 3 files changed, 137 insertions(+), 214 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index dcd7528660..844b575843 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -1,4 +1,3 @@ -import numpy as np from drizzle.resample import Drizzle from . import resample_utils @@ -33,6 +32,9 @@ def __init__(self, product, outwcs=None, wt_scl=None, provided, the WCS is taken from product. wt_scl : str, optional + .. note:: + This parameter is no longer used. + How each input image should be scaled. The choices are `exptime`, which scales each image by its exposure time, or `expsq`, which scales each image by the exposure time squared. If not set, then each @@ -55,8 +57,6 @@ def __init__(self, product, outwcs=None, wt_scl=None, The value a pixel is set to in the output if the input image does not overlap it. The default value of NAN sets NaN values. """ - - # Initialize the object fields self._product = product @@ -100,17 +100,18 @@ def __init__(self, product, outwcs=None, wt_scl=None, disable_ctx=disable_ctx ) + # Since the context array is dynamic, it must be re-assigned + # back to the product's `con` attribute. product.con = self.out_ctx - # Since the context array is dynamic, it must be re-assigned - # back to the product's `con` attribute. @property def outcon(self): """Return the context array""" return self.out_ctx - def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, - expin=1.0, in_units="cps", wt_scl=1.0, iscale=1.0): + def add_image(self, insci, inwcs, pixmap=None, inwht=None, + xmin=0, xmax=0, ymin=0, ymax=0, + expin=1.0, in_units="cps", iscale=1.0): """ Combine an input image with the output drizzled image. @@ -134,6 +135,10 @@ def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, Must have the same dimensions as insci. If none is supplied, the weighting is set to one. + pixmap : array, optional + A 3D array that maps input image coordinates onto the output + array. If provided, it can improve performance. + xmin : int, optional This and the following three parameters set a bounding rectangle on the input image. Only pixels on the input image inside this @@ -176,11 +181,12 @@ def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, """ # Compute the mapping between the input and output pixel coordinates # for use in drizzle.cdrizzle.tdriz - pixmap = resample_utils.calc_gwcs_pixmap( - inwcs, - self.outwcs, - insci.shape - ) + if pixmap is None: + pixmap = resample_utils.calc_gwcs_pixmap( + inwcs, + self.outwcs, + insci.shape + ) log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") log.debug(f"Input Sci shape: {insci.shape}") @@ -195,7 +201,7 @@ def add_image(self, insci, inwcs, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, pixmap=pixmap, scale=iscale, weight_map=inwht, - wht_scale=wt_scl, # hard-coded for JWST count-rate data + wht_scale=1.0, # hard-coded for JWST count-rate data pixfrac=self.pixfrac, in_units=in_units, xmin=xmin, diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index 2c5a7c314c..c43defe850 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -5,7 +5,7 @@ import numpy as np import psutil -from drizzle import cdrizzle +from drizzle.resample import Drizzle from spherical_geometry.polygon import SphericalPolygon from stdatamodels.jwst import datamodels @@ -284,8 +284,13 @@ def resample_group(self, input_models, indices, compute_error=False): del example_image # Initialize the output with the wcs - driz = gwcs_drizzle.GWCSDrizzle(output_model, pixfrac=self.pixfrac, - kernel=self.kernel, fillval=self.fillval) + driz = gwcs_drizzle.GWCSDrizzle( + output_model, + pixfrac=self.pixfrac, + kernel=self.kernel, + fillval=self.fillval, + disable_ctx=True, + ) # Also make a temporary model to hold error data error_model, driz_error = None, None @@ -322,10 +327,15 @@ def resample_group(self, input_models, indices, compute_error=False): data.shape, img.meta.wcs.bounding_box ) - + pixmap = resample_utils.calc_gwcs_pixmap( + img.meta.wcs, + output_model.meta.wcs, + img.data.shape, + ) driz.add_image( data, img.meta.wcs, + pixmap=pixmap, iscale=iscale, inwht=inwht, xmin=xmin, @@ -341,6 +351,7 @@ def resample_group(self, input_models, indices, compute_error=False): driz_error.add_image( img.err, img.meta.wcs, + pixmap=pixmap, iscale=iscale, inwht=inwht, xmin=xmin, @@ -354,8 +365,12 @@ def resample_group(self, input_models, indices, compute_error=False): # copy the drizzled error into the output model if compute_error: - output_model.err = error_model.data + output_model.err[:, :] = driz.out_img del error_model + # Since the context array is dynamic, it must be re-assigned + # back to the product's `con` attribute. + output_model.data[:, :] = driz.out_img + output_model.wht[:, :] = driz.out_wht return output_model @@ -417,8 +432,12 @@ def resample_many_to_one(self, input_models): ) # Initialize the output with the wcs - driz = gwcs_drizzle.GWCSDrizzle(output_model, pixfrac=self.pixfrac, - kernel=self.kernel, fillval=self.fillval) + driz = gwcs_drizzle.GWCSDrizzle( + output_model, + pixfrac=self.pixfrac, + kernel=self.kernel, + fillval=self.fillval, + ) log.info("Resampling science data") with input_models: @@ -459,6 +478,12 @@ def resample_many_to_one(self, input_models): del data, inwht input_models.shelve(img) + # Since the context array is dynamic, it must be re-assigned + # back to the product's `con` attribute. + output_model.data[:, :] = driz.out_img + output_model.wht[:, :] = driz.out_wht + output_model.con = driz.out_ctx + if self.blendheaders: blender.finalize_model(output_model) @@ -499,10 +524,35 @@ def resample_variance_arrays(self, output_model, input_models): total_weight_flat_var = np.zeros_like(output_model.data) with input_models: for i, model in enumerate(input_models): + # compute pixmap and inwht arrays common to all variance arrays: + inwht = resample_utils.build_driz_weight( + model, + weight_type=self.weight_type, # weights match science + good_bits=self.good_bits, + ) + pixmap = resample_utils.calc_gwcs_pixmap( + model.meta.wcs, + output_model.meta.wcs, + model.data.shape, + ) + xmin, xmax, ymin, ymax = resample_utils._resample_range( + model.data.shape, + model.meta.wcs.bounding_box, + ) + # Do the read noise variance first, so it can be # used for weights if needed rn_var = self._resample_one_variance_array( - "var_rnoise", model, output_model) + "var_rnoise", + input_model=model, + output_shape=output_model.data.shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) # Find valid weighting values in the variance if rn_var is not None: @@ -537,7 +587,16 @@ def resample_variance_arrays(self, output_model, input_models): # Now do poisson and flat variance, updating only valid new values # (zero is a valid value; negative, inf, or NaN are not) pn_var = self._resample_one_variance_array( - "var_poisson", model, output_model) + "var_poisson", + input_model=model, + output_shape=output_model.data.shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) if pn_var is not None: mask = (pn_var >= 0) & np.isfinite(pn_var) & (weight > 0) weighted_pn_var[mask] = np.nansum( @@ -550,7 +609,16 @@ def resample_variance_arrays(self, output_model, input_models): total_weight_pn_var[mask] += weight[mask] flat_var = self._resample_one_variance_array( - "var_flat", model, output_model) + "var_flat", + input_model=model, + output_shape=output_model.data.shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) if flat_var is not None: mask = (flat_var >= 0) & np.isfinite(flat_var) & (weight > 0) weighted_flat_var[mask] = np.nansum( @@ -588,7 +656,9 @@ def resample_variance_arrays(self, output_model, input_models): del weighted_rn_var, weighted_pn_var, weighted_flat_var del total_weight_rn_var, total_weight_pn_var, total_weight_flat_var - def _resample_one_variance_array(self, name, input_model, output_model): + def _resample_one_variance_array(self, name, input_model, output_shape, + inwht, pixmap, + xmin=None, xmax=None, ymin=None, ymax=None): """Resample one variance image from an input model. The error image is passed to drizzle instead of the variance, to @@ -611,43 +681,45 @@ def _resample_one_variance_array(self, name, input_model, output_model): ) return - # Make input weight map - inwht = resample_utils.build_driz_weight( - input_model, - weight_type=self.weight_type, # weights match science - good_bits=self.good_bits - ) - - resampled_error = np.zeros_like(output_model.data) - outwht = np.zeros_like(output_model.data) - outcon = np.zeros_like(output_model.con) + iscale = input_model.meta.iscale - xmin, xmax, ymin, ymax = resample_utils._resample_range( - variance.shape, - input_model.meta.wcs.bounding_box + # Resample the error array. Fill "unpopulated" pixels with NaNs. + driz = Drizzle( + kernel=self.kernel, + fillval=np.nan, + out_shape=output_shape, + out_img=None, + out_wht=None, + out_ctx=None, + exptime=0, + begin_ctx_id=0, + max_ctx_id=1, + disable_ctx=True ) - iscale = input_model.meta.iscale + log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") + log.debug(f"Input Sci shape: {variance.shape}") + log.debug(f"Output Sci shape: {output_shape}") - # Resample the error array. Fill "unpopulated" pixels with NaNs. - self.drizzle_arrays( - np.sqrt(variance), - inwht, - input_model.meta.wcs, - output_model.meta.wcs, - resampled_error, - outwht, - outcon, - iscale=iscale, + # Call 'drizzle' to perform image combination + log.info(f"Drizzling {variance.shape} --> {output_shape}") + + driz.add_image( + data=np.sqrt(variance), + exptime=input_model.meta.exposure.exposure_time, + pixmap=pixmap, + scale=iscale, + weight_map=inwht, + wht_scale=1.0, # hard-coded for JWST count-rate data pixfrac=self.pixfrac, - kernel=self.kernel, - fillval='NAN', + in_units="cps", xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) - return resampled_error ** 2 + + return driz.out_img ** 2 def update_exposure_times(self, output_model, input_models): """Modify exposure time metadata in-place""" @@ -685,168 +757,6 @@ def update_exposure_times(self, output_model, input_models): output_model.meta.exposure.duration = duration output_model.meta.exposure.elapsed_exposure_time = duration - @staticmethod - def drizzle_arrays(insci, inwht, input_wcs, output_wcs, outsci, outwht, - outcon, uniqid=1, xmin=0, xmax=0, ymin=0, ymax=0, - iscale=1.0, pixfrac=1.0, kernel='square', - fillval="NAN", wtscale=1.0): - """ - Low level routine for performing 'drizzle' operation on one image. - - The interface is compatible with STScI code. All images are Python - ndarrays, instead of filenames. File handling (input and output) is - performed by the calling routine. - - Parameters - ---------- - - insci : 2d array - A 2d numpy array containing the input image to be drizzled. - - inwht : 2d array - A 2d numpy array containing the pixel by pixel weighting. - Must have the same dimensions as insci. If none is supplied, - the weighting is set to one. - - input_wcs : gwcs.WCS object - The world coordinate system of the input image. - - output_wcs : gwcs.WCS object - The world coordinate system of the output image. - - outsci : 2d array - A 2d numpy array containing the output image produced by - drizzling. On the first call it should be set to zero. - Subsequent calls it will hold the intermediate results. This - is modified in-place. - - outwht : 2d array - A 2d numpy array containing the output counts. On the first - call it should be set to zero. On subsequent calls it will - hold the intermediate results. This is modified in-place. - - outcon : 2d or 3d array, optional - A 2d or 3d numpy array holding a bitmap of which image was an input - for each output pixel. Should be integer zero on first call. - Subsequent calls hold intermediate results. This is modified - in-place. - - uniqid : int, optional - The id number of the input image. Should be one the first time - this function is called and incremented by one on each subsequent - call. - - xmin : int, optional - This and the following three parameters set a bounding rectangle - on the input image. Only pixels on the input image inside this - rectangle will have their flux added to the output image. Xmin - sets the minimum value of the x dimension. The x dimension is the - dimension that varies quickest on the image. All four parameters - are zero based, counting starts at zero. - - xmax : int, optional - Sets the maximum value of the x dimension on the bounding box - of the input image. If ``xmax = 0``, no maximum will - be set in the x dimension (all pixels in a row of the input image - will be resampled). - - ymin : int, optional - Sets the minimum value in the y dimension on the bounding box. The - y dimension varies less rapidly than the x and represents the line - index on the input image. - - ymax : int, optional - Sets the maximum value in the y dimension. If ``ymax = 0``, - no maximum will be set in the y dimension (all pixels in a column - of the input image will be resampled). - - iscale : float, optional - A scale factor to be applied to pixel intensities of the - input image before resampling. - - pixfrac : float, optional - The fraction of a pixel that the pixel flux is confined to. The - default value of 1 has the pixel flux evenly spread across the image. - A value of 0.5 confines it to half a pixel in the linear dimension, - so the flux is confined to a quarter of the pixel area when the square - kernel is used. - - kernel: str, optional - The name of the kernel used to combine the input. The choice of - kernel controls the distribution of flux over the kernel. The kernel - names are: "square", "gaussian", "point", "turbo", "lanczos2", - and "lanczos3". The square kernel is the default. - - fillval: str, optional - The value a pixel is set to in the output if the input image does - not overlap it. The default value of NAN sets NaN values. - - Returns - ------- - A tuple with three values: a version string, the number of pixels - on the input image that do not overlap the output image, and the - number of complete lines on the input image that do not overlap the - output input image. - - """ - - # Insure that the fillval parameter gets properly interpreted for use with tdriz - if str(fillval).strip() == '': - fillval = 'NAN' - else: - fillval = str(fillval) - - if insci.dtype > np.float32: - insci = insci.astype(np.float32) - - # Add input weight image if it was not passed in - if inwht is None: - inwht = np.ones_like(insci) - - # Compute what plane of the context image this input would - # correspond to: - planeid = int((uniqid - 1) / 32) - - # Check if the context image has this many planes - if outcon.ndim == 3: - nplanes = outcon.shape[0] - elif outcon.ndim == 2: - nplanes = 1 - else: - nplanes = 0 - - if nplanes <= planeid: - raise IndexError("Not enough planes in drizzle context image") - - # Alias context image to the requested plane if 3d - if outcon.ndim == 3: - outcon = outcon[planeid] - - # Compute the mapping between the input and output pixel coordinates - # for use in drizzle.cdrizzle.tdriz - pixmap = resample_utils.calc_gwcs_pixmap(input_wcs, output_wcs, insci.shape) - - log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") - log.debug(f"Input Sci shape: {insci.shape}") - log.debug(f"Output Sci shape: {outsci.shape}") - - log.info(f"Drizzling {insci.shape} --> {outsci.shape}") - - _vers, _nmiss, _nskip = cdrizzle.tdriz( - insci, inwht, pixmap, - outsci, outwht, outcon, - uniqid=uniqid, - xmin=xmin, xmax=xmax, - ymin=ymin, ymax=ymax, - scale=iscale, - pixfrac=pixfrac, - kernel=kernel, - in_units="cps", - expscale=1.0, - wtscale=wtscale, - fillstr=fillval - ) - def _get_boundary_points(xmin, xmax, ymin, ymax, dx=None, dy=None, shrink=0): """ diff --git a/jwst/resample/resample_utils.py b/jwst/resample/resample_utils.py index c527b450f1..330f5fe454 100644 --- a/jwst/resample/resample_utils.py +++ b/jwst/resample/resample_utils.py @@ -124,6 +124,13 @@ def shape_from_bounding_box(bounding_box): def calc_gwcs_pixmap(in_wcs, out_wcs, shape=None): """ Return a pixel grid map from input frame to output frame. """ + # from drizzle.utils import calc_pixmap + # if shape is not None and not np.array_equiv(shape, in_wcs.array_shape): + # in_wcs = deepcopy(in_wcs) + # in_wcs.array_shape = shape + + # return calc_pixmap(wcs_from=in_wcs, wcs_to=out_wcs) + if shape: bb = wcs_bbox_from_shape(shape) log.debug("Bounding box from data shape: {}".format(bb)) From bc65dc204505972515e1c0626f8433b69b1914f1 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Thu, 17 Oct 2024 00:45:57 -0400 Subject: [PATCH 07/19] Performance improvements --- jwst/resample/resample.py | 375 +++++++++++++++++++++++--------------- 1 file changed, 225 insertions(+), 150 deletions(-) diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index c43defe850..f6df5a670c 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -439,9 +439,17 @@ def resample_many_to_one(self, input_models): fillval=self.fillval, ) - log.info("Resampling science data") + self._init_variance_arrays(output_model) + self._init_exptime_counters() + + log.info("Resampling science and variance data") + + leading_group_idx = [v[0] for v in input_models.group_indices.values()] + with input_models: - for img in input_models: + for idx, img in enumerate(input_models): + if idx in leading_group_idx: + self._update_exptime(img) if self.blendheaders: blender.accumulate(img) iscale = self._get_intensity_scale(img) @@ -460,22 +468,42 @@ def resample_many_to_one(self, input_models): else: data = img.data.copy() - xmin, xmax, ymin, ymax = resample_utils._resample_range( + in_image_limits = resample_utils._resample_range( data.shape, img.meta.wcs.bounding_box ) + xmin, xmax, ymin, ymax = in_image_limits + + pixmap = resample_utils.calc_gwcs_pixmap( + img.meta.wcs, + output_model.meta.wcs, + data.shape, + ) driz.add_image( data, img.meta.wcs, + pixmap=pixmap, iscale=iscale, inwht=inwht, xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) + + # Resample variance arrays in input_models to output_model + self._resample_variance_arrays( + model=img, + inwht=inwht, + pixmap=pixmap, + in_image_limits=in_image_limits, + output_shape=output_model.meta.wcs.array_shape, + ) + del data, inwht + del img.meta.iscale + input_models.shelve(img) # Since the context array is dynamic, it must be re-assigned @@ -487,8 +515,9 @@ def resample_many_to_one(self, input_models): if self.blendheaders: blender.finalize_model(output_model) - # Resample variance arrays in input_models to output_model - self.resample_variance_arrays(output_model, input_models) + # compute final variances: + self._compute_resample_variance_totals(output_model) + var_components = [ output_model.var_rnoise, output_model.var_poisson, @@ -501,10 +530,146 @@ def resample_many_to_one(self, input_models): all_nan = np.all(np.isnan(var_components), axis=0) output_model.err[all_nan] = np.nan - self.update_exposure_times(output_model, input_models) + self._get_exptime_totals(output_model) return ModelLibrary([output_model,], on_disk=False) + def _init_variance_arrays(self, output_model): + self._weighted_rn_var = np.full_like(output_model.data, np.nan) + self._weighted_pn_var = np.full_like(output_model.data, np.nan) + self._weighted_flat_var = np.full_like(output_model.data, np.nan) + self._total_weight_rn_var = np.zeros_like(output_model.data) + self._total_weight_pn_var = np.zeros_like(output_model.data) + self._total_weight_flat_var = np.zeros_like(output_model.data) + + def _resample_variance_arrays(self, model, inwht, pixmap, in_image_limits, + output_shape): + xmin, xmax, ymin, ymax = in_image_limits + + # Do the read noise variance first, so it can be + # used for weights if needed + rn_var = self._resample_one_variance_array( + "var_rnoise", + input_model=model, + output_shape=output_shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) + + # Find valid weighting values in the variance + if rn_var is not None: + mask = (rn_var > 0) & np.isfinite(rn_var) + else: + mask = np.full_like(rn_var, False) + + # Set the weight for the image from the weight type + weight = np.ones(output_shape) + if self.weight_type == "ivm" and rn_var is not None: + weight[mask] = rn_var[mask] ** -1 + elif self.weight_type == "exptime": + if resample_utils.check_for_tmeasure(model): + weight[:] = model.meta.exposure.measurement_time + else: + weight[:] = model.meta.exposure.exposure_time + + # Weight and add the readnoise variance + # Note: floating point overflow is an issue if variance weights + # are used - it can't be squared before multiplication + if rn_var is not None: + mask = (rn_var >= 0) & np.isfinite(rn_var) & (weight > 0) + self._weighted_rn_var[mask] = np.nansum( + [ + self._weighted_rn_var[mask], + rn_var[mask] * weight[mask] * weight[mask] + ], + axis=0 + ) + self._total_weight_rn_var[mask] += weight[mask] + + # Now do poisson and flat variance, updating only valid new values + # (zero is a valid value; negative, inf, or NaN are not) + pn_var = self._resample_one_variance_array( + "var_poisson", + input_model=model, + output_shape=output_shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) + if pn_var is not None: + mask = (pn_var >= 0) & np.isfinite(pn_var) & (weight > 0) + self._weighted_pn_var[mask] = np.nansum( + [ + self._weighted_pn_var[mask], + pn_var[mask] * weight[mask] * weight[mask] + ], + axis=0 + ) + self._total_weight_pn_var[mask] += weight[mask] + + flat_var = self._resample_one_variance_array( + "var_flat", + input_model=model, + output_shape=output_shape, + inwht=inwht, + pixmap=pixmap, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax + ) + if flat_var is not None: + mask = (flat_var >= 0) & np.isfinite(flat_var) & (weight > 0) + self._weighted_flat_var[mask] = np.nansum( + [ + self._weighted_flat_var[mask], + flat_var[mask] * weight[mask] * weight[mask] + ], + axis=0 + ) + self._total_weight_flat_var[mask] += weight[mask] + + def _compute_resample_variance_totals(self, output_model): + # Divide by the total weights, squared, and set in the output model. + # Zero weight and missing values are NaN in the output. + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", "invalid value*", RuntimeWarning) + warnings.filterwarnings("ignore", "divide by zero*", RuntimeWarning) + + output_variance = ( + self._weighted_rn_var / self._total_weight_rn_var / + self._total_weight_rn_var + ) + setattr(output_model, "var_rnoise", output_variance) + + output_variance = ( + self._weighted_pn_var / self._total_weight_pn_var / + self._total_weight_pn_var + ) + setattr(output_model, "var_poisson", output_variance) + + output_variance = ( + self._weighted_flat_var / self._total_weight_flat_var / + self._total_weight_flat_var + ) + setattr(output_model, "var_flat", output_variance) + + del ( + self._weighted_rn_var, + self._weighted_pn_var, + self._weighted_flat_var, + self._total_weight_rn_var, + self._total_weight_pn_var, + self._total_weight_flat_var, + ) + def resample_variance_arrays(self, output_model, input_models): """Resample variance arrays from input_models to the output_model. @@ -515,13 +680,10 @@ def resample_variance_arrays(self, output_model, input_models): The output_model is modified in place. """ + self._init_variance_arrays(output_model) + output_shape = output_model.meta.wcs.array_shape log.info("Resampling variance components") - weighted_rn_var = np.full_like(output_model.data, np.nan) - weighted_pn_var = np.full_like(output_model.data, np.nan) - weighted_flat_var = np.full_like(output_model.data, np.nan) - total_weight_rn_var = np.zeros_like(output_model.data) - total_weight_pn_var = np.zeros_like(output_model.data) - total_weight_flat_var = np.zeros_like(output_model.data) + with input_models: for i, model in enumerate(input_models): # compute pixmap and inwht arrays common to all variance arrays: @@ -535,126 +697,24 @@ def resample_variance_arrays(self, output_model, input_models): output_model.meta.wcs, model.data.shape, ) - xmin, xmax, ymin, ymax = resample_utils._resample_range( + in_image_limits = resample_utils._resample_range( model.data.shape, model.meta.wcs.bounding_box, ) - - # Do the read noise variance first, so it can be - # used for weights if needed - rn_var = self._resample_one_variance_array( - "var_rnoise", - input_model=model, - output_shape=output_model.data.shape, - inwht=inwht, - pixmap=pixmap, - xmin=xmin, - xmax=xmax, - ymin=ymin, - ymax=ymax - ) - - # Find valid weighting values in the variance - if rn_var is not None: - mask = (rn_var > 0) & np.isfinite(rn_var) - else: - mask = np.full_like(rn_var, False) - - # Set the weight for the image from the weight type - weight = np.ones(output_model.data.shape) - if self.weight_type == "ivm" and rn_var is not None: - weight[mask] = rn_var[mask] ** -1 - elif self.weight_type == "exptime": - if resample_utils.check_for_tmeasure(model): - weight[:] = model.meta.exposure.measurement_time - else: - weight[:] = model.meta.exposure.exposure_time - - # Weight and add the readnoise variance - # Note: floating point overflow is an issue if variance weights - # are used - it can't be squared before multiplication - if rn_var is not None: - mask = (rn_var >= 0) & np.isfinite(rn_var) & (weight > 0) - weighted_rn_var[mask] = np.nansum( - [ - weighted_rn_var[mask], - rn_var[mask] * weight[mask] * weight[mask] - ], - axis=0 - ) - total_weight_rn_var[mask] += weight[mask] - - # Now do poisson and flat variance, updating only valid new values - # (zero is a valid value; negative, inf, or NaN are not) - pn_var = self._resample_one_variance_array( - "var_poisson", - input_model=model, - output_shape=output_model.data.shape, - inwht=inwht, - pixmap=pixmap, - xmin=xmin, - xmax=xmax, - ymin=ymin, - ymax=ymax - ) - if pn_var is not None: - mask = (pn_var >= 0) & np.isfinite(pn_var) & (weight > 0) - weighted_pn_var[mask] = np.nansum( - [ - weighted_pn_var[mask], - pn_var[mask] * weight[mask] * weight[mask] - ], - axis=0 - ) - total_weight_pn_var[mask] += weight[mask] - - flat_var = self._resample_one_variance_array( - "var_flat", - input_model=model, - output_shape=output_model.data.shape, + self._resample_variance_arrays( + model=model, inwht=inwht, pixmap=pixmap, - xmin=xmin, - xmax=xmax, - ymin=ymin, - ymax=ymax + in_image_limits=in_image_limits, + output_shape=output_shape, ) - if flat_var is not None: - mask = (flat_var >= 0) & np.isfinite(flat_var) & (weight > 0) - weighted_flat_var[mask] = np.nansum( - [ - weighted_flat_var[mask], - flat_var[mask] * weight[mask] * weight[mask] - ], - axis=0 - ) - total_weight_flat_var[mask] += weight[mask] del model.meta.iscale - del weight + del inwht input_models.shelve(model, i) # We now have a sum of the weighted resampled variances. - # Divide by the total weights, squared, and set in the output model. - # Zero weight and missing values are NaN in the output. - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", "invalid value*", RuntimeWarning) - warnings.filterwarnings("ignore", "divide by zero*", RuntimeWarning) - - output_variance = (weighted_rn_var - / total_weight_rn_var / total_weight_rn_var) - setattr(output_model, "var_rnoise", output_variance) - - output_variance = (weighted_pn_var - / total_weight_pn_var / total_weight_pn_var) - setattr(output_model, "var_poisson", output_variance) - - output_variance = (weighted_flat_var - / total_weight_flat_var / total_weight_flat_var) - setattr(output_model, "var_flat", output_variance) - - del weighted_rn_var, weighted_pn_var, weighted_flat_var - del total_weight_rn_var, total_weight_pn_var, total_weight_flat_var + self._compute_resample_variance_totals(output_model) def _resample_one_variance_array(self, name, input_model, output_shape, inwht, pixmap, @@ -721,41 +781,56 @@ def _resample_one_variance_array(self, name, input_model, output_shape, return driz.out_img ** 2 - def update_exposure_times(self, output_model, input_models): - """Modify exposure time metadata in-place""" - total_exposure_time = 0. - exposure_times = {'start': [], 'end': []} - duration = 0.0 - total_measurement_time = 0.0 - measurement_time_failures = [] - with input_models: - for _, indices in input_models.group_indices.items(): - model = input_models.borrow(indices[0]) - total_exposure_time += model.meta.exposure.exposure_time - if not resample_utils.check_for_tmeasure(model): - measurement_time_failures.append(1) - else: - total_measurement_time += model.meta.exposure.measurement_time - measurement_time_failures.append(0) - exposure_times['start'].append(model.meta.exposure.start_time) - exposure_times['end'].append(model.meta.exposure.end_time) - duration += model.meta.exposure.duration - input_models.shelve(model, indices[0], modify=False) + def _init_exptime_counters(self): + self._total_exposure_time = 0. + self._exposure_times = {'start': [], 'end': []} + self._duration = 0.0 + self._total_measurement_time = 0.0 + self._measurement_time_failures = [] + + def _update_exptime(self, model): + self._total_exposure_time += model.meta.exposure.exposure_time + if not resample_utils.check_for_tmeasure(model): + self._measurement_time_failures.append(1) + else: + self._total_measurement_time += model.meta.exposure.measurement_time + self._measurement_time_failures.append(0) + self._exposure_times['start'].append(model.meta.exposure.start_time) + self._exposure_times['end'].append(model.meta.exposure.end_time) + self._duration += model.meta.exposure.duration + def _get_exptime_totals(self, output_model): # Update some basic exposure time values based on output_model - output_model.meta.exposure.exposure_time = total_exposure_time - if not any(measurement_time_failures): - output_model.meta.exposure.measurement_time = total_measurement_time - output_model.meta.exposure.start_time = min(exposure_times['start']) - output_model.meta.exposure.end_time = max(exposure_times['end']) + output_model.meta.exposure.exposure_time = self._total_exposure_time + if not any(self._measurement_time_failures): + output_model.meta.exposure.measurement_time = self._total_measurement_time + output_model.meta.exposure.start_time = min(self._exposure_times['start']) + output_model.meta.exposure.end_time = max(self._exposure_times['end']) # Update other exposure time keywords: # XPOSURE (identical to the total effective exposure time, EFFEXPTM) - xposure = total_exposure_time + xposure = self._total_exposure_time output_model.meta.exposure.effective_exposure_time = xposure # DURATION (identical to TELAPSE, elapsed time) - output_model.meta.exposure.duration = duration - output_model.meta.exposure.elapsed_exposure_time = duration + output_model.meta.exposure.duration = self._duration + output_model.meta.exposure.elapsed_exposure_time = self._duration + + del ( + self._total_exposure_time, + self._exposure_times, + self._duration, + self._total_measurement_time, + self._measurement_time_failures, + ) + + def update_exposure_times(self, output_model, input_models): + """Modify exposure time metadata in-place""" + self._init_exptime_counters() + with input_models: + for _, indices in input_models.group_indices.items(): + model = input_models.borrow(indices[0]) + self._update_exptime(model) + self._get_exptime_totals(output_model) def _get_boundary_points(xmin, xmax, ymin, ymax, dx=None, dy=None, shrink=0): From 7f02ed5610eb7c5bbe3f3679ec9981c7456a5a0c Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Fri, 18 Oct 2024 01:35:52 -0400 Subject: [PATCH 08/19] Fix a faling unit test for spec --- jwst/resample/resample.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index f6df5a670c..f2e12000c2 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -297,8 +297,12 @@ def resample_group(self, input_models, indices, compute_error=False): if compute_error: error_model = output_model.copy() driz_error = gwcs_drizzle.GWCSDrizzle( - error_model, pixfrac=self.pixfrac, - kernel=self.kernel, fillval=self.fillval) + error_model, + pixfrac=self.pixfrac, + kernel=self.kernel, + fillval=self.fillval, + disable_ctx=True, + ) log.info(f"{len(indices)} exposures to drizzle together") for index in indices: @@ -341,7 +345,7 @@ def resample_group(self, input_models, indices, compute_error=False): xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) del data @@ -357,20 +361,20 @@ def resample_group(self, input_models, indices, compute_error=False): xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) input_models.shelve(img, index, modify=False) del img - # copy the drizzled error into the output model - if compute_error: - output_model.err[:, :] = driz.out_img - del error_model # Since the context array is dynamic, it must be re-assigned # back to the product's `con` attribute. output_model.data[:, :] = driz.out_img output_model.wht[:, :] = driz.out_wht + # copy the drizzled error into the output model + if compute_error: + output_model.err[:, :] = driz_error.out_img + del error_model return output_model From a4f75ec10f226334812cfd0d2b8e86769b7e603b Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Fri, 18 Oct 2024 00:42:25 -0400 Subject: [PATCH 09/19] Address reviewers comments + remove no longer used stuff --- jwst/resample/gwcs_drizzle.py | 16 +--------------- jwst/resample/resample.py | 3 ++- jwst/resample/resample_utils.py | 7 ------- 3 files changed, 3 insertions(+), 23 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 844b575843..6d5de3ab9c 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -12,7 +12,7 @@ class GWCSDrizzle(Drizzle): Combine images using the drizzle algorithm """ - def __init__(self, product, outwcs=None, wt_scl=None, + def __init__(self, product, outwcs=None, pixfrac=1.0, kernel="square", fillval="NAN", disable_ctx=False, n_max_images=None): """ @@ -31,15 +31,6 @@ def __init__(self, product, outwcs=None, wt_scl=None, The world coordinate system (WCS) of the resampled image. If not provided, the WCS is taken from product. - wt_scl : str, optional - .. note:: - This parameter is no longer used. - - How each input image should be scaled. The choices are `exptime`, - which scales each image by its exposure time, or `expsq`, which scales - each image by the exposure time squared. If not set, then each - input image is scaled by its own weight map. - pixfrac : float, optional The fraction of a pixel that the pixel flux is confined to. The default value of 1 has the pixel flux evenly spread across the image. @@ -104,11 +95,6 @@ def __init__(self, product, outwcs=None, wt_scl=None, # back to the product's `con` attribute. product.con = self.out_ctx - @property - def outcon(self): - """Return the context array""" - return self.out_ctx - def add_image(self, insci, inwcs, pixmap=None, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, expin=1.0, in_units="cps", iscale=1.0): diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index f2e12000c2..92e550375f 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -514,7 +514,8 @@ def resample_many_to_one(self, input_models): # back to the product's `con` attribute. output_model.data[:, :] = driz.out_img output_model.wht[:, :] = driz.out_wht - output_model.con = driz.out_ctx + if driz.out_ctx is not None: + output_model.con = driz.out_ctx if self.blendheaders: blender.finalize_model(output_model) diff --git a/jwst/resample/resample_utils.py b/jwst/resample/resample_utils.py index 330f5fe454..c527b450f1 100644 --- a/jwst/resample/resample_utils.py +++ b/jwst/resample/resample_utils.py @@ -124,13 +124,6 @@ def shape_from_bounding_box(bounding_box): def calc_gwcs_pixmap(in_wcs, out_wcs, shape=None): """ Return a pixel grid map from input frame to output frame. """ - # from drizzle.utils import calc_pixmap - # if shape is not None and not np.array_equiv(shape, in_wcs.array_shape): - # in_wcs = deepcopy(in_wcs) - # in_wcs.array_shape = shape - - # return calc_pixmap(wcs_from=in_wcs, wcs_to=out_wcs) - if shape: bb = wcs_bbox_from_shape(shape) log.debug("Bounding box from data shape: {}".format(bb)) From 1a352c3c3a7661062b2ff70737924ed1d12b20b7 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Fri, 18 Oct 2024 12:02:35 -0400 Subject: [PATCH 10/19] don't update .con when disable_ctx is True --- jwst/resample/gwcs_drizzle.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 6d5de3ab9c..8d6977a270 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -67,16 +67,19 @@ def __init__(self, product, outwcs=None, if not self.outwcs: raise ValueError("Either an existing file or wcs must be supplied") - ctx = product.con begin_ctx_id = 0 - if ctx is not None: - if ctx.size == 0: - ctx = None - elif ctx.ndim not in [2, 3]: - # TODO: this message seems odd - raise ValueError( - f"Drizzle context image has wrong dimensions: {product}" - ) + if disable_ctx: + ctx = None + else: + ctx = product.con + if ctx is not None: + if ctx.size == 0: + ctx = None + elif ctx.ndim not in [2, 3]: + # TODO: this message seems odd + raise ValueError( + f"Drizzle context image has wrong dimensions: {product}" + ) super().__init__( kernel=kernel, @@ -93,7 +96,8 @@ def __init__(self, product, outwcs=None, # Since the context array is dynamic, it must be re-assigned # back to the product's `con` attribute. - product.con = self.out_ctx + if not disable_ctx: + product.con = self.out_ctx def add_image(self, insci, inwcs, pixmap=None, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, From ba2de2f90227ef273eeac5202d6ef4e01575de88 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Fri, 18 Oct 2024 15:48:38 -0400 Subject: [PATCH 11/19] improve error message --- jwst/resample/gwcs_drizzle.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 8d6977a270..0a7059665f 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -76,9 +76,10 @@ def __init__(self, product, outwcs=None, if ctx.size == 0: ctx = None elif ctx.ndim not in [2, 3]: - # TODO: this message seems odd raise ValueError( - f"Drizzle context image has wrong dimensions: {product}" + "Provided (via product.con) context image has wrong " + "dimensions: expected a 2D or 3D array but got a " + f"{ctx.ndim}D array." ) super().__init__( From 676f4fae4947352570788249a25b5bbc68515910 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Sat, 19 Oct 2024 04:41:49 -0400 Subject: [PATCH 12/19] Don't store iscale in model. Create output model on demand --- jwst/resample/gwcs_drizzle.py | 61 ++----- jwst/resample/resample.py | 309 ++++++++++++++++++--------------- jwst/resample/resample_spec.py | 10 +- 3 files changed, 185 insertions(+), 195 deletions(-) diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py index 0a7059665f..51b5633022 100644 --- a/jwst/resample/gwcs_drizzle.py +++ b/jwst/resample/gwcs_drizzle.py @@ -12,9 +12,8 @@ class GWCSDrizzle(Drizzle): Combine images using the drizzle algorithm """ - def __init__(self, product, outwcs=None, - pixfrac=1.0, kernel="square", fillval="NAN", - disable_ctx=False, n_max_images=None): + def __init__(self, output_wcs=None, pixfrac=1.0, kernel="square", + fillval="NAN", disable_ctx=False, n_max_images=None): """ Create a new Drizzle output object and set the drizzle parameters. @@ -48,58 +47,22 @@ def __init__(self, product, outwcs=None, The value a pixel is set to in the output if the input image does not overlap it. The default value of NAN sets NaN values. """ - # Initialize the object fields - self._product = product - + self.output_wcs = output_wcs self.pixfrac = pixfrac - self.outexptime = product.meta.exposure.measurement_time or 0.0 - - self.outsci = product.data - self.outwht = product.wht - - if outwcs: - self.outwcs = outwcs - else: - self.outwcs = product.meta.wcs - - # Check field values - if not self.outwcs: - raise ValueError("Either an existing file or wcs must be supplied") - - begin_ctx_id = 0 - if disable_ctx: - ctx = None - else: - ctx = product.con - if ctx is not None: - if ctx.size == 0: - ctx = None - elif ctx.ndim not in [2, 3]: - raise ValueError( - "Provided (via product.con) context image has wrong " - "dimensions: expected a 2D or 3D array but got a " - f"{ctx.ndim}D array." - ) - super().__init__( kernel=kernel, fillval=fillval, - out_shape=None, - out_img=product.data, - out_wht=product.wht, - out_ctx=ctx, - exptime=self.outexptime, - begin_ctx_id=begin_ctx_id, + out_shape=output_wcs.array_shape, + out_img=None, + out_wht=None, + out_ctx=None, + exptime=0.0, + begin_ctx_id=0, max_ctx_id=n_max_images, disable_ctx=disable_ctx ) - # Since the context array is dynamic, it must be re-assigned - # back to the product's `con` attribute. - if not disable_ctx: - product.con = self.out_ctx - def add_image(self, insci, inwcs, pixmap=None, inwht=None, xmin=0, xmax=0, ymin=0, ymax=0, expin=1.0, in_units="cps", iscale=1.0): @@ -175,16 +138,16 @@ def add_image(self, insci, inwcs, pixmap=None, inwht=None, if pixmap is None: pixmap = resample_utils.calc_gwcs_pixmap( inwcs, - self.outwcs, + self.output_wcs, insci.shape ) log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") log.debug(f"Input Sci shape: {insci.shape}") - log.debug(f"Output Sci shape: {self.outsci.shape}") + log.debug(f"Output Sci shape: {self.out_img.shape}") # Call 'drizzle' to perform image combination - log.info(f"Drizzling {insci.shape} --> {self.outsci.shape}") + log.info(f"Drizzling {insci.shape} --> {self.out_img.shape}") super().add_image( data=insci, diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index 92e550375f..488057bc0f 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -14,9 +14,9 @@ from jwst.datamodels import ModelLibrary from jwst.associations.asn_from_list import asn_from_list -from . import gwcs_drizzle from jwst.model_blender.blender import ModelBlender from jwst.resample import resample_utils +from jwst.resample.gwcs_drizzle import GWCSDrizzle log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) @@ -89,6 +89,12 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, self.input_pixscale0 = None # computed pixel scale of the first image (deg) self._recalc_pscale_ratio = pscale is not None + # TODO: self._iscales is used to store computed iscale for each + # input model. + # It is used only by self.resample_variance_arrays(). + # if that method no longer needs to be supported, we can remove this. + self._iscales = {} + log.info(f"Driz parameter kernel: {self.kernel}") log.info(f"Driz parameter pixfrac: {self.pixfrac}") log.info(f"Driz parameter fillval: {self.fillval}") @@ -115,13 +121,13 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, self.output_wcs.array_shape = output_shape[::-1] if output_wcs.pixel_area is None: - output_pix_area = compute_image_pixel_area(self.output_wcs) - if output_pix_area is None: + self.output_pix_area = compute_image_pixel_area(self.output_wcs) + if self.output_pix_area is None: raise ValueError( "Unable to compute output pixel area from 'output_wcs'." ) else: - output_pix_area = output_wcs.pixel_area + self.output_pix_area = output_wcs.pixel_area else: # Define output WCS based on all inputs, including a reference WCS: @@ -139,18 +145,18 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, # use the same algorithm as for when output_wcs is provided by the # user. tr = self.output_wcs.pipeline[0].transform - output_pix_area = ( + self.output_pix_area = ( np.deg2rad(tr['cdelt1'].factor.value) * np.deg2rad(tr['cdelt2'].factor.value) ) if pscale is None: - pscale = np.rad2deg(np.sqrt(output_pix_area)) + pscale = np.rad2deg(np.sqrt(self.output_pix_area)) log.info(f'Computed output pixel scale: {3600.0 * pscale} arcsec.') self.pscale = pscale # in deg - log.debug('Output mosaic size: {}'.format(self.output_wcs.array_shape)) + log.debug(f"Output mosaic size: {tuple(self.output_wcs.pixel_shape)}") allowed_memory = kwargs['allowed_memory'] if allowed_memory is None: @@ -174,23 +180,25 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, if not can_allocate: raise OutputTooLargeError( - f'Combined ImageModel size {self.output_wcs.array_shape} ' + f'Combined ImageModel size {self.output_array_shape} ' f'requires {bytes2human(required_memory)}. ' f'Model cannot be instantiated.' ) - self.blank_output = datamodels.ImageModel(tuple(self.output_wcs.array_shape)) + + def _create_output_model(self, ref_input_model=None): + """ Create a new blank model and update it's meta with info from ``ref_input_model``. """ + output_model = datamodels.ImageModel(None) # tuple(self.output_wcs.array_shape)) # update meta data and wcs - with input_models: - example_model = input_models.borrow(0) - self.blank_output.update(example_model) - input_models.shelve(example_model, 0, modify=False) - del example_model - self.blank_output.meta.wcs = self.output_wcs - self.blank_output.meta.photometry.pixelarea_steradians = output_pix_area - self.blank_output.meta.photometry.pixelarea_arcsecsq = ( - output_pix_area * np.rad2deg(3600)**2 - ) + if ref_input_model is not None: + output_model.update(ref_input_model) + output_model.meta.wcs = self.output_wcs + if self.output_pix_area is not None: + output_model.meta.photometry.pixelarea_steradians = self.output_pix_area + output_model.meta.photometry.pixelarea_arcsecsq = ( + self.output_pix_area * np.rad2deg(3600)**2 + ) + return output_model def do_drizzle(self, input_models): """Pick the correct drizzling mode based on self.single @@ -264,80 +272,96 @@ def resample_group(self, input_models, indices, compute_error=False): If True, an approximate error image will be resampled alongside the science image. """ - output_model = self.blank_output.copy() + output_model = None - copy_asn_info_from_library(input_models, output_model) + # Initialize the output with the wcs + driz = GWCSDrizzle( + self.output_wcs, + kernel=self.kernel, + fillval=self.fillval, + disable_ctx=True, + ) - with input_models: - example_image = input_models.borrow(indices[0]) - - # Determine output file type from input exposure filenames - # Use this for defining the output filename - indx = example_image.meta.filename.rfind('.') - output_type = example_image.meta.filename[indx:] - output_root = '_'.join(example_image.meta.filename.replace( - output_type, '').split('_')[:-1]) - output_model.meta.filename = ( - f'{output_root}_' - f'{self.intermediate_suffix}{output_type}') - input_models.shelve(example_image, indices[0], modify=False) - del example_image - - # Initialize the output with the wcs - driz = gwcs_drizzle.GWCSDrizzle( - output_model, + # Also make a temporary model to hold error data + if compute_error: + driz_error = GWCSDrizzle( + self.output_wcs, pixfrac=self.pixfrac, kernel=self.kernel, fillval=self.fillval, disable_ctx=True, ) - # Also make a temporary model to hold error data - error_model, driz_error = None, None - if compute_error: - error_model = output_model.copy() - driz_error = gwcs_drizzle.GWCSDrizzle( - error_model, - pixfrac=self.pixfrac, - kernel=self.kernel, - fillval=self.fillval, - disable_ctx=True, + log.info(f"{len(indices)} exposures to drizzle together") + for index in indices: + img = input_models.borrow(index) + if output_model is None: + output_model = self._create_output_model( + ref_input_model=img + ) + # Determine output file type from input exposure filenames + # Use this for defining the output filename + indx = img.meta.filename.rfind('.') + output_type = img.meta.filename[indx:] + output_root = '_'.join(img.meta.filename.replace( + output_type, + '' + ).split('_')[:-1]) + output_model.meta.filename = ( + f'{output_root}_' + f'{self.intermediate_suffix}{output_type}' ) + copy_asn_info_from_library(input_models, output_model) - log.info(f"{len(indices)} exposures to drizzle together") - for index in indices: - img = input_models.borrow(index) - if isinstance(img, datamodels.SlitModel): - # must call this explicitly to populate area extension - # although the existence of this extension may not be necessary - img.area = img.area - iscale = self._get_intensity_scale(img) - log.debug(f'Using intensity scale iscale={iscale}') + if isinstance(img, datamodels.SlitModel): + # must call this explicitly to populate area extension + # although the existence of this extension may not be necessary + img.area = img.area - inwht = resample_utils.build_driz_weight( - img, - weight_type=self.weight_type, - good_bits=self.good_bits - ) + iscale = self._get_intensity_scale(img) + self._iscales[index] = iscale + log.debug(f'Using intensity scale iscale={iscale}') - # apply sky subtraction - blevel = img.meta.background.level - if not img.meta.background.subtracted and blevel is not None: - data = img.data - blevel - else: - data = img.data + inwht = resample_utils.build_driz_weight( + img, + weight_type=self.weight_type, + good_bits=self.good_bits + ) - xmin, xmax, ymin, ymax = resample_utils._resample_range( - data.shape, - img.meta.wcs.bounding_box - ) - pixmap = resample_utils.calc_gwcs_pixmap( - img.meta.wcs, - output_model.meta.wcs, - img.data.shape, - ) - driz.add_image( - data, + # apply sky subtraction + blevel = img.meta.background.level + if not img.meta.background.subtracted and blevel is not None: + data = img.data - blevel + else: + data = img.data + + xmin, xmax, ymin, ymax = resample_utils._resample_range( + data.shape, + img.meta.wcs.bounding_box + ) + pixmap = resample_utils.calc_gwcs_pixmap( + img.meta.wcs, + self.output_wcs, + img.data.shape, + ) + driz.add_image( + data, + img.meta.wcs, + pixmap=pixmap, + iscale=iscale, + inwht=inwht, + xmin=xmin, + xmax=xmax, + ymin=ymin, + ymax=ymax, + ) + del data + + # make an approximate error image by drizzling it + # in the same way the image is handled + if compute_error: + driz_error.add_image( + img.err, img.meta.wcs, pixmap=pixmap, iscale=iscale, @@ -347,34 +371,17 @@ def resample_group(self, input_models, indices, compute_error=False): ymin=ymin, ymax=ymax, ) - del data - - # make an approximate error image by drizzling it - # in the same way the image is handled - if compute_error: - driz_error.add_image( - img.err, - img.meta.wcs, - pixmap=pixmap, - iscale=iscale, - inwht=inwht, - xmin=xmin, - xmax=xmax, - ymin=ymin, - ymax=ymax, - ) - input_models.shelve(img, index, modify=False) - del img + input_models.shelve(img, index, modify=False) + del img - # Since the context array is dynamic, it must be re-assigned - # back to the product's `con` attribute. - output_model.data[:, :] = driz.out_img - output_model.wht[:, :] = driz.out_wht + output_model.data = driz.out_img + output_model.wht = driz.out_wht + del driz # copy the drizzled error into the output model if compute_error: - output_model.err[:, :] = driz_error.out_img - del error_model + output_model.err = driz_error.out_img + del driz_error return output_model @@ -403,28 +410,22 @@ def resample_many_to_many(self, input_models): else: output_models.append(output_model) - if not self.in_memory: + if self.in_memory: + # build ModelLibrary as a list of in-memory models + return ModelLibrary(output_models, on_disk=False) + else: # build ModelLibrary as an association from the output files # this saves memory if there are multiple groups asn = asn_from_list(output_models, product_name='outlier_i2d') asn_dict = json.loads(asn.dump()[1]) # serializes the asn and converts to dict return ModelLibrary(asn_dict, on_disk=True) - # otherwise just build it as a list of in-memory models - return ModelLibrary(output_models, on_disk=False) - def resample_many_to_one(self, input_models): """Resample and coadd many inputs to a single output. Used for stage 3 resampling """ - output_model = self.blank_output.copy() - output_model.meta.filename = self.output_filename - output_model.meta.resample.weight_type = self.weight_type - output_model.meta.resample.pointings = len(input_models.group_names) - - # copy over asn information - copy_asn_info_from_library(input_models, output_model) + output_model = None if self.blendheaders: blender = ModelBlender( @@ -436,14 +437,14 @@ def resample_many_to_one(self, input_models): ) # Initialize the output with the wcs - driz = gwcs_drizzle.GWCSDrizzle( - output_model, + driz = GWCSDrizzle( + self.output_wcs, pixfrac=self.pixfrac, kernel=self.kernel, fillval=self.fillval, ) - self._init_variance_arrays(output_model) + self._init_variance_arrays() self._init_exptime_counters() log.info("Resampling science and variance data") @@ -452,19 +453,35 @@ def resample_many_to_one(self, input_models): with input_models: for idx, img in enumerate(input_models): + if output_model is None: + output_model = self._create_output_model( + ref_input_model=img + ) + # Determine output file type from input exposure filenames + # Use this for defining the output filename + output_model.meta.filename = self.output_filename + output_model.meta.resample.weight_type = self.weight_type + output_model.meta.resample.pointings = len(input_models.group_names) + + # copy over asn information + copy_asn_info_from_library(input_models, output_model) + if idx in leading_group_idx: self._update_exptime(img) + if self.blendheaders: blender.accumulate(img) + iscale = self._get_intensity_scale(img) + self._iscales[idx] = iscale log.debug(f'Using intensity scale iscale={iscale}') - img.meta.iscale = iscale inwht = resample_utils.build_driz_weight( img, weight_type=self.weight_type, good_bits=self.good_bits, ) + # apply sky subtraction blevel = img.meta.background.level if not img.meta.background.subtracted and blevel is not None: @@ -499,6 +516,7 @@ def resample_many_to_one(self, input_models): # Resample variance arrays in input_models to output_model self._resample_variance_arrays( model=img, + iscale=iscale, inwht=inwht, pixmap=pixmap, in_image_limits=in_image_limits, @@ -506,19 +524,17 @@ def resample_many_to_one(self, input_models): ) del data, inwht - del img.meta.iscale input_models.shelve(img) # Since the context array is dynamic, it must be re-assigned # back to the product's `con` attribute. - output_model.data[:, :] = driz.out_img - output_model.wht[:, :] = driz.out_wht + output_model.data = driz.out_img + output_model.wht = driz.out_wht if driz.out_ctx is not None: output_model.con = driz.out_ctx - if self.blendheaders: - blender.finalize_model(output_model) + del driz # compute final variances: self._compute_resample_variance_totals(output_model) @@ -535,20 +551,24 @@ def resample_many_to_one(self, input_models): all_nan = np.all(np.isnan(var_components), axis=0) output_model.err[all_nan] = np.nan + if self.blendheaders: + blender.finalize_model(output_model) + self._get_exptime_totals(output_model) return ModelLibrary([output_model,], on_disk=False) - def _init_variance_arrays(self, output_model): - self._weighted_rn_var = np.full_like(output_model.data, np.nan) - self._weighted_pn_var = np.full_like(output_model.data, np.nan) - self._weighted_flat_var = np.full_like(output_model.data, np.nan) - self._total_weight_rn_var = np.zeros_like(output_model.data) - self._total_weight_pn_var = np.zeros_like(output_model.data) - self._total_weight_flat_var = np.zeros_like(output_model.data) - - def _resample_variance_arrays(self, model, inwht, pixmap, in_image_limits, - output_shape): + def _init_variance_arrays(self): + shape = tuple(self.output_wcs.array_shape) + self._weighted_rn_var = np.full(shape, np.nan, dtype=np.float32) + self._weighted_pn_var = np.full(shape, np.nan, dtype=np.float32) + self._weighted_flat_var = np.full(shape, np.nan, dtype=np.float32) + self._total_weight_rn_var = np.zeros(shape, dtype=np.float32) + self._total_weight_pn_var = np.zeros(shape, dtype=np.float32) + self._total_weight_flat_var = np.zeros(shape, dtype=np.float32) + + def _resample_variance_arrays(self, model, iscale, inwht, pixmap, + in_image_limits, output_shape): xmin, xmax, ymin, ymax = in_image_limits # Do the read noise variance first, so it can be @@ -556,13 +576,13 @@ def _resample_variance_arrays(self, model, inwht, pixmap, in_image_limits, rn_var = self._resample_one_variance_array( "var_rnoise", input_model=model, - output_shape=output_shape, + iscale=iscale, inwht=inwht, pixmap=pixmap, xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) # Find valid weighting values in the variance @@ -600,13 +620,13 @@ def _resample_variance_arrays(self, model, inwht, pixmap, in_image_limits, pn_var = self._resample_one_variance_array( "var_poisson", input_model=model, - output_shape=output_shape, + iscale=iscale, inwht=inwht, pixmap=pixmap, xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) if pn_var is not None: mask = (pn_var >= 0) & np.isfinite(pn_var) & (weight > 0) @@ -622,13 +642,13 @@ def _resample_variance_arrays(self, model, inwht, pixmap, in_image_limits, flat_var = self._resample_one_variance_array( "var_flat", input_model=model, - output_shape=output_shape, + iscale=iscale, inwht=inwht, pixmap=pixmap, xmin=xmin, xmax=xmax, ymin=ymin, - ymax=ymax + ymax=ymax, ) if flat_var is not None: mask = (flat_var >= 0) & np.isfinite(flat_var) & (weight > 0) @@ -685,7 +705,7 @@ def resample_variance_arrays(self, output_model, input_models): The output_model is modified in place. """ - self._init_variance_arrays(output_model) + self._init_variance_arrays() output_shape = output_model.meta.wcs.array_shape log.info("Resampling variance components") @@ -706,22 +726,23 @@ def resample_variance_arrays(self, output_model, input_models): model.data.shape, model.meta.wcs.bounding_box, ) + iscale = self._iscales[i] self._resample_variance_arrays( model=model, + iscale=iscale, inwht=inwht, pixmap=pixmap, in_image_limits=in_image_limits, output_shape=output_shape, ) - del model.meta.iscale del inwht input_models.shelve(model, i) # We now have a sum of the weighted resampled variances. self._compute_resample_variance_totals(output_model) - def _resample_one_variance_array(self, name, input_model, output_shape, + def _resample_one_variance_array(self, name, input_model, iscale, inwht, pixmap, xmin=None, xmax=None, ymin=None, ymax=None): """Resample one variance image from an input model. @@ -746,7 +767,7 @@ def _resample_one_variance_array(self, name, input_model, output_shape, ) return - iscale = input_model.meta.iscale + output_shape = tuple(self.output_wcs.array_shape) # Resample the error array. Fill "unpopulated" pixels with NaNs. driz = Drizzle( diff --git a/jwst/resample/resample_spec.py b/jwst/resample/resample_spec.py index 0d90d0d891..e55f821767 100644 --- a/jwst/resample/resample_spec.py +++ b/jwst/resample/resample_spec.py @@ -78,6 +78,12 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, self.in_memory = kwargs.get('in_memory', True) self._recalc_pscale_ratio = False + # TODO: self._iscales is used to store computed iscale for each + # input model. + # It is used only by self.resample_variance_arrays(). + # if that method no longer needs to be supported, we can remove this. + self._iscales = {} + log.info(f"Driz parameter kernel: {self.kernel}") log.info(f"Driz parameter pixfrac: {self.pixfrac}") log.info(f"Driz parameter fillval: {self.fillval}") @@ -136,7 +142,6 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, self.pscale_ratio = nominal_area / output_pix_area else: self.pscale_ratio = 1.0 - # Set the output shape if specified if output_shape is not None: self.output_wcs.array_shape = output_shape[::-1] @@ -186,12 +191,13 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, # update meta data and wcs self.blank_output.update(input_models[0]) self.blank_output.meta.wcs = self.output_wcs + + self.output_pix_area = output_pix_area if output_pix_area is not None: self.blank_output.meta.photometry.pixelarea_steradians = output_pix_area self.blank_output.meta.photometry.pixelarea_arcsecsq = ( output_pix_area * np.rad2deg(3600)**2) - def build_nirspec_output_wcs(self, input_models, refmodel=None): """ Create a spatial/spectral WCS covering the footprint of the input. From f13840961564a177a12389c4ec9cfa13385233e5 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Sat, 19 Oct 2024 12:23:36 -0400 Subject: [PATCH 13/19] fix output model creation for spec --- jwst/resample/resample_spec.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/jwst/resample/resample_spec.py b/jwst/resample/resample_spec.py index e55f821767..a4f11d481e 100644 --- a/jwst/resample/resample_spec.py +++ b/jwst/resample/resample_spec.py @@ -179,24 +179,28 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, else: output_pix_area = None + self.output_pix_area = output_pix_area + if pscale is None: log.info(f'Specified output pixel scale ratio: {self.pscale_ratio}.') pscale = compute_spectral_pixel_scale( self.output_wcs, disp_axis=disp_axis) log.info(f'Computed output pixel scale: {3600.0 * pscale:.5g} arcsec.') - # Output model - self.blank_output = datamodels.SlitModel(tuple(self.output_wcs.array_shape)) + def _create_output_model(self, ref_input_model=None): + """ Create a new blank model and update it's meta with info from ``ref_input_model``. """ + output_model = datamodels.SlitModel(None) # update meta data and wcs - self.blank_output.update(input_models[0]) - self.blank_output.meta.wcs = self.output_wcs - - self.output_pix_area = output_pix_area - if output_pix_area is not None: - self.blank_output.meta.photometry.pixelarea_steradians = output_pix_area - self.blank_output.meta.photometry.pixelarea_arcsecsq = ( - output_pix_area * np.rad2deg(3600)**2) + if ref_input_model is not None: + output_model.update(ref_input_model) + output_model.meta.wcs = self.output_wcs + if self.output_pix_area is not None: + output_model.meta.photometry.pixelarea_steradians = self.output_pix_area + output_model.meta.photometry.pixelarea_arcsecsq = ( + self.output_pix_area * np.rad2deg(3600)**2 + ) + return output_model def build_nirspec_output_wcs(self, input_models, refmodel=None): """ From 78247b81973bb3d57b1925cbf86cc23afed8a47d Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Sun, 20 Oct 2024 18:01:56 -0400 Subject: [PATCH 14/19] Replace GWCSDrizzle with Drizzle --- jwst/resample/gwcs_drizzle.py | 165 --------------------------------- jwst/resample/resample.py | 78 ++++++++-------- jwst/resample/resample_spec.py | 2 + 3 files changed, 41 insertions(+), 204 deletions(-) delete mode 100644 jwst/resample/gwcs_drizzle.py diff --git a/jwst/resample/gwcs_drizzle.py b/jwst/resample/gwcs_drizzle.py deleted file mode 100644 index 51b5633022..0000000000 --- a/jwst/resample/gwcs_drizzle.py +++ /dev/null @@ -1,165 +0,0 @@ - -from drizzle.resample import Drizzle -from . import resample_utils - -import logging -log = logging.getLogger(__name__) -log.setLevel(logging.DEBUG) - - -class GWCSDrizzle(Drizzle): - """ - Combine images using the drizzle algorithm - - """ - def __init__(self, output_wcs=None, pixfrac=1.0, kernel="square", - fillval="NAN", disable_ctx=False, n_max_images=None): - """ - Create a new Drizzle output object and set the drizzle parameters. - - Parameters - ---------- - - product : DataModel - A data model containing results from a previous run. The three - extensions SCI, WHT, and CTX contain the combined image, total counts - and image id bitmap, respectively. The WCS of the combined image is - also read from the SCI extension. - - outwcs : `gwcs.WCS` - The world coordinate system (WCS) of the resampled image. If not - provided, the WCS is taken from product. - - pixfrac : float, optional - The fraction of a pixel that the pixel flux is confined to. The - default value of 1 has the pixel flux evenly spread across the image. - A value of 0.5 confines it to half a pixel in the linear dimension, - so the flux is confined to a quarter of the pixel area when the square - kernel is used. - - kernel : str, optional - The name of the kernel used to combine the inputs. The choice of - kernel controls the distribution of flux over the kernel. The kernel - names are: "square", "gaussian", "point", "turbo", "lanczos2", - and "lanczos3". The square kernel is the default. - - fillval : str, optional - The value a pixel is set to in the output if the input image does - not overlap it. The default value of NAN sets NaN values. - """ - self.output_wcs = output_wcs - self.pixfrac = pixfrac - - super().__init__( - kernel=kernel, - fillval=fillval, - out_shape=output_wcs.array_shape, - out_img=None, - out_wht=None, - out_ctx=None, - exptime=0.0, - begin_ctx_id=0, - max_ctx_id=n_max_images, - disable_ctx=disable_ctx - ) - - def add_image(self, insci, inwcs, pixmap=None, inwht=None, - xmin=0, xmax=0, ymin=0, ymax=0, - expin=1.0, in_units="cps", iscale=1.0): - """ - Combine an input image with the output drizzled image. - - Instead of reading the parameters from a fits file, you can set - them by calling this lower level method. `Add_fits_file` calls - this method after doing its setup. - - Parameters - ---------- - - insci : array - A 2d numpy array containing the input image to be drizzled. - it is an error to not supply an image. - - inwcs : wcs - The world coordinate system of the input image. This is - used to convert the pixels to the output coordinate system. - - inwht : array, optional - A 2d numpy array containing the pixel by pixel weighting. - Must have the same dimensions as insci. If none is supplied, - the weighting is set to one. - - pixmap : array, optional - A 3D array that maps input image coordinates onto the output - array. If provided, it can improve performance. - - xmin : int, optional - This and the following three parameters set a bounding rectangle - on the input image. Only pixels on the input image inside this - rectangle will have their flux added to the output image. Xmin - sets the minimum value of the x dimension. The x dimension is the - dimension that varies quickest on the image. All four parameters - are zero based, counting starts at zero. - - xmax : int, optional - Sets the maximum value of the x dimension on the bounding box - of the input image. If ``xmax = 0``, no maximum will - be set in the x dimension (all pixels in a row of the input image - will be resampled). - - ymin : int, optional - Sets the minimum value in the y dimension on the bounding box. The - y dimension varies less rapidly than the x and represents the line - index on the input image. - - ymax : int, optional - Sets the maximum value in the y dimension. If ``ymax = 0``, - no maximum will be set in the y dimension (all pixels in a column - of the input image will be resampled). - - expin : float, optional - The exposure time of the input image, a positive number. The - exposure time is used to scale the image if the units are counts and - to scale the image weighting if the drizzle was initialized with - wt_scl equal to "exptime" or "expsq." - - in_units : str, optional - The units of the input image. The units can either be "counts" - or "cps" (counts per second.) If the value is counts, before using - the input image it is scaled by dividing it by the exposure time. - - iscale : float, optional - A scale factor to be applied to pixel intensities of the - input image before resampling. - - """ - # Compute the mapping between the input and output pixel coordinates - # for use in drizzle.cdrizzle.tdriz - if pixmap is None: - pixmap = resample_utils.calc_gwcs_pixmap( - inwcs, - self.output_wcs, - insci.shape - ) - - log.debug(f"Pixmap shape: {pixmap[:,:,0].shape}") - log.debug(f"Input Sci shape: {insci.shape}") - log.debug(f"Output Sci shape: {self.out_img.shape}") - - # Call 'drizzle' to perform image combination - log.info(f"Drizzling {insci.shape} --> {self.out_img.shape}") - - super().add_image( - data=insci, - exptime=expin, - pixmap=pixmap, - scale=iscale, - weight_map=inwht, - wht_scale=1.0, # hard-coded for JWST count-rate data - pixfrac=self.pixfrac, - in_units=in_units, - xmin=xmin, - xmax=xmax, - ymin=ymin, - ymax=ymax, - ) diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index 488057bc0f..7df2217412 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -16,7 +16,6 @@ from jwst.model_blender.blender import ModelBlender from jwst.resample import resample_utils -from jwst.resample.gwcs_drizzle import GWCSDrizzle log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) @@ -150,6 +149,9 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, np.deg2rad(tr['cdelt2'].factor.value) ) + self.out_arr_shape = tuple(self.output_wcs.array_shape) + log.debug(f"Output mosaic size: {self.out_arr_shape}") + if pscale is None: pscale = np.rad2deg(np.sqrt(self.output_pix_area)) log.info(f'Computed output pixel scale: {3600.0 * pscale} arcsec.') @@ -170,7 +172,7 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, available_memory = psutil.virtual_memory().available + psutil.swap_memory().total # compute the output array size - required_memory = np.prod(self.output_wcs.array_shape) * dtype.itemsize + required_memory = np.prod(self.out_arr_shape) * dtype.itemsize # compare used to available used_fraction = required_memory / available_memory @@ -275,23 +277,20 @@ def resample_group(self, input_models, indices, compute_error=False): output_model = None # Initialize the output with the wcs - driz = GWCSDrizzle( - self.output_wcs, + driz = Drizzle( + out_shape=self.out_arr_shape, kernel=self.kernel, fillval=self.fillval, disable_ctx=True, ) - # Also make a temporary model to hold error data if compute_error: - driz_error = GWCSDrizzle( - self.output_wcs, - pixfrac=self.pixfrac, + driz_error = Drizzle( + out_shape=self.out_arr_shape, kernel=self.kernel, fillval=self.fillval, disable_ctx=True, ) - log.info(f"{len(indices)} exposures to drizzle together") for index in indices: img = input_models.borrow(index) @@ -344,12 +343,16 @@ def resample_group(self, input_models, indices, compute_error=False): self.output_wcs, img.data.shape, ) + driz.add_image( - data, - img.meta.wcs, + data=data, + exptime=img.meta.exposure.exposure_time, # GWCSDrizzle.add_image param default was 1.0 pixmap=pixmap, - iscale=iscale, - inwht=inwht, + scale=iscale, + weight_map=inwht, + wht_scale=1.0, # hard-coded for JWST count-rate data + pixfrac=self.pixfrac, + in_units="cps", # GWCSDrizzle.add_image param default xmin=xmin, xmax=xmax, ymin=ymin, @@ -361,17 +364,19 @@ def resample_group(self, input_models, indices, compute_error=False): # in the same way the image is handled if compute_error: driz_error.add_image( - img.err, - img.meta.wcs, + data=img.err, + exptime=img.meta.exposure.exposure_time, # GWCSDrizzle.add_image param default pixmap=pixmap, - iscale=iscale, - inwht=inwht, + scale=iscale, + weight_map=inwht, + wht_scale=1.0, # hard-coded for JWST count-rate data + pixfrac=self.pixfrac, + in_units="cps", # GWCSDrizzle.add_image param default xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, ) - input_models.shelve(img, index, modify=False) del img @@ -436,14 +441,13 @@ def resample_many_to_one(self, input_models): ] ) - # Initialize the output with the wcs - driz = GWCSDrizzle( - self.output_wcs, - pixfrac=self.pixfrac, + driz = Drizzle( + out_shape=self.out_arr_shape, kernel=self.kernel, fillval=self.fillval, + max_ctx_id=len(input_models), + disable_ctx=False, ) - self._init_variance_arrays() self._init_exptime_counters() @@ -502,17 +506,19 @@ def resample_many_to_one(self, input_models): ) driz.add_image( - data, - img.meta.wcs, + data=data, + exptime=img.meta.exposure.exposure_time, # GWCSDrizzle.add_image param default pixmap=pixmap, - iscale=iscale, - inwht=inwht, + scale=iscale, + weight_map=inwht, + wht_scale=1.0, # hard-coded for JWST count-rate data + pixfrac=self.pixfrac, + in_units="cps", # GWCSDrizzle.add_image param default xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax, ) - # Resample variance arrays in input_models to output_model self._resample_variance_arrays( model=img, @@ -520,7 +526,7 @@ def resample_many_to_one(self, input_models): inwht=inwht, pixmap=pixmap, in_image_limits=in_image_limits, - output_shape=output_model.meta.wcs.array_shape, + output_shape=self.out_arr_shape, ) del data, inwht @@ -559,7 +565,7 @@ def resample_many_to_one(self, input_models): return ModelLibrary([output_model,], on_disk=False) def _init_variance_arrays(self): - shape = tuple(self.output_wcs.array_shape) + shape = self.out_arr_shape self._weighted_rn_var = np.full(shape, np.nan, dtype=np.float32) self._weighted_pn_var = np.full(shape, np.nan, dtype=np.float32) self._weighted_flat_var = np.full(shape, np.nan, dtype=np.float32) @@ -706,7 +712,7 @@ def resample_variance_arrays(self, output_model, input_models): The output_model is modified in place. """ self._init_variance_arrays() - output_shape = output_model.meta.wcs.array_shape + output_shape = self.out_arr_shape log.info("Resampling variance components") with input_models: @@ -767,19 +773,13 @@ def _resample_one_variance_array(self, name, input_model, iscale, ) return - output_shape = tuple(self.output_wcs.array_shape) + output_shape = self.out_arr_shape # Resample the error array. Fill "unpopulated" pixels with NaNs. driz = Drizzle( + out_shape=output_shape, kernel=self.kernel, fillval=np.nan, - out_shape=output_shape, - out_img=None, - out_wht=None, - out_ctx=None, - exptime=0, - begin_ctx_id=0, - max_ctx_id=1, disable_ctx=True ) diff --git a/jwst/resample/resample_spec.py b/jwst/resample/resample_spec.py index a4f11d481e..e53d233839 100644 --- a/jwst/resample/resample_spec.py +++ b/jwst/resample/resample_spec.py @@ -180,6 +180,8 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, output_pix_area = None self.output_pix_area = output_pix_area + self.out_arr_shape = tuple(self.output_wcs.array_shape) + log.debug(f"Output mosaic size: {self.out_arr_shape}") if pscale is None: log.info(f'Specified output pixel scale ratio: {self.pscale_ratio}.') From 334168a4e4231012b758978976f0144ed5febd7d Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Mon, 21 Oct 2024 02:56:00 -0400 Subject: [PATCH 15/19] fix a bug in the name of output array shape variable --- jwst/resample/resample.py | 20 ++++++++++---------- jwst/resample/resample_spec.py | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index 7df2217412..c1f6271214 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -149,8 +149,8 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, np.deg2rad(tr['cdelt2'].factor.value) ) - self.out_arr_shape = tuple(self.output_wcs.array_shape) - log.debug(f"Output mosaic size: {self.out_arr_shape}") + self.output_array_shape = tuple(self.output_wcs.array_shape) + log.debug(f"Output mosaic size: {self.output_array_shape}") if pscale is None: pscale = np.rad2deg(np.sqrt(self.output_pix_area)) @@ -172,7 +172,7 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, available_memory = psutil.virtual_memory().available + psutil.swap_memory().total # compute the output array size - required_memory = np.prod(self.out_arr_shape) * dtype.itemsize + required_memory = np.prod(self.output_array_shape) * dtype.itemsize # compare used to available used_fraction = required_memory / available_memory @@ -278,7 +278,7 @@ def resample_group(self, input_models, indices, compute_error=False): # Initialize the output with the wcs driz = Drizzle( - out_shape=self.out_arr_shape, + out_shape=self.output_array_shape, kernel=self.kernel, fillval=self.fillval, disable_ctx=True, @@ -286,7 +286,7 @@ def resample_group(self, input_models, indices, compute_error=False): # Also make a temporary model to hold error data if compute_error: driz_error = Drizzle( - out_shape=self.out_arr_shape, + out_shape=self.output_array_shape, kernel=self.kernel, fillval=self.fillval, disable_ctx=True, @@ -442,7 +442,7 @@ def resample_many_to_one(self, input_models): ) driz = Drizzle( - out_shape=self.out_arr_shape, + out_shape=self.output_array_shape, kernel=self.kernel, fillval=self.fillval, max_ctx_id=len(input_models), @@ -526,7 +526,7 @@ def resample_many_to_one(self, input_models): inwht=inwht, pixmap=pixmap, in_image_limits=in_image_limits, - output_shape=self.out_arr_shape, + output_shape=self.output_array_shape, ) del data, inwht @@ -565,7 +565,7 @@ def resample_many_to_one(self, input_models): return ModelLibrary([output_model,], on_disk=False) def _init_variance_arrays(self): - shape = self.out_arr_shape + shape = self.output_array_shape self._weighted_rn_var = np.full(shape, np.nan, dtype=np.float32) self._weighted_pn_var = np.full(shape, np.nan, dtype=np.float32) self._weighted_flat_var = np.full(shape, np.nan, dtype=np.float32) @@ -712,7 +712,7 @@ def resample_variance_arrays(self, output_model, input_models): The output_model is modified in place. """ self._init_variance_arrays() - output_shape = self.out_arr_shape + output_shape = self.output_array_shape log.info("Resampling variance components") with input_models: @@ -773,7 +773,7 @@ def _resample_one_variance_array(self, name, input_model, iscale, ) return - output_shape = self.out_arr_shape + output_shape = self.output_array_shape # Resample the error array. Fill "unpopulated" pixels with NaNs. driz = Drizzle( diff --git a/jwst/resample/resample_spec.py b/jwst/resample/resample_spec.py index e53d233839..66c39618ac 100644 --- a/jwst/resample/resample_spec.py +++ b/jwst/resample/resample_spec.py @@ -180,8 +180,8 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, output_pix_area = None self.output_pix_area = output_pix_area - self.out_arr_shape = tuple(self.output_wcs.array_shape) - log.debug(f"Output mosaic size: {self.out_arr_shape}") + self.output_array_shape = tuple(self.output_wcs.array_shape) + log.debug(f"Output mosaic size: {self.output_array_shape}") if pscale is None: log.info(f'Specified output pixel scale ratio: {self.pscale_ratio}.') From 4349f05e0f3f18339c82ab398f46f570c1efa6df Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Mon, 21 Oct 2024 15:50:57 -0400 Subject: [PATCH 16/19] delete drizzled_model and update median model meta from the first image --- jwst/outlier_detection/utils.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/jwst/outlier_detection/utils.py b/jwst/outlier_detection/utils.py index d0de3bf8b8..e137dc3251 100644 --- a/jwst/outlier_detection/utils.py +++ b/jwst/outlier_detection/utils.py @@ -107,6 +107,10 @@ def median_without_resampling(input_models, in_memory = not input_models._on_disk ngroups = len(input_models) + if save_intermediate_results: + # create an empty image model for the median data + median_model = datamodels.ImageModel(None) + with input_models: for i in range(len(input_models)): @@ -128,6 +132,10 @@ def median_without_resampling(input_models, err_computer = MedianComputer(input_shape, in_memory, buffer_size, dtype) else: err_computer = None + if save_intermediate_results: + # update median model's meta with meta from the first model: + median_model.update(drizzled_model) + median_model.meta.wcs = median_wcs weight_threshold = compute_weight_threshold(weight, maskpt) drizzled_data[weight < weight_threshold] = np.nan @@ -137,6 +145,7 @@ def median_without_resampling(input_models, err_computer.append(drizzled_err, i) input_models.shelve(drizzled_model, i, modify=False) + del drizzled_model # Perform median combination on set of drizzled mosaics median_data = computer.evaluate() @@ -147,13 +156,11 @@ def median_without_resampling(input_models, if save_intermediate_results: # Save median model to fits - median_model = datamodels.ImageModel(median_data) + median_model.data = median_data if return_error: median_model.err = median_err - median_model.update(drizzled_model) - median_model.meta.wcs = median_wcs _fileio.save_median(median_model, make_output_path) - del drizzled_model + if return_error: return median_data, median_wcs, median_err @@ -219,6 +226,10 @@ def median_with_resampling(input_models, indices_by_group = list(input_models.group_indices.values()) ngroups = len(indices_by_group) + if save_intermediate_results: + # create an empty image model for the median data + median_model = datamodels.ImageModel(None) + with input_models: for i, indices in enumerate(indices_by_group): @@ -238,6 +249,10 @@ def median_with_resampling(input_models, err_computer = MedianComputer(input_shape, in_memory, buffer_size, dtype) else: err_computer = None + if save_intermediate_results: + # update median model's meta with meta from the first model: + median_model.update(drizzled_model) + median_model.meta.wcs = median_wcs weight_threshold = compute_weight_threshold(drizzled_model.wht, maskpt) drizzled_model.data[drizzled_model.wht < weight_threshold] = np.nan @@ -245,6 +260,7 @@ def median_with_resampling(input_models, if return_error: drizzled_model.err[drizzled_model.wht < weight_threshold] = np.nan err_computer.append(drizzled_model.err, i) + del drizzled_model # Perform median combination on set of drizzled mosaics median_data = computer.evaluate() @@ -255,15 +271,12 @@ def median_with_resampling(input_models, if save_intermediate_results: # Save median model to fits - median_model = datamodels.ImageModel(median_data) + median_model.data = median_data if return_error: median_model.err = median_err - median_model.update(drizzled_model) - median_model.meta.wcs = median_wcs # drizzled model already contains asn_id make_output_path = partial(make_output_path, asn_id=None) _fileio.save_median(median_model, make_output_path) - del drizzled_model if return_error: return median_data, median_wcs, median_err From 60d0c754899da83be5c8491b7a6ba7649e3bd075 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Mon, 21 Oct 2024 16:41:01 -0400 Subject: [PATCH 17/19] fix nirspec_fs_spec3 file names --- jwst/outlier_detection/utils.py | 1 - jwst/regtest/test_nirspec_fs_spec3.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/jwst/outlier_detection/utils.py b/jwst/outlier_detection/utils.py index e137dc3251..ac0d2b6192 100644 --- a/jwst/outlier_detection/utils.py +++ b/jwst/outlier_detection/utils.py @@ -161,7 +161,6 @@ def median_without_resampling(input_models, median_model.err = median_err _fileio.save_median(median_model, make_output_path) - if return_error: return median_data, median_wcs, median_err else: diff --git a/jwst/regtest/test_nirspec_fs_spec3.py b/jwst/regtest/test_nirspec_fs_spec3.py index 4fd2d62cf5..997b121961 100644 --- a/jwst/regtest/test_nirspec_fs_spec3.py +++ b/jwst/regtest/test_nirspec_fs_spec3.py @@ -36,7 +36,7 @@ def test_nirspec_fs_spec3(run_pipeline, rtdata_module, fitsdiff_default_kwargs, rtdata = rtdata_module if suffix == "median": - output = f"jw01309022001_04102_00002_nrs2_{slit_name}_{suffix}.fits" + output = f"jw01309022001_04102_00001_nrs2_{slit_name}_{suffix}.fits" # also ensure drizzled and blot models were created with the correct names assert os.path.isfile(output.replace("median", "outlier_s2d")) assert os.path.isfile(output.replace("median", "blot")) From a73900cb7f88a9dff31e4f276803ed7e385f6aaa Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Thu, 24 Oct 2024 02:54:26 -0400 Subject: [PATCH 18/19] update requirements. pin drizzle --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 584a785efd..8efd2dd897 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "astropy>=5.3", "BayesicFitting>=3.0.1", "crds>=11.17.14", - "drizzle @ git+https://github.com/mcara/drizzle.git@redesign-api", + "drizzle>=2.0.0", "gwcs>=0.21.0,<0.23.0", "numpy>=1.22,<2.0", "opencv-python-headless>=4.6.0.66", @@ -266,7 +266,7 @@ lint.ignore = [ "E741", # ambiguous variable name (O/0, l/I, etc.) "E722", # Do not use bare `except` "ANN101", # Missing type annotation for self in method - "ANN102", # Missing type annotation for cls in classmethod + "ANN102", # Missing type annotation for cls in classmethod ] [tool.ruff.lint.per-file-ignores] From 085265ec0de0c39395863721eaad09a3f5cd8fa9 Mon Sep 17 00:00:00 2001 From: Mihai Cara Date: Thu, 24 Oct 2024 03:04:03 -0400 Subject: [PATCH 19/19] remove _iscales and resample_variance_arrays --- jwst/resample/resample.py | 55 ---------------------------------- jwst/resample/resample_spec.py | 6 ---- 2 files changed, 61 deletions(-) diff --git a/jwst/resample/resample.py b/jwst/resample/resample.py index c1f6271214..72b325cbd3 100644 --- a/jwst/resample/resample.py +++ b/jwst/resample/resample.py @@ -88,12 +88,6 @@ def __init__(self, input_models, output=None, single=False, blendheaders=True, self.input_pixscale0 = None # computed pixel scale of the first image (deg) self._recalc_pscale_ratio = pscale is not None - # TODO: self._iscales is used to store computed iscale for each - # input model. - # It is used only by self.resample_variance_arrays(). - # if that method no longer needs to be supported, we can remove this. - self._iscales = {} - log.info(f"Driz parameter kernel: {self.kernel}") log.info(f"Driz parameter pixfrac: {self.pixfrac}") log.info(f"Driz parameter fillval: {self.fillval}") @@ -318,7 +312,6 @@ def resample_group(self, input_models, indices, compute_error=False): img.area = img.area iscale = self._get_intensity_scale(img) - self._iscales[index] = iscale log.debug(f'Using intensity scale iscale={iscale}') inwht = resample_utils.build_driz_weight( @@ -477,7 +470,6 @@ def resample_many_to_one(self, input_models): blender.accumulate(img) iscale = self._get_intensity_scale(img) - self._iscales[idx] = iscale log.debug(f'Using intensity scale iscale={iscale}') inwht = resample_utils.build_driz_weight( @@ -701,53 +693,6 @@ def _compute_resample_variance_totals(self, output_model): self._total_weight_flat_var, ) - def resample_variance_arrays(self, output_model, input_models): - """Resample variance arrays from input_models to the output_model. - - Variance images from each input model are resampled individually and - added to a weighted sum. If ``weight_type`` is 'ivm', the inverse of the - resampled read noise variance is used as the weight for all the variance - components. If ``weight_type`` is 'exptime', the exposure time is used. - - The output_model is modified in place. - """ - self._init_variance_arrays() - output_shape = self.output_array_shape - log.info("Resampling variance components") - - with input_models: - for i, model in enumerate(input_models): - # compute pixmap and inwht arrays common to all variance arrays: - inwht = resample_utils.build_driz_weight( - model, - weight_type=self.weight_type, # weights match science - good_bits=self.good_bits, - ) - pixmap = resample_utils.calc_gwcs_pixmap( - model.meta.wcs, - output_model.meta.wcs, - model.data.shape, - ) - in_image_limits = resample_utils._resample_range( - model.data.shape, - model.meta.wcs.bounding_box, - ) - iscale = self._iscales[i] - self._resample_variance_arrays( - model=model, - iscale=iscale, - inwht=inwht, - pixmap=pixmap, - in_image_limits=in_image_limits, - output_shape=output_shape, - ) - - del inwht - input_models.shelve(model, i) - - # We now have a sum of the weighted resampled variances. - self._compute_resample_variance_totals(output_model) - def _resample_one_variance_array(self, name, input_model, iscale, inwht, pixmap, xmin=None, xmax=None, ymin=None, ymax=None): diff --git a/jwst/resample/resample_spec.py b/jwst/resample/resample_spec.py index 66c39618ac..5160034ddb 100644 --- a/jwst/resample/resample_spec.py +++ b/jwst/resample/resample_spec.py @@ -78,12 +78,6 @@ def __init__(self, input_models, output=None, single=False, blendheaders=False, self.in_memory = kwargs.get('in_memory', True) self._recalc_pscale_ratio = False - # TODO: self._iscales is used to store computed iscale for each - # input model. - # It is used only by self.resample_variance_arrays(). - # if that method no longer needs to be supported, we can remove this. - self._iscales = {} - log.info(f"Driz parameter kernel: {self.kernel}") log.info(f"Driz parameter pixfrac: {self.pixfrac}") log.info(f"Driz parameter fillval: {self.fillval}")