Coverage for tests/assemble_coadd_test_utils.py: 96%

189 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-09-14 02:49 -0700

1# This file is part of drp_tasks. 

2# 

3# LSST Data Management System 

4# This product includes software developed by the 

5# LSST Project (http://www.lsst.org/). 

6# See COPYRIGHT file at the top of the source tree. 

7# 

8# This program is free software: you can redistribute it and/or modify 

9# it under the terms of the GNU General Public License as published by 

10# the Free Software Foundation, either version 3 of the License, or 

11# (at your option) any later version. 

12# 

13# This program is distributed in the hope that it will be useful, 

14# but WITHOUT ANY WARRANTY; without even the implied warranty of 

15# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

16# GNU General Public License for more details. 

17# 

18# You should have received a copy of the LSST License Statement and 

19# the GNU General Public License along with this program. If not, 

20# see <https://www.lsstcorp.org/LegalNotices/>. 

21# 

22"""Set up simulated test data and simplified APIs for AssembleCoaddTask 

23and its derived classes. 

24 

25This is not intended to test accessing data with the Butler and instead uses 

26mock Butler data references to pass in the simulated data. 

27""" 

28 

29import numpy as np 

30from astro_metadata_translator import makeObservationInfo 

31from astropy import units as u 

32from astropy.coordinates import Angle, EarthLocation, SkyCoord 

33from astropy.time import Time 

34 

35import lsst.afw.geom as afwGeom 

36import lsst.afw.image as afwImage 

37import lsst.afw.math as afwMath 

38import lsst.afw.table as afwTable 

39import lsst.geom as geom 

40import lsst.pipe.base as pipeBase 

41from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

42from lsst.cell_coadds.test_utils import generate_data_id 

43from lsst.geom import arcseconds, degrees 

44from lsst.meas.algorithms.testUtils import plantSources 

45from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

46from lsst.pipe.base import InMemoryDatasetHandle 

47from lsst.pipe.tasks.coaddBase import growValidPolygons 

48from lsst.pipe.tasks.coaddInputRecorder import CoaddInputRecorderConfig, CoaddInputRecorderTask 

49from lsst.skymap import Index2D, PatchInfo 

50 

51__all__ = ["makeMockSkyInfo", "MockCoaddTestData"] 

52 

53 

54class _MockTractInfo: 

55 

56 def __init__(self, patch_info: PatchInfo): 

57 self._patch_info = patch_info 

58 

59 def __getitem__(self, _) -> PatchInfo: 

60 return self._patch_info 

61 

62 def getBBox(self) -> geom.Box2I: 

63 return self._patch_info.getOuterBBox() 

64 

65 

66def makeMockSkyInfo(bbox, wcs, patch): 

67 """Construct a `Struct` containing the geometry of the patch to be coadded. 

68 

69 Parameters 

70 ---------- 

71 bbox : `lsst.geom.Box` 

72 Bounding box of the patch to be coadded. 

73 wcs : `lsst.afw.geom.SkyWcs` 

74 Coordinate system definition (wcs) for the exposure. 

75 

76 Returns 

77 ------- 

78 skyInfo : `lsst.pipe.base.Struct` 

79 Patch geometry information. 

80 """ 

81 patchInfo = PatchInfo( 

82 index=Index2D(0, 0), 

83 sequentialIndex=patch, 

84 innerBBox=bbox, 

85 outerBBox=bbox, 

86 tractWcs=wcs, 

87 numCellsPerPatchInner=1, 

88 cellInnerDimensions=(bbox.width, bbox.height), 

89 ) 

90 skyInfo = pipeBase.Struct(bbox=bbox, wcs=wcs, tractInfo=_MockTractInfo(patchInfo), patchInfo=patchInfo) 

91 return skyInfo 

92 

93 

94class MockCoaddTestData: 

