Merge pull request #54 from lsst/tickets/DM-53132

DM-53132: Add APIs to stitch noise images

Author: arunkannawadi

python/lsst/cell_coadds/_image_planes.py

@@ -89,11 +89,42 @@ def noise_realizations(self) -> Sequence[ImageLike]:
         """
         raise NotImplementedError()
 
-    def asMaskedImage(self) -> MaskedImageF:
+    def asMaskedImage(
+        self,
+        *,
+        noise_index: int | None = None,
+    ) -> MaskedImageF:
         """Return an `lsst.afw.image.MaskedImage` view of the image, mask, and
         variance planes.
+
+        Parameters
+        ----------
+        noise_index : `int` or `None`, optional
+            If `None`, return the masked image formed from the
+            main image, mask, and variance planes.  If an integer index is
+            provided, return the masked image formed from the specified noise
+            realization, along with the mask and variance planes.
+
+        Returns
+        -------
+        masked_image : `lsst.afw.image.MaskedImageF`
+            The masked image formed from the specified planes.
+
+        Raises
+        ------
+        ValueError
+            Raised if ``noise_index`` is out of range for the available noise
+            realizations.
         """
-        return MaskedImageF(self.image, self.mask, self.variance)
+        if noise_index is None:
+            return MaskedImageF(self.image, self.mask, self.variance)
+        elif 0 <= noise_index < len(self.noise_realizations):
+            return MaskedImageF(self.noise_realizations[noise_index], self.mask, self.variance)
+        else:
+            raise ValueError(
+                f"noise_index {noise_index} is out of range for "
+                f"{len(self.noise_realizations)} noise realizations"
+            )
 
 
 class OwnedImagePlanes(ImagePlanes):

python/lsst/cell_coadds/_stitched_coadd.py

@@ -166,9 +166,25 @@ def set_cell_edges(self, *, edge_width: int = 1, edge_mask_name: str = "CELL_EDG
         for cell in self._cell_coadd.cells.values():
             cell._set_cell_edges(edge_width=edge_width, edge_mask_name=edge_mask_name)
 
-    def asExposure(self) -> ExposureF:
-        """Return an `lsst.afw.image.Exposure` view of this piecewise image."""
-        result = ExposureF(self.asMaskedImage())
+    def asExposure(
+        self,
+        *,
+        noise_index: int | None = None,
+    ) -> ExposureF:
+        """Return an `lsst.afw.image.Exposure` view of this piecewise image.
+
+        Parameters
+        ----------
+        noise_index : `int` or `None`, optional
+            If specified, return an exposure containing the specified noise
+            realization instead of the data image.
+
+        Returns
+        -------
+        exposure : `lsst.afw.image.ExposureF`
+            The stitched exposure.
+        """
+        result = ExposureF(self.asMaskedImage(noise_index=noise_index))
         # Exposure components derived from "common" components are all simple.
         result.setWcs(self._cell_coadd.wcs)
         result.setFilter(FilterLabel(band=self.band))

tests/test_coadds.py

@@ -561,9 +561,29 @@ def test_asExposure(self):
                     geom.Extent2I(self.inner_size_x, self.inner_size_y),
                 )
                 index = Index2D(x=x, y=y)
-                self.assertImagesEqual(exposure.image[bbox], self.exposures[index].image[bbox])
-                self.assertImagesEqual(exposure.variance[bbox], self.exposures[index].variance[bbox])
-                self.assertImagesEqual(exposure.mask[bbox], self.exposures[index].mask[bbox])
+                self.assertMaskedImagesEqual(exposure[bbox], self.exposures[index][bbox])
+
+        self.assertMaskedImagesEqual(self.stitched_coadd.asExposure(noise_index=None), exposure)
+        for noise_index in range(self.n_noise_realizations):
+            noise_exposure = self.stitched_coadd.asExposure(noise_index=noise_index)
+
+            self.assertImagesEqual(noise_exposure.variance, exposure.variance)
+            self.assertImagesEqual(noise_exposure.mask, exposure.mask)
+
+            for y in range(self.ny):
+                for x in range(self.nx):
+                    bbox = geom.Box2I(
+                        geom.Point2I(self.x0 + x * self.inner_size_x, self.y0 + y * self.inner_size_y),
+                        geom.Extent2I(self.inner_size_x, self.inner_size_y),
+                    )
+                    index = Index2D(x=x, y=y)
+                    self.assertImagesEqual(
+                        noise_exposure.image[bbox],
+                        self.multiple_cell_coadd.cells[index].outer.noise_realizations[noise_index][bbox],
+                    )
+
+        with self.assertRaises(ValueError):
+            self.stitched_coadd.asExposure(noise_index=self.n_noise_realizations)
 
     def test_aperture_correction(self):
         """Test the aperture correction values are what we expect."""
@@ -613,9 +633,7 @@ def test_fits(self):
             read_exposure = ExposureF.readFits(filename)  # Test the readFits method.
 
         # Test that the image planes are identical.
-        self.assertImagesEqual(read_exposure.image, write_exposure.image)
-        self.assertImagesEqual(read_exposure.variance, write_exposure.variance)
-        self.assertImagesEqual(read_exposure.mask, write_exposure.mask)
+        self.assertMaskedImagesEqual(read_exposure, write_exposure)
 
         # Test the PSF images in the StitchedPsf.
         for index in write_exposure.psf.images.indices():