Coverage for tests/test_deblend.py: 12%

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 

34import lsst.daf.butler 

35from lsst.daf.butler import Butler, Config, DatasetType, StorageClass, FileDataset, DatasetRef 

36from lsst.daf.butler.tests import makeTestRepo, makeTestCollection 

37from lsst.geom import Extent2I, Point2D, Point2I 

38from lsst.meas.algorithms import SourceDetectionTask 

39from lsst.meas.extensions.scarlet.scarletDeblendTask import ScarletDeblendTask 

40from lsst.meas.extensions.scarlet.deconvolveExposureTask import DeconvolveExposureTask 

41from lsst.pipe.base import NoWorkFound, Struct 

42from utils import initData, SersicModel, PsfModel 

43 

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

45 

46 

47class TestDeblend(lsst.utils.tests.TestCase): 

48 def setUp(self): 

49 self.modelPsf = scl.utils.integrated_circular_gaussian(sigma=0.8).astype(np.float32) 

50 psfRadius = 20 

51 psfShape = (2 * psfRadius + 1, 2 * psfRadius + 1) 

52 self.psfs = [ 

53 GaussianPsf(psfShape[1], psfShape[0], 1.0), 

54 GaussianPsf(psfShape[1], psfShape[0], 1.2), 

55 GaussianPsf(psfShape[1], psfShape[0], 1.4), 

56 ] 

57 self.imagePsf = np.asarray( 

58 [psf.computeImage(psf.getAveragePosition()).array for psf in self.psfs] 

59 ).astype(np.float32) 

60 self.imagePsf /= self.imagePsf.sum(axis=(1, 2))[:, None, None] 

61 self.bands = tuple("gri") 

62 

63 self.models = [ 

64 # Isolated source 

65 PsfModel( 

66 center=(30, 15), 

67 spectrum=np.array([8, 2, 1]), 

68 bands=self.bands, 

69 ), 

70 # Two source blend 

71 SersicModel( 

72 center=(40, 20), 

73 major=5, 

74 minor=2, 

75 radius=15, 

76 theta=-np.pi/4, 

77 n=1, 

78 spectrum=np.array([2, 4, 8]), 

79 bands=self.bands, 

80 ), 

81 PsfModel( 

82 center=(12, 20), 

83 spectrum=np.array([1, 2, 8]), 

84 bands=self.bands, 

85 ), 

86 # 3 source blend 

87 SersicModel( 

88 center=(25, 70), 

89 major=5, 

90 minor=2, 

91 radius=20, 

92 theta=np.pi/48, 

93 n=1, 

94 spectrum=np.array([2, 5, 8]), 

95 bands=self.bands, 

96 ), 

97 PsfModel( 

98 center=(32, 60), 

99 spectrum=np.array([1, 2, 8]), 

100 bands=self.bands, 

101 ), 

102 PsfModel( 

103 center=(16, 80), 

104 spectrum=np.array([8, 2, 1]), 

105 bands=self.bands, 

106 ), 

107 # Large blend 

108 SersicModel( 

109 center=(70, 70), 

110 major=5, 

111 minor=2, 

112 radius=25, 

113 theta=0, 

114 n=1, 

115 spectrum=np.array([2, 10, 18]), 

116 bands=self.bands, 

117 ), 

118 SersicModel( 

119 center=(85, 85), 

120 major=5, 

121 minor=2, 

122 radius=25, 

123 theta=np.pi/2, 

124 n=1, 

125 spectrum=np.array([5, 10, 20]), 

126 bands=self.bands, 

127 ), 

128 ] 

129 

130 def scarlet_image_to_exposure( 

131 self, 

132 image: scl.Image, 

133 noise: np.ndarray, 

134 ) -> afwImage.MultibandExposure: 

135 masked_image = afwImage.MultibandMaskedImage.fromArrays( 

136 image.bands, image.data, None, noise**2 

137 ) 

138 coadds = [ 

139 afwImage.Exposure(img, dtype=img.image.array.dtype) for img in masked_image 

140 ] 

