DM-42529: Accelerate the utility function(s) needed for CloughTocher2DInterpolator

include/lsst/meas/algorithms.h

@@ -31,3 +31,4 @@
 #include "lsst/meas/algorithms/CoaddPsf.h"
 #include "lsst/meas/algorithms/WarpedPsf.h"
 #include "lsst/meas/algorithms/CoaddBoundedField.h"
+#include "lsst/meas/algorithms/CloughTocher2DInterpolatorUtils.h"

include/lsst/meas/algorithms/CloughTocher2DInterpolatorUtils.h

@@ -0,0 +1,98 @@
+// -*- LSST-C++ -*-
+/*
+ * This file is part of meas_algorithms.
+ *
+ * Developed for the LSST Data Management System.
+ * This product includes software developed by the LSST Project
+ * (https://www.lsst.org).
+ * See the COPYRIGHT file at the top-level directory of this distribution
+ * for details of code ownership.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#ifndef LSST_MEAS_ALGORITHMS_CLOUGHTOCHER2DINTERPOLATORUTILS_H
+#define LSST_MEAS_ALGORITHMS_CLOUGHTOCHER2DINTERPOLATORUTILS_H
+
+/** \file
+ * \brief Provide pixel-level utilities to support CloughTocher2DInterpolator.
+ */
+
+#include <vector>
+#include "lsst/afw/image/Image.h"
+#include "lsst/afw/image/MaskedImage.h"
+
+namespace lsst {
+namespace meas {
+namespace algorithms {
+
+/**
+@brief This class contains static utility methods that are used by the CloughTocher2DInterpolatorTask
+and exists solely to provide a namespace.
+*/
+class CloughTocher2DInterpolatorUtils {
+    typedef float PixelT;
+
+public:
+    /**
+     * @brief Find the positions of bad pixels as defined by the masks, and also
+     * find the location and pixel value of good pixels around the bad pixels.
+     *
+     * @param[in] mimage MaskedImage to find the pixels from.
+     * @param[in] maskPlanes List of mask planes to consider as bad pixels.
+     * @param[in] buffer Number of pixels to dilate the bad pixels by.
+     * @return A pair of arrays, the first containing the bad pixels and the second containing the good
+     * pixels.
+     *
+     * @note The bad pixels array has shape (N, 3) where N is the number of bad pixels. The first column is
+     * the x coordinate, the second column is the y coordinate, and the third column is the pixel value.
+     * The good pixels array has shape (M, 3) where M is the number of good pixels and has the same format as
+     * the previous one.
+     */
+    static std::pair<ndarray::Array<float, 2, 2>, ndarray::Array<float, 2, 2>> findGoodPixelsAroundBadPixels(
+            afw::image::MaskedImage<PixelT, afw::image::MaskPixel, afw::image::VariancePixel> const &mimage,
+            std::vector<std::string> const &maskPlanes, int const buffer);
+
+    /**
+     * @brief Update the values (third column) of the pixelArray from the image.
+     *
+     * @param[out] pixelArray The three-column array to update.
+     * @param[in] image The image to update the pixel values from.
+     *
+     * @note The pixelArray is typically one of the arrays returned by findGoodPixelsAroundBadPixels.
+     *
+     * @see updateImageFromArray
+     */
+    static void updateArrayFromImage(ndarray::Array<float, 2, 2> const &pixelArray,
+                                     afw::image::Image<PixelT> const &image);
+
+    /**
+     * @brief Update the pixel values of the image from the pixelArray
+     *
+     * @param[out] image The image to update.
+     * @param[in] pixelArray The three-column array to update the pixel values from.
+     *
+     * @note The pixelArray is typically one of the arrays returned by findGoodPixelsAroundBadPixels.
+     *
+     * @see updateArrayFromImage
+     */
+    static void updateImageFromArray(afw::image::Image<PixelT> &image,
+                                     ndarray::Array<float const, 2, 2> const &pixelArray);
+};
+
+}  // namespace algorithms
+}  // namespace meas
+}  // namespace lsst
+
+#endif

python/lsst/meas/algorithms/SConscript

