Add CoaddPsf construction in default reducer subtask.

- Add tests of newly-generated PSFs
- Fix a few docstrings
- Add some fixes that arose during testing

python/lsst/ip/diffim/imageMapReduce.py

@@ -30,6 +30,8 @@
 import lsst.afw.image as afwImage
 import lsst.afw.math as afwMath
 import lsst.afw.geom as afwGeom
+import lsst.afw.table as afwTable
+import lsst.meas.algorithms as measAlg
 import lsst.pex.config as pexConfig
 import lsst.pipe.base as pipeBase
 
@@ -201,7 +203,7 @@ def run(self, mapperResults, exposure, **kwargs):
             basis for the resulting exposure (if
             self.config.mapperSubtask.reduceOperation is not 'none')
         kwargs :
-            additional keyword arguments propagaged from
+            additional keyword arguments propagated from
             `ImageMapReduceTask.run`.
 
         Returns
@@ -212,12 +214,11 @@ def run(self, mapperResults, exposure, **kwargs):
         Notes
         -----
         And currently known issues:
-        1. This currently should correctly handle overlapping sub-exposures.
+        1. This currently correctly handles overlapping sub-exposures.
            For overlapping sub-exposures, use `config.reduceOperation='average'`.
-        2. This currently does not correctly handle varying PSFs (in fact,
-           it just copies over the PSF from the original exposure). To be
-           investigated in DM-9629.
-        3. To be done: correct handling of masks as well.
+        2. This correctly handles varying PSFs, constructing the resulting
+           exposure's PSF via CoaddPsf (DM-9629).
+        3. To be done: correct handling of masks
         4. This logic currently makes *two* copies of the original exposure
            (one here and one in `mapperSubtask.run()`). Possibly of concern
            for large images on memory-constrained systems.
@@ -230,14 +231,15 @@ def run(self, mapperResults, exposure, **kwargs):
         newMI = newExp.getMaskedImage()
 
         reduceOp = self.config.reduceOperation
+        weights = None
         if reduceOp == 'copy':
             newMI.getImage()[:, :] = np.nan
             newMI.getVariance()[:, :] = np.nan
         else:
             newMI.getImage()[:, :] = 0.
             newMI.getVariance()[:, :] = 0.
             if reduceOp == 'average':  # make an array to keep track of weights
-                weights = afwImage.ImageF(newMI.getBBox())  # must be a float to divide later.
+                weights = afwImage.ImageI(newMI.getBBox())
 
         for item in mapperResults:
             item = item.subExposure  # Expected named value in the pipeBase.Struct
@@ -248,38 +250,76 @@ def run(self, mapperResults, exposure, **kwargs):
             subExp = afwImage.ExposureF(newExp, item.getBBox())
             subMI = subExp.getMaskedImage()
             patchMI = item.getMaskedImage()
-            isNan = (np.isnan(patchMI.getImage().getArray()) |
-                     np.isnan(patchMI.getVariance().getArray()))
-            noNans = np.sum(isNan) <= 0
+            isNotNan = ~(np.isnan(patchMI.getImage().getArray()) |
+                         np.isnan(patchMI.getVariance().getArray()))
             if reduceOp == 'copy':
-                if noNans:
-                    subMI[:, :] = patchMI
-                else:
-                    isNotNan = ~isNan
-                    subMI.getImage().getArray()[isNotNan] = patchMI.getImage().getArray()[isNotNan]
-                    subMI.getVariance().getArray()[isNotNan] = patchMI.getVariance().getArray()[isNotNan]
+                subMI.getImage().getArray()[isNotNan] = patchMI.getImage().getArray()[isNotNan]
+                subMI.getVariance().getArray()[isNotNan] = patchMI.getVariance().getArray()[isNotNan]
 
             if reduceOp == 'sum' or reduceOp == 'average':  # much of these two options is the same
-                if noNans:
-                    subMI += patchMI
-                    if reduceOp == 'average':
-                        # wsubim is a view into the `weights` Image, so here we simply add one to
-                        # the region of `weights` confined by `item.getBBox()`.
-                        wsubim = afwImage.ImageF(weights, item.getBBox())
-                        wsubim += 1.
-                else:
-                    isNotNan = ~isNan
-                    subMI.getImage().getArray()[isNotNan] += patchMI.getImage().getArray()[isNotNan]
-                    subMI.getVariance().getArray()[isNotNan] += patchMI.getVariance().getArray()[isNotNan]
-                    if reduceOp == 'average':
-                        wsubim = afwImage.ImageF(weights, item.getBBox())
-                        wsubim.getArray()[isNotNan] += 1.
+                subMI.getImage().getArray()[isNotNan] += patchMI.getImage().getArray()[isNotNan]
+                subMI.getVariance().getArray()[isNotNan] += patchMI.getVariance().getArray()[isNotNan]
+                if reduceOp == 'average':
+                    # wsubim is a view into the `weights` Image
+                    wsubim = afwImage.ImageI(weights, item.getBBox())
+                    wsubim.getArray()[isNotNan] += 1
 
         if reduceOp == 'average':
