Merge pull request #387 from astrofrog/coadd-output-array

Add ability to specify output array and footprint in reproject_and_coadd

reproject/mosaicking/coadd.py

@@ -22,6 +22,8 @@ def reproject_and_coadd(
     combine_function="mean",
     match_background=False,
     background_reference=None,
+    output_array=None,
+    output_footprint=None,
     **kwargs,
 ):
     """
@@ -85,6 +87,14 @@ 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
+        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
+        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
         Keyword arguments to be passed to the reprojection function.
     """
@@ -93,8 +103,6 @@ def reproject_and_coadd(
     # up memory usage. We could probably still have references to array
     # objects, but we'd just make sure these were memory mapped
 
-    # TODO: add support for specifying output array
-
     # Validate inputs
 
     if combine_function not in ("mean", "sum", "median", "first", "last", "min", "max"):
@@ -109,6 +117,17 @@ 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}"
+        )
+
     # Start off by reprojecting individual images to the final projection
 
     arrays = []
@@ -223,15 +242,15 @@ def reproject_and_coadd(
 
     # At this point, the images are now ready to be co-added.
 
-    # TODO: provide control over final dtype
-
-    final_array = np.zeros(shape_out)
-    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":
-        final_array[...] = np.inf
+        output_array[...] = np.inf
     elif combine_function == "max":
-        final_array[...] = -np.inf
+        output_array[...] = -np.inf
 
     if combine_function in ("mean", "sum"):
         for array in arrays:
@@ -240,33 +259,33 @@ def reproject_and_coadd(
             # means/sums.
             array.array[array.footprint == 0] = 0
 
-            final_array[array.view_in_original_array] += array.array * array.footprint
-            final_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"):
-                final_array /= final_footprint
+                output_array /= output_footprint
 
     elif combine_function in ("first", "last", "min", "max"):
         for array in arrays:
             if combine_function == "first":
-                mask = final_footprint[array.view_in_original_array] == 0
+                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 < final_array[array.view_in_original_array]
+                    array.array < output_array[array.view_in_original_array]
                 )
             elif combine_function == "max":
                 mask = (array.footprint > 0) & (
-                    array.array > final_array[array.view_in_original_array]
+                    array.array > output_array[array.view_in_original_array]
                 )
 
-            final_footprint[array.view_in_original_array] = np.where(
-                mask, array.footprint, final_footprint[array.view_in_original_array]
+            output_footprint[array.view_in_original_array] = np.where(
+                mask, array.footprint, output_footprint[array.view_in_original_array]
             )
-            final_array[array.view_in_original_array] = np.where(
-                mask, array.array, final_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 == "median":
@@ -276,6 +295,6 @@ def reproject_and_coadd(
         raise NotImplementedError("combine_function='median' is not yet implemented")
 
     if combine_function in ("min", "max"):
-        final_array[final_footprint == 0] = 0.0
+        output_array[output_footprint == 0] = 0.0
 
-    return final_array, final_footprint
+    return output_array, output_footprint

reproject/mosaicking/tests/test_coadd.py

@@ -108,6 +108,31 @@ def test_coadd_with_overlap(self, reproject_function):
 
         assert_allclose(array, self.array, atol=ATOL)
 
+    def test_coadd_with_outputs(self, tmp_path, reproject_function):
+        # Test the options to specify output array/footprint
+
+        input_data = self._get_tiles(self._overlapping_views)
+
+        output_array = np.memmap(
+            tmp_path / "array.np", mode="w+", dtype=float, shape=self.array.shape
+        )
+        output_footprint = np.memmap(
+            tmp_path / "footprint.np", mode="w+", dtype=float, shape=self.array.shape
+        )
+
+        array, footprint = reproject_and_coadd(
+            input_data,
+            self.wcs,
+            shape_out=self.array.shape,
+            combine_function="mean",
+            reproject_function=reproject_function,
+            output_array=output_array,
+            output_footprint=output_footprint,
+        )
+
+        assert_allclose(output_array, self.array, atol=ATOL)
+        assert_allclose(output_footprint, footprint, atol=ATOL)
+
     @pytest.mark.parametrize("combine_function", ["first", "last", "min", "max"])
     def test_coadd_with_overlap_first_last(self, reproject_function, combine_function):
         views = self._overlapping_views