@@ -3,6 +3,7 @@ from lsst.sconsUtils import scripts
 
 scripts.BasicSConscript.python(['_algorithmsLib'], [
     '_algorithmsLib.cc',
+    'cloughTocher2DInterpolatorUtils.cc',
     'cr.cc',
     'coaddBoundedField.cc',
     'coaddPsf/coaddPsf.cc',

python/lsst/meas/algorithms/__init__.py

@@ -26,7 +26,7 @@
 import lsst.afw.image
 import lsst.afw.math
 
-from._algorithmsLib import *
+from ._algorithmsLib import *
 from .psfCandidate import * #python
 from .coaddPsf import *
 from .simple_curve import *

python/lsst/meas/algorithms/_algorithmsLib.cc

@@ -38,12 +38,13 @@ void wrapDoubleGaussianPsf(WrapperCollection &wrappers);
 void wrapInterp(WrapperCollection &wrappers);
 void wrapPcaPsf(WrapperCollection &wrappers);
 void wrapWarpedPsf(WrapperCollection &wrappers);
-void wapPsfCandidate(WrapperCollection &wrappers);
+void wrapPsfCandidate(WrapperCollection &wrappers);
 void wrapImagePsf(WrapperCollection &wrappers);
 void wrapSpatialModelPsf(WrapperCollection &wrappers);
 void wrapCoaddPsf(WrapperCollection &wrappers);
 void wrapCoaddBoundedField(WrapperCollection &wrappers);
 void wrapCr(WrapperCollection &wrappers);
+void wrapCloughTocher2DInterpolatorUtils(WrapperCollection &wrappers);
 
 PYBIND11_MODULE(_algorithmsLib, mod) {
 WrapperCollection wrappers(mod, "lsst.meas.algorithms");
@@ -60,11 +61,12 @@ WrapperCollection wrappers(mod, "lsst.meas.algorithms");
     wrapInterp(wrappers);
     wrapPcaPsf(wrappers);
     wrapWarpedPsf(wrappers);
-    wapPsfCandidate(wrappers);
+    wrapPsfCandidate(wrappers);
     wrapSpatialModelPsf(wrappers);
     wrapCoaddPsf(wrappers);
     wrapCoaddBoundedField(wrappers);
     wrapCr(wrappers);
+    wrapCloughTocher2DInterpolatorUtils(wrappers);
     wrappers.finish();
 }
 

python/lsst/meas/algorithms/cloughTocher2DInterpolator.py

@@ -0,0 +1,144 @@
+# This file is part of meas_algorithms.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+__all__ = (
+    "CloughTocher2DInterpolateConfig",
+    "CloughTocher2DInterpolateTask",
+)
+
+
+from lsst.pex.config import Config, Field, ListField
+from lsst.pipe.base import Task
+from scipy.interpolate import CloughTocher2DInterpolator
+
+from . import CloughTocher2DInterpolatorUtils as ctUtils
+
+
+class CloughTocher2DInterpolateConfig(Config):
+    """Config for CloughTocher2DInterpolateTask."""
+
+    badMaskPlanes = ListField[str](
+        doc="List of mask planes to interpolate over.",
+        default=["BAD", "SAT", "CR"],
+    )
+    fillValue = Field[float](
+        doc="Constant value to fill outside of the convex hull of the good "
+        "pixels. A long (longer than twice the ``interpLength``) streak of "
+        "bad pixels at an edge will be set to this value.",
+        default=0.0,
+    )
+    interpLength = Field[int](
+        doc="Maximum number of pixels away from a bad pixel to include in "
+        "building the interpolant. Must be greater than or equal to 1.",
+        default=4,
+        check=lambda x: x >= 1,
+    )
+
+
+class CloughTocher2DInterpolateTask(Task):
+    """Interpolated over bad pixels using CloughTocher2DInterpolator.
+
+    Pixels with mask bits set to any of those listed ``badMaskPlanes`` config
+    are considered bad and are interpolated over. All good (non-bad) pixels
+    within ``interpLength`` pixels of a bad pixel in either direction are used
+    to construct the interpolant.  An extended streak of bad pixels at an edge,
+    longer than ``interpLength``, is set to `fillValue`` specified in config.
+    """
+
+    ConfigClass = CloughTocher2DInterpolateConfig
+    _DefaultName = "cloughTocher2DInterpolate"
+
+    def run(
+        self,
+        maskedImage,
+        badpix=None,
+        goodpix=None,
+    ):
+        """Interpolate over bad pixels in a masked image.
+
+        This modifies the ``image`` attribute of the ``maskedImage`` in place.
+        This method returns, and accepts, the coordinates of the bad pixels
+        that were interpolated over, and the coordinates and values of the
+        good pixels that were used to construct the interpolant. This avoids
+        having to search for the bad and the good pixels repeatedly when the
+        mask plane is shared among many images, as would be the case with
+        noise realizations.
+
+        Parameters
+        ----------
+        maskedImage : `~lsst.afw.image.MaskedImageF`
+            Image on which to perform interpolation (and modify in-place).
+        badpix: `numpy.ndarray`, optional
+            N x 3 numpy array, where N is the number of bad pixels.
+            The coordinates of the bad pixels to interpolate over in the first
+            two columns, and the pixel values (unused) in the third column.
+            If None, then the coordinates of the bad pixels are determined by
+            an exhaustive search over the image. If ``goodpix`` is not
+            provided, then this parameter is ignored.
+        goodpix: `numpy.ndarray`, optional
+            M x 3 numpy array, where M is the number of good pixels.
+            The first two columns are the coordinates of the good pixels around
+            ``badpix`` that must be included when constructing the interpolant.
+            the interpolant. The values are populated from the image plane of
+            the ``maskedImage`` and returned (provided values will be ignored).
+            If ``badpix`` is not provided, then this parameter is ignored.
+
+        Returns
+        -------
+        badpix: `numpy.ndarray`
+            N x 3 numpy array, where N is the number of bad pixels.
+            The coordinates of the bad pixels that were interpolated over are
+            in the first two columns, and the corresponding pixel values in the
+            third. If ``badpix`` was provided, this is the same as the input.
+        goodpix: `numpy.ndarray`
+            M x 3 numpy array, where M is the number of bad pixels.
+            The coordinates of the good pixels that were used to construct the
+            interpolant arein the first two columns, and the corresponding
+            pixel values in the third. If ``goodpix`` was provided, the first
+            two columns are same as the input, with the third column updated
+            with the pixel values from the image plane of the ``maskedImage``.
+        """
+
+        if badpix is None or goodpix is None:
+            badpix, goodpix = ctUtils.findGoodPixelsAroundBadPixels(
+                maskedImage,
+                self.config.badMaskPlanes,
+                buffer=self.config.interpLength,
+            )
+        else:
+            # Even if badpix and goodpix is provided, make sure to update
+            # the values of goodpix.
+            ctUtils.updateArrayFromImage(goodpix, maskedImage.image)
+
+        # Construct the interpolant with goodpix.
+        interpolator = CloughTocher2DInterpolator(
+            list(zip(goodpix[:, 0], goodpix[:, 1])),
+            goodpix[:, 2],
+            fill_value=self.config.fillValue,
+        )
+
+        # Compute the interpolated values at bad pixel locations.
+        badpix[:, 2] = interpolator(badpix[:, :2])
+
+        # Fill in the bad pixels.
+        ctUtils.updateImageFromArray(maskedImage.image, badpix)
+
+        return badpix, goodpix

python/lsst/meas/algorithms/cloughTocher2DInterpolatorUtils.cc

@@ -0,0 +1,64 @@
+// -*- LSST-C++ -*-
+/*
+ * This file is part of meas_algorithms.
+ *
+ * Developed for the LSST Data Management System.
+ * This product includes software developed by the LSST Project
+ * (https://www.lsst.org).
+ * See the COPYRIGHT file at the top-level directory of this distribution
+ * for details of code ownership.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include "pybind11/pybind11.h"
+#include "ndarray/pybind11.h"
+#include "lsst/cpputils/python.h"
+#include "pybind11/stl.h"
+
+#include "lsst/geom/Box.h"
+#include "lsst/meas/algorithms/CloughTocher2DInterpolatorUtils.h"
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+namespace lsst {
+namespace meas {
+namespace algorithms {
+namespace {
+
+void declareCloughTocher2DInterpolatorUtils(lsst::cpputils::python::WrapperCollection &wrappers) {
+    using PyCloughTocher2DInterpolatorUtils =
+            py::class_<CloughTocher2DInterpolatorUtils, std::shared_ptr<CloughTocher2DInterpolatorUtils>>;
+    auto clsCloughTocher2DInterpolatorUtils = wrappers.wrapType(
+            PyCloughTocher2DInterpolatorUtils(wrappers.module, "CloughTocher2DInterpolatorUtils"),
+            [](auto &mod, auto &cls) {
+                cls.def_static("findGoodPixelsAroundBadPixels",
+                               &CloughTocher2DInterpolatorUtils::findGoodPixelsAroundBadPixels, "mimage"_a,
+                               "badList"_a, "buffer"_a = 4);
+                cls.def_static("updateArrayFromImage", &CloughTocher2DInterpolatorUtils::updateArrayFromImage,
+                               "array"_a, "image"_a);
+                cls.def_static("updateImageFromArray", &CloughTocher2DInterpolatorUtils::updateImageFromArray,
+                               "image"_a, "array"_a);
+            });
+}
+}  // namespace
+
+void wrapCloughTocher2DInterpolatorUtils(lsst::cpputils::python::WrapperCollection &wrappers) {
+    declareCloughTocher2DInterpolatorUtils(wrappers);
+}
+
+}  // namespace algorithms
+}  // namespace meas
+}  // namespace lsst

python/lsst/meas/algorithms/psfCandidate/psfCandidate.cc

@@ -76,7 +76,7 @@ void declarePsfCandidate(lsst::cpputils::python::WrapperCollection &wrappers, st
 
 }  // namespace
 
-void wapPsfCandidate(lsst::cpputils::python::WrapperCollection &wrappers) {
+void wrapPsfCandidate(lsst::cpputils::python::WrapperCollection &wrappers) {
     declarePsfCandidate<float>(wrappers, "F");
 }