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 mDeconvolved = self.deconvolve(data) 

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

256 

257 catalog = result.deblendedCatalog 

258 objectParents = result.objectParents 

259 modelData = result.scarletModelData 

260 observedPsf = modelData.metadata["psf"] 

261 modelPsf = modelData.metadata["model_psf"] 

262 

263 # Check that isolated sources are handled correctly 

264 isolated = catalog[(catalog["parent"] == 0)] 

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

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

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

268 isolated_array = catalog_footprint.spans.asArray() 

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

270 

271 # Check that the origin is correct 

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

273 

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

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

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

277 # options of loading catalog footprints. 

278 for useFlux in [False, True]: 

279 for band in self.bands: 

280 bandIndex = self.bands.index(band) 

281 coadd = data.mCoadd[band] 

282 

283 if useFlux: 

284 imageForRedistribution = coadd 

285 else: 

286 imageForRedistribution = None 

287 

288 mes.io.updateCatalogFootprints( 

289 modelData, 

290 catalog, 

291 band=band, 

292 imageForRedistribution=imageForRedistribution, 

293 removeScarletData=False, 

294 updateFluxColumns=True, 

295 ) 

296 

297 # Check that the number of deblended children is consistent 

298 parents = objectParents[ 

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

300 self.assertEqual( 

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

302 ) 

303 

304 for parent in parents: 

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

306 

307 # Extract the parent blend data 

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

309 parentFootprint = parent.getFootprint() 

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

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

312 yx0 = (y0, x0) 

313 

314 for child in children: 

315 fp = child.getFootprint() 

316 img = fp.extractImage(fill=0.0) 

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

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

319 # detected peak is in the same location as the 

320 # scarlet peak. 

321 # If the peak is shifted, 

322 # the flux value will be correct but 

323 # deblend_peak_center is not the correct location. 

324 px = child.get("deblend_peak_center_x") 

325 py = child.get("deblend_peak_center_y") 

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

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

328 

329 # Check that the peak positions match the catalog entry 

330 peaks = fp.getPeaks() 

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

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

333 

334 # Load the data to check against the HeavyFootprint 

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

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

337 # the model. 

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

339 observation = scl.Observation.empty( 

340 bands=("dummy",), 

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

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

343 bbox=modelBox, 

344 dtype=np.float32, 

345 ) 

346 blend = mes.io.monochromaticDataToScarlet( 

347 blendData=blendData, 

348 bandIndex=bandIndex, 

349 observation=observation, 

350 ) 

351 

352 # Get the scarlet model for the source 

353 source = next( 

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

355 ) 

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

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

358 

359 if useFlux: 

360 assert imageForRedistribution is not None 

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

362 # the HeavyFootprint. 

363 # The HeavyFootprint needs to be projected onto 

364 # the image of the flux-redistributed model, 

365 # since the HeavyFootprint 

366 # may trim rows or columns. 

367 _images = imageForRedistribution[ 

368 parentFootprint.getBBox() 

369 ].image.array 

370 blend.observation.images = scl.Image( 

371 _images[None, :, :], 

372 yx0=yx0, 

373 bands=("dummy",), 

374 ) 

375 blend.observation.weights = scl.Image( 

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

377 yx0=yx0, 

378 bands=("dummy",), 

379 ) 

380 blend.conserve_flux() 

381 model = source.flux_weighted_image.data[0] 

382 my0, mx0 = source.flux_weighted_image.yx0 

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

384 fp.insert(image) 

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

386 else: 

387 # Get the model for the source and test 

388 # against the HeavyFootprint 

389 bbox = fp.getBBox() 

390 bbox = mes.utils.bboxToScarletBox(bbox) 

391 model = blend.observation.convolve( 

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

393 ).data[0] 

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

395 

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

397 maxId = np.max(objectParents["id"]) 

398 for src in catalog: 

399 fp = src.getFootprint() 

400 self.assertEqual(len(fp.peaks), 1) 

401 if src["parent"] > 0: 

402 # Check that source IDs are greater than the max parent ID 

403 self.assertGreater(src["id"], maxId) 

404 

405 # Ensure that sources are sorted by parent ID 

406 np.testing.assert_array_equal(sorted(catalog["parent"]), catalog["parent"]) 

407 

408 # Check that the catalog matches the expected results 

409 self.assertEqual(len(catalog), len(self.models)) 

410 

411 def test_skipped(self): 

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

413 # and "large" footprint 

414 config = ScarletDeblendTask.ConfigClass() 

415 config.maxFootprintArea = 2000 

416 config.maxNumberOfPeaks = 2 

417 config.catchFailures = False 

418 

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

420 mDeconvolved = self.deconvolve(data) 

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

422 

423 catalog = result.objectParents 

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

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

426 self.assertEqual(np.sum(parents["deblend_skipped_parentTooBig"]), 1) 

427 self.assertEqual(np.sum(parents["deblend_skipped_tooManyPeaks"]), 1) 

428 

429 def test_persistence(self): 

430 # Test that the model data is persisted correctly 

431 data = self.initialize_data(self.models) 

432 repo = self._setup_butler() 

433 mDeconvolved = self.deconvolve(data) 

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

435 modelData = result.scarletModelData 

436 bands = modelData.metadata["bands"] 

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

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

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

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

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

442 np.testing.assert_almost_equal(model_psf2, model_psf) 

443 psf = modelData.metadata["psf"] 

444 psf2 = modelData2.metadata["psf"] 

445 np.testing.assert_almost_equal(psf2, psf) 

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

447 

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

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

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

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

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

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

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

455 

456 for sourceId in modelData.isolated.keys(): 

457 isolatedData1 = modelData.isolated[sourceId] 

458 isolatedData2 = modelData2.isolated[sourceId] 

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

460 np.testing.assert_array_equal( 

461 isolatedData1.span_array, 

462 isolatedData2.span_array, 

463 ) 

464 

465 # Test extracting a single blend 

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

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

468 

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

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

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

472 

473 # Test extracting two blends 

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

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

476 for parentId in [2, 3]: 

477 parentData1 = modelData.blends[parentId] 

478 parentData2 = modelData2.blends[parentId] 

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

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

481 blendData1 = parentData1.children[blendId] 

482 blendData2 = parentData2.children[blendId] 

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

484 

485 def test_legacy_model(self): 

486 repo = self._setup_butler() 

487 storageClass = StorageClass( 

488 "LsstScarletModelData", 

489 pytype=mes.io.LsstScarletModelData, 

490 ) 

491 datasetType = DatasetType( 

492 "old_scarlet_model_data", 

493 dimensions=(), 

494 storageClass=storageClass, 

495 universe=repo.dimensions, 

496 ) 

497 ref = DatasetRef( 

498 datasetType, 

499 run="test_ingestion", 

500 dataId={}, 

501 ) 

502 dataset = FileDataset( 

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

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

505 refs=[ref], 

506 ) 

507 

508 # Ingest the legacy model into the butler 

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

510 repo.registry.registerDatasetType(datasetType) 

511 butler.ingest(dataset) 

512 

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

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

515 

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

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

518 

519 def test_older_legacy_model(self): 

520 repo = self._setup_butler() 

521 oldStorageClass = StorageClass( 

522 "ScarletModelData", 

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

524 ) 

525 oldDatasetType = DatasetType( 

526 "old_scarlet_model_data", 

527 dimensions=(), 

528 storageClass=oldStorageClass, 

529 universe=repo.dimensions, 

530 ) 

531 ref = DatasetRef( 

532 oldDatasetType, 

533 run="test_ingestion", 

534 dataId={}, 

535 ) 

536 dataset = FileDataset( 

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

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

539 refs=[ref], 

540 ) 

541 

542 # Ingest the legacy model into the butler 

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

544 repo.registry.registerDatasetType(oldDatasetType) 

545 butler.ingest(dataset) 

546 

547 # Load the base repo config from the repository 

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

549 

550 # Load the storage class override config 

551 override_path = os.path.join( 

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

553 "configs", 

554 "storageClasses.yaml" 

555 ) 

556 override_config = Config(override_path) 

557 

558 # Merge the configs (update base with override) 

559 base_config.update(override_config) 

560 

561 # Create Butler with the merged config 

562 # The config now contains both the repo info and 

563 # the storage class overrides 

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

565 

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

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

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

569 

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

571 # Test that two ScarletBlendData objects are equal 

572 # up to machine precision. 

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

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

575 

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

577 # once converted into scarlet lite Blend objects. 

578 blend1 = blendData1.minimal_data_to_blend( 

579 model_psf, 

580 psf, 

581 bands, 

582 dtype=np.float32, 

583 ) 

584 blend2 = blendData2.minimal_data_to_blend( 

585 model_psf, 

586 psf, 

587 bands, 

588 dtype=np.float32, 

589 ) 

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

591 

592 def _setup_butler(self): 

593 # Initialize a Butler to test persistence 

594 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) 