141 mCoadd = afwImage.MultibandExposure.fromExposures(image.bands, coadds) 

142 for b, coadd in enumerate(mCoadd): 

143 coadd.setPsf(self.psfs[b]) 

144 return mCoadd 

145 

146 def initialize_data( 

147 self, 

148 models, 

149 deconvolveConfig=None, 

150 deblendConfig=None, 

151 doDetect: bool = True, 

152 ): 

153 if deconvolveConfig is None: 

154 deconvolveConfig = DeconvolveExposureTask.ConfigClass() 

155 if deblendConfig is None: 

156 deblendConfig = ScarletDeblendTask.ConfigClass() 

157 # Generate the data for the test 

158 deconvolved, convolved = initData(models, self.modelPsf, self.imagePsf) 

159 # Set the random seed so that the noise field is unaffected 

160 # and add noise to the image 

161 np.random.seed(0) 

162 noise = 0.05 * (np.random.rand(*convolved.shape).astype(np.float32) - 0.5) 

163 noisyImage = convolved.copy() 

164 noisyImage._data += noise 

165 # Create the multiband coadd 

166 mCoadd = self.scarlet_image_to_exposure(noisyImage, noise) 

167 # Initialze tasks 

168 inputSchema = SourceTable.makeMinimalSchema() 

169 table = SourceTable.make(inputSchema) 

170 detectionTask = SourceDetectionTask(schema=inputSchema) 

171 schemaMapper = SchemaMapper(inputSchema) 

172 schemaMapper.addMinimalSchema(inputSchema) 

173 schema = schemaMapper.getOutputSchema() 

174 deconvolveTask = DeconvolveExposureTask(config=deconvolveConfig) 

175 deblendTask = ScarletDeblendTask(schema=schema, config=deblendConfig) 

176 

177 result = Struct( 

178 deconvolved=deconvolved, 

179 convolved=convolved, 

180 noise=noise, 

181 noisyImage=noisyImage, 

182 mCoadd=mCoadd, 

183 detectionTask=detectionTask, 

184 deconvolveTask=deconvolveTask, 

185 deblendTask=deblendTask, 

186 ) 

187 

188 if doDetect: 

189 # Generate a detection catalog 

190 detectionResult = detectionTask.run(table, mCoadd["r"]) 

191 table = SourceCatalog.Table.make(schema) 

192 catalog = SourceCatalog(table) 

193 catalog.extend(detectionResult.sources, schemaMapper) 

194 result.catalog = catalog 

195 

196 return result 

197 

198 def deconvolve(self, data: Struct): 

199 deconvolvedCoadds = [] 

200 deconvolveTask = data.deconvolveTask 

201 if deconvolveTask.config.useFootprints: 

202 catalog = data.catalog 

203 else: 

204 catalog = None 

205 for coadd in data.mCoadd: 

206 deconvolvedCoadd = deconvolveTask.run(coadd, catalog).deconvolved 

207 deconvolvedCoadds.append(deconvolvedCoadd) 

208 mDeconvolved = afwImage.MultibandExposure.fromExposures(self.bands, deconvolvedCoadds) 

209 return mDeconvolved 

210 

211 def test_default_deconvolve(self): 

212 data = self.initialize_data(self.models) 

213 deconvolved = self.deconvolve(data) 

214 

215 diff = data.deconvolved.data - deconvolved.image.array 

216 # Due to peakiness of Sersic models the center has a sharp peak, 

217 # so we ignore a 3x3 region around each source center 

218 for model in self.models: 

219 yc, xc = model.center 

220 for x in (-1, 0, 1): 

221 for y in (-1, 0, 1): 

222 diff[:, yc+y, xc+x] = 0 

223 self.assertTrue(np.max(diff[:2]) < 10*np.std(data.noise)) 

224 self.assertTrue(np.max(diff[2]) < 20*np.std(data.noise)) 

225 