95 """Generate repeatable simulated exposures with consistent metadata that 

96 are realistic enough to test the image coaddition algorithms. 

97 

98 Notes 

99 ----- 

100 The simple GaussianPsf used by lsst.meas.algorithms.testUtils.plantSources 

101 will always return an average position of (0, 0). 

102 The bounding box of the exposures MUST include (0, 0), or else the PSF will 

103 not be valid and `AssembleCoaddTask` will fail with the error 

104 'Could not find a valid average position for CoaddPsf'. 

105 

106 Parameters 

107 ---------- 

108 shape : `lsst.geom.Extent2I`, optional 

109 Size of the bounding box of the exposures to be simulated, in pixels. 

110 offset : `lsst.geom.Point2I`, optional 

111 Pixel coordinate of the lower left corner of the bounding box. 

112 backgroundLevel : `float`, optional 

113 Background value added to all pixels in the simulated images. 

114 seed : `int`, optional 

115 Seed value to initialize the random number generator. 

116 nSrc : `int`, optional 

117 Number of sources to simulate. 

118 fluxRange : `float`, optional 

119 Range in flux amplitude of the simulated sources. 

120 noiseLevel : `float`, optional 

121 Standard deviation of the noise to add to each pixel. 

122 sourceSigma : `float`, optional 

123 Average amplitude of the simulated sources, 

124 relative to ``noiseLevel`` 

125 minPsfSize : `float`, optional 

126 The smallest PSF width (sigma) to use, in pixels. 

127 maxPsfSize : `float`, optional 

128 The largest PSF width (sigma) to use, in pixels. 

129 pixelScale : `lsst.geom.Angle`, optional 

130 The plate scale of the simulated images. 

131 ra : `lsst.geom.Angle`, optional 

132 Right Ascension of the boresight of the camera for the observation. 

133 dec : `lsst.geom.Angle`, optional 

134 Declination of the boresight of the camera for the observation. 

135 ccd : `int`, optional 

136 CCD number to put in the metadata of the exposure. 

137 patch : `int`, optional 

138 Unique identifier for a subdivision of a tract. 

139 tract : `int`, optional 

140 Unique identifier for a tract of a skyMap. 

141 

142 Raises 

143 ------ 

144 ValueError 

145 If the bounding box does not contain the pixel coordinate (0, 0). 

146 This is due to `GaussianPsf` that is used by 

147 `~lsst.meas.algorithms.testUtils.plantSources` 

148 lacking the option to specify the pixel origin. 

149 """ 

150 

151 rotAngle = 0.0 * degrees 

152 """Rotation of the pixel grid on the sky, East from North 

153 (`lsst.geom.Angle`). 

154 """ 

155 filterLabel = None 

156 """The filter definition, usually set in the current instruments' obs 

157 package. For these tests, a simple filter is defined without using an obs 

158 package (`lsst.afw.image.FilterLabel`). 

159 """ 

160 rngData = None 

161 """Pre-initialized random number generator for constructing the test images 

162 repeatably (`numpy.random.Generator`). 

163 """ 

164 rngMods = None 

165 """Pre-initialized random number generator for applying modifications to 

166 the test images for only some test cases (`numpy.random.Generator`). 

167 """ 

168 kernelSize = None 

169 "Width of the kernel used for simulating sources, in pixels." 

170 exposures = {} 

171 """The simulated test data, with variable PSF size 

172 (`dict` of `lsst.afw.image.Exposure`) 

173 """ 

174 matchedExposures = {} 

175 """The simulated exposures, all with PSF width set to ``maxPsfSize`` 

176 (`dict` of `lsst.afw.image.Exposure`). 

177 """ 

178 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

179 """The photometric zero point to use for converting counts to flux units 

180 (`lsst.afw.image.PhotoCalib`). 

181 """ 

182 badMaskPlanes = ["NO_DATA", "BAD"] 

183 """Mask planes that, if set, the associated pixel should not be included in 

184 the coaddTempExp. 

185 """ 

186 detector = None 

187 "Properties of the CCD for the exposure (`lsst.afw.cameraGeom.Detector`)." 

188 