-            newMI /= weights
+            wts = weights.getArray().astype(np.float)
+            newMI.getImage().getArray()[:, :] /= wts
+            newMI.getVariance().getArray()[:, :] /= wts
+
+        # Not sure how to construct a PSF when reduceOp=='copy'...
+        if reduceOp == 'sum' or reduceOp == 'average':
+            psf = self._constructPsf(mapperResults, exposure)
+            newExp.setPsf(psf)
 
         return pipeBase.Struct(exposure=newExp)
 
+    def _constructPsf(self, mapperResults, exposure):
+        """Construct a CoaddPsf based on PSFs from individual subExposures
+
+        Currently uses (and returns) a CoaddPsf. TBD if we want to
+        create a custom subclass of CoaddPsf to differentiate it.
+
+        Parameters
+        ----------
+        mapperResults : list
+            list of `pipeBase.Struct` returned by `ImageMapperSubtask.run`.
+            For this to work, each element of `mapperResults` must contain
+            a `subExposure` element, from which the component Psfs are
+            extracted (thus the reducerTask cannot have
+            `reduceOperation = 'none'`.
+        exposure : afwImage.Exposure
+            the original exposure which is used here solely for its
+            bounding-box and WCS.
+
+        Returns
+        -------
+        A `measAlg.CoaddPsf` constructed from the PSFs of the individual
+        subExposures.
+        """
+        schema = afwTable.ExposureTable.makeMinimalSchema()
+        schema.addField("weight", type="D", doc="Coadd weight")
+        mycatalog = afwTable.ExposureCatalog(schema)
+
+        # We're just using the exposure's WCS (assuming that the subExposures'
+        # WCSs are the same, which they better be!).
+        wcsref = exposure.getWcs()
+        for i, res in enumerate(mapperResults):
+            subExp = res.subExposure
+            record = mycatalog.getTable().makeRecord()
+            record.setPsf(subExp.getPsf())
+            record.setWcs(subExp.getWcs())
+            record.setBBox(subExp.getBBox())
+            record['weight'] = 1.0
+            record['id'] = i
+            mycatalog.append(record)
+
+        # create the coaddpsf
+        psf = measAlg.CoaddPsf(catalog=mycatalog, coaddWcs=wcsref, weightFieldName='weight')
+        return psf
+
 
 class ImageMapReduceConfig(pexConfig.Config):
     """Configuration parameters for the ImageMapReduceTask
@@ -483,7 +523,7 @@ def _runMapper(self, exposure, doClone=False, **kwargs):
         a list of `pipeBase.Struct`s as returned by `mapperSubtask.run`.
         """
         if self.boxes0 is None:
-            self._generateGrid(exposure)
+            self._generateGrid(exposure, **kwargs)  # possibly pass `forceEvenSized`
         if len(self.boxes0) != len(self.boxes1):
             raise ValueError('Bounding boxes list and expanded bounding boxes list are of different lengths')
 
@@ -523,7 +563,7 @@ def _reduceImage(self, mapperResults, exposure, **kwargs):
         result = self.reducerSubtask.run(mapperResults, exposure, **kwargs)
         return result
 