226 context = mes.scarletDeblendTask.ScarletDeblendContext.build( 

227 data.mCoadd, 

228 deconvolved, 

229 data.catalog, 

230 data.deblendTask.ConfigClass() 

231 ) 

232 

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

234 

235 def test_catalog_free_deconvolve(self): 

236 config = DeconvolveExposureTask.ConfigClass() 

237 config.useFootprints = False 

238 data = self.initialize_data(self.models, deconvolveConfig=config) 

239 deconvolved = self.deconvolve(data) 

240 

241 diff = data.deconvolved.data - deconvolved.image.array 

242 # Due to peakiness of Sersic models the center has a sharp peak, 

243 # so we ignore a 3x3 region around each source center 

244 for model in self.models: 

245 yc, xc = model.center 

246 for x in (-1, 0, 1): 

247 for y in (-1, 0, 1): 

248 diff[:, yc+y, xc+x] = 0 

249 self.assertTrue(np.max(diff[:2]) < 10*np.std(data.noise)) 

250 self.assertTrue(np.max(diff[2]) < 20*np.std(data.noise)) 

251 

252 def test_footprints(self): 

253 data = self.initialize_data(self.models) 

254 nParents = len(data.catalog) 

255 mDeconvolved = self.deconvolve(data) 

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

257 

258 config = data.deblendTask.config 

259 catalog = result.deblendedCatalog 

260 modelData = result.scarletModelData 

261 observedPsf = modelData.metadata["psf"] 

262 modelPsf = modelData.metadata["model_psf"] 

263 

264 # Check that isolated sources are handled correctly 

265 isolated = catalog[(catalog["parent"] == 0) & (catalog["deblend_nPeaks"] == 1)] 

266 self.assertEqual(len(isolated), len(modelData.isolated)) 

267 for sid, source in modelData.isolated.items(): 

268 catalog_footprint = catalog.find(sid).getFootprint() 

269 isolated_array = catalog_footprint.spans.asArray() 

270 np.testing.assert_array_equal(source.span_array, isolated_array) 

271 

272 # Check that the origin is correct 

273 self.assertTupleEqual(source.origin[::-1], tuple(catalog_footprint.getBBox().getMin())) 

274 

275 # Attach the footprints in each band and compare to the full 

276 # data model. This is done in each band, both with and without 

277 # flux re-distribution to test all of the different possible 

278 # options of loading catalog footprints. 

279 for useFlux in [False, True]: 

280 for band in self.bands: 

281 bandIndex = self.bands.index(band) 

282 coadd = data.mCoadd[band] 

283 

284 if useFlux: 

285 imageForRedistribution = coadd 

286 else: 

287 imageForRedistribution = None 

288 

289 mes.io.updateCatalogFootprints( 

290 modelData, 

291 catalog, 

292 band=band, 

293 imageForRedistribution=imageForRedistribution, 

294 removeScarletData=False, 

295 updateFluxColumns=True, 

296 ) 

297 

298 # Check that the number of deblended children is consistent 

299 parents = catalog[(catalog["parent"] == 0) & (catalog["deblend_nPeaks"] > 1)] 

300 self.assertEqual( 

301 np.sum(catalog["deblend_nChild"]), len(catalog) - len(parents) - len(isolated) 

302 ) 

303 

304 for parent in parents: 

305 children = catalog[catalog["parent"] == parent.get("id")] 

306 # Check that nChild is set correctly 

307 self.assertEqual(len(children), parent.get("deblend_nChild")) 

308 for parentCol, childCol in config.columnInheritance.items(): 

309 np.testing.assert_array_equal( 

310 parent.get(parentCol), children[childCol] 

311 ) 

312 

313 # Extract the parent blend data 

314 parentBlendData = modelData.blends[parent.getId()] 

315 parentFootprint = parent.getFootprint() 

316 x0, y0 = parentFootprint.getBBox().getMin() 

317 width, height = parentFootprint.getBBox().getDimensions() 

318 yx0 = (y0, x0) 

319 

