From 50bb23f4140af985bff8ee0c9f7b02fc7e0322a6 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 17:13:14 -0400
Subject: [PATCH 01/13] add ability to specify output array in coadd, and allow
 specification of block size making

---
 reproject/mosaicking/coadd.py | 46 +++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 21 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 8de178e35..649eb78b8 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -22,8 +22,9 @@ def reproject_and_coadd(
     combine_function="mean",
     match_background=False,
     background_reference=None,
-    output_array=None,
-    output_footprint=None,
+    final_array=None,
+    final_footprint=None,
+    block_sizes=None,
     **kwargs,
 ):
     """
@@ -89,15 +90,19 @@ def reproject_and_coadd(
         If `None`, the background matching will make it so that the average of
         the corrections for all images is zero. If an integer, this specifies
         the index of the image to use as a reference.
-    output_array : array or None
+    final_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.
-    output_footprint : array or None
+        specified with `shape_out` or derived from the output
+        projection.
+    final_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 cube (optional; meant for dask use)
+
+    kwargs
         Keyword arguments to be passed to the reprojection function.
 
     Returns
@@ -128,16 +133,12 @@ def reproject_and_coadd(
 
     wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)
 
-    if output_array is not None and output_array.shape != shape_out:
-        raise ValueError(
-            "If you specify an output array, it must have a shape matching "
-            f"the output shape {shape_out}"
-        )
-    if output_footprint is not None and output_footprint.shape != shape_out:
-        raise ValueError(
-            "If you specify an output footprint array, it must have a shape matching "
-            f"the output shape {shape_out}"
-        )
+    if final_array is not None and final_array.shape != shape_out:
+        raise ValueError("If you specify an output array, it must have a shape matching "
+                         f"the output shape {shape_out}")
+    if final_footprint is not None and final_footprint.shape != shape_out:
+        raise ValueError("If you specify an output footprint array, it must have a shape matching "
+                         f"the output shape {shape_out}")
 
     # Start off by reprojecting individual images to the final projection
 
@@ -203,6 +204,9 @@ def reproject_and_coadd(
 
         shape_out_indiv = (jmax - jmin, imax - imin)
 
+        if block_sizes is not None:
+            kwargs['block_size'] = block_sizes[idata]
+
         # TODO: optimize handling of weights by making reprojection functions
         # able to handle weights, and make the footprint become the combined
         # footprint + weight map
@@ -253,10 +257,10 @@ def reproject_and_coadd(
 
     # 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 final_array is None:
+        final_array = np.zeros(shape_out)
+    if final_footprint is None:
+        final_footprint = np.zeros(shape_out)
 
     if combine_function == "min":
         output_array[...] = np.inf

From c012c5b1d60df9d2b5e24173bd4f1a7d768e58d7 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 17:59:58 -0400
Subject: [PATCH 02/13] allow reproject_and_coadd to operate on cubes

---
 reproject/mosaicking/coadd.py | 57 ++++++++++++++++++++++++-----------
 1 file changed, 39 insertions(+), 18 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 649eb78b8..c56dc6dc9 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -167,13 +167,23 @@ def reproject_and_coadd(
         # 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))
+        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])
+            xc_out, yc_out = wcs_out.celestial.world_to_pixel(wcs_in.celestial.pixel_to_world(xc, yc))
+            zc_out = wcs_out.spectral.world_to_pixel(wcs_in.spectral.pixel_to_world(zc))
+            shape_out_cel = shape_out[1:]
 
         # Determine the cutout parameters
 
@@ -183,26 +193,37 @@ def reproject_and_coadd(
 
         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)
+        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:
             kwargs['block_size'] = block_sizes[idata]

From 71c85a599918adafd9021b02e0aa9d189c1e185f Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 18:04:21 -0400
Subject: [PATCH 03/13] implement astrofrog suggestions

---
 reproject/mosaicking/coadd.py | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index c56dc6dc9..bac8421e8 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -22,8 +22,8 @@ def reproject_and_coadd(
     combine_function="mean",
     match_background=False,
     background_reference=None,
-    final_array=None,
-    final_footprint=None,
+    output_array=None,
+    output_footprint=None,
     block_sizes=None,
     **kwargs,
 ):
@@ -90,17 +90,18 @@ def reproject_and_coadd(
         If `None`, the background matching will make it so that the average of
         the corrections for all images is zero. If an integer, this specifies
         the index of the image to use as a reference.
-    final_array : array or None
+    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.
-    final_footprint : array or None
+    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.
     block_sizes : list of tuples or None
-        The block size to use for each cube (optional; meant for dask use)
+        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
 
     kwargs
         Keyword arguments to be passed to the reprojection function.
@@ -133,10 +134,10 @@ def reproject_and_coadd(
 
     wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)
 
-    if final_array is not None and final_array.shape != shape_out:
+    if output_array is not None and output_array.shape != shape_out:
         raise ValueError("If you specify an output array, it must have a shape matching "
                          f"the output shape {shape_out}")
-    if final_footprint is not None and final_footprint.shape != shape_out:
+    if output_footprint is not None and output_footprint.shape != shape_out:
         raise ValueError("If you specify an output footprint array, it must have a shape matching "
                          f"the output shape {shape_out}")
 
@@ -226,7 +227,10 @@ def reproject_and_coadd(
 
 
         if block_sizes is not None:
-            kwargs['block_size'] = block_sizes[idata]
+            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
@@ -278,10 +282,10 @@ def reproject_and_coadd(
 
     # At this point, the images are now ready to be co-added.
 
-    if final_array is None:
-        final_array = np.zeros(shape_out)
-    if final_footprint is None:
-        final_footprint = np.zeros(shape_out)
+    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

From e40477c06c473951f56c8576ee104e948dae163d Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 18:28:46 -0400
Subject: [PATCH 04/13] extend dimensionality of reprojectedarraysubset

---
 reproject/mosaicking/coadd.py        |   3 +-
 reproject/mosaicking/subset_array.py | 125 +++++++++++++++++++--------
 2 files changed, 92 insertions(+), 36 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index bac8421e8..7c58d4a59 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -214,6 +214,7 @@ def reproject_and_coadd(
                     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)))
@@ -264,7 +265,7 @@ def reproject_and_coadd(
             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).
diff --git a/reproject/mosaicking/subset_array.py b/reproject/mosaicking/subset_array.py
index 37bde7ebb..9e9334a0e 100644
--- a/reproject/mosaicking/subset_array.py
+++ b/reproject/mosaicking/subset_array.py
@@ -15,35 +15,56 @@ class ReprojectedArraySubset:
     # rather than the center, which is not well defined for even-sized
     # cutouts.
 
-    def __init__(self, array, footprint, imin, imax, jmin, jmax):
+    def __init__(self, array, footprint, imin, imax, jmin, jmax, kmin=None, kmax=None):
         self.array = array
         self.footprint = footprint
         self.imin = imin
         self.imax = imax
         self.jmin = jmin
         self.jmax = jmax
+        self.kmin = kmin
+        self.kmax = kmax
 
     def __repr__(self):
-        return f"<ReprojectedArraySubset at [{self.jmin}:{self.jmax},{self.imin}:{self.imax}]>"
+        if self.kmin is not None:
+            return f"<ReprojectedArraySubset at [{self.kmin}:{self.kmax},{self.jmin}:{self.jmax},{self.imin}:{self.imax}]>"
+        else:
+            return f"<ReprojectedArraySubset at [{self.jmin}:{self.jmax},{self.imin}:{self.imax}]>"
 
     @property
     def view_in_original_array(self):
-        return (slice(self.jmin, self.jmax), slice(self.imin, self.imax))
+        if self.kmin is not None:
+            return (slice(self.kmin, self.kmax), slice(self.jmin, self.jmax), slice(self.imin, self.imax))
+        else
+            return (slice(self.jmin, self.jmax), slice(self.imin, self.imax))
 
     @property
     def shape(self):
-        return (self.jmax - self.jmin, self.imax - self.imin)
+        if self.kmin is not None:
+            return (self.kmax - self.kmin, self.jmax - self.jmin, self.imax - self.imin)
+        else:
+            return (self.jmax - self.jmin, self.imax - self.imin)
 
     def overlaps(self, other):
         # Note that the use of <= or >= instead of < and > is due to
         # the fact that the max values are exclusive (so +1 above the
         # last value).
-        return not (
-            self.imax <= other.imin
-            or other.imax <= self.imin
-            or self.jmax <= other.jmin
-            or other.jmax <= self.jmin
-        )
+        if self.kmin is not None:
+            return not (
+                self.imax <= other.imin
+                or other.imax <= self.imin
+                or self.jmax <= other.jmin
+                or other.jmax <= self.jmin
+                or self.kmax <= other.kmin
+                or other.kmax <= self.kmin
+            )
+        else:
+            return not (
+                self.imax <= other.imin
+                or other.imax <= self.imin
+                or self.jmax <= other.jmin
+                or other.jmax <= self.jmin
+            )
 
     def __add__(self, other):
         return self._operation(other, operator.add)
@@ -71,29 +92,63 @@ def _operation(self, other, op):
         if jmax < jmin:
             jmax = jmin
 
-        # Extract cutout from each
 
-        self_array = self.array[
-            jmin - self.jmin : jmax - self.jmin,
-            imin - self.imin : imax - self.imin,
-        ]
-        self_footprint = self.footprint[
-            jmin - self.jmin : jmax - self.jmin,
-            imin - self.imin : imax - self.imin,
-        ]
-
-        other_array = other.array[
-            jmin - other.jmin : jmax - other.jmin,
-            imin - other.imin : imax - other.imin,
-        ]
-        other_footprint = other.footprint[
-            jmin - other.jmin : jmax - other.jmin,
-            imin - other.imin : imax - other.imin,
-        ]
-
-        # Carry out operator and store result in ReprojectedArraySubset
-
-        array = op(self_array, other_array)
-        footprint = (self_footprint > 0) & (other_footprint > 0)
-
-        return ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax)
+        if self.kmin is None:
+            # Extract cutout from each
+
+            self_array = self.array[
+                jmin - self.jmin : jmax - self.jmin,
+                imin - self.imin : imax - self.imin,
+            ]
+            self_footprint = self.footprint[
+                jmin - self.jmin : jmax - self.jmin,
+                imin - self.imin : imax - self.imin,
+            ]
+
+            other_array = other.array[
+                jmin - other.jmin : jmax - other.jmin,
+                imin - other.imin : imax - other.imin,
+            ]
+            other_footprint = other.footprint[
+                jmin - other.jmin : jmax - other.jmin,
+                imin - other.imin : imax - other.imin,
+            ]
+
+            # Carry out operator and store result in ReprojectedArraySubset
+
+            array = op(self_array, other_array)
+            footprint = (self_footprint > 0) & (other_footprint > 0)
+
+            return ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax)
+
+        else:
+            # Extract cutout from each
+
+            self_array = self.array[
+                kmin - self.kmin : kmax - self.kmin,
+                jmin - self.jmin : jmax - self.jmin,
+                imin - self.imin : imax - self.imin,
+            ]
+            self_footprint = self.footprint[
+                kmin - self.kmin : kmax - self.kmin,
+                jmin - self.jmin : jmax - self.jmin,
+                imin - self.imin : imax - self.imin,
+            ]
+
+            other_array = other.array[
+                kmin - other.kmin : kmax - other.kmin,
+                jmin - other.jmin : jmax - other.jmin,
+                imin - other.imin : imax - other.imin,
+            ]
+            other_footprint = other.footprint[
+                kmin - other.kmin : kmax - other.kmin,
+                jmin - other.jmin : jmax - other.jmin,
+                imin - other.imin : imax - other.imin,
+            ]
+
+            # Carry out operator and store result in ReprojectedArraySubset
+
+            array = op(self_array, other_array)
+            footprint = (self_footprint > 0) & (other_footprint > 0)
+
+            return ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax, kmin, kmax)

From a35c415dd99b9c354746c419a311128900cac26a Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 18:31:05 -0400
Subject: [PATCH 05/13] fix typo

---
 reproject/mosaicking/subset_array.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/reproject/mosaicking/subset_array.py b/reproject/mosaicking/subset_array.py
index 9e9334a0e..bf0d299a9 100644
--- a/reproject/mosaicking/subset_array.py
+++ b/reproject/mosaicking/subset_array.py
@@ -35,7 +35,7 @@ def __repr__(self):
     def view_in_original_array(self):
         if self.kmin is not None:
             return (slice(self.kmin, self.kmax), slice(self.jmin, self.jmax), slice(self.imin, self.imax))
-        else
+        else:
             return (slice(self.jmin, self.jmax), slice(self.imin, self.imax))
 
     @property

From e177cbff46fe60eeed226062228b262b45eb82ba Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 18:55:01 -0400
Subject: [PATCH 06/13] only store all the arrays in memory if matching
 backgrounds; otherwise very wasteful

---
 reproject/mosaicking/coadd.py | 33 ++++++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 7c58d4a59..be74e5f85 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -141,9 +141,15 @@ def reproject_and_coadd(
         raise ValueError("If you specify an output footprint array, it must have a shape matching "
                          f"the output shape {shape_out}")
 
+    if output_array is None:
+        output_array = np.zeros(shape_out)
+    if output_footprint is None:
+        output_footprint = np.zeros(shape_out)
+
     # Start off by reprojecting individual images to the final projection
 
-    arrays = []
+    if match_background:
+        arrays = []
 
     for idata in range(len(input_data)):
         # We need to pre-parse the data here since we need to figure out how to
@@ -270,7 +276,18 @@ def reproject_and_coadd(
         # 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:
@@ -281,12 +298,14 @@ def reproject_and_coadd(
         for array, correction in zip(arrays, corrections, strict=True):
             array.array -= correction
 
-    # At this point, the images are now ready to be co-added.
+        # 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 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 combine_function == "min":
         output_array[...] = np.inf

From 32080b42791e2ed8b17b0c12ab372dc3e5c1f202 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 21:38:54 -0400
Subject: [PATCH 07/13] add a progressbar option

---
 reproject/mosaicking/coadd.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index be74e5f85..4ec50e8ee 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -102,6 +102,9 @@ def reproject_and_coadd(
     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.
@@ -130,6 +133,9 @@ def reproject_and_coadd(
             "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)
@@ -151,7 +157,7 @@ def reproject_and_coadd(
     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)

From 600783524ccdb06ae9d23809a4a9d57c7ed2a0f6 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 23:42:34 -0400
Subject: [PATCH 08/13] add progressbar as kwd

---
 reproject/mosaicking/coadd.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 4ec50e8ee..5ca605b65 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -25,6 +25,7 @@ def reproject_and_coadd(
     output_array=None,
     output_footprint=None,
     block_sizes=None,
+    progressbar=False,
     **kwargs,
 ):
     """

