Coverage for tests/test_isPrimaryFlag.py: 94%
176 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-24 09:42 +0000
1# This file is part of pipe_tasks.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22import os
23import unittest
24import numpy as np
26from lsst.geom import Point2I, Box2I, Extent2I
27from lsst.skymap import TractInfo
28from lsst.skymap.patchInfo import PatchInfo
29import lsst.afw.detection as afwDetect
30import lsst.afw.image as afwImage
31import lsst.utils.tests
32from lsst.pipe.tasks.characterizeImage import CharacterizeImageTask, CharacterizeImageConfig
33from lsst.pipe.tasks.calibrate import CalibrateTask, CalibrateConfig
34from lsst.meas.algorithms import SourceDetectionTask, SkyObjectsTask, SetPrimaryFlagsTask
35import lsst.meas.extensions.scarlet as mes
36from lsst.meas.extensions.scarlet.scarletDeblendTask import ScarletDeblendTask
37from lsst.meas.extensions.scarlet.deconvolveExposureTask import DeconvolveExposureTask
38from lsst.meas.base import SingleFrameMeasurementTask
39from lsst.afw.table import SourceCatalog
41TESTDIR = os.path.abspath(os.path.dirname(__file__))
44class NullTract(TractInfo):
45 """A Tract not contained in the MockSkyMap.
47 BaseSkyMap.findTract(coord) will always return a Tract,
48 even if the coord isn't located in the Tract.
49 In order to mimick this functionality we create a
50 NullTract for regions of the MockSkyMap that
51 aren't contained in any of the tracts.
52 """
53 def __init__(self):
54 pass
56 def getId(self):
57 return None
60class MockTractInfo:
61 """A Tract based on a bounding box and WCS.
63 Testing is made easier when we can specifically define
64 a Tract in terms of its bounding box in pixel coordinates
65 along with a WCS for the exposure.
67 Only the relevant methods from `TractInfo` needed to make
68 test pass are implemented here. Since this is just for
69 testing, it isn't sophisticated and requires developers to
70 ensure that the size of the bounding box is evenly divisible
71 by the number of patches in the Tract.
72 """
73 def __init__(self, name, bbox, wcs, numPatches):
74 self.name = name
75 self.bbox = bbox
76 self.wcs = wcs
77 self._numPatches = numPatches
78 assert bbox.getWidth()%numPatches[0] == 0
79 assert bbox.getHeight()%numPatches[1] == 0
80 self.patchWidth = bbox.getWidth()//numPatches[0]
81 self.patchHeight = bbox.getHeight()//numPatches[1]
83 def contains(self, coord):
84 pixel = self.wcs.skyToPixel(coord)
85 return self.bbox.contains(Point2I(pixel))
87 def getId(self):
88 return self.name
90 def getNumPatches(self):
91 return self._numPatches
93 def getPatchInfo(self, index):
94 x, y = index
95 width = self.patchWidth
96 height = self.patchHeight
98 x = x*self.patchWidth
99 y = y*self.patchHeight
101 bbox = Box2I(Point2I(x, y), Extent2I(width, height))
103 nx, ny = self._numPatches
104 sequentialIndex = nx*y + x
106 patchInfo = PatchInfo(
107 index=index,
108 innerBBox=bbox,
109 outerBBox=bbox,
110 sequentialIndex=sequentialIndex,
111 tractWcs=self.wcs
112 )
113 return patchInfo
115 def __getitem__(self, index):
116 return self.getPatchInfo(index)
118 def __iter__(self):
119 xNum, yNum = self.getNumPatches()
120 for y in range(yNum):
121 for x in range(xNum):
122 yield self.getPatchInfo((x, y))
125class MockSkyMap:
126 """A SkyMap based on a list of bounding boxes.
128 Testing is made easier when we can specifically define
129 a Tract in terms of its bounding box in pixel coordinates
130 along with a WCS for the exposure. This class allows us
131 to define the tract(s) in the SkyMap and create
132 them.
133 """
134 def __init__(self, bboxes, wcs, numPatches):
135 self.bboxes = bboxes
136 self.wcs = wcs
137 self.numPatches = numPatches
139 def __iter__(self):
140 for b, bbox in enumerate(self.bboxes):
141 yield self.generateTract(b)
143 def __getitem__(self, index):
144 return self.generateTract(index)
146 def generateTract(self, index):
147 return MockTractInfo(index, self.bboxes[index], self.wcs, self.numPatches)
149 def findTract(self, coord):
150 for tractInfo in self:
151 if tractInfo.contains(coord):
152 return tractInfo
154 return NullTract()
157class IsPrimaryTestCase(lsst.utils.tests.TestCase):
159 def setUp(self):
160 # Load sample input from disk
161 expPath = os.path.join(TESTDIR, "data", "v695833-e0-c000-a00.sci.fits")
162 self.exposure = afwImage.ExposureF(expPath)
164 # Characterize the image (create PSF, etc.)
165 charImConfig = CharacterizeImageConfig()
166 charImConfig.measureApCorr.sourceSelector["science"].doSignalToNoise = False
167 charImTask = CharacterizeImageTask(config=charImConfig)
168 self.charImResults = charImTask.run(self.exposure)
170 def tearDown(self):
171 del self.exposure
172 self.charImResults
174 def testIsSinglePrimaryFlag(self):
175 """Tests detect_isPrimary column gets added when run, and that sources
176 labelled as detect_isPrimary are not sky sources and have no children.
177 """
178 calibConfig = CalibrateConfig()
179 calibConfig.doAstrometry = False
180 calibConfig.doPhotoCal = False
181 calibConfig.doComputeSummaryStats = False
182 calibTask = CalibrateTask(config=calibConfig)
183 calibResults = calibTask.run(self.charImResults.exposure)
184 outputCat = calibResults.outputCat
185 self.assertTrue("detect_isPrimary" in outputCat.schema.getNames())
186 # make sure all sky sources are flagged as not primary
187 self.assertEqual(sum((outputCat["detect_isPrimary"]) & (outputCat["sky_source"])), 0)
188 # make sure all parent sources are flagged as not primary
189 self.assertEqual(sum((outputCat["detect_isPrimary"]) & (outputCat["deblend_nChild"] > 0)), 0)
191 with self.assertRaises(KeyError):
192 outputCat.getSchema().find("detect_isDelendedModelPrimary")
194 def testIsScarletPrimaryFlag(self):
195 """Test detect_isPrimary column when scarlet is used as the deblender
196 """
197 # We need a multiband coadd for scarlet,
198 # even though there is only one band
199 coadds = afwImage.MultibandExposure.fromExposures(["test"], [self.exposure])
201 # Create a SkyMap with a tract that contains a portion of the image,
202 # subdivided into 3x3 patches
203 wcs = self.exposure.getWcs()
204 tractBBox = Box2I(Point2I(100, 100), Extent2I(900, 900))
205 skyMap = MockSkyMap([tractBBox], wcs, (3, 3))
206 tractInfo = skyMap[0]
207 patchInfo = tractInfo[0, 0]
208 patchBBox = patchInfo.getInnerBBox()
210 schema = SourceCatalog.Table.makeMinimalSchema()
211 # Initialize the detection task
212 detectionTask = SourceDetectionTask(schema=schema)
214 # Initialize the fake source injection task
215 skyConfig = SkyObjectsTask.ConfigClass()
216 skySourcesTask = SkyObjectsTask(name="skySources", config=skyConfig)
218 # Build the Footprint peak schema that the coadd merge step would
219 # provide. The deblender reads the pseudo-source flag from each
220 # Footprint's peak (not the source record) and transfers any extra
221 # peak fields onto the source schema, so `merge_peak_sky` must live on
222 # the peak schema and must not be pre-added to the source schema here.
223 peakSchema = afwDetect.PeakTable.makeMinimalSchema()
224 skyPeakKey = peakSchema.addField("merge_peak_sky", type="Flag")
226 # Initialize the deconvolution task
227 deconvolveConfig = DeconvolveExposureTask.ConfigClass()
228 deconvolveTask = DeconvolveExposureTask(config=deconvolveConfig)
230 # Initialize the deblender task
231 scarletConfig = ScarletDeblendTask.ConfigClass()
232 scarletConfig.maxIter = 20
233 scarletConfig.columnInheritance["merge_peak_sky"] = "merge_peak_sky"
234 scarletConfig.processSingles = True
235 deblendTask = ScarletDeblendTask(schema=schema, peakSchema=peakSchema, config=scarletConfig)
237 # We'll customize the configuration of measurement to just run the
238 # minimal number of plugins to make setPrimaryFlags work.
239 # As of DM-51670 we also include `base_PsfFlux` to ensure that
240 # the measurement plugins run correctly with the split between
241 # parent and child catalogs.
242 measureConfig = SingleFrameMeasurementTask.ConfigClass()
243 measureConfig.plugins.names = ["base_SdssCentroid", "base_SkyCoord", "base_PsfFlux"]
244 measureConfig.slots.apFlux = None
245 measureConfig.slots.shape = None
246 measureConfig.slots.modelFlux = None
247 measureConfig.slots.calibFlux = None
248 measureConfig.slots.gaussianFlux = None
249 measureTask = SingleFrameMeasurementTask(config=measureConfig, schema=schema)
250 setPrimaryTask = SetPrimaryFlagsTask(schema=schema, isSingleFrame=False)
252 table = SourceCatalog.Table.make(schema)
253 # detect sources
254 detectionResult = detectionTask.run(table, coadds["test"])
255 catalog = detectionResult.sources
256 # The deblender requires every Footprint's peaks to share the peak
257 # schema passed to it, so rebuild the detected Footprints with the
258 # merge peak schema (their peaks carry no pseudo-source flag).
259 for src in catalog:
260 oldFoot = src.getFootprint()
261 newFoot = afwDetect.Footprint(oldFoot.spans, peakSchema)
262 for peak in oldFoot.peaks:
263 newFoot.addPeak(peak.getFx(), peak.getFy(), peak.getPeakValue())
264 src.setFootprint(newFoot)
265 # add fake sources, flagging the sky peak as the coadd merge step does
266 # in `MergeDetectionsTask.getSkySourceFootprints`
267 skySources = skySourcesTask.run(mask=self.exposure.mask, seed=0)
268 for foot in skySources[:5]:
269 src = catalog.addNew()
270 peak = foot.getPeaks()[0]
271 newFoot = afwDetect.Footprint(foot.spans, peakSchema)
272 newFoot.addPeak(peak.getFx(), peak.getFy(), peak.getPeakValue())
273 newFoot.getPeaks()[0].set(skyPeakKey, True)
274 src.setFootprint(newFoot)
275 # deconvolve the images
276 deconvolved = deconvolveTask.run(coadds["test"], catalog).deconvolved
277 mDeconvolved = afwImage.MultibandExposure.fromExposures(["test"], [deconvolved])
278 # deblend
279 # This is a hack because the variance is not calibrated properly
280 # (it is 3 orders of magnitude too high), which causes the deblender
281 # to improperly deblend most sources due to the sparsity constraint.
282 coadds.variance.array[:] = 2e-1
283 mDeconvolved.variance.array[:] = 2e-1
284 result = deblendTask.run(coadds, mDeconvolved, catalog)
285 modelData = result.scarletModelData
286 catalog = result.deblendedCatalog
287 # Attach footprints to the catalog
288 mes.io.updateCatalogFootprints(
289 modelData=modelData,
290 catalog=catalog,
291 band="test",
292 imageForRedistribution=coadds["test"],
293 removeScarletData=True,
294 updateFluxColumns=True,
295 )
297 # measure
298 measureTask.run(catalog, self.exposure)
299 outputCat = catalog
300 # Set the primary flags
301 setPrimaryTask.run(outputCat, skyMap=skyMap, tractInfo=tractInfo, patchInfo=patchInfo)
303 # There should be the same number of deblenedPrimary and
304 # deblendedModelPrimary sources,
305 # since they both have the same blended sources and only differ
306 # over which model to use for the isolated sources.
307 isPseudo = outputCat["merge_peak_sky"]
309 # Check that all 5 pseudo-sources were created
310 self.assertEqual(np.sum(isPseudo), 5)
312 self.assertEqual(
313 np.sum(outputCat["detect_isDeblendedSource"] & ~isPseudo),
314 np.sum(outputCat["detect_isDeblendedModelSource"]))
316 # Check that the sources contained in a tract are all marked appropriately
317 x = outputCat["slot_Centroid_x"]
318 y = outputCat["slot_Centroid_y"]
319 tractInner = tractBBox.contains(x, y)
320 np.testing.assert_array_equal(outputCat["detect_isTractInner"], tractInner)
322 # Check that the sources contained in a patch are all marked appropriately
323 patchInner = patchBBox.contains(x, y)
324 np.testing.assert_array_equal(outputCat["detect_isPatchInner"], patchInner)
326 # make sure all sky sources are flagged as not primary
327 self.assertEqual(sum((outputCat["detect_isPrimary"]) & (outputCat["merge_peak_sky"])), 0)
329 # Check that sky objects have not been deblended
330 # (deblended sources have parent > 0)
331 np.testing.assert_array_equal(
332 isPseudo,
333 isPseudo & (outputCat["parent"] == 0)
334 )
336 # Check that measurements were performed on all of the children
337 self.assertTrue(np.all(outputCat["base_PsfFlux_instFlux"] != 0) and np.all(np.isfinite(
338 outputCat["base_PsfFlux_instFlux"])))
341class MemoryTester(lsst.utils.tests.MemoryTestCase):
342 pass
345def setup_module(module):
346 lsst.utils.tests.init()
349if __name__ == "__main__": 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 lsst.utils.tests.init()
351 unittest.main()