320 for child in children: 

321 fp = child.getFootprint() 

322 img = fp.extractImage(fill=0.0) 

323 # Check that the flux at the center is correct. 

324 # Note: this only works in this test image because the 

325 # detected peak is in the same location as the 

326 # scarlet peak. 

327 # If the peak is shifted, 

328 # the flux value will be correct but 

329 # deblend_peak_center is not the correct location. 

330 px = child.get("deblend_peak_center_x") 

331 py = child.get("deblend_peak_center_y") 

332 flux = img[Point2I(px, py)] 

333 self.assertEqual(flux, child.get("deblend_peak_instFlux")) 

334 

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

336 

337 # Check that the peak positions match the catalog entry 

338 peaks = fp.getPeaks() 

339 self.assertEqual(px, peaks[0].getIx()) 

340 self.assertEqual(py, peaks[0].getIy()) 

341 

342 # Load the data to check against the HeavyFootprint 

343 blendData = parentBlendData.children[child["deblend_blendId"]] 

344 # We need to set an observation in order to convolve 

345 # the model. 

346 modelBox = scl.Box((height, width), origin=(y0, x0)) 

347 observation = scl.Observation.empty( 

348 bands=("dummy",), 

349 psfs=observedPsf[bandIndex][None, :, :], 

350 model_psf=modelPsf[None, :, :], 

351 bbox=modelBox, 

352 dtype=np.float32, 

353 ) 

354 blend = mes.io.monochromaticDataToScarlet( 

355 blendData=blendData, 

356 bandIndex=bandIndex, 

357 observation=observation, 

358 ) 

359 

360 # Get the scarlet model for the source 

361 source = next( 

362 src for src in blend.sources if src.metadata["id"] == child.getId() 

363 ) 

364 self.assertEqual(source.center[1], px) 

365 self.assertEqual(source.center[0], py) 

366 

367 if useFlux: 

368 assert imageForRedistribution is not None 

369 # Get the flux re-weighted model and test against 

370 # the HeavyFootprint. 

371 # The HeavyFootprint needs to be projected onto 

372 # the image of the flux-redistributed model, 

373 # since the HeavyFootprint 

374 # may trim rows or columns. 

375 _images = imageForRedistribution[ 

376 parentFootprint.getBBox() 

377 ].image.array 

378 blend.observation.images = scl.Image( 

379 _images[None, :, :], 

380 yx0=yx0, 

381 bands=("dummy",), 

382 ) 

383 blend.observation.weights = scl.Image( 

384 parentFootprint.spans.asArray()[None, :, :], 

385 yx0=yx0, 

386 bands=("dummy",), 

387 ) 

388 blend.conserve_flux() 

389 model = source.flux_weighted_image.data[0] 

390 my0, mx0 = source.flux_weighted_image.yx0 

391 image = afwImage.ImageF(model, xy0=Point2I(mx0, my0)) 

392 fp.insert(image) 

393 np.testing.assert_almost_equal(image.array, model) 

394 else: 

395 # Get the model for the source and test 

396 # against the HeavyFootprint 

397 bbox = fp.getBBox() 

398 bbox = mes.utils.bboxToScarletBox(bbox) 

399 model = blend.observation.convolve( 

400 source.get_model().project(bbox=bbox), mode="real" 

401 ).data[0] 

402 np.testing.assert_almost_equal(img.array, model) 

403 

404 # Check that all sources have the correct number of peaks 

405 for src in catalog: 

406 fp = src.getFootprint() 

407 self.assertEqual(len(fp.peaks), src.get("deblend_nPeaks")) 

408 

409 # Check that the catalog matches the expected results 

410 nModels = len(self.models) 

411 self.assertEqual(len(catalog), nParents+nModels-len(isolated)) 

412 

413 def test_skipped(self): 

414 # Use tight configs to force skipping a 3 source footprint 

415 # and "large" footprint 

416 config = ScarletDeblendTask.ConfigClass() 

417 config.maxFootprintArea = 2000 

