Coverage for tests/test_deblend.py: 12%

392 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-10-11 01:54 -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/>. 

21 

22import os 

23import unittest 

24import tempfile 

25 

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 

42 

43TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

44 

45 

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") 

61 

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 ] 

128 

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 

144 

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) 

175 

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 ) 

186 

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 

194 

195 return result 

196 

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 

209 

210 def test_default_deconvolve(self): 

211 data = self.initialize_data(self.models) 

212 deconvolved = self.deconvolve(data) 

213 

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)) 

224 

225 context = mes.scarletDeblendTask.ScarletDeblendContext.build( 

226 data.mCoadd, 

227 deconvolved, 

228 data.catalog, 

229 data.deblendTask.ConfigClass() 

230 ) 

231 

232 self.assertEqual(len(context.footprints), 4) 

233 

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) 

239 

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)) 

250 

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) 

256 

257 config = data.deblendTask.config 

258 catalog = result.deblendedCatalog 

259 modelData = result.scarletModelData 

260 observedPsf = modelData.metadata["psf"] 

261 modelPsf = modelData.metadata["model_psf"] 

262 

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] 

271 

272 if useFlux: 

273 imageForRedistribution = coadd 

274 else: 

275 imageForRedistribution = None 

276 

277 mes.io.updateCatalogFootprints( 

278 modelData, 

279 catalog, 

280 band=band, 

281 imageForRedistribution=imageForRedistribution, 

282 removeScarletData=False, 

283 updateFluxColumns=True, 

284 ) 

285 

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 ) 

291 

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 ) 

300 

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) 

307 

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")) 

322 

323 self.assertEqual(child.get("deblend_nPeaks"), len(fp.peaks)) 

324 

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()) 

329 

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 ) 

347 

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) 

354 

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) 

391 

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")) 

396 

397 # Check that the catalog matches the expected results 

398 nModels = len(self.models) 

399 self.assertEqual(len(catalog), nParents+nModels) 

400 

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 

408 

409 data = self.initialize_data(self.models, deblendConfig=config) 

410 mDeconvolved = self.deconvolve(data) 

411 result = data.deblendTask.run(data.mCoadd, mDeconvolved, data.catalog) 

412 

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) 

418 

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)) 

437 

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) 

445 

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) 

449 

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) 

453 

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) 

465 

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 ) 

488 

489 # Ingest the legacy model into the butler 

490 butler = makeTestCollection(repo, uniqueId="ingestion") 

491 repo.registry.registerDatasetType(datasetType) 

492 butler.ingest(dataset) 

493 

494 model = butler.get("old_scarlet_model_data", dataId={}) 

495 self.assertEqual(len(model.blends), 2) 

496 

497 test = butler.get("old_scarlet_model_data", dataId={}, parameters={"blend_id": 3495976385350991873}) 

498 self.assertEqual(len(test.blends), 1) 

499 

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)) 

505 

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) 

521 

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 

543 

544 

545class BadPsf(Psf): 

546 def __init__(self, validPoint: Point2D, psf: GaussianPsf): 

547 self.validPoint = validPoint 

548 self.psf = psf 

549 super().__init__() 

550 

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}") 

555 

556 

557class TestUtils(lsst.utils.tests.TestCase): 

558 def setUp(self): 

559 self.bands = tuple("gri") 

560 

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) 

570 

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) 

575 

576 def test_computeNearestPsfGood(self): 

577 # Test that using a valid PSF works normally 

578 psf, psfImage = self._generateGoodPsf() 

579 coadd = self._generateCoadd(psf) 

580 

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) 

586 

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)]]) 

593 

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) 

601 

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) 

607 

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) 

613 

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) 

619 

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) 

624 

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)]]) 

632 

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) 

638 

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) 

647 

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")) 

653 

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) 

661 

662 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), catalog) 

663 self.assertIsNone(psfArray) 

664 self.assertIsNone(newCoadd) 

665 

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) 

675 

676 # Test that building the observation fails without a catalog 

677 with self.assertRaises(NoWorkFound): 

678 mes.utils.buildObservation(modelPsf, Point2I(25, 25), mCoadd) 

679 

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) 

683 

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 

691 

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) 

701 

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 

708 

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) 

715 

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") 

725 

726 table = SourceTable.make(schema) 

727 catalog = SourceCatalog(table) 

728 

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) 

735 

736 for band in bands: 

737 src[f"merge_footprint_{band}"] = True 

738 footprint.peaks[f"merge_peak_{band}"] = True 

739 return catalog 

740 

741 

742class MemoryTester(lsst.utils.tests.MemoryTestCase): 

743 pass 

744 

745 

746def setup_module(module): 

747 lsst.utils.tests.init() 

748 

749 

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()