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 for sourceId in modelData.isolated.keys(): 

459 isolatedData1 = modelData.isolated[sourceId] 

460 isolatedData2 = modelData2.isolated[sourceId] 

461 self.assertTupleEqual(isolatedData1.origin, isolatedData2.origin) 

462 np.testing.assert_array_equal( 

463 isolatedData1.span_array, 

464 isolatedData2.span_array, 

465 ) 

466 

467 # Test extracting a single blend 

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

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

470 

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

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

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

474 

475 # Test extracting two blends 

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

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

478 for parentId in [2, 3]: 

479 parentData1 = modelData.blends[parentId] 

480 parentData2 = modelData2.blends[parentId] 

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

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

483 blendData1 = parentData1.children[blendId] 

484 blendData2 = parentData2.children[blendId] 

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

486 

487 def test_legacy_model(self): 

488 repo = self._setup_butler() 

489 storageClass = StorageClass( 

490 "LsstScarletModelData", 

491 pytype=mes.io.LsstScarletModelData, 

492 ) 

493 datasetType = DatasetType( 

494 "old_scarlet_model_data", 

495 dimensions=(), 

496 storageClass=storageClass, 

497 universe=repo.dimensions, 

498 ) 

499 ref = DatasetRef( 

500 datasetType, 

501 run="test_ingestion", 

502 dataId={}, 

503 ) 

504 dataset = FileDataset( 

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

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

507 refs=[ref], 

508 ) 

509 

510 # Ingest the legacy model into the butler 

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

512 repo.registry.registerDatasetType(datasetType) 

513 butler.ingest(dataset) 

514 

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

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

517 

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

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

520 

521 def test_older_legacy_model(self): 

522 repo = self._setup_butler() 

523 oldStorageClass = StorageClass( 

524 "ScarletModelData", 

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

526 ) 

527 oldDatasetType = DatasetType( 

528 "old_scarlet_model_data", 

529 dimensions=(), 

530 storageClass=oldStorageClass, 

531 universe=repo.dimensions, 

532 ) 

533 ref = DatasetRef( 

534 oldDatasetType, 

535 run="test_ingestion", 

536 dataId={}, 

537 ) 

538 dataset = FileDataset( 

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

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

541 refs=[ref], 

542 ) 

543 

544 # Ingest the legacy model into the butler 

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

546 repo.registry.registerDatasetType(oldDatasetType) 

547 butler.ingest(dataset) 

548 

549 # Load the base repo config from the repository 

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

551 

552 # Load the storage class override config 

553 override_path = os.path.join( 

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

555 "configs", 

556 "storageClasses.yaml" 

557 ) 

558 override_config = Config(override_path) 

559 

560 # Merge the configs (update base with override) 

561 base_config.update(override_config) 

562 

563 # Create Butler with the merged config 

564 # The config now contains both the repo info and 

565 # the storage class overrides 

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

567 

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

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

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

571 

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

573 # Test that two ScarletBlendData objects are equal 

574 # up to machine precision. 

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

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

577 

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

579 # once converted into scarlet lite Blend objects. 

580 blend1 = blendData1.minimal_data_to_blend( 

581 model_psf, 

582 psf, 

583 bands, 

584 dtype=np.float32, 

585 ) 

586 blend2 = blendData2.minimal_data_to_blend( 

587 model_psf, 

588 psf, 

589 bands, 

590 dtype=np.float32, 

591 ) 

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

593 

594 def _setup_butler(self): 

595 # Initialize a Butler to test persistence 

596 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) 

597 self.repo_dir = repo_dir.name 

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

599 config = Config() 

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

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

602 storageClass = StorageClass( 

603 "LsstScarletModelData", 

604 pytype=mes.io.LsstScarletModelData, 

605 parameters=('blend_id',), 

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

607 ) 

608 datasetType = DatasetType( 

609 "scarlet_model_data", 

610 dimensions=(), 

611 storageClass=storageClass, 

612 universe=repo.dimensions, 

613 ) 

614 repo.registry.registerDatasetType(datasetType) 

615 return repo 

616 

617 

618class BadPsf(Psf): 

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

620 self.validPoint = validPoint 

621 self.psf = psf 

622 super().__init__() 

623 

624 def computeKernelImage(self, location: Point2D): 

625 if location == self.validPoint: 

626 return self.psf.computeKernelImage(location) 

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

628 

629 

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

631 def setUp(self): 

632 self.bands = tuple("gri") 

633 

634 def test_box_transforms(self): 

635 # Test scarlet Box to BBox 

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

637 bbox = mes.utils.scarletBoxToBBox(box) 

638 x0, y0 = bbox.getMin() 

639 width = bbox.getWidth() 

640 height = bbox.getHeight() 

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

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

643 

644 # Test back to scarlet Box 

645 newBox = mes.utils.bboxToScarletBox(bbox) 

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

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

648 

649 def test_computeNearestPsfGood(self): 

650 # Test that using a valid PSF works normally 

651 psf, psfImage = self._generateGoodPsf() 

652 coadd = self._generateCoadd(psf) 

653 

654 # Test that computing the PSF works 

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

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

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

658 self.assertEqual(dist, 0) 

659 

660 def test_computeNearestPsfRecoverable(self): 

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

662 # will fallback to a valid location. 

663 psf, psfImage = self._generateGoodPsf() 

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

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

666 

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

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

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

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

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

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

673 self.assertEqual(dist, 5) 

674 

675 def test_computeNearestPsfBad(self): 

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

677 psf = self._generateGoodPsf() 

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

679 catalog = self._generateCatalog(self.bands) 

680 

681 # Test that computing the PSF cannot generate a PSF 

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

683 self.assertIsNone(derivedPsf) 

684 self.assertIsNone(center) 

685 self.assertIsNone(dist) 

686 

687 def test_computeNearestPsfMultiBandGood(self): 

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

689 bands = tuple("gri") 

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

691 mCoadd = self._generateMultibandCoadd(psfs, bands) 

692 

693 # Test that computing the PSF works 

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

695 np.testing.assert_array_equal(psfArray, psfImage) 

696 self.assertTupleEqual(newCoadd.bands, bands) 

697 

698 def test_computeNearestPsfMultiBandRecoverable(self): 

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

700 bands = tuple("gri") 

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