From 89754902c143c729d3bb5ba3f10f55bbce8a8113 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Mon, 13 Mar 2023 23:46:41 -0400
Subject: [PATCH 09/13] change dimensional subsetting to operate on low-level
 wcs

---
 reproject/mosaicking/coadd.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 5ca605b65..908010c34 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -195,8 +195,9 @@ def reproject_and_coadd(
             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])
-            xc_out, yc_out = wcs_out.celestial.world_to_pixel(wcs_in.celestial.pixel_to_world(xc, yc))
-            zc_out = wcs_out.spectral.world_to_pixel(wcs_in.spectral.pixel_to_world(zc))
+            # 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:]
 
         # Determine the cutout parameters

From 9a98b4261420f1d7be8e61c21af8dbbfcee2cfd3 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Wed, 15 Mar 2023 11:48:53 -0400
Subject: [PATCH 10/13] add a helpful error msg

---
 reproject/mosaicking/coadd.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 908010c34..22823022c 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -199,6 +199,8 @@ def reproject_and_coadd(
             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
 

From 22cfce97f104039a923f77ab572ad0925bcac668 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Sat, 9 Sep 2023 11:02:50 -0400
Subject: [PATCH 11/13] Fix several issues

---
 reproject/mosaicking/coadd.py        | 102 ++++++++++++++-------------
 reproject/mosaicking/subset_array.py |  19 +++--
 reproject/utils.py                   |   5 +-
 3 files changed, 66 insertions(+), 60 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 22823022c..62faef59c 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -26,6 +26,7 @@ def reproject_and_coadd(
     output_footprint=None,
     block_sizes=None,
     progressbar=False,
+    blank_pixel_value=np.nan,
     **kwargs,
 ):
     """
@@ -128,6 +129,9 @@ def reproject_and_coadd(
 
     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(
@@ -142,11 +146,15 @@ def reproject_and_coadd(
     wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)
 
     if output_array is not None and output_array.shape != shape_out:
-        raise ValueError("If you specify an output array, it must have a shape matching "
-                         f"the output shape {shape_out}")
+        raise ValueError(
+            "If you specify an output array, it must have a shape matching "
+            f"the output shape {shape_out}"
+        )
     if output_footprint is not None and output_footprint.shape != shape_out:
-        raise ValueError("If you specify an output footprint array, it must have a shape matching "
-                         f"the output shape {shape_out}")
+        raise ValueError(
+            "If you specify an output footprint array, it must have a shape matching "
+            f"the output shape {shape_out}"
+        )
 
     if output_array is None:
         output_array = np.zeros(shape_out)
@@ -154,7 +162,6 @@ def reproject_and_coadd(
         output_footprint = np.zeros(shape_out)
 
     # Start off by reprojecting individual images to the final projection
-
     if match_background:
         arrays = []
 
@@ -189,6 +196,7 @@ def reproject_and_coadd(
             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
@@ -196,8 +204,12 @@ def reproject_and_coadd(
             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))
+            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}")
@@ -230,7 +242,7 @@ def reproject_and_coadd(
                     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
+            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)))
@@ -242,12 +254,11 @@ def reproject_and_coadd(
                 )
             shape_out_indiv = (kmax - kmin, 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]
+                kwargs["block_size"] = block_sizes[idata]
             else:
-                kwargs['block_size'] = block_sizes
+                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
@@ -298,7 +309,6 @@ def reproject_and_coadd(
                 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:
         offset_matrix = determine_offset_matrix(arrays)
@@ -308,64 +318,58 @@ def reproject_and_coadd(
         for array, correction in zip(arrays, corrections, strict=True):
             array.array -= correction
 
-        # At this point, the images are now ready to be co-added.
-
-        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 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
diff --git a/reproject/mosaicking/subset_array.py b/reproject/mosaicking/subset_array.py
index bf0d299a9..5133644c0 100644
--- a/reproject/mosaicking/subset_array.py
+++ b/reproject/mosaicking/subset_array.py
@@ -34,7 +34,11 @@ def __repr__(self):
     @property
     def view_in_original_array(self):
         if self.kmin is not None:
-            return (slice(self.kmin, self.kmax), slice(self.jmin, self.jmax), slice(self.imin, self.imax))
+            return (
+                slice(self.kmin, self.kmax),
+                slice(self.jmin, self.jmax),
+                slice(self.imin, self.imax),
+            )
         else:
             return (slice(self.jmin, self.jmax), slice(self.imin, self.imax))
 
@@ -92,26 +96,21 @@ def _operation(self, other, op):
         if jmax < jmin:
             jmax = jmin
 
-
         if self.kmin is None:
             # Extract cutout from each
 
             self_array = self.array[
-                jmin - self.jmin : jmax - self.jmin,
-                imin - self.imin : imax - self.imin,
+                jmin - self.jmin : jmax - self.jmin, imin - self.imin : imax - self.imin
             ]
             self_footprint = self.footprint[
-                jmin - self.jmin : jmax - self.jmin,
-                imin - self.imin : imax - self.imin,
+                jmin - self.jmin : jmax - self.jmin, imin - self.imin : imax - self.imin
             ]
 
             other_array = other.array[
-                jmin - other.jmin : jmax - other.jmin,
-                imin - other.imin : imax - other.imin,
+                jmin - other.jmin : jmax - other.jmin, imin - other.imin : imax - other.imin
             ]
             other_footprint = other.footprint[
-                jmin - other.jmin : jmax - other.jmin,
-                imin - other.imin : imax - other.imin,
+                jmin - other.jmin : jmax - other.jmin, imin - other.imin : imax - other.imin
             ]
 
             # Carry out operator and store result in ReprojectedArraySubset
diff --git a/reproject/utils.py b/reproject/utils.py
index 3ff46555c..9f72eabbe 100644
--- a/reproject/utils.py
+++ b/reproject/utils.py
@@ -38,7 +38,10 @@ def _dask_to_numpy_memmap(dask_array, tmp_dir):
     with tempfile.TemporaryDirectory() as zarr_tmp:
         # First compute and store the dask array to zarr using whatever
         # the default scheduler is at this point
-        dask_array.to_zarr(zarr_tmp)
+        try:
+            dask_array.to_zarr(zarr_tmp)
+        except ValueError:
+            dask_array.rechunk().to_zarr(zarr_tmp)
 
         # Load the array back to dask
         zarr_array = da.from_zarr(zarr_tmp)

From bc503a47efaebd09f6ef8ea2985e95f0bef383ad Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Tue, 4 Jun 2024 23:49:19 +0100
Subject: [PATCH 12/13] Rename progressbar to progress_bar, document
 blank_pixel_value, and remove 'median' as an advertised option since it isn't
 implemented.

---
 reproject/mosaicking/coadd.py | 39 ++++++++++++++++-------------------
 1 file changed, 18 insertions(+), 21 deletions(-)

diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 62faef59c..0b3e4d4ee 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -11,6 +11,10 @@
 __all__ = ["reproject_and_coadd"]
 
 
+def _noop(iterable):
+    return iterable
+
+
 def reproject_and_coadd(
     input_data,
     output_projection,
@@ -25,7 +29,7 @@ def reproject_and_coadd(
     output_array=None,
     output_footprint=None,
     block_sizes=None,
-    progressbar=False,
+    progress_bar=None,
     blank_pixel_value=np.nan,
     **kwargs,
 ):
@@ -80,7 +84,7 @@ def reproject_and_coadd(
         `~astropy.io.fits.HDUList` instance, specifies the HDU to use.
     reproject_function : callable
         The function to use for the reprojection.
-    combine_function : { 'mean', 'sum', 'median', 'first', 'last', 'min', 'max' }
+    combine_function : { 'mean', 'sum', 'first', 'last', 'min', 'max' }
         The type of function to use for combining the values into the final
         image. For 'first' and 'last', respectively, the reprojected images are
         simply overlaid on top of each other. With respect to the order of the
@@ -103,12 +107,15 @@ def reproject_and_coadd(
         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
+        if you want the sample block size for all data sets.
+    progress_bar : callable, optional
+        If specified, use this as a progress_bar to track loop iterations over
         data sets.
+    blank_pixel_value : float, optional
+        Value to use for areas of the resulting mosaic that do not have input
+        data.
 
-    kwargs
+    **kwargs
         Keyword arguments to be passed to the reprojection function.
 
     Returns
@@ -127,19 +134,16 @@ def reproject_and_coadd(
 
     # Validate inputs
 
-    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 combine_function not in ("mean", "sum", "first", "last", "min", "max"):
+        raise ValueError("combine_function should be one of mean/sum/first/last/min/max")
 
     if reproject_function is None:
         raise ValueError(
             "reprojection function should be specified with the reproject_function argument"
         )
 
-    if not progressbar:
-        progressbar = lambda x: x
+    if progress_bar is None:
+        progress_bar = _noop
 
     # Parse the output projection to avoid having to do it for each
 
@@ -165,7 +169,7 @@ def reproject_and_coadd(
     if match_background:
         arrays = []
 
-    for idata in progressbar(range(len(input_data))):
+    for idata in progress_bar(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)
@@ -363,13 +367,6 @@ def reproject_and_coadd(
                     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")
-
     output_array[output_footprint == 0] = blank_pixel_value
 
     return output_array, output_footprint

From f5e3016d5cbea5fbd648b08b9ccb913b87356ee3 Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Wed, 5 Jun 2024 11:54:25 +0100
Subject: [PATCH 13/13] Generalize reproject_and_coadd to N-dimensions and fix
 test failures

---
 reproject/array_utils.py                      |  21 +-
 reproject/mosaicking/coadd.py                 | 180 ++++++++----------
 reproject/mosaicking/subset_array.py          | 152 +++++----------
 reproject/mosaicking/tests/test_coadd.py      |   1 -
 .../mosaicking/tests/test_subset_array.py     |  42 ++--
 5 files changed, 176 insertions(+), 220 deletions(-)

diff --git a/reproject/array_utils.py b/reproject/array_utils.py
index ec3a39a2e..1d5aae562 100644
--- a/reproject/array_utils.py
+++ b/reproject/array_utils.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-__all__ = ["map_coordinates"]
+__all__ = ["map_coordinates", "sample_array_edges"]
 
 
 def map_coordinates(image, coords, **kwargs):
@@ -35,3 +35,22 @@ def map_coordinates(image, coords, **kwargs):
     values[reset] = kwargs.get("cval", 0.0)
 
     return values
+
+
+def sample_array_edges(shape, *, n_samples):
+    # Given an N-dimensional array shape, sample each edge of the array using
+    # the requested number of samples (which will include vertices). To do this
+    # we iterate through the dimensions and for each one we sample the points
+    # in that dimension and iterate over the combination of other vertices.
+    # Returns an array with dimensions (N, n_samples)
+    all_positions = []
+    ndim = len(shape)
+    shape = np.array(shape)
+    for idim in range(ndim):
+        for vertex in range(2**ndim):
+            positions = -0.5 + shape * ((vertex & (2 ** np.arange(ndim))) > 0).astype(int)
+            positions = np.broadcast_to(positions, (n_samples, ndim)).copy()
+            positions[:, idim] = np.linspace(-0.5, shape[idim] - 0.5, n_samples)
+            all_positions.append(positions)
+    positions = np.unique(np.vstack(all_positions), axis=0).T
+    return positions
diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py
index 0b3e4d4ee..bbeb652d9 100644
--- a/reproject/mosaicking/coadd.py
+++ b/reproject/mosaicking/coadd.py
@@ -4,6 +4,7 @@
 from astropy.wcs import WCS
 from astropy.wcs.wcsapi import SlicedLowLevelWCS
 
+from ..array_utils import sample_array_edges
 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
@@ -30,15 +31,13 @@ def reproject_and_coadd(
     output_footprint=None,
     block_sizes=None,
     progress_bar=None,
-    blank_pixel_value=np.nan,
+    blank_pixel_value=0,
     **kwargs,
 ):
     """
-    Given a set of input images, reproject and co-add these to a single
+    Given a set of input data, reproject and co-add these to a single
     final image.
 
-    This currently only works with 2-d images with celestial WCS.
-
     Parameters
     ----------
     input_data : iterable
@@ -149,24 +148,31 @@ def reproject_and_coadd(
 
     wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)
 
-    if output_array is not None and output_array.shape != shape_out:
+    if output_array is None:
+        output_array = np.zeros(shape_out)
+    elif output_array.shape != shape_out:
         raise ValueError(
             "If you specify an output array, it must have a shape matching "
             f"the output shape {shape_out}"
         )
-    if output_footprint is not None and output_footprint.shape != shape_out:
+
+    if output_footprint is None:
+        output_footprint = np.zeros(shape_out)
+    elif output_footprint.shape != shape_out:
         raise ValueError(
             "If you specify an output footprint array, it must have a shape matching "
             f"the output shape {shape_out}"
         )
 
-    if output_array is None:
-        output_array = np.zeros(shape_out)
-    if output_footprint is None:
-        output_footprint = np.zeros(shape_out)
+    # Define 'on-the-fly' mode: in the case where we don't need to match
+    # the backgrounds and we are combining with 'mean' or 'sum', we don't
+    # have to keep track of the intermediate arrays and can just modify
+    # the output array on-the-fly
+    on_the_fly = not match_background and combine_function in ("mean", "sum")
 
     # Start off by reprojecting individual images to the final projection
-    if match_background:
+
+    if not on_the_fly:
         arrays = []
 
     for idata in progress_bar(range(len(input_data))):
@@ -192,31 +198,9 @@ def reproject_and_coadd(
         # 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.
-        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}")
+
+        edges = sample_array_edges(array_in.shape, n_samples=11)[::-1]
+        edges_out = wcs_out.world_to_pixel(wcs_in.pixel_to_world(*edges))[::-1]
 
         # Determine the cutout parameters
 
@@ -224,39 +208,32 @@ def reproject_and_coadd(
         # such as all-sky images or full solar disk views. In this case we skip
         # this step and just use the full output WCS for reprojection.
 
-        if np.any(np.isnan(xc_out)) or np.any(np.isnan(yc_out)):
-            imin = 0
-            imax = shape_out_cel[1]
-            jmin = 0
-            jmax = shape_out_cel[0]
-        else:
-            imin = max(0, int(np.floor(xc_out.min() + 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_cel[0], int(np.ceil(yc_out.max() + 0.5)))
+        ndim_out = len(shape_out)
 
-        if imax < imin or jmax < jmin:
+        skip_data = False
+        if np.any(np.isnan(edges_out)):
+            bounds = list(zip([0] * ndim_out, shape_out))
+        else:
+            bounds = []
+            for idim in range(ndim_out):
+                imin = max(0, int(np.floor(edges_out[idim].min() + 0.5)))
+                imax = min(shape_out[idim], int(np.ceil(edges_out[idim].max() + 0.5)))
+                bounds.append((imin, imax))
+                if imax < imin:
+                    skip_data = True
+                    break
+
+        if skip_data:
             continue
 
-        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)
+        slice_out = tuple([slice(imin, imax) for (imin, imax) in bounds])
+
+        if isinstance(wcs_out, WCS):
+            wcs_out_indiv = wcs_out[slice_out]
+        else:
+            wcs_out_indiv = SlicedLowLevelWCS(wcs_out.low_level_wcs, slice_out)
+
+        shape_out_indiv = [imax - imin for (imin, imax) in bounds]
 
         if block_sizes is not None:
             if len(block_sizes) == len(input_data) and len(block_sizes[idata]) == len(shape_out):
@@ -296,22 +273,20 @@ def reproject_and_coadd(
             weights[reset] = 0.0
             footprint *= weights
 
-        array = ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax, kmin, kmax)
+        array = ReprojectedArraySubset(array, footprint, bounds)
 
         # TODO: make sure we gracefully handle the case where the
         # output image is empty (due e.g. to no overlap).
 
-        if match_background:
-            arrays.append(array)
+        if on_the_fly:
+            # 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
         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
+            arrays.append(array)
 
     # If requested, try and match the backgrounds.
     if match_background and len(arrays) > 1:
@@ -322,11 +297,6 @@ def reproject_and_coadd(
         for array, correction in zip(arrays, corrections, strict=True):
             array.array -= correction
 
-    if combine_function == "min":
-        output_array[...] = np.inf
-    elif combine_function == "max":
-        output_array[...] = -np.inf
-
     if combine_function in ("mean", "sum"):
         if match_background:
             # if we're not matching the background, this part has already been done
@@ -336,8 +306,8 @@ def reproject_and_coadd(
                 # 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"):
@@ -345,28 +315,32 @@ def reproject_and_coadd(
                 output_array[output_footprint == 0] = blank_pixel_value
 
     elif combine_function in ("first", "last", "min", "max"):
-        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]
+        if combine_function == "min":
+            output_array[...] = np.inf
+        elif combine_function == "max":
+            output_array[...] = -np.inf
+
+        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]
                 )
-                output_array[array.view_in_original_array] = np.where(
-                    mask, 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]
+            )
+            output_array[array.view_in_original_array] = np.where(
+                mask, array.array, output_array[array.view_in_original_array]
+            )
+
     output_array[output_footprint == 0] = blank_pixel_value
 
     return output_array, output_footprint
diff --git a/reproject/mosaicking/subset_array.py b/reproject/mosaicking/subset_array.py
index 5133644c0..010114e0a 100644
--- a/reproject/mosaicking/subset_array.py
+++ b/reproject/mosaicking/subset_array.py
@@ -15,60 +15,36 @@ class ReprojectedArraySubset:
     # rather than the center, which is not well defined for even-sized
     # cutouts.
 
-    def __init__(self, array, footprint, imin, imax, jmin, jmax, kmin=None, kmax=None):
+    def __init__(self, array, footprint, bounds):
         self.array = array
         self.footprint = footprint
-        self.imin = imin
-        self.imax = imax
-        self.jmin = jmin
-        self.jmax = jmax
-        self.kmin = kmin
-        self.kmax = kmax
+        self.bounds = bounds
 
     def __repr__(self):
-        if self.kmin is not None:
-            return f"<ReprojectedArraySubset at [{self.kmin}:{self.kmax},{self.jmin}:{self.jmax},{self.imin}:{self.imax}]>"
-        else:
-            return f"<ReprojectedArraySubset at [{self.jmin}:{self.jmax},{self.imin}:{self.imax}]>"
+        bounds_str = "[" + ",".join(f"{imin}:{imax}" for (imin, imax) in self.bounds) + "]"
+        return f"<ReprojectedArraySubset at {bounds_str}>"
 
     @property
     def view_in_original_array(self):
-        if self.kmin is not None:
-            return (
-                slice(self.kmin, self.kmax),
-                slice(self.jmin, self.jmax),
-                slice(self.imin, self.imax),
-            )
-        else:
-            return (slice(self.jmin, self.jmax), slice(self.imin, self.imax))
+        return tuple([slice(imin, imax) for (imin, imax) in self.bounds])
 
     @property
     def shape(self):
-        if self.kmin is not None:
-            return (self.kmax - self.kmin, self.jmax - self.jmin, self.imax - self.imin)
-        else:
-            return (self.jmax - self.jmin, self.imax - self.imin)
+        return tuple((imax - imin) for (imin, imax) in self.bounds)
 
     def overlaps(self, other):
         # Note that the use of <= or >= instead of < and > is due to
         # the fact that the max values are exclusive (so +1 above the
         # last value).
-        if self.kmin is not None:
-            return not (
-                self.imax <= other.imin
-                or other.imax <= self.imin
-                or self.jmax <= other.jmin
-                or other.jmax <= self.jmin
-                or self.kmax <= other.kmin
-                or other.kmax <= self.kmin
-            )
-        else:
-            return not (
-                self.imax <= other.imin
-                or other.imax <= self.imin
-                or self.jmax <= other.jmin
-                or other.jmax <= self.jmin
+        if len(self.bounds) != len(other.bounds):
+            raise ValueError(
+                f"Mismatch in number of dimensions, expected "
+                f"{len(self.bounds)} dimensions and got {len(other.bounds)}"
             )
+        for (imin, imax), (imin_other, imax_other) in zip(self.bounds, other.bounds, strict=False):
+            if imax <= imin_other or imax_other <= imin:
+                return False
+        return True
 
     def __add__(self, other):
         return self._operation(other, operator.add)
@@ -83,71 +59,39 @@ def __truediv__(self, other):
         return self._operation(other, operator.truediv)
 
     def _operation(self, other, op):
+        if len(self.bounds) != len(other.bounds):
+            raise ValueError(
+                f"Mismatch in number of dimensions, expected "
+                f"{len(self.bounds)} dimensions and got {len(other.bounds)}"
+            )
+
         # Determine cutout parameters for overlap region
 
-        imin = max(self.imin, other.imin)
-        imax = min(self.imax, other.imax)
-        jmin = max(self.jmin, other.jmin)
-        jmax = min(self.jmax, other.jmax)
-
-        if imax < imin:
-            imax = imin
-
-        if jmax < jmin:
-            jmax = jmin
-
-        if self.kmin is None:
-            # Extract cutout from each
-
-            self_array = self.array[
-                jmin - self.jmin : jmax - self.jmin, imin - self.imin : imax - self.imin
-            ]
-            self_footprint = self.footprint[
-                jmin - self.jmin : jmax - self.jmin, imin - self.imin : imax - self.imin
-            ]
-
-            other_array = other.array[
-                jmin - other.jmin : jmax - other.jmin, imin - other.imin : imax - other.imin
-            ]
-            other_footprint = other.footprint[
-                jmin - other.jmin : jmax - other.jmin, imin - other.imin : imax - other.imin
-            ]
-
-            # Carry out operator and store result in ReprojectedArraySubset
-
-            array = op(self_array, other_array)
-            footprint = (self_footprint > 0) & (other_footprint > 0)
-
-            return ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax)
-
-        else:
-            # Extract cutout from each
-
-            self_array = self.array[
-                kmin - self.kmin : kmax - self.kmin,
-                jmin - self.jmin : jmax - self.jmin,
-                imin - self.imin : imax - self.imin,
-            ]
-            self_footprint = self.footprint[
-                kmin - self.kmin : kmax - self.kmin,
-                jmin - self.jmin : jmax - self.jmin,
-                imin - self.imin : imax - self.imin,
-            ]
-
-            other_array = other.array[
-                kmin - other.kmin : kmax - other.kmin,
-                jmin - other.jmin : jmax - other.jmin,
-                imin - other.imin : imax - other.imin,
-            ]
-            other_footprint = other.footprint[
-                kmin - other.kmin : kmax - other.kmin,
-                jmin - other.jmin : jmax - other.jmin,
-                imin - other.imin : imax - other.imin,
-            ]
-
-            # Carry out operator and store result in ReprojectedArraySubset
-
-            array = op(self_array, other_array)
-            footprint = (self_footprint > 0) & (other_footprint > 0)
-
-            return ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax, kmin, kmax)
+        overlap_bounds = []
+        self_slices = []
+        other_slices = []
+        for (imin, imax), (imin_other, imax_other) in zip(self.bounds, other.bounds, strict=False):
+            imin_overlap = max(imin, imin_other)
+            imax_overlap = min(imax, imax_other)
+            if imax_overlap < imin_overlap:
+                imax_overlap = imin_overlap
+            overlap_bounds.append((imin_overlap, imax_overlap))
+            self_slices.append(slice(imin_overlap - imin, imax_overlap - imin))
+            other_slices.append(slice(imin_overlap - imin_other, imax_overlap - imin_other))
+
+        self_slices = tuple(self_slices)
+
+        self_array = self.array[self_slices]
+        self_footprint = self.footprint[self_slices]
+
+        other_slices = tuple(other_slices)
+
+        other_array = other.array[other_slices]
+        other_footprint = other.footprint[other_slices]
+
+        # Carry out operator and store result in ReprojectedArraySubset
+
+        array = op(self_array, other_array)
+        footprint = (self_footprint > 0) & (other_footprint > 0)
+
+        return ReprojectedArraySubset(array, footprint, overlap_bounds)
diff --git a/reproject/mosaicking/tests/test_coadd.py b/reproject/mosaicking/tests/test_coadd.py
index d9e3de683..04b7c271c 100644
--- a/reproject/mosaicking/tests/test_coadd.py
+++ b/reproject/mosaicking/tests/test_coadd.py
@@ -80,7 +80,6 @@ def test_coadd_no_overlap(self, combine_function, reproject_function):
 
         input_data = self._get_tiles(self._nonoverlapping_views)
 
-        input_data = [(self.array, self.wcs)]
         array, footprint = reproject_and_coadd(
             input_data,
             self.wcs,
diff --git a/reproject/mosaicking/tests/test_subset_array.py b/reproject/mosaicking/tests/test_subset_array.py
index 3ecde69f9..dc05ef76c 100644
--- a/reproject/mosaicking/tests/test_subset_array.py
+++ b/reproject/mosaicking/tests/test_subset_array.py
@@ -14,21 +14,35 @@ def setup_method(self, method):
         self.array1 = np.random.random((123, 87))
         self.array2 = np.random.random((123, 87))
         self.array3 = np.random.random((123, 87))
+        self.array4 = np.random.random((123, 87, 16))
 
         self.footprint1 = (self.array1 > 0.5).astype(int)
         self.footprint2 = (self.array2 > 0.5).astype(int)
         self.footprint3 = (self.array3 > 0.5).astype(int)
+        self.footprint4 = (self.array4 > 0.5).astype(int)
 
         self.subset1 = ReprojectedArraySubset(
-            self.array1[20:88, 34:40], self.footprint1[20:88, 34:40], 34, 40, 20, 88
+            self.array1[20:88, 34:40],
+            self.footprint1[20:88, 34:40],
+            [(20, 88), (34, 40)],
         )
 
         self.subset2 = ReprojectedArraySubset(
-            self.array2[50:123, 37:42], self.footprint2[50:123, 37:42], 37, 42, 50, 123
+            self.array2[50:123, 37:42],
+            self.footprint2[50:123, 37:42],
+            [(50, 123), (37, 42)],
         )
 
         self.subset3 = ReprojectedArraySubset(
-            self.array3[40:50, 11:19], self.footprint3[40:50, 11:19], 11, 19, 40, 50
+            self.array3[40:50, 11:19],
+            self.footprint3[40:50, 11:19],
+            [(40, 50), (11, 19)],
+        )
+
+        self.subset4 = ReprojectedArraySubset(
+            self.array4[30:35, 40:45, 1:4],
+            self.footprint4[30:35, 40:45, 1:4],
+            [(30, 35), (40, 45), (1, 4)],
         )
 
     def test_repr(self):
@@ -55,17 +69,23 @@ def test_overlaps(self):
     @pytest.mark.parametrize("op", [operator.add, operator.sub, operator.mul, operator.truediv])
     def test_arithmetic(self, op):
         subset = op(self.subset1, self.subset2)
-        assert subset.imin == 37
-        assert subset.imax == 40
-        assert subset.jmin == 50
-        assert subset.jmax == 88
+        assert subset.bounds == [(50, 88), (37, 40)]
         expected = op(self.array1[50:88, 37:40], self.array2[50:88, 37:40])
         assert_equal(subset.array, expected)
 
     def test_arithmetic_nooverlap(self):
         subset = self.subset1 - self.subset3
-        assert subset.imin == 34
-        assert subset.imax == 34
-        assert subset.jmin == 40
-        assert subset.jmax == 50
+        assert subset.bounds == [(40, 50), (34, 34)]
         assert subset.shape == (10, 0)
+
+    def test_overlaps_dimension_mismatch(self):
+        with pytest.raises(
+            ValueError, match=("Mismatch in number of dimensions, expected 2 dimensions and got 3")
+        ):
+            self.subset1.overlaps(self.subset4)
+
+    def test_arithmetic_dimension_mismatch(self):
+        with pytest.raises(
+            ValueError, match=("Mismatch in number of dimensions, expected 2 dimensions and got 3")
+        ):
+            self.subset1 - self.subset4