Add some more flexibility to coadd: output array specification

reproject/mosaicking/coadd.py

@@ -24,6 +24,9 @@
     background_reference=None,
     output_array=None,
     output_footprint=None,
+    block_sizes=None,
+    progressbar=False,
+    blank_pixel_value=np.nan,
     **kwargs,
 ):
     """
@@ -92,12 +95,20 @@
     output_array : array or None
         The final output array.  Specify this if you already have an
         appropriately-shaped array to store the data in.  Must match shape
-        specified with ``shape_out`` or derived from the output projection.
+        specified with `shape_out` or derived from the output
+        projection.
     output_footprint : array or None
         The final output footprint array.  Specify this if you already have an
         appropriately-shaped array to store the data in.  Must match shape
-        specified with ``shape_out`` or derived from the output projection.
-    **kwargs
+        specified with `shape_out` or derived from the output projection.
+    block_sizes : list of tuples or None
+        The block size to use for each dataset.  Could also be a single tuple
+        if you want the sample block size for all data sets
+    progressbar : False
+        If specified, use this as a progressbar to track loop iterations over
+        data sets.
+
+    kwargs
         Keyword arguments to be passed to the reprojection function.
 
     Returns
@@ -118,12 +129,18 @@
 
     if combine_function not in ("mean", "sum", "median", "first", "last", "min", "max"):
         raise ValueError("combine_function should be one of mean/sum/median/first/last/min/max")
+    elif combine_function == "median":
+        # Note to devs: the exception shoudl be raised as early as possible
+        raise NotImplementedError("combine_function='median' is not yet implemented")
 
     if reproject_function is None:
         raise ValueError(
             "reprojection function should be specified with the reproject_function argument"
         )
 
+    if not progressbar:
+        progressbar = lambda x: x
+
     # Parse the output projection to avoid having to do it for each
 
     wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)
@@ -139,11 +156,16 @@
             f"the output shape {shape_out}"
         )
 
-    # Start off by reprojecting individual images to the final projection
+    if output_array is None:
+        output_array = np.zeros(shape_out)
+    if output_footprint is None:
+        output_footprint = np.zeros(shape_out)
 
-    arrays = []
+    # Start off by reprojecting individual images to the final projection
+    if match_background:
+        arrays = []
 
-    for idata in range(len(input_data)):
+    for idata in progressbar(range(len(input_data))):
         # We need to pre-parse the data here since we need to figure out how to
         # optimize/minimize the size of each output tile (see below).
         array_in, wcs_in = parse_input_data(input_data[idata], hdu_in=hdu_in)
@@ -166,13 +188,31 @@
         # significant distortion (when the edges of the input image become
         # convex in the output projection), and transforming every edge pixel,
         # which provides a lot of redundant information.
-        ny, nx = array_in.shape
-        n_per_edge = 11
-        xs = np.linspace(-0.5, nx - 0.5, n_per_edge)
-        ys = np.linspace(-0.5, ny - 0.5, n_per_edge)
-        xs = np.concatenate((xs, np.full(n_per_edge, xs[-1]), xs, np.full(n_per_edge, xs[0])))
-        ys = np.concatenate((np.full(n_per_edge, ys[0]), ys, np.full(n_per_edge, ys[-1]), ys))
-        xc_out, yc_out = wcs_out.world_to_pixel(wcs_in.pixel_to_world(xs, ys))
+        if array_in.ndim == 2:
+            ny, nx = array_in.shape
+            n_per_edge = 11
+            xs = np.linspace(-0.5, nx - 0.5, n_per_edge)
+            ys = np.linspace(-0.5, ny - 0.5, n_per_edge)
+            xs = np.concatenate((xs, np.full(n_per_edge, xs[-1]), xs, np.full(n_per_edge, xs[0])))
+            ys = np.concatenate((np.full(n_per_edge, ys[0]), ys, np.full(n_per_edge, ys[-1]), ys))
+            xc_out, yc_out = wcs_out.world_to_pixel(wcs_in.pixel_to_world(xs, ys))
+            shape_out_cel = shape_out
+        elif array_in.ndim == 3:
+            # for cubes, we only handle single corners now
+            nz, ny, nx = array_in.shape
+            xc = np.array([-0.5, nx - 0.5, nx - 0.5, -0.5])
+            yc = np.array([-0.5, -0.5, ny - 0.5, ny - 0.5])
+            zc = np.array([-0.5, nz - 0.5])
+            # TODO: figure out what to do here if the low_level_wcs doesn't  support subsetting
+            xc_out, yc_out = wcs_out.low_level_wcs.celestial.world_to_pixel(
+                wcs_in.celestial.pixel_to_world(xc, yc)
+            )
+            zc_out = wcs_out.low_level_wcs.spectral.world_to_pixel(
+                wcs_in.spectral.pixel_to_world(zc)
+            )
+            shape_out_cel = shape_out[1:]
+        else:
+            raise ValueError(f"Wrong number of dimensions: {array_in.ndim}")
 
         # Determine the cutout parameters
 
@@ -182,26 +222,43 @@
 
         if np.any(np.isnan(xc_out)) or np.any(np.isnan(yc_out)):
             imin = 0
-            imax = shape_out[1]
+            imax = shape_out_cel[1]
             jmin = 0
-            jmax = shape_out[0]
+            jmax = shape_out_cel[0]
         else:
             imin = max(0, int(np.floor(xc_out.min() + 0.5)))
-            imax = min(shape_out[1], int(np.ceil(xc_out.max() + 0.5)))
+            imax = min(shape_out_cel[1], int(np.ceil(xc_out.max() + 0.5)))
             jmin = max(0, int(np.floor(yc_out.min() + 0.5)))
-            jmax = min(shape_out[0], int(np.ceil(yc_out.max() + 0.5)))
+            jmax = min(shape_out_cel[0], int(np.ceil(yc_out.max() + 0.5)))
 
         if imax < imin or jmax < jmin:
             continue
 
-        if isinstance(wcs_out, WCS):
-            wcs_out_indiv = wcs_out[jmin:jmax, imin:imax]
-        else:
-            wcs_out_indiv = SlicedLowLevelWCS(
-                wcs_out.low_level_wcs, (slice(jmin, jmax), slice(imin, imax))
-            )
+        if array_in.ndim == 2:
+            if isinstance(wcs_out, WCS):
+                wcs_out_indiv = wcs_out[jmin:jmax, imin:imax]
+            else:
+                wcs_out_indiv = SlicedLowLevelWCS(
+                    wcs_out.low_level_wcs, (slice(jmin, jmax), slice(imin, imax))
+                )
+            shape_out_indiv = (jmax - jmin, imax - imin)
+            kmin, kmax = None, None  # for reprojectedarraysubset below
+        elif array_in.ndim == 3:
+            kmin = max(0, int(np.floor(zc_out.min() + 0.5)))
+            kmax = min(shape_out[0], int(np.ceil(zc_out.max() + 0.5)))
+            if isinstance(wcs_out, WCS):
+                wcs_out_indiv = wcs_out[kmin:kmax, jmin:jmax, imin:imax]
+            else:
+                wcs_out_indiv = SlicedLowLevelWCS(
+                    wcs_out.low_level_wcs, (slice(kmin, kmax), slice(jmin, jmax), slice(imin, imax))
+                )
+            shape_out_indiv = (kmax - kmin, jmax - jmin, imax - imin)
 
-        shape_out_indiv = (jmax - jmin, imax - imin)
+        if block_sizes is not None:
+            if len(block_sizes) == len(input_data) and len(block_sizes[idata]) == len(shape_out):
+                kwargs["block_size"] = block_sizes[idata]
+            else:
+                kwargs["block_size"] = block_sizes
 
         # TODO: optimize handling of weights by making reprojection functions
         # able to handle weights, and make the footprint become the combined
@@ -235,12 +292,22 @@
             weights[reset] = 0.0
             footprint *= weights
 
-        array = ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax)
+        array = ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax, kmin, kmax)
 
         # TODO: make sure we gracefully handle the case where the
         # output image is empty (due e.g. to no overlap).
 
-        arrays.append(array)
+        if match_background:
+            arrays.append(array)
+        else:
+            if combine_function in ("mean", "sum"):
+                # By default, values outside of the footprint are set to NaN
+                # but we set these to 0 here to avoid getting NaNs in the
+                # means/sums.
+                array.array[array.footprint == 0] = 0
+
+                output_array[array.view_in_original_array] += array.array * array.footprint
+                output_footprint[array.view_in_original_array] += array.footprint
 
     # If requested, try and match the backgrounds.
     if match_background and len(arrays) > 1:
@@ -251,62 +318,58 @@
         for array, correction in zip(arrays, corrections):
             array.array -= correction
 
-    # At this point, the images are now ready to be co-added.
-
-    if output_array is None:
-        output_array = np.zeros(shape_out)
-    if output_footprint is None:
-        output_footprint = np.zeros(shape_out)
-
     if combine_function == "min":
         output_array[...] = np.inf
     elif combine_function == "max":
         output_array[...] = -np.inf
 
     if combine_function in ("mean", "sum"):
-        for array in arrays:
-            # By default, values outside of the footprint are set to NaN
-            # but we set these to 0 here to avoid getting NaNs in the
-            # means/sums.
-            array.array[array.footprint == 0] = 0
+        if match_background:
+            # if we're not matching the background, this part has already been done
+            for array in arrays:
+                # By default, values outside of the footprint are set to NaN
+                # but we set these to 0 here to avoid getting NaNs in the
+                # means/sums.
+                array.array[array.footprint == 0] = 0
 
-            output_array[array.view_in_original_array] += array.array * array.footprint
-            output_footprint[array.view_in_original_array] += array.footprint
+        output_array[array.view_in_original_array] += array.array * array.footprint
+        output_footprint[array.view_in_original_array] += array.footprint
 
         if combine_function == "mean":
             with np.errstate(invalid="ignore"):
                 output_array /= output_footprint
-                output_array[output_footprint == 0] = 0
+                output_array[output_footprint == 0] = blank_pixel_value
 
     elif combine_function in ("first", "last", "min", "max"):
-        for array in arrays:
-            if combine_function == "first":
-                mask = (output_footprint[array.view_in_original_array] == 0) & (array.footprint > 0)
-            elif combine_function == "last":
-                mask = array.footprint > 0
-            elif combine_function == "min":
-                mask = (array.footprint > 0) & (
-                    array.array < output_array[array.view_in_original_array]
+        if match_background:
+            for array in arrays:
+                if combine_function == "first":
+                    mask = output_footprint[array.view_in_original_array] == 0
+                elif combine_function == "last":
+                    mask = array.footprint > 0
+                elif combine_function == "min":
+                    mask = (array.footprint > 0) & (
+                        array.array < output_array[array.view_in_original_array]
+                    )
+                elif combine_function == "max":
+                    mask = (array.footprint > 0) & (
+                        array.array > output_array[array.view_in_original_array]
+                    )
+
+                output_footprint[array.view_in_original_array] = np.where(
+                    mask, array.footprint, output_footprint[array.view_in_original_array]
                 )
-            elif combine_function == "max":
-                mask = (array.footprint > 0) & (
-                    array.array > output_array[array.view_in_original_array]
+                output_array[array.view_in_original_array] = np.where(
+                    mask, array.array, output_array[array.view_in_original_array]
                 )
 
-            output_footprint[array.view_in_original_array] = np.where(
-                mask, array.footprint, output_footprint[array.view_in_original_array]
-            )
-            output_array[array.view_in_original_array] = np.where(
-                mask, array.array, output_array[array.view_in_original_array]
-            )
-
     elif combine_function == "median":
         # Here we need to operate in chunks since we could otherwise run
         # into memory issues
 
+        # this is redundant, but left as a note-to-devs about where such an implementation belongs
         raise NotImplementedError("combine_function='median' is not yet implemented")
 
-    if combine_function in ("min", "max"):
-        output_array[output_footprint == 0] = 0.0
+    output_array[output_footprint == 0] = blank_pixel_value
 
     return output_array, output_footprint