418 config.maxNumberOfPeaks = 2 

419 config.catchFailures = False 

420 

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

422 mDeconvolved = self.deconvolve(data) 

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

424 

425 catalog = result.deblendedCatalog 

426 parents = catalog[catalog["parent"] == 0] 

427 self.assertEqual(np.sum(parents["deblend_skipped"]), 2) 

428 self.assertEqual(np.sum(parents["deblend_parentTooBig"]), 1) 

429 self.assertEqual(np.sum(parents["deblend_tooManyPeaks"]), 1) 

430 

431 def test_persistence(self): 

432 # Test that the model data is persisted correctly 

433 data = self.initialize_data(self.models) 

434 repo = self._setup_butler() 

435 mDeconvolved = self.deconvolve(data) 

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

437 modelData = result.scarletModelData 

438 bands = modelData.metadata["bands"] 

439 butler = makeTestCollection(repo, uniqueId="test_run1") 

440 butler.put(modelData, "scarlet_model_data", dataId={}) 

441 modelData2 = butler.get("scarlet_model_data", dataId={}) 

442 model_psf = modelData.metadata["model_psf"][None, :, :] 

443 model_psf2 = modelData2.metadata["model_psf"][None, :, :] 

444 np.testing.assert_almost_equal(model_psf2, model_psf) 

445 psf = modelData.metadata["psf"] 

446 psf2 = modelData2.metadata["psf"] 

447 np.testing.assert_almost_equal(psf2, psf) 

448 self.assertEqual(len(modelData2.blends), len(modelData.blends)) 

449 

450 for parentId in modelData.blends.keys(): 

451 nChildren = len(modelData.blends[parentId].children) 

452 self.assertEqual(nChildren, len(modelData2.blends[parentId].children)) 

453 for blendId in modelData.blends[parentId].children: 

454 blendData1 = modelData.blends[parentId].children[blendId] 

455 blendData2 = modelData2.blends[parentId].children[blendId] 

456 self._test_blend(blendData1, blendData2, model_psf, psf, bands) 

457 

458 # Test extracting a single blend 

459 modelData2 = butler.get("scarlet_model_data", dataId={}, parameters={"blend_id": parentId}) 

460 self.assertEqual(len(modelData2.blends), 1) 

461 

462 for blendId, blendData1 in modelData.blends[parentId].children.items(): 

463 blendData2 = modelData2.blends[parentId].children[blendId] 

464 self._test_blend(blendData1, blendData2, model_psf, psf, bands) 

465 

466 # Test extracting two blends 

467 modelData2 = butler.get("scarlet_model_data", dataId={}, parameters={"blend_id": [2, 3]}) 

468 self.assertEqual(len(modelData2.blends), 2) 

469 for parentId in [2, 3]: 

470 parentData1 = modelData.blends[parentId] 

471 parentData2 = modelData2.blends[parentId] 

472 self.assertEqual(len(parentData1.children), len(parentData2.children)) 

473 for blendId in parentData1.children.keys(): 

474 blendData1 = parentData1.children[blendId] 

475 blendData2 = parentData2.children[blendId] 

476 self._test_blend(blendData1, blendData2, model_psf, psf, bands) 

477 

478 def test_legacy_model(self): 

479 repo = self._setup_butler() 

480 storageClass = StorageClass( 

481 "LsstScarletModelData", 

482 pytype=mes.io.LsstScarletModelData, 

483 ) 

484 datasetType = DatasetType( 

485 "old_scarlet_model_data", 

486 dimensions=(), 

487 storageClass=storageClass, 

488 universe=repo.dimensions, 

489 ) 

490 ref = DatasetRef( 

491 datasetType, 

492 run="test_ingestion", 

493 dataId={}, 

494 ) 

495 dataset = FileDataset( 

496 path=os.path.join(TESTDIR, "data", "v29_models.json"), 

497 formatter="lsst.daf.butler.formatters.json.JsonFormatter", 

498 refs=[ref], 

499 ) 