-    def _generateGrid(self, exposure, forceEvenSized=False):
+    def _generateGrid(self, exposure, forceEvenSized=False, **kwargs):
         """Generate two lists of bounding boxes that evenly grid `exposure`
 
         Unless the config was provided with `centroidCoordsX` and
@@ -644,8 +684,32 @@ def offsetAndClipBoxes(bbox0, bbox1, xoff, yoff, bbox):
             bb1.shift(afwGeom.Extent2I(xoff, yoff))
             bb1.clip(bbox)
             if forceEvenSized:
-                bb0.grow(afwGeom.Extent2I(bb0.getWidth() % 2, bb0.getHeight() % 2))
-                bb1.grow(afwGeom.Extent2I(bb1.getWidth() % 2, bb1.getHeight() % 2))
+                if bb0.getWidth() % 2 == 1:  # grow to the right
+                    bb0.include(afwGeom.Point2I(bb0.getMaxX()+1, bb0.getMaxY())) # Expand by 1 pixel!
+                    bb0.clip(bbox)
+                    if bb0.getWidth() % 2 == 1:  # clipped at right -- so grow to the left
+                        bb0.include(afwGeom.Point2I(bb0.getMinX()-1, bb0.getMaxY()))
+                        bb0.clip(bbox)
+                if bb0.getHeight() % 2 == 1: # grow upwards
+                    bb0.include(afwGeom.Point2I(bb0.getMaxX(), bb0.getMaxY()+1)) # Expand by 1 pixel!
+                    bb0.clip(bbox)
+                    if bb0.getHeight() % 2 == 1: # clipped upwards -- so grow down
+                        bb0.include(afwGeom.Point2I(bb0.getMaxX(), bb0.getMinY()-1))
+                        bb0.clip(bbox)
+
+                if bb1.getWidth() % 2 == 1:  # grow to the left
+                    bb1.include(afwGeom.Point2I(bb1.getMaxX()+1, bb1.getMaxY())) # Expand by 1 pixel!
+                    bb1.clip(bbox)
+                    if bb1.getWidth() % 2 == 1:  # clipped at right -- so grow to the left
+                        bb1.include(afwGeom.Point2I(bb1.getMinX()-1, bb1.getMaxY()))
+                        bb1.clip(bbox)
+                if bb1.getHeight() % 2 == 1: # grow downwards
+                    bb1.include(afwGeom.Point2I(bb1.getMaxX(), bb1.getMaxY()+1)) # Expand by 1 pixel!
+                    bb1.clip(bbox)
+                    if bb1.getHeight() % 2 == 1: # clipped upwards -- so grow down
+                        bb1.include(afwGeom.Point2I(bb1.getMaxX(), bb1.getMinY()-1))
+                        bb1.clip(bbox)
+
             return bb0, bb1
 
         xoff = 0
@@ -659,7 +723,7 @@ def offsetAndClipBoxes(bbox0, bbox1, xoff, yoff, bbox):
                     self.boxes1.append(bb1)
             xoff += gridStepX
 
-    def _plotBoxes(self, fullBBox, skip=3):
+    def plotBoxes(self, fullBBox, skip=3):
         """Plot both grids of boxes using matplotlib.
 
         Will compute the grid via `_generateGrid` if

tests/testImageMapReduce.py

@@ -27,9 +27,11 @@
 import lsst.utils.tests
 import lsst.afw.image as afwImage
 import lsst.afw.math as afwMath
+import lsst.afw.geom as afwGeom
 import lsst.pex.config as pexConfig
 import lsst.meas.algorithms as measAlg
 import lsst.pipe.base as pipeBase