595 self.repo_dir = repo_dir.name 

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

597 config = Config() 

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

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

600 storageClass = StorageClass( 

601 "LsstScarletModelData", 

602 pytype=mes.io.LsstScarletModelData, 

603 parameters=('blend_id',), 

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

605 ) 

606 datasetType = DatasetType( 

607 "scarlet_model_data", 

608 dimensions=(), 

609 storageClass=storageClass, 

610 universe=repo.dimensions, 

611 ) 

612 repo.registry.registerDatasetType(datasetType) 

613 return repo 

614 

615 

616class BadPsf(Psf): 

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

618 self.validPoint = validPoint 

619 self.psf = psf 

620 super().__init__() 

621 

622 def computeKernelImage(self, location: Point2D): 

623 if location == self.validPoint: 

624 return self.psf.computeKernelImage(location) 

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

626 

627 

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

629 def setUp(self): 

630 self.bands = tuple("gri") 

631 

632 def test_box_transforms(self): 

633 # Test scarlet Box to BBox 

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

635 bbox = mes.utils.scarletBoxToBBox(box) 

636 x0, y0 = bbox.getMin() 

637 width = bbox.getWidth() 

638 height = bbox.getHeight() 

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

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

641 

642 # Test back to scarlet Box 

643 newBox = mes.utils.bboxToScarletBox(bbox) 

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

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