189 def __init__( 

190 self, 

191 shape=geom.Extent2I(201, 301), 

192 offset=geom.Point2I(-123, -45), 

193 backgroundLevel=314.592, 

194 seed=42, 

195 nSrc=37, 

196 fluxRange=2.0, 

197 noiseLevel=5, 

198 sourceSigma=200.0, 

199 minPsfSize=1.5, 

200 maxPsfSize=3.0, 

201 pixelScale=0.2 * arcseconds, 

202 ra=209.0 * degrees, 

203 dec=-20.25 * degrees, 

204 ccd=37, 

205 patch=42, 

206 tract=0, 

207 ): 

208 self.ra = ra 

209 self.dec = dec 

210 self.pixelScale = pixelScale 

211 self.patch = patch 

212 self.tract = tract 

213 self.filterLabel = afwImage.FilterLabel(band="gTest", physical="gTest") 

214 self.rngData = np.random.default_rng(seed) 

215 self.rngMods = np.random.default_rng(seed + 1) 

216 self.bbox = geom.Box2I(offset, shape) 

217 if not self.bbox.contains(0, 0): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true

218 raise ValueError(f"The bounding box must contain the coordinate (0, 0). {repr(self.bbox)}") 

219 self.wcs = self.makeDummyWcs() 

220 

221 # Set up properties of the simulations 

222 nSigmaForKernel = 5 