500 

501 # Ingest the legacy model into the butler 

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

503 repo.registry.registerDatasetType(datasetType) 

504 butler.ingest(dataset) 

505 

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

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

508 

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

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

511 

512 def test_older_legacy_model(self): 

513 repo = self._setup_butler() 

514 oldStorageClass = StorageClass( 

515 "ScarletModelData", 

516 pytype=lsst.scarlet.lite.io.ScarletModelData, 

517 ) 

518 oldDatasetType = DatasetType( 

519 "old_scarlet_model_data", 

520 dimensions=(), 

521 storageClass=oldStorageClass, 

522 universe=repo.dimensions, 

523 ) 

524 ref = DatasetRef( 

525 oldDatasetType, 

526 run="test_ingestion", 

527 dataId={}, 

528 ) 

529 dataset = FileDataset( 

530 path=os.path.join(TESTDIR, "data", "v29_models.json"), 

531 formatter="lsst.daf.butler.formatters.json.JsonFormatter", 

532 refs=[ref], 

533 ) 

534 

535 # Ingest the legacy model into the butler 

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

537 repo.registry.registerDatasetType(oldDatasetType) 

538 butler.ingest(dataset) 

539 

540 # Load the base repo config from the repository 

541 base_config = Config(os.path.join(self.repo_dir, "butler.yaml")) 

542 

543 # Load the storage class override config 

544 override_path = os.path.join( 

545 os.path.dirname(lsst.daf.butler.__file__), 

546 "configs", 

547 "storageClasses.yaml" 

548 ) 

549 override_config = Config(override_path) 

550 

551 # Merge the configs (update base with override) 

552 base_config.update(override_config) 

553 

554 # Create Butler with the merged config 

555 # The config now contains both the repo info and 

556 # the storage class overrides 

557 newButler = Butler.from_config(base_config, collections=butler.collections) 

558 

559 model = newButler.get("old_scarlet_model_data", dataId={}, storageClass="LsstScarletModelData") 

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

561 self.assertEqual(len(model.isolated), 0) 

562 

563 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands): 

564 # Test that two ScarletBlendData objects are equal 

565 # up to machine precision. 

566 self.assertTupleEqual(blendData1.origin, blendData2.origin) 

567 self.assertEqual(len(blendData1.sources), len(blendData2.sources)) 

568 

569 # Test that the two blends are equal up to machine precision 

570 # once converted into scarlet lite Blend objects. 

571 blend1 = blendData1.minimal_data_to_blend( 

572 model_psf, 

573 psf, 

574 bands, 

575 dtype=np.float32, 

576 ) 

577 blend2 = blendData2.minimal_data_to_blend( 

578 model_psf, 

579 psf, 

580 bands, 

581 dtype=np.float32, 

582 ) 

583 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data) 

584 

585 def _setup_butler(self): 

586 # Initialize a Butler to test persistence 

587 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) 

588 self.repo_dir = repo_dir.name 

589 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir) 

590 config = Config() 

591 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore" 

592 repo = makeTestRepo(repo_dir.name, config=config) 

593 storageClass = StorageClass( 

594 "LsstScarletModelData", 

595 pytype=mes.io.LsstScarletModelData, 

596 parameters=('blend_id',), 

597 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate", 

598 ) 

599 datasetType = DatasetType( 

600 "scarlet_model_data", 

601 dimensions=(), 

602 storageClass=storageClass, 

603 universe=repo.dimensions, 

604 ) 

605 repo.registry.registerDatasetType(datasetType) 

606 return repo 

607 

608 

609class BadPsf(Psf): 

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

611 self.validPoint = validPoint 

612 self.psf = psf 

613 super().__init__() 

614 

615 def computeKernelImage(self, location: Point2D): 

616 if location == self.validPoint: 

617 return self.psf.computeKernelImage(location) 

618 raise InvalidPsfError(f"Invalid PSF at location {location}") 

619 

620 

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

622 def setUp(self): 

