update_dq improvement in findsat_mrt

acstools/findsat_mrt.py

@@ -1360,6 +1360,12 @@ def rebin(self):
 
         LOG.info('Rebinning the data by {}'.format(self.binsize))
 
+        # trigger a warning if the image dimensions are not even
+        if ((self.image.shape[0] % self.binsize != 0)) | \
+           ((self.image.shape[1] % self.binsize != 0)):
+            LOG.warning('Warning: image dimensions do not evenly divide by bin'
+                        'size.\n Extra rows/columns will be trimmed first')
+
         # setting a warning filter for the nansum calculations. These
         # warnings are inconsequential and already accounted for in the code
         with warnings.catch_warnings():
@@ -1406,4 +1412,4 @@ def update_dq(self, dqval=16384, verbose=True):
             raise ValueError(f'DQ array can only be updated for FLC/FLT images, not {self.image_type}')  # noqa
 
         update_dq(self.image_file, self.extension + 2, self.mask, dqval=dqval,
-                  verbose=verbose)
+                  verbose=verbose, expand_mask=True)

acstools/utils_findsat_mrt.py

@@ -10,7 +10,7 @@
 from astropy.convolution import convolve, Gaussian2DKernel
 from astropy.io import fits
 from astropy.modeling import models, fitting
-from astropy.nddata import Cutout2D
+from astropy.nddata import Cutout2D, block_replicate
 from astropy.stats import sigma_clip
 from astropy.table import Table, vstack
 from astropy.utils.exceptions import AstropyUserWarning
@@ -1179,7 +1179,8 @@ def radon(image, theta=None, circle=False, *, preserve_range=False,
         return radon_image
 
 
-def update_dq(filename, ext, mask, dqval=16384, verbose=True):
+def update_dq(filename, ext, mask, dqval=16384, verbose=True,
+              expand_mask=False):
     """Update the given image and DQ extension with the given
     satellite trails mask and flag.
 
@@ -1199,10 +1200,54 @@ def update_dq(filename, ext, mask, dqval=16384, verbose=True):
         tailored for ACS/WFC.
     verbose : bool, optional
         Print extra information to the terminal.
+    expand_mask: int, optional
+        Allows the mask to be expanded to match the size of the original mask.
+        This is relevant for masks that were generated used binned versions of
+        the original image
 
     """
     with fits.open(filename, mode='update') as pf:
         dqarr = pf[ext].data
+
+        # check that the old mask and new mask are the same size
+        same_size = mask.shape == dqarr.shape
+
+        if not same_size:
+            LOG.info('Inconsistent mask sizes: \n'
+                     'Input mask: {} \n'
+                     'DQ array: {}'.format(mask.shape, dqarr.shape))
+            if not expand_mask:
+                LOG.warning('Cannot proceed due to size mismatch.\n'
+                            'Set expand_mask=True if the mask was generated '
+                            'with binned data and needs to be enlarged to the '
+                            'original size')
+                return
+            else:
+                LOG.info('Enlarging mask to original size')
+                factor_x = np.floor(dqarr.shape[1]/mask.shape[1]).astype(int)
+                factor_y = np.floor(dqarr.shape[0]/mask.shape[0]).astype(int)
+                mask = block_replicate(mask, [factor_y, factor_x],
+                                       conserve_sum=False)
+
+                # for original data that had odd dimensions before binning, the
+                # mask still will not match the original. Duplicate end
+                # rows/columns in this case
+
+                same_size = mask.shape == dqarr.shape
+                if not same_size:
+                    LOG.info('Still inconsistent mask sizes: \n'
+                             'Input mask: {} \n'
+                             'DQ array: {}'.format(mask.shape, dqarr.shape))
+                    LOG.info('Duplicating end rows/columns to match orig size')
+                    while mask.shape[0] != dqarr.shape[0]:
+                        mask = np.vstack([mask,
+                                          mask[-1, :]])
+
+                    while mask.shape[1] != dqarr.shape[1]:
+                        mask = np.hstack([mask,
+                                          mask[:, -1].reshape(mask.shape[0],
+                                                              1)])
+
         old_mask = (dqval & dqarr) != 0  # Existing flagged trails
         new_mask = mask & ~old_mask  # Only flag previously unflagged trails
         npix_updated = np.count_nonzero(new_mask)