diff --git a/reproject/mosaicking/coadd.py b/reproject/mosaicking/coadd.py index cece83407..66c06ddcb 100644 --- a/reproject/mosaicking/coadd.py +++ b/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 diff --git a/reproject/mosaicking/tests/test_coadd.py b/reproject/mosaicking/tests/test_coadd.py index ad7373e52..6e03ab6e9 100644 --- a/reproject/mosaicking/tests/test_coadd.py +++ b/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