623 self.bands = tuple("gri") 

624 

625 def test_box_transforms(self): 

626 # Test scarlet Box to BBox 

627 box = scl.Box((25, 30), (17, 5)) 

628 bbox = mes.utils.scarletBoxToBBox(box) 

629 x0, y0 = bbox.getMin() 

630 width = bbox.getWidth() 

631 height = bbox.getHeight() 

632 self.assertTupleEqual((y0, x0), box.origin) 

633 self.assertTupleEqual((height, width), box.shape) 

634 

635 # Test back to scarlet Box 

636 newBox = mes.utils.bboxToScarletBox(bbox) 

637 self.assertTupleEqual(newBox.origin, box.origin) 

638 self.assertTupleEqual(newBox.shape, box.shape) 

639 

640 def test_computeNearestPsfGood(self): 

641 # Test that using a valid PSF works normally 

642 psf, psfImage = self._generateGoodPsf() 

643 coadd = self._generateCoadd(psf) 

644 

645 # Test that computing the PSF works 

646 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, None, "g", Point2D(25, 25)) 

647 np.testing.assert_array_equal(derivedPsf.array, psfImage) 

648 self.assertEqual(center, Point2D(25, 25)) 

649 self.assertEqual(dist, 0) 

650 

651 def test_computeNearestPsfRecoverable(self): 

652 # Test that using a PSF not defined at the initial location 

653 # will fallback to a valid location. 

654 psf, psfImage = self._generateGoodPsf() 

655 coadd = self._generateCoadd(BadPsf(Point2D(1, 1), psf)) 

656 catalog = self._generateCatalog(self.bands, [[(1, 1, 10)]]) 

657 

658 # Test that computing the PSF works after finding a new location. 

659 # Since the PSF above is *only* defined at (1, 1) it will fail to 

660 # compute a PSF image at (4, 5) but should fall back to (1, 1). 

661 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, catalog, None, Point2D(4, 5)) 

662 np.testing.assert_array_equal(derivedPsf.array, psfImage) 

663 self.assertEqual(center, Point2I(1, 1)) 

664 self.assertEqual(dist, 5) 

665 

666 def test_computeNearestPsfBad(self): 

667 # Test that a PSF that cannot find a matching location returns None 

668 psf = self._generateGoodPsf() 

669 coadd = self._generateCoadd(BadPsf(Point2D(1, 1), psf)) 

670 catalog = self._generateCatalog(self.bands) 

671 

672 # Test that computing the PSF cannot generate a PSF 

673 derivedPsf, center, dist = mes.utils.computeNearestPsf(coadd, catalog, None, Point2D(4, 5)) 

674 self.assertIsNone(derivedPsf) 

675 self.assertIsNone(center) 

676 self.assertIsNone(dist) 

677 

678 def test_computeNearestPsfMultiBandGood(self): 

679 # Test that a valid PSF in every band works normally 

680 bands = tuple("gri") 

681 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4]) 

682 mCoadd = self._generateMultibandCoadd(psfs, bands) 

683 

684 # Test that computing the PSF works 

685 psfArray, newCoadd = mes.utils.computeNearestPsfMultiBand(mCoadd, Point2D(25, 25), None) 

686 np.testing.assert_array_equal(psfArray, psfImage) 

687 self.assertTupleEqual(newCoadd.bands, bands) 

688 

689 def test_computeNearestPsfMultiBandRecoverable(self): 

690 # Test that a Psf at a different location is still recoverable 

691 bands = tuple("gri") 

692 psfs, psfImage = self._generateMultibandPsf([1.0, 1.2, 1.4]) 

693 psfs[1] = BadPsf(Point2D(1, 1), psfs[1]) 

694 mCoadd = self._generateMultibandCoadd(psfs, bands) 

695 catalog = self._generateCatalog(self.bands, [[(1, 1, 10)]]) 

696 

697 # Test that computing the PSF works because the catalog has a peak 

698 # at the location of the BadPsf. 

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