Merge pull request #612 from lsst/tickets/DM-32219

DM-32219: Add HeavyFootprint addTo and subtractFrom methods.

python/lsst/afw/detection/__init__.py

@@ -23,3 +23,4 @@
 from ._footprintMerge import *
 from ._peak import *
 from .multiband import *
+from ._heavyFootprintContinued import *

python/lsst/afw/detection/_heavyFootprintContinued.py

@@ -0,0 +1,55 @@
+# This file is part of afw.
+#
+# 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.
+#
+import numpy as np
+
+from lsst.utils import TemplateMeta
+
+from ._detection import HeavyFootprintI, HeavyFootprintF, HeavyFootprintD, HeavyFootprintU
+
+__all__ = []  # import this module only for its side effects
+
+
+class HeavyFootprint(metaclass=TemplateMeta):  # noqa: F811
+    def addTo(self, image):
+        """Add this heavy footprint to an image.
+
+        Parameters
+        ----------
+        image : `lsst.afw.image`
+        """
+        indices = self.spans.indices()
+        image.array[indices[0, :] - image.getY0(),
+                    indices[1, :] - image.getX0()] += self.getImageArray()
+
+    def subtractFrom(self, image):
+        """Subtract this heavy footprint from an image.
+
+        Parameters
+        ----------
+        image : `lsst.afw.image`
+        """
+        indices = self.spans.indices()
+        image.array[indices[0, :] - image.getY0(),
+                    indices[1, :] - image.getX0()] -= self.getImageArray()
+
+
+HeavyFootprint.register(np.int32, HeavyFootprintI)
+HeavyFootprint.register(np.float32, HeavyFootprintF)
+HeavyFootprint.register(np.float64, HeavyFootprintD)
+HeavyFootprint.register(np.uint16, HeavyFootprintU)

python/lsst/afw/geom/_spanSet.cc

@@ -31,6 +31,7 @@
 #include <iostream>
 
 #include "ndarray/pybind11.h"
+#include "ndarray/Array.h"
 
 #include "lsst/afw/geom/ellipses/Quadrupole.h"
 #include "lsst/afw/geom/SpanSet.h"
@@ -259,19 +260,19 @@ void declareSpanSet(lsst::utils::python::WrapperCollection &wrappers) {
                        (std::shared_ptr<SpanSet>(*)(geom::ellipses::Ellipse const &)) & SpanSet::fromShape);
         cls.def("split", &SpanSet::split);
         cls.def("findEdgePixels", &SpanSet::findEdgePixels);
-        cls.def("indices", [](SpanSet const &self) -> std::pair<std::vector<int>, std::vector<int>> {
-            std::vector<int> yind;
-            std::vector<int> xind;
-            yind.reserve(self.getArea());
-            xind.reserve(self.getArea());
+        cls.def("indices", [](SpanSet const &self) -> ndarray::Array<int, 2, 2> {
+            unsigned long dims = 2;
+            ndarray::Array<int, 2, 2> inds = ndarray::allocate(ndarray::makeVector(dims, self.getArea()));
+            int element = 0;
             for (auto const &span : self) {
                 auto y = span.getY();
                 for (int x = span.getX0(); x <= span.getX1(); ++x) {
-                    yind.push_back(y);
-                    xind.push_back(x);
+                    inds[0][element] = y;
+                    inds[1][element] = x;
+                    element++;
                 }
             }
-            return std::make_pair(yind, xind);
+            return inds;
         });
 
         /* SpanSet Operators */

tests/test_heavyFootprint.py

@@ -301,6 +301,43 @@ def testDot(self):
             self.assertEqual(hfp1.dot(hfp2), dot)
             self.assertEqual(hfp2.dot(hfp1), dot)
 
+    def testAddSubtract(self):
+        """Test that we can add and subtract a HeavyFootprint."""
+        fs = afwDetect.FootprintSet(self.mi, afwDetect.Threshold(1))
+
+        fs.makeHeavy(self.mi)
+
+        bbox = lsst.geom.BoxI(lsst.geom.PointI(9, 1), lsst.geom.ExtentI(7, 4))
+        compImage = self.mi.Factory(self.mi, bbox, afwImage.LOCAL, True)
+        compImage.set((10, 0x0, 0))
+
+        blankImage = compImage.clone()
+        blankImage.set((0, 0x0, 0))
+
+        testImage = compImage.clone()
+
+        for foot in fs.getFootprints():
+            # Add via addTo
+            foot.addTo(testImage.image)
+
+            # Add by using a temporary image and insert
+            foot.insert(blankImage.image)
+            inserted = blankImage.image.array > 0
+            compImage.image.array[inserted] += blankImage.image.array[inserted]
+
+            np.testing.assert_array_almost_equal(testImage.image.array,
+                                                 compImage.image.array)
+
+            # Subtract via subtractFrom
+            foot.subtractFrom(testImage.image)
+
+            compImage.image.array[inserted] -= blankImage.image.array[inserted]
+
+            np.testing.assert_array_almost_equal(testImage.image.array,
+                                                 compImage.image.array)
+
+            blankImage.set((0, 0x0, 0))
+
 
 class TestMemory(lsst.utils.tests.MemoryTestCase):
     pass