+import lsst.daf.base as dafBase
 
 from lsst.ip.diffim.imageMapReduce import (ImageMapReduceTask, ImageMapReduceConfig,
                                            ImageMapperSubtask, ImageMapperSubtaskConfig,
@@ -40,6 +42,52 @@ def setup_module(module):
     lsst.utils.tests.init()
 
 
+def makeWcs(offset=0):
+    # taken from $AFW_DIR/tests/testMakeWcs.py
+    metadata = dafBase.PropertySet()
+    metadata.set("SIMPLE", "T")
+    metadata.set("BITPIX", -32)
+    metadata.set("NAXIS", 2)
+    metadata.set("NAXIS1", 1024)
+    metadata.set("NAXIS2", 1153)
+    metadata.set("RADECSYS", 'FK5')
+    metadata.set("EQUINOX", 2000.)
+    metadata.setDouble("CRVAL1", 215.604025685476)
+    metadata.setDouble("CRVAL2", 53.1595451514076)
+    metadata.setDouble("CRPIX1", 1109.99981456774 + offset)
+    metadata.setDouble("CRPIX2", 560.018167811613 + offset)
+    metadata.set("CTYPE1", 'RA---SIN')
+    metadata.set("CTYPE2", 'DEC--SIN')
+    metadata.setDouble("CD1_1", 5.10808596133527E-05)
+    metadata.setDouble("CD1_2", 1.85579539217196E-07)
+    metadata.setDouble("CD2_2", -5.10281493481982E-05)
+    metadata.setDouble("CD2_1", -8.27440751733828E-07)
+    return afwImage.makeWcs(metadata)
+
+
+def getPsfMoments(psfArray):
+    # Borrowed and modified from meas_algorithms/testCoaddPsf
+    sumx2 = sumy2 = sumy = sumx = sum = 0.0
+    for x in range(psfArray.shape[0]):
+        for y in range(psfArray.shape[1]):
+            f = psfArray[x, y]
+            sumx2 += x*x*f
+            sumy2 += y*y*f
+            sumx += x*f
+            sumy += y*f
+            sum += f
+    xbar = sumx/sum
+    ybar = sumy/sum
+    mxx = sumx2 - 2*xbar*sumx + xbar*xbar*sum
+    myy = sumy2 - 2*ybar*sumy + ybar*ybar*sum
+    return sum, xbar, ybar, mxx, myy
+
+
+def getPsfSecondMoments(psfArray):
+    sum, xbar, ybar, mxx, myy = getPsfMoments(psfArray)
+    return mxx, myy
+
+
 class AddAmountImageMapperSubtaskConfig(ImageMapperSubtaskConfig):
     """Configuration parameters for the AddAmountImageMapperSubtask
     """
@@ -153,6 +201,7 @@ def _makeImage(self):
         self.exposure.setPsf(measAlg.DoubleGaussianPsf(11, 11, 2.0, 3.7))
         mi = self.exposure.getMaskedImage()
         mi.set(0.)
+        self.exposure.setWcs(makeWcs())  # required for PSF construction via CoaddPsf
 
     def testCopySumNoOverlaps(self):
         self._testCopySumNoOverlaps(reduceOp='copy', withNaNs=False)
@@ -162,7 +211,7 @@ def testCopySumNoOverlaps(self):
 
     def _testCopySumNoOverlaps(self, reduceOp='copy', withNaNs=False):
         """Test sample grid task that adds 5.0 to input image and uses
-        default 'copy' `reduceOperation`. Optionally add NaNs to subimages.
+        `reduceOperation = 'copy'`. Optionally add NaNs to subimages.
         """
         config = AddAmountImageMapReduceConfig()
         task = ImageMapReduceTask(config)
@@ -180,6 +229,8 @@ def _testCopySumNoOverlaps(self, reduceOp='copy', withNaNs=False):
         if reduceOp != 'sum':
             self.assertFloatsAlmostEqual(mi[~isnan], newArr[~isnan] - 5.,
                                          msg='Failed on withNaNs: %s' % str(withNaNs))
+        else:  # We don't construct a new PSF if reduceOperation == 'copy'.
+            self._testCoaddPsf(newExp)
 
     def testAverageWithOverlaps(self):
         self._testAverageWithOverlaps(withNaNs=False)
@@ -206,6 +257,37 @@ def _testAverageWithOverlaps(self, withNaNs=False):
         mi = self.exposure.getMaskedImage().getImage().getArray()
         self.assertFloatsAlmostEqual(mi[~isnan], newArr[~isnan] - 5.,
                                      msg='Failed on withNaNs: %s' % str(withNaNs))
+        self._testCoaddPsf(newExp)
+
+    def _testCoaddPsf(self, newExposure):
+        """Test that the new CoaddPsf of the `newExposure` returns PSF images
+        ~identical to the input PSF of `self.exposure` across a grid
+        covering the entire exposure bounding box.
+        """
+        origPsf = self.exposure.getPsf()
+        newPsf = newExposure.getPsf()
+        self.assertTrue(isinstance(newPsf, measAlg.CoaddPsf))
+        extentX = int(self.exposure.getWidth()*0.05)
+        extentY = int(self.exposure.getHeight()*0.05)
+        for x in np.linspace(extentX, self.exposure.getWidth()-extentX, 10):
+            for y in np.linspace(extentY, self.exposure.getHeight()-extentY, 10):
+                point = afwGeom.Point2D(np.rint(x), np.rint(y))
+                oPsf = origPsf.computeImage(point).getArray()
+                nPsf = newPsf.computeImage(point).getArray()
+                if oPsf.shape[0] < nPsf.shape[0]:  # sometimes CoaddPsf does this.
+                    oPsf = np.pad(oPsf, ((1, 1), (0, 0)), mode='constant')
+                elif oPsf.shape[0] > nPsf.shape[0]:
+                    nPsf = np.pad(nPsf, ((1, 1), (0, 0)), mode='constant')
+                if oPsf.shape[1] < nPsf.shape[1]:  # sometimes CoaddPsf does this.
+                    oPsf = np.pad(oPsf, ((0, 0), (1, 1)), mode='constant')
+                elif oPsf.shape[1] > nPsf.shape[1]:
+                    nPsf = np.pad(nPsf, ((0, 0), (1, 1)), mode='constant')
+                # pixel-wise comparison -- pretty stringent
+                self.assertFloatsAlmostEqual(oPsf, nPsf, atol=1e-4, msg='Failed on Psf')
+
+                origMmts = np.array(getPsfSecondMoments(oPsf))
+                newMmts = np.array(getPsfSecondMoments(nPsf))
+                self.assertFloatsAlmostEqual(origMmts, newMmts, atol=1e-4, msg='Failed on Psf')
 
     def testAverageVersusCopy(self):
         self._testAverageVersusCopy(withNaNs=False)