Flag out NaN weights in coadd & allow for weights to have different WCS than data

reproject/mosaicking/coadd.py

@@ -11,6 +11,7 @@
 from astropy.wcs.wcsapi import SlicedLowLevelWCS
 
 from ..array_utils import iterate_chunks, sample_array_edges
+from ..interpolation.core import _validate_wcs
 from ..utils import parse_input_data, parse_input_weights, parse_output_projection
 from .background import determine_offset_matrix, solve_corrections_sgd
 from .subset_array import ReprojectedArraySubset
@@ -211,7 +212,18 @@ def reproject_and_coadd(
             if input_weights is None:
                 weights_in = None
             else:
-                weights_in = parse_input_weights(input_weights[idata], hdu_weights=hdu_weights)
+                weights_in, weights_wcs = parse_input_weights(
+                    input_weights[idata], hdu_weights=hdu_weights
+                )
+                if weights_wcs is None:
+                    # if weights are passed as an array
+                    weights_wcs = wcs_in
+                else:
+                    try:
+                        _validate_wcs(weights_wcs, wcs_in, weights_in.shape, shape_out)
+                    except ValueError:
+                        # WCS is not valid (most likely, it is blank?)
+                        weights_wcs = wcs_in
                 if np.any(np.isnan(weights_in)):
                     weights_in = np.nan_to_num(weights_in)
 
@@ -347,18 +359,20 @@ def reproject_and_coadd(
                 )
 
                 weights = reproject_function(
-                    (weights_in, wcs_in),
+                    (weights_in, weights_wcs),
                     output_projection=wcs_out_indiv,
                     shape_out=shape_out_indiv,
                     hdu_in=hdu_in,
                     output_array=weights,
                     return_footprint=False,
                     **kwargs,
                 )
+                reset = np.isnan(array) | np.isnan(weights)
+            else:
+                reset = np.isnan(array)
 
             # For the purposes of mosaicking, we mask out NaN values from the array
             # and set the footprint to 0 at these locations.
-            reset = np.isnan(array)
             array[reset] = 0.0
             footprint[reset] = 0.0
 

reproject/mosaicking/tests/test_coadd.py

@@ -304,7 +304,7 @@ def test_coadd_background_matching_with_nan(self, reproject_function, intermedia
         assert_allclose(array - np.mean(array), self.array - np.mean(self.array), atol=ATOL)
 
     @pytest.mark.filterwarnings("ignore:unclosed file:ResourceWarning")
-    @pytest.mark.parametrize("mode", ["arrays", "filenames", "hdus"])
+    @pytest.mark.parametrize("mode", ["arrays", "filenames", "hdus", "hdulist"])
     def test_coadd_with_weights(self, tmpdir, reproject_function, mode, intermediate_memmap):
         # Make sure that things work properly when specifying weights
 
@@ -328,6 +328,10 @@ def test_coadd_with_weights(self, tmpdir, reproject_function, mode, intermediate
             hdu1 = fits.ImageHDU(weight1)
             hdu2 = fits.ImageHDU(weight2)
             input_weights = [hdu1, hdu2]
+        elif mode == "hdulist":
+            hdu1 = fits.HDUList([fits.ImageHDU(weight1, header=self.wcs.to_header())])
+            hdu2 = fits.HDUList([fits.ImageHDU(weight2, header=self.wcs.to_header())])
+            input_weights = [hdu1, hdu2]
 
         array, footprint = reproject_and_coadd(
             input_data,
@@ -343,6 +347,54 @@ def test_coadd_with_weights(self, tmpdir, reproject_function, mode, intermediate
 
         assert_allclose(array, expected, atol=ATOL)
 
+    @pytest.mark.filterwarnings("ignore:unclosed file:ResourceWarning")
+    def test_coadd_with_weights_with_wcs(self, tmpdir, reproject_function, intermediate_memmap):
+        # Make sure that things work properly when specifying weights that have offset WCS
+
+        array1 = self.array + 1
+        array2 = self.array - 1
+
+        weight1 = np.cumsum(np.ones_like(self.array), axis=1) - 1
+        weight2 = weight1[:, ::-1]
+
+        input_data = [(array1, self.wcs), (array2, self.wcs)]
+
+        # make weight WCS pixel scale bigger so that weights encompass data
+        weightwcs = self.wcs.copy()
+        weightwcs.wcs.cdelt *= 1.1
+
+        hdu1 = fits.ImageHDU(weight1, header=weightwcs.to_header())
+        hdu2 = fits.ImageHDU(weight2, header=weightwcs.to_header())
+        input_weights = [hdu1, hdu2]
+
+        array, footprint = reproject_and_coadd(
+            input_data,
+            self.wcs,
+            shape_out=self.array.shape,
+            combine_function="mean",
+            input_weights=input_weights,
+            reproject_function=reproject_function,
+            match_background=False,
+        )
+
+        weights1_reprojected = reproject_function(
+            hdu1, self.wcs, shape_out=self.array.shape, return_footprint=False
+        )
+        weights2_reprojected = reproject_function(
+            hdu2, self.wcs, shape_out=self.array.shape, return_footprint=False
+        )
+        array1_reprojected = reproject_function(
+            input_data[0], self.wcs, shape_out=self.array.shape, return_footprint=False
+        )
+        array2_reprojected = reproject_function(
+            input_data[1], self.wcs, shape_out=self.array.shape, return_footprint=False
+        )
+        expected = (
+            array1_reprojected * weights1_reprojected + array2_reprojected * weights2_reprojected
+        ) / (weights1_reprojected + weights2_reprojected)
+
+        assert_allclose(array, expected, atol=ATOL)
+
 
 HEADER_SOLAR_OUT = """
 WCSAXES =                    2

reproject/utils.py

@@ -198,7 +198,7 @@ def parse_input_weights(input_weights, hdu_weights=None):
     """
 
     if isinstance(input_weights, str):
-        return parse_input_data(fits.open(input_weights), hdu_in=hdu_weights)[0]
+        return parse_input_data(fits.open(input_weights), hdu_in=hdu_weights)
     elif isinstance(input_weights, HDUList):
         if hdu_weights is None:
             if len(input_weights) > 1:
@@ -208,11 +208,16 @@ def parse_input_weights(input_weights, hdu_weights=None):
                 )
             else:
                 hdu_weights = 0
-        return parse_input_data(input_weights[hdu_weights])[0]
+        return parse_input_data(input_weights[hdu_weights])
     elif isinstance(input_weights, PrimaryHDU | ImageHDU | CompImageHDU):
-        return input_weights.data
+        if "CTYPE1" in input_weights.header:
+            # all valid WCSes have CTYPE1 specified, at least
+            ww = WCS(input_weights.header)
+        else:
+            ww = None
+        return input_weights.data, ww
     elif isinstance(input_weights, np.ndarray):
-        return input_weights
+        return input_weights, None
     else:
         raise TypeError("input_weights should either be an HDU object or a Numpy array")