223 self.kernelSize = (int(maxPsfSize * nSigmaForKernel + 0.5) // 2) * 2 + 1 # make sure it is odd 

224 

225 bufferSize = self.kernelSize // 2 

226 x0, y0 = self.bbox.getBegin() 

227 xSize, ySize = self.bbox.getDimensions() 

228 # Set the pixel coordinates and fluxes of the simulated sources. 

229 self.xLoc = self.rngData.random(nSrc) * (xSize - 2 * bufferSize) + bufferSize + x0 

230 self.yLoc = self.rngData.random(nSrc) * (ySize - 2 * bufferSize) + bufferSize + y0 

231 self.flux = (self.rngData.random(nSrc) * (fluxRange - 1.0) + 1.0) * sourceSigma * noiseLevel 

232 

233 self.backgroundLevel = backgroundLevel 

234 self.noiseLevel = noiseLevel 

235 self.minPsfSize = minPsfSize 

236 self.maxPsfSize = maxPsfSize 

237 self.detector = DetectorWrapper(name=f"detector {ccd}", id=ccd).detector 

238 

239 def setDummyCoaddInputs(self, exposure, expId): 

240 """Generate an `ExposureCatalog` as though the exposures had been 

241 processed using `make_direct_warp`. 

242 

243 Parameters 

244 ---------- 

245 exposure : `lsst.afw.image.Exposure` 

246 The exposure to construct a `CoaddInputs` `ExposureCatalog` for. 

247 expId : `int` 

248 A unique identifier for the visit. 

249 """ 

250 badPixelMask = afwImage.Mask.getPlaneBitMask(self.badMaskPlanes) 

251 nGoodPix = np.sum(exposure.getMask().getArray() & badPixelMask == 0) 

252 

253 config = CoaddInputRecorderConfig() 

254 inputRecorder = CoaddInputRecorderTask(config=config, name="inputRecorder") 

255 tempExpInputRecorder = inputRecorder.makeCoaddTempExpRecorder(expId, num=1) 

256 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

257 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

258 growValidPolygons(exposure.getInfo().getCoaddInputs(), growBy=0) 

259 

260 def makeCoaddTempExp(self, rawExposure, visitInfo, expId, apCorrMap=None): 

261 """Add the metadata required by `AssembleCoaddTask` to an exposure. 

262 

263 Parameters 

264 ---------- 

265 rawExposure : `lsst.afw.image.Exposure` 

266 The simulated exposure. 

267 visitInfo : `lsst.afw.image.VisitInfo` 

268 VisitInfo containing metadata for the exposure. 

269 expId : `int` 

270 A unique identifier for the visit. 

271 

272 Returns 

273 ------- 

274 tempExp : `lsst.afw.image.Exposure` 

275 The exposure, with all of the metadata needed for coaddition. 

276 """ 

277 tempExp = rawExposure.clone() 

278 tempExp.setWcs(self.wcs) 

279 

280 tempExp.setFilter(self.filterLabel) 

281 tempExp.setPhotoCalib(self.photoCalib) 

282 tempExp.setApCorrMap(apCorrMap) 

283 tempExp.getInfo().setVisitInfo(visitInfo) 

284 tempExp.getInfo().setDetector(self.detector) 

285 self.setDummyCoaddInputs(tempExp, expId) 

286 return tempExp 

287 

288 def makeDummyWcs(self, rotAngle=None, pixelScale=None, crval=None, flipX=True): 

289 """Make a World Coordinate System object for testing. 

290 

291 Parameters 

292 ---------- 

293 rotAngle : `lsst.geom.Angle` 

294 Rotation of the CD matrix, East from North 

295 pixelScale : `lsst.geom.Angle` 

296 Pixel scale of the projection. 

297 crval : `lsst.afw.geom.SpherePoint` 

298 Coordinates of the reference pixel of the wcs. 

299 flipX : `bool`, optional 

300 Flip the direction of increasing Right Ascension. 

301 

302 Returns 

303 ------- 

304 wcs : `lsst.afw.geom.skyWcs.SkyWcs` 

305 A wcs that matches the inputs. 

306 """ 

307 if rotAngle is None: 307 ↛ 309line 307 didn't jump to line 309 because the condition on line 307 was always true

308 rotAngle = self.rotAngle 

309 if pixelScale is None: 309 ↛ 311line 309 didn't jump to line 311 because the condition on line 309 was always true

310 pixelScale = self.pixelScale 

311 if crval is None: 311 ↛ 313line 311 didn't jump to line 313 because the condition on line 311 was always true

312 crval = geom.SpherePoint(self.ra, self.dec) 

313 crpix = geom.Box2D(self.bbox).getCenter() 

314 cdMatrix = afwGeom.makeCdMatrix(scale=pixelScale, orientation=rotAngle, flipX=flipX) 

315 wcs = afwGeom.makeSkyWcs(crpix=crpix, crval=crval, cdMatrix=cdMatrix) 

316 return wcs 

317 

318 def makeDummyVisitInfo(self, exposureId, randomizeTime=False): 

319 """Make a self-consistent visitInfo object for testing. 

320 

321 Parameters 

322 ---------- 

323 exposureId : `int`, optional 

324 Unique integer identifier for this observation. 

325 randomizeTime : `bool`, optional 

326 Add a random offset within a 6 hour window to the observation time. 

327 

328 Returns 

329 ------- 

330 visitInfo : `lsst.afw.image.VisitInfo` 

331 VisitInfo for the exposure. 

332 """ 

333 lsstLat = -30.244639 * u.degree 

334 lsstLon = -70.749417 * u.degree 

335 lsstAlt = 2663.0 * u.m 

336 lsstTemperature = 20.0 * u.Celsius 

337 lsstHumidity = 40.0 # in percent 

338 lsstPressure = 73892.0 * u.pascal 

339 loc = EarthLocation(lat=lsstLat, lon=lsstLon, height=lsstAlt) 

340 

341 time = Time(2000.0, format="jyear", scale="tt") 

342 if randomizeTime: 342 ↛ 345line 342 didn't jump to line 345 because the condition on line 342 was always true

343 # Pick a random time within a 6 hour window 

344 time += 6 * u.hour * (self.rngMods.random() - 0.5) 

345 radec = SkyCoord( 

346 dec=self.dec.asDegrees(), 

347 ra=self.ra.asDegrees(), 

348 unit="deg", 

349 obstime=time, 

350 frame="icrs", 

351 location=loc, 

352 ) 

353 airmass = float(1.0 / np.sin(radec.altaz.alt)) 

354 obsInfo = makeObservationInfo( 

355 location=loc, 

356 detector_exposure_id=exposureId, 

357 datetime_begin=time, 

358 datetime_end=time, 

359 boresight_airmass=airmass, 

360 boresight_rotation_angle=Angle(0.0 * u.degree), 

361 boresight_rotation_coord="sky", 

362 temperature=lsstTemperature, 

363 pressure=lsstPressure, 

364 relative_humidity=lsstHumidity, 

365 tracking_radec=radec, 

366 altaz_begin=radec.altaz, 

367 observation_type="science", 

368 ) 

369 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

370 return visitInfo 

371 

372 def makeDummyApCorrMap(self): 

373 """Make a dummy aperture correction map for testing. 

374 

375 This method returns an `~lsst.afw.image.ApCorrMap` instance with 

376 random values for the "algo1_instFlux", "algo1_instFluxErr", 

377 "algo2_instFlux", and "algo2_instFluxErr" fields that are spatially 

378 constant. 

379 

380 Returns 

381 ------- 

382 apCorrMap : `lsst.afw.image.ApCorrMap` 

383 Aperture correction map for the exposure. 

384 """ 

385 apCorrMap = afwImage.ApCorrMap() 

386 apCorrMap["algo1_instFlux"] = afwMath.ChebyshevBoundedField( 

387 self.bbox, 

388 np.array([[self.rngMods.random()]]), 

389 ) 

390 apCorrMap["algo1_instFluxErr"] = afwMath.ChebyshevBoundedField( 

391 self.bbox, 

392 np.array([[self.rngMods.random()]]), 

393 ) 

394 apCorrMap["algo2_instFlux"] = afwMath.ChebyshevBoundedField( 

395 self.bbox, 

396 np.array([[self.rngMods.random()]]), 

397 ) 

398 apCorrMap["algo2_instFluxErr"] = afwMath.ChebyshevBoundedField( 

399 self.bbox, 

400 np.array([[self.rngMods.random()]]), 

401 ) 

402 

403 return apCorrMap 

404 

405 def makeTestImage( 

406 self, 

407 expId, 

408 noiseLevel=None, 

409 psfSize=None, 

410 backgroundLevel=None, 

411 detectionSigma=5.0, 

412 badRegionBox=None, 

413 ): 

414 """Make a reproduceable PSF-convolved masked image for testing. 

415 

416 Parameters 

417 ---------- 

418 expId : `int` 

419 A unique identifier to use to refer to the visit. 

420 noiseLevel : `float`, optional 

421 Standard deviation of the noise to add to each pixel. 

422 psfSize : `float`, optional 

423 Width of the PSF of the simulated sources, in pixels. 

424 backgroundLevel : `float`, optional 

425 Background value added to all pixels in the simulated images. 

426 detectionSigma : `float`, optional 

427 Threshold amplitude of the image to set the "DETECTED" mask. 

428 badRegionBox : `lsst.geom.Box2I`, optional 

429 Add a bad region bounding box (set to "BAD"). 

430 """ 

431 if backgroundLevel is None: 431 ↛ 433line 431 didn't jump to line 433 because the condition on line 431 was always true

432 backgroundLevel = self.backgroundLevel 

433 if noiseLevel is None: 433 ↛ 435line 433 didn't jump to line 435 because the condition on line 433 was always true

434 noiseLevel = 5.0 

435 visitInfo = self.makeDummyVisitInfo(expId, randomizeTime=True) 

436 

437 if psfSize is None: 437 ↛ 439line 437 didn't jump to line 439 because the condition on line 437 was always true

438 psfSize = self.rngMods.random() * (self.maxPsfSize - self.minPsfSize) + self.minPsfSize 

439 nSrc = len(self.flux) 

440 sigmas = [psfSize for src in range(nSrc)] 

441 sigmasPsfMatched = [self.maxPsfSize for src in range(nSrc)] 

442 coordList = list(zip(self.xLoc, self.yLoc, self.flux, sigmas)) 

443 coordListPsfMatched = list(zip(self.xLoc, self.yLoc, self.flux, sigmasPsfMatched)) 

444 xSize, ySize = self.bbox.getDimensions() 

445 model = plantSources( 

446 self.bbox, self.kernelSize, self.backgroundLevel, coordList, addPoissonNoise=False 

447 ) 

448 modelPsfMatched = plantSources( 

449 self.bbox, self.kernelSize, self.backgroundLevel, coordListPsfMatched, addPoissonNoise=False 

450 ) 

451 model.variance.array = np.abs(model.image.array) + noiseLevel 

452 modelPsfMatched.variance.array = np.abs(modelPsfMatched.image.array) + noiseLevel 

453 noise = self.rngData.random((ySize, xSize)) * noiseLevel 

454 noise -= np.median(noise) 

455 model.image.array += noise 

456 modelPsfMatched.image.array += noise 

457 detectedMask = afwImage.Mask.getPlaneBitMask("DETECTED") 

458 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel 

459 model.mask.array[model.image.array > detectionThreshold] += detectedMask 

460 

461 if badRegionBox is not None: 

462 model.mask[badRegionBox] = afwImage.Mask.getPlaneBitMask("BAD") 

463 

464 apCorrMap = self.makeDummyApCorrMap() 

465 exposure = self.makeCoaddTempExp(model, visitInfo, expId, apCorrMap) 

466 matchedExposure = self.makeCoaddTempExp(modelPsfMatched, visitInfo, expId, apCorrMap) 

467 exposure.metadata["BUNIT"] = "nJy" 

468 matchedExposure.metadata["BUNIT"] = "nJy" 

469 return exposure, matchedExposure 

470 

471 def makeVisitSummaryTableHandle(self, warpHandle): 

472 schema = afwTable.ExposureTable.makeMinimalSchema() 

473 schema.addField("visit", type=np.int64, doc="Visit ID") 

474 schema.addField("ccd", type=np.int64, doc="CCD ID") 

475 schema.addField("meanVar", type=float, units="nJy**2", doc="Mean variance") 

476 visitSummaryTable = afwTable.ExposureCatalog(schema) 

477 

478 warp = warpHandle.get() 

479 for ccd in warp.getInfo().getCoaddInputs().ccds: 

480 record = visitSummaryTable.addNew() 

481 record.set("id", ccd["ccd"]) 

482 record.set("ccd", ccd["ccd"]) 

483 record.set("visit", ccd["visit"]) 

484 record.set("meanVar", 0.3 + self.rngMods.random()) 

485 record.setPhotoCalib(afwImage.PhotoCalib(calibrationMean=10.0)) 

486 

487 handle = InMemoryDatasetHandle(visitSummaryTable, dataId=warpHandle.dataId) 

488 return handle 

489 

490 @staticmethod 

491 def makeDataRefList(exposures, tract=0, patch=42): 

492 """Make data references from the simulated exposures that can be 

493 retrieved using the Gen 3 Butler API. 

494 

495 Parameters 

496 ---------- 

497 exposures : `Mapping` [`Any`, `~lsst.afw.image.ExposureF`] 

498 A mapping of exposure IDs to ExposureF objects. 

499 tract : `int`, optional 

500 Unique identifier for a tract of a skyMap. 

501 patch : `int`, optional 

502 Unique identifier for a subdivision of a tract. 

503 

504 Returns 

505 ------- 

506 dataRefList : `list` [`~lsst.pipe.base.InMemoryDatasetHandle`] 

507 The data references. 

508 

509 Raises 

510 ------ 

511 ValueError 

512 If an unknown `warpType` is supplied. 

513 """ 

514 dataRefList = [] 

515 for expId in exposures: 

516 exposure = exposures[expId] 

517 dataRef = pipeBase.InMemoryDatasetHandle( 

518 exposure, 

519 storageClass="ExposureF", 

520 copy=True, 

521 dataId=generate_data_id( 

522 tract=tract, 

523 patch=patch, 

524 visit_id=expId, 

525 ), 

526 ) 

527 dataRefList.append(dataRef) 

528 return dataRefList