646 

647 def test_computeNearestPsfGood(self): 

648 # Test that using a valid PSF works normally 

649 psf, psfImage = self._generateGoodPsf() 

650 coadd = self._generateCoadd(psf) 

651 

652 # Test that computing the PSF works 

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

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

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

656 self.assertEqual(dist, 0) 

657 

658 def test_computeNearestPsfRecoverable(self): 

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

660 # will fallback to a valid location. 

661 psf, psfImage = self._generateGoodPsf() 

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

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

664 

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

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

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

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

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

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

671 self.assertEqual(dist, 5) 

672 

673 def test_computeNearestPsfBad(self): 

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

675 psf = self._generateGoodPsf() 

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

677 catalog = self._generateCatalog(self.bands) 

678 

679 # Test that computing the PSF cannot generate a PSF 

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

681 self.assertIsNone(derivedPsf) 

682 self.assertIsNone(center) 

683 self.assertIsNone(dist) 

684 

685 def test_computeNearestPsfMultiBandGood(self): 

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

687 bands = tuple("gri") 

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

689 mCoadd = self._generateMultibandCoadd(psfs, bands) 

690 

691 # Test that computing the PSF works 

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

693 np.testing.assert_array_equal(psfArray, psfImage) 

694 self.assertTupleEqual(newCoadd.bands, bands) 

695 

696 def test_computeNearestPsfMultiBandRecoverable(self): 

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

698 bands = tuple("gri") 

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

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

701 mCoadd = self._generateMultibandCoadd(psfs, bands) 

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