Coverage for tests/test_deblend.py: 12%
392 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 04:06 -0700
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 04:06 -0700
1# This file is part of meas_extensions_scarlet.
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 tempfile
26import lsst.afw.image as afwImage
27import lsst.meas.extensions.scarlet as mes
28import lsst.scarlet.lite as scl
29import lsst.utils.tests
30import numpy as np
31from lsst.afw.detection import Footprint, GaussianPsf, InvalidPsfError, PeakTable, Psf
32from lsst.afw.geom import SpanSet
33from lsst.afw.table import SourceCatalog, SourceTable, SchemaMapper
34from lsst.daf.butler import Config, DatasetType, StorageClass, FileDataset, DatasetRef
35from lsst.daf.butler.tests import makeTestRepo, makeTestCollection
36from lsst.geom import Extent2I, Point2D, Point2I
37from lsst.meas.algorithms import SourceDetectionTask
38from lsst.meas.extensions.scarlet.scarletDeblendTask import ScarletDeblendTask
39from lsst.meas.extensions.scarlet.deconvolveExposureTask import DeconvolveExposureTask
40from lsst.pipe.base import NoWorkFound, Struct
41from utils import initData, SersicModel, PsfModel
43TESTDIR = os.path.abspath(os.path.dirname(__file__))
46class TestDeblend(lsst.utils.tests.TestCase):
47 def setUp(self):
48 self.modelPsf = scl.utils.integrated_circular_gaussian(sigma=0.8).astype(np.float32)
49 psfRadius = 20
50 psfShape = (2 * psfRadius + 1, 2 * psfRadius + 1)
51 self.psfs = [
52 GaussianPsf(psfShape[1], psfShape[0], 1.0),
53 GaussianPsf(psfShape[1], psfShape[0], 1.2),
54 GaussianPsf(psfShape[1], psfShape[0], 1.4),
55 ]
56 self.imagePsf = np.asarray(
57 [psf.computeImage(psf.getAveragePosition()).array for psf in self.psfs]
58 ).astype(np.float32)
59 self.imagePsf /= self.imagePsf.sum(axis=(1, 2))[:, None, None]
60 self.bands = tuple("gri")
62 self.models = [
63 # Isolated source
64 PsfModel(
65 center=(30, 15),
66 spectrum=np.array([8, 2, 1]),
67 bands=self.bands,
68 ),
69 # Two source blend
70 SersicModel(
71 center=(40, 20),
72 major=5,
73 minor=2,
74 radius=15,
75 theta=-np.pi/4,
76 n=1,
77 spectrum=np.array([2, 4, 8]),
78 bands=self.bands,
79 ),
80 PsfModel(
81 center=(12, 20),
82 spectrum=np.array([1, 2, 8]),
83 bands=self.bands,
84 ),
85 # 3 source blend
86 SersicModel(
87 center=(25, 70),
88 major=5,
89 minor=2,
90 radius=20,
91 theta=np.pi/48,
92 n=1,
93 spectrum=np.array([2, 5, 8]),
94 bands=self.bands,
95 ),
96 PsfModel(
97 center=(32, 60),
98 spectrum=np.array([1, 2, 8]),
99 bands=self.bands,
100 ),
101 PsfModel(
102 center=(16, 80),
103 spectrum=np.array([8, 2, 1]),
104 bands=self.bands,
105 ),
106 # Large blend
107 SersicModel(
108 center=(70, 70),
109 major=5,
110 minor=2,
111 radius=25,
112 theta=0,
113 n=1,
114 spectrum=np.array([2, 10, 18]),
115 bands=self.bands,
116 ),
117 SersicModel(
118 center=(85, 85),
119 major=5,
120 minor=2,
121 radius=25,
122 theta=np.pi/2,
123 n=1,
124 spectrum=np.array([5, 10, 20]),
125 bands=self.bands,
126 ),
127 ]
129 def scarlet_image_to_exposure(
130 self,
131 image: scl.Image,
132 noise: np.ndarray,
133 ) -> afwImage.MultibandExposure:
134 masked_image = afwImage.MultibandMaskedImage.fromArrays(
135 image.bands, image.data, None, noise**2
136 )
137 coadds = [
138 afwImage.Exposure(img, dtype=img.image.array.dtype) for img in masked_image
139 ]
140 mCoadd = afwImage.MultibandExposure.fromExposures(image.bands, coadds)
141 for b, coadd in enumerate(mCoadd):
142 coadd.setPsf(self.psfs[b])
143 return mCoadd
145 def initialize_data(
146 self,
147 models,
148 deconvolveConfig=None,
149 deblendConfig=None,
150 doDetect: bool = True,
151 ):
152 if deconvolveConfig is None:
153 deconvolveConfig = DeconvolveExposureTask.ConfigClass()
154 if deblendConfig is None:
155 deblendConfig = ScarletDeblendTask.ConfigClass()
156 # Generate the data for the test
157 deconvolved, convolved = initData(models, self.modelPsf, self.imagePsf)
158 # Set the random seed so that the noise field is unaffected
159 # and add noise to the image
160 np.random.seed(0)
161 noise = 0.05 * (np.random.rand(*convolved.shape).astype(np.float32) - 0.5)
162 noisyImage = convolved.copy()
163 noisyImage._data += noise
164 # Create the multiband coadd
165 mCoadd = self.scarlet_image_to_exposure(noisyImage, noise)
166 # Initialze tasks
167 inputSchema = SourceTable.makeMinimalSchema()
168 table = SourceTable.make(inputSchema)
169 detectionTask = SourceDetectionTask(schema=inputSchema)
170 schemaMapper = SchemaMapper(inputSchema)
171 schemaMapper.addMinimalSchema(inputSchema)
172 schema = schemaMapper.getOutputSchema()
173 deconvolveTask = DeconvolveExposureTask(config=deconvolveConfig)
174 deblendTask = ScarletDeblendTask(schema=schema, config=deblendConfig)
176 result = Struct(
177 deconvolved=deconvolved,
178 convolved=convolved,
179 noise=noise,
180 noisyImage=noisyImage,
181 mCoadd=mCoadd,
182 detectionTask=detectionTask,
183 deconvolveTask=deconvolveTask,
184 deblendTask=deblendTask,
185 )
187 if doDetect:
188 # Generate a detection catalog
189 detectionResult = detectionTask.run(table, mCoadd["r"])
190 table = SourceCatalog.Table.make(schema)
191 catalog = SourceCatalog(table)
192 catalog.extend(detectionResult.sources, schemaMapper)
193 result.catalog = catalog
195 return result
197 def deconvolve(self, data: Struct):
198 deconvolvedCoadds = []
199 deconvolveTask = data.deconvolveTask
200 if deconvolveTask.config.useFootprints:
201 catalog = data.catalog
202 else:
203 catalog = None
204 for coadd in data.mCoadd:
205 deconvolvedCoadd = deconvolveTask.run(coadd, catalog).deconvolved
206 deconvolvedCoadds.append(deconvolvedCoadd)
207 mDeconvolved = afwImage.MultibandExposure.fromExposures(self.bands, deconvolvedCoadds)
208 return mDeconvolved
210 def test_default_deconvolve(self):
211 data = self.initialize_data(self.models)
212 deconvolved = self.deconvolve(data)
214 diff = data.deconvolved.data - deconvolved.image.array
215 # Due to peakiness of Sersic models the center has a sharp peak,
216 # so we ignore a 3x3 region around each source center
217 for model in self.models:
218 yc, xc = model.center
219 for x in (-1, 0, 1):
220 for y in (-1, 0, 1):
221 diff[:, yc+y, xc+x] = 0
222 self.assertTrue(np.max(diff[:2]) < 10*np.std(data.noise))
223 self.assertTrue(np.max(diff[2]) < 20*np.std(data.noise))
225 context = mes.scarletDeblendTask.ScarletDeblendContext.build(
226 data.mCoadd,
227 deconvolved,
228 data.catalog,
229 data.deblendTask.ConfigClass()
230 )
232 self.assertEqual(len(context.footprints), 4)
234 def test_catalog_free_deconvolve(self):
235 config = DeconvolveExposureTask.ConfigClass()
236 config.useFootprints = False
237 data = self.initialize_data(self.models, deconvolveConfig=config)
238 deconvolved = self.deconvolve(data)
240 diff = data.deconvolved.data - deconvolved.image.array
241 # Due to peakiness of Sersic models the center has a sharp peak,
242 # so we ignore a 3x3 region around each source center
243 for model in self.models:
244 yc, xc = model.center
245 for x in (-1, 0, 1):
246 for y in (-1, 0, 1):
247 diff[:, yc+y, xc+x] = 0
248 self.assertTrue(np.max(diff[:2]) < 10*np.std(data.noise))
249 self.assertTrue(np.max(diff[2]) < 20*np.std(data.noise))
251 def test_footprints(self):
252 data = self.initialize_data(self.models)
253 nParents = len(data.catalog)
254 mDeconvolved = self.deconvolve(data)
255 result = data.deblendTask.run(data.mCoadd, mDeconvolved, data.catalog)
257 config = data.deblendTask.config
258 catalog = result.deblendedCatalog
259 modelData = result.scarletModelData
260 observedPsf = modelData.metadata["psf"]
261 modelPsf = modelData.metadata["model_psf"]
263 # Attach the footprints in each band and compare to the full
264 # data model. This is done in each band, both with and without
265 # flux re-distribution to test all of the different possible
266 # options of loading catalog footprints.
267 for useFlux in [False, True]:
268 for band in self.bands:
269 bandIndex = self.bands.index(band)
270 coadd = data.mCoadd[band]
272 if useFlux:
273 imageForRedistribution = coadd
274 else:
275 imageForRedistribution = None
277 mes.io.updateCatalogFootprints(
278 modelData,
279 catalog,
280 band=band,
281 imageForRedistribution=imageForRedistribution,
282 removeScarletData=False,
283 updateFluxColumns=True,
284 )
286 # Check that the number of deblended children is consistent
287 parents = catalog[catalog["parent"] == 0]
288 self.assertEqual(
289 np.sum(catalog["deblend_nChild"]), len(catalog) - len(parents)
290 )
292 for parent in parents:
293 children = catalog[catalog["parent"] == parent.get("id")]
294 # Check that nChild is set correctly
295 self.assertEqual(len(children), parent.get("deblend_nChild"))
296 for parentCol, childCol in config.columnInheritance.items():
297 np.testing.assert_array_equal(
298 parent.get(parentCol), children[childCol]
299 )
301 # Extract the parent blend data
302 parentBlendData = modelData.blends[parent.getId()]
303 parentFootprint = parent.getFootprint()
304 x0, y0 = parentFootprint.getBBox().getMin()
305 width, height = parentFootprint.getBBox().getDimensions()
306 yx0 = (y0, x0)
308 for child in children:
309 fp = child.getFootprint()
310 img = fp.extractImage(fill=0.0)
311 # Check that the flux at the center is correct.
312 # Note: this only works in this test image because the
313 # detected peak is in the same location as the
314 # scarlet peak.
315 # If the peak is shifted,
316 # the flux value will be correct but
317 # deblend_peak_center is not the correct location.
318 px = child.get("deblend_peak_center_x")
319 py = child.get("deblend_peak_center_y")
320 flux = img[Point2I(px, py)]
321 self.assertEqual(flux, child.get("deblend_peak_instFlux"))
323 self.assertEqual(child.get("deblend_nPeaks"), len(fp.peaks))
325 # Check that the peak positions match the catalog entry
326 peaks = fp.getPeaks()
327 self.assertEqual(px, peaks[0].getIx())
328 self.assertEqual(py, peaks[0].getIy())
330 # Load the data to check against the HeavyFootprint
331 blendData = parentBlendData.children[child["deblend_blendId"]]
332 # We need to set an observation in order to convolve
333 # the model.
334 modelBox = scl.Box((height, width), origin=(y0, x0))
335 observation = scl.Observation.empty(
336 bands=("dummy",),
337 psfs=observedPsf[bandIndex][None, :, :],
338 model_psf=modelPsf[None, :, :],
339 bbox=modelBox,
340 dtype=np.float32,
341 )
342 blend = mes.io.monochromaticDataToScarlet(
343 blendData=blendData,
344 bandIndex=bandIndex,
345 observation=observation,
346 )
348 # Get the scarlet model for the source
349 source = next(
350 src for src in blend.sources if src.record_id == child.getId()
351 )
352 self.assertEqual(source.center[1], px)
353 self.assertEqual(source.center[0], py)
355 if useFlux:
356 assert imageForRedistribution is not None
357 # Get the flux re-weighted model and test against
358 # the HeavyFootprint.
359 # The HeavyFootprint needs to be projected onto
360 # the image of the flux-redistributed model,
361 # since the HeavyFootprint
362 # may trim rows or columns.
363 _images = imageForRedistribution[
364 parentFootprint.getBBox()
365 ].image.array
366 blend.observation.images = scl.Image(
367 _images[None, :, :],
368 yx0=yx0,
369 bands=("dummy",),
370 )
371 blend.observation.weights = scl.Image(
372 parentFootprint.spans.asArray()[None, :, :],
373 yx0=yx0,
374 bands=("dummy",),
375 )
376 blend.conserve_flux()
377 model = source.flux_weighted_image.data[0]
378 my0, mx0 = source.flux_weighted_image.yx0
379 image = afwImage.ImageF(model, xy0=Point2I(mx0, my0))
380 fp.insert(image)
381 np.testing.assert_almost_equal(image.array, model)
382 else:
383 # Get the model for the source and test
384 # against the HeavyFootprint
385 bbox = fp.getBBox()
386 bbox = mes.utils.bboxToScarletBox(bbox)
387 model = blend.observation.convolve(
388 source.get_model().project(bbox=bbox), mode="real"
389 ).data[0]
390 np.testing.assert_almost_equal(img.array, model)
392 # Check that all sources have the correct number of peaks
393 for src in catalog:
394 fp = src.getFootprint()
395 self.assertEqual(len(fp.peaks), src.get("deblend_nPeaks"))
397 # Check that the catalog matches the expected results
398 nModels = len(self.models)
399 self.assertEqual(len(catalog), nParents+nModels)
401 def test_skipped(self):
402 # Use tight configs to force skipping a 3 source footprint
403 # and "large" footprint
404 config = ScarletDeblendTask.ConfigClass()
405 config.maxFootprintArea = 1000
406 config.maxNumberOfPeaks = 2
407 config.catchFailures = False
409 data = self.initialize_data(self.models, deblendConfig=config)
410 mDeconvolved = self.deconvolve(data)
411 result = data.deblendTask.run(data.mCoadd, mDeconvolved, data.catalog)
413 catalog = result.deblendedCatalog
414 parents = catalog[catalog["parent"] == 0]
415 self.assertEqual(np.sum(parents["deblend_skipped"]), 2)
416 self.assertEqual(np.sum(parents["deblend_parentTooBig"]), 1)
417 self.assertEqual(np.sum(parents["deblend_tooManyPeaks"]), 1)
419 def test_persistence(self):
420 # Test that the model data is persisted correctly
421 data = self.initialize_data(self.models)
422 repo = self._setup_butler()
423 mDeconvolved = self.deconvolve(data)
424 result = data.deblendTask.run(data.mCoadd, mDeconvolved, data.catalog)
425 modelData = result.scarletModelData
426 bands = modelData.metadata["bands"]
427 butler = makeTestCollection(repo, uniqueId="test_run1")
428 butler.put(modelData, "scarlet_model_data", dataId={})
429 modelData2 = butler.get("scarlet_model_data", dataId={})
430 model_psf = modelData.metadata["model_psf"][None, :, :]
431 model_psf2 = modelData2.metadata["model_psf"][None, :, :]
432 np.testing.assert_almost_equal(model_psf2, model_psf)
433 psf = modelData.metadata["psf"]
434 psf2 = modelData2.metadata["psf"]
435 np.testing.assert_almost_equal(psf2, psf)
436 self.assertEqual(len(modelData2.blends), len(modelData.blends))
438 for parentId in modelData.blends.keys():
439 nChildren = len(modelData.blends[parentId].children)
440 self.assertEqual(nChildren, len(modelData2.blends[parentId].children))
441 for blendId in modelData.blends[parentId].children:
442 blendData1 = modelData.blends[parentId].children[blendId]
443 blendData2 = modelData2.blends[parentId].children[blendId]
444 self._test_blend(blendData1, blendData2, model_psf, psf, bands)
446 # Test extracting a single blend
447 modelData2 = butler.get("scarlet_model_data", dataId={}, parameters={"blend_id": parentId})
448 self.assertEqual(len(modelData2.blends), 1)
450 for blendId, blendData1 in modelData.blends[parentId].children.items():
451 blendData2 = modelData2.blends[parentId].children[blendId]
452 self._test_blend(blendData1, blendData2, model_psf, psf, bands)
454 # Test extracting two blends
455 modelData2 = butler.get("scarlet_model_data", dataId={}, parameters={"blend_id": [1, 2]})
456 self.assertEqual(len(modelData2.blends), 2)
457 for parentId in [1, 2]:
458 parentData1 = modelData.blends[parentId]
459 parentData2 = modelData2.blends[parentId]
460 self.assertEqual(len(parentData1.children), len(parentData2.children))
461 for blendId in parentData1.children.keys():
462 blendData1 = parentData1.children[blendId]
463 blendData2 = parentData2.children[blendId]
464 self._test_blend(blendData1, blendData2, model_psf, psf, bands)
466 def test_legacy_model(self):
467 repo = self._setup_butler()
468 storageClass = StorageClass(
469 "ScarletModelData",
470 pytype=scl.io.ScarletModelData,
471 )
472 datasetType = DatasetType(
473 "old_scarlet_model_data",
474 dimensions=(),
475 storageClass=storageClass,
476 universe=repo.dimensions,
477 )
478 ref = DatasetRef(
479 datasetType,
480 run="test_ingestion",
481 dataId={},
482 )
483 dataset = FileDataset(
484 path=os.path.join(TESTDIR, "data", "v29_models.json"),
485 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
486 refs=[ref],
487 )
489 # Ingest the legacy model into the butler
490 butler = makeTestCollection(repo, uniqueId="ingestion")
491 repo.registry.registerDatasetType(datasetType)
492 butler.ingest(dataset)
494 model = butler.get("old_scarlet_model_data", dataId={})
495 self.assertEqual(len(model.blends), 2)
497 test = butler.get("old_scarlet_model_data", dataId={}, parameters={"blend_id": 3495976385350991873})
498 self.assertEqual(len(test.blends), 1)
500 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands):
501 # Test that two ScarletBlendData objects are equal
502 # up to machine precision.
503 self.assertTupleEqual(blendData1.origin, blendData2.origin)
504 self.assertEqual(len(blendData1.sources), len(blendData2.sources))
506 # Test that the two blends are equal up to machine precision
507 # once converted into scarlet lite Blend objects.
508 blend1 = blendData1.minimal_data_to_blend(
509 model_psf,
510 psf,
511 bands,
512 dtype=np.float32,
513 )
514 blend2 = blendData2.minimal_data_to_blend(
515 model_psf,
516 psf,
517 bands,
518 dtype=np.float32,
519 )
520 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data)
522 def _setup_butler(self):
523 # Initialize a Butler to test persistence
524 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
525 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir)
526 config = Config()
527 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
528 repo = makeTestRepo(repo_dir.name, config=config)
529 storageClass = StorageClass(
530 "ScarletModelData",
531 pytype=scl.io.ScarletModelData,
532 parameters=('blend_id',),
533 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate",
534 )
535 datasetType = DatasetType(
536 "scarlet_model_data",
537 dimensions=(),
538 storageClass=storageClass,
539 universe=repo.dimensions,
540 )
541 repo.registry.registerDatasetType(datasetType)
542 return repo
545class BadPsf(Psf):
546 def __init__(self, validPoint: Point2D, psf: GaussianPsf):
547 self.validPoint = validPoint
548 self.psf = psf
549 super().__init__()
551 def computeKernelImage(self, location: Point2D):
552 if location == self.validPoint:
553 return self.psf.computeKernelImage(location)
554 raise InvalidPsfError(f"Invalid PSF at location {location}")
557class TestUtils(lsst.utils.tests.TestCase):
558 def setUp(self):
559 self.bands = tuple("gri")
561 def test_box_transforms(self):
562 # Test scarlet Box to BBox
563 box = scl.Box((25, 30), (17, 5))
564 bbox = mes.utils.scarletBoxToBBox(box)
565 x0, y0 = bbox.getMin()
566 width = bbox.getWidth()
567 height = bbox.getHeight()
568 self.assertTupleEqual((y0, x0), box.origin)
569 self.assertTupleEqual((height, width), box.shape)
571 # Test back to scarlet Box
572 newBox = mes.utils.bboxToScarletBox(bbox)
573 self.assertTupleEqual(newBox.origin, box.origin)
574 self.assertTupleEqual(newBox.shape, box.shape)
576 def test_computeNearestPsfGood(self):
577 # Test that using a valid PSF works normally
578 psf, psfImage = self._generateGoodPsf()
579 coadd = self._generateCoadd(psf)
581 # Test that computing the PSF works
582 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, None, "g", Point2D(25, 25))
583 np.testing.assert_array_equal(derivedPsf.array, psfImage)
584 self.assertEqual(center, Point2D(25, 25))
585 self.assertEqual(dist, 0)
587 def test_computeNearestPsfRecoverable(self):
588 # Test that using a PSF not defined at the initial location
589 # will fallback to a valid location.
590 psf, psfImage = self._generateGoodPsf()
591 coadd = self._generateCoadd(BadPsf(Point2D(1, 1), psf))
592 catalog = self._generateCatalog(self.bands, [[(1, 1, 10)]])
594 # Test that computing the PSF works after finding a new location.
595 # Since the PSF above is *only* defined at (1, 1) it will fail to
596 # compute a PSF image at (4, 5) but should fall back to (1, 1).
597 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, catalog, None, Point2D(4, 5))
598 np.testing.assert_array_equal(derivedPsf.array, psfImage)
599 self.assertEqual(center, Point2I(1, 1))
600 self.assertEqual(dist, 5)
602 def test_computeNearestPsfBad(self):
603 # Test that a PSF that cannot find a matching location returns None
604 psf = self._generateGoodPsf()
605 coadd = self._generateCoadd(BadPsf(Point2D(1, 1), psf))
606 catalog = self._generateCatalog(self.bands)
608 # Test that computing the PSF cannot generate a PSF
609 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, catalog, None, Point2D(4, 5))
610 self.assertIsNone(derivedPsf)
611 self.assertIsNone(center)
612 self.assertIsNone(dist)
614 def test_computeNearestPsfMultiBandGood(self):
615 # Test that a valid PSF in every band works normally
616 bands = tuple("gri")
617 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4])
618 mCoadd = self._generateMultibandCoadd(psfs, bands)
620 # Test that computing the PSF works
621 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), None)
622 np.testing.assert_array_equal(psfArray, psfImage)
623 self.assertTupleEqual(newCoadd.bands, bands)
625 def test_computeNearestPsfMultiBandRecoverable(self):
626 # Test that a Psf at a different location is still recoverable
627 bands = tuple("gri")
628 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4])
629 psfs[1] = BadPsf(Point2D(1, 1), psfs[1])
630 mCoadd = self._generateMultibandCoadd(psfs, bands)
631 catalog = self._generateCatalog(self.bands, [[(1, 1, 10)]])
633 # Test that computing the PSF works because the catalog has a peak
634 # at the location of the BadPsf.
635 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), catalog)
636 np.testing.assert_array_equal(psfArray, psfImage)
637 self.assertTupleEqual(newCoadd.bands, bands)
639 def test_computeNearestPsfMultiBandIncomplete(self):
640 # Test that missing a PSF in one band returns a PSF and
641 # an exposure that are missing bands.
642 bands = tuple("gri")
643 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4])
644 psfs[1] = BadPsf(Point2D(1, 1), psfs[1])
645 mCoadd = self._generateMultibandCoadd(psfs, bands)
646 catalog = self._generateCatalog(self.bands)
648 # Test that computing the PSF works for the g- and i-band PSFs that
649 # are not BadPsf.
650 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), catalog)
651 np.testing.assert_array_equal(psfArray, np.delete(psfImage, 1, axis=0))
652 self.assertTupleEqual(newCoadd.bands, tuple("gi"))
654 def test_computeNearestPsfMultiBandBad(self):
655 # Test that None is returned if none of the PSFs can be computed
656 bands = tuple("gri")
657 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4])
658 psfs = [BadPsf(Point2D(1, 1), psfs) for psf in psfs]
659 mCoadd = self._generateMultibandCoadd(psfs, bands)
660 catalog = self._generateCatalog(self.bands)
662 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), catalog)
663 self.assertIsNone(psfArray)
664 self.assertIsNone(newCoadd)
666 def test_buildObservationBadPsfs(self):
667 # Test that creating an observation with all bad PSFs
668 # raises NoWorkFound
669 modelPsf = scl.utils.integrated_circular_gaussian(sigma=0.8).astype(np.float32)
670 bands = tuple("gri")
671 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4])
672 psfs = [BadPsf(Point2D(1, 1), psf) for psf in psfs]
673 mCoadd = self._generateMultibandCoadd(psfs, bands)
674 catalog = self._generateCatalog(self.bands)
676 # Test that building the observation fails without a catalog
677 with self.assertRaises(NoWorkFound):
678 mes.utils.buildObservation(modelPsf, Point2I(25, 25), mCoadd)
680 # Test that building the observation fails even with a catalog
681 with self.assertRaises(NoWorkFound):
682 mes.utils.buildObservation(modelPsf, Point2I(25, 25), mCoadd, catalog=catalog)
684 def _generateGoodPsf(self, sigma: float = 1.0):
685 # Generate a PSF and Image of the PSF
686 psfRadius = 20
687 psfShape = (2 * psfRadius + 1, 2 * psfRadius + 1)
688 psf = GaussianPsf(psfShape[1], psfShape[0], sigma)
689 psfImage = psf.computeImage(psf.getAveragePosition()).array
690 return psf, psfImage
692 def _generateMultibandPsf(self, sigmas: list[float]):
693 # Generate a multiband PSF with a BadPsf for each None value in sigmas
694 psfs = []
695 psfImages = []
696 for sigma in sigmas:
697 psf, psfImage = self._generateGoodPsf(sigma)
698 psfs.append(psf)
699 psfImages.append(psfImage)
700 return psfs, np.asarray(psfImages)
702 def _generateCoadd(self, psf: Psf):
703 # Create an empty exposure
704 masked_image = afwImage.MaskedImage(Extent2I(50, 50), dtype=np.float32)
705 coadd = afwImage.Exposure(masked_image, dtype=np.float32)
706 coadd.setPsf(psf)
707 return coadd
709 def _generateMultibandCoadd(self, psfs: Psf, bands: list[str]):
710 # Create an empty multi-band exposure
711 coadds = []
712 for psf in psfs:
713 coadds.append(self._generateCoadd(psf))
714 return afwImage.MultibandExposure.fromExposures(bands, coadds)
716 def _generateCatalog(self, bands, footprints: list[list[tuple[int, int, int]]] | None = None):
717 # Generate a catalog with a source for each footprint
718 if footprints is None:
719 footprints = []
720 schema = SourceTable.makeMinimalSchema()
721 peakSchema = PeakTable.makeMinimalSchema()
722 for band in bands:
723 schema.addField(f"merge_footprint_{band}", type="Flag")
724 peakSchema.addField(f"merge_peak_{band}", type="Flag")
726 table = SourceTable.make(schema)
727 catalog = SourceCatalog(table)
729 for peaks in footprints:
730 src = catalog.addNew()
731 footprint = Footprint(SpanSet(), peakSchema)
732 for peak in peaks:
733 footprint.addPeak(*peak)
734 src.setFootprint(footprint)
736 for band in bands:
737 src[f"merge_footprint_{band}"] = True
738 footprint.peaks[f"merge_peak_{band}"] = True
739 return catalog
742class MemoryTester(lsst.utils.tests.MemoryTestCase):
743 pass
746def setup_module(module):
747 lsst.utils.tests.init()
750if __name__ == "__main__": 750 ↛ 751line 750 didn't jump to line 751 because the condition on line 750 was never true
751 lsst.utils.tests.init()
752 unittest.main()