Coverage for tests / assemble_coadd_test_utils.py: 26%

182 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-28 09:16 +0000

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

28import numpy as np 

29from astro_metadata_translator import makeObservationInfo 

30from astropy import units as u 

31from astropy.coordinates import Angle, EarthLocation, SkyCoord 

32from astropy.time import Time 

33 

34import lsst.afw.geom as afwGeom 

35import lsst.afw.image as afwImage 

36import lsst.afw.math as afwMath 

37import lsst.afw.table as afwTable 

38import lsst.geom as geom 

39import lsst.pipe.base as pipeBase 

40from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

41from lsst.cell_coadds.test_utils import generate_data_id 

42from lsst.geom import arcseconds, degrees 

43from lsst.meas.algorithms.testUtils import plantSources 

44from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

45from lsst.pipe.base import InMemoryDatasetHandle 

46from lsst.pipe.tasks.coaddBase import growValidPolygons 

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

48from lsst.skymap import Index2D, PatchInfo 

49 

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

51 

52 

53def makeMockSkyInfo(bbox, wcs, patch): 

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

55 

56 Parameters 

57 ---------- 

58 bbox : `lsst.geom.Box` 

59 Bounding box of the patch to be coadded. 

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

61 Coordinate system definition (wcs) for the exposure. 

62 

63 Returns 

64 ------- 

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

66 Patch geometry information. 

67 """ 

68 

69 patchInfo = PatchInfo( 

70 index=Index2D(0, 0), 

71 sequentialIndex=patch, 

72 innerBBox=bbox, 

73 outerBBox=bbox, 

74 tractWcs=wcs, 

75 numCellsPerPatchInner=1, 

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

77 ) 

78 skyInfo = pipeBase.Struct(bbox=bbox, wcs=wcs, patchInfo=patchInfo) 

79 return skyInfo 

80 

81 

82class MockCoaddTestData: 

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

84 are realistic enough to test the image coaddition algorithms. 

85 

86 Notes 

87 ----- 

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

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

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

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

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

93 

94 Parameters 

95 ---------- 

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

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

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

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

100 backgroundLevel : `float`, optional 

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

102 seed : `int`, optional 

103 Seed value to initialize the random number generator. 

104 nSrc : `int`, optional 

105 Number of sources to simulate. 

106 fluxRange : `float`, optional 

107 Range in flux amplitude of the simulated sources. 

108 noiseLevel : `float`, optional 

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

110 sourceSigma : `float`, optional 

111 Average amplitude of the simulated sources, 

112 relative to ``noiseLevel`` 

113 minPsfSize : `float`, optional 

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

115 maxPsfSize : `float`, optional 

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

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

118 The plate scale of the simulated images. 

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

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

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

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

123 ccd : `int`, optional 

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

125 patch : `int`, optional 

126 Unique identifier for a subdivision of a tract. 

127 tract : `int`, optional 

128 Unique identifier for a tract of a skyMap. 

129 

130 Raises 

131 ------ 

132 ValueError 

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

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

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

136 lacking the option to specify the pixel origin. 

137 """ 

138 

139 rotAngle = 0.0 * degrees 

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

141 (`lsst.geom.Angle`). 

142 """ 

143 filterLabel = None 

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

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

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

147 """ 

148 rngData = None 

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

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

151 """ 

152 rngMods = None 

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

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

155 """ 

156 kernelSize = None 

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

158 exposures = {} 

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

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

161 """ 

162 matchedExposures = {} 

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

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

165 """ 

166 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

169 """ 

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

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

172 the coaddTempExp. 

173 """ 

174 detector = None 

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

176 

177 def __init__( 

178 self, 

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

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

181 backgroundLevel=314.592, 

182 seed=42, 

183 nSrc=37, 

184 fluxRange=2.0, 

185 noiseLevel=5, 

186 sourceSigma=200.0, 

187 minPsfSize=1.5, 

188 maxPsfSize=3.0, 

189 pixelScale=0.2 * arcseconds, 

190 ra=209.0 * degrees, 

191 dec=-20.25 * degrees, 

192 ccd=37, 

193 patch=42, 

194 tract=0, 

195 ): 

196 self.ra = ra 

197 self.dec = dec 

198 self.pixelScale = pixelScale 

199 self.patch = patch 

200 self.tract = tract 

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

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

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

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

205 if not self.bbox.contains(0, 0): 

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

207 self.wcs = self.makeDummyWcs() 

208 

209 # Set up properties of the simulations 

210 nSigmaForKernel = 5 

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

212 

213 bufferSize = self.kernelSize // 2 

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

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

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

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

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

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

220 

221 self.backgroundLevel = backgroundLevel 

222 self.noiseLevel = noiseLevel 

223 self.minPsfSize = minPsfSize 

224 self.maxPsfSize = maxPsfSize 

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

226 

227 def setDummyCoaddInputs(self, exposure, expId): 

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

229 processed using `make_direct_warp`. 

230 

231 Parameters 

232 ---------- 

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

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

235 expId : `int` 

236 A unique identifier for the visit. 

237 """ 

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

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

240 

241 config = CoaddInputRecorderConfig() 

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

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

244 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

245 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

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

247 

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

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

250 

251 Parameters 

252 ---------- 

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

254 The simulated exposure. 

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

256 VisitInfo containing metadata for the exposure. 

257 expId : `int` 

258 A unique identifier for the visit. 

259 

260 Returns 

261 ------- 

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

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

264 """ 

265 tempExp = rawExposure.clone() 

266 tempExp.setWcs(self.wcs) 

267 

268 tempExp.setFilter(self.filterLabel) 

269 tempExp.setPhotoCalib(self.photoCalib) 

270 tempExp.setApCorrMap(apCorrMap) 

271 tempExp.getInfo().setVisitInfo(visitInfo) 

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

273 self.setDummyCoaddInputs(tempExp, expId) 

274 return tempExp 

275 

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

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

278 

279 Parameters 

280 ---------- 

281 rotAngle : `lsst.geom.Angle` 

282 Rotation of the CD matrix, East from North 

283 pixelScale : `lsst.geom.Angle` 

284 Pixel scale of the projection. 

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

286 Coordinates of the reference pixel of the wcs. 

287 flipX : `bool`, optional 

288 Flip the direction of increasing Right Ascension. 

289 

290 Returns 

291 ------- 

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

293 A wcs that matches the inputs. 

294 """ 

295 if rotAngle is None: 

296 rotAngle = self.rotAngle 

297 if pixelScale is None: 

298 pixelScale = self.pixelScale 

299 if crval is None: 

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

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

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

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

304 return wcs 

305 

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

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

308 

309 Parameters 

310 ---------- 

311 exposureId : `int`, optional 

312 Unique integer identifier for this observation. 

313 randomizeTime : `bool`, optional 

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

315 

316 Returns 

317 ------- 

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

319 VisitInfo for the exposure. 

320 """ 

321 lsstLat = -30.244639 * u.degree 

322 lsstLon = -70.749417 * u.degree 

323 lsstAlt = 2663.0 * u.m 

324 lsstTemperature = 20.0 * u.Celsius 

325 lsstHumidity = 40.0 # in percent 

326 lsstPressure = 73892.0 * u.pascal 

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

328 

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

330 if randomizeTime: 

331 # Pick a random time within a 6 hour window 

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

333 radec = SkyCoord( 

334 dec=self.dec.asDegrees(), 

335 ra=self.ra.asDegrees(), 

336 unit="deg", 

337 obstime=time, 

338 frame="icrs", 

339 location=loc, 

340 ) 

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

342 obsInfo = makeObservationInfo( 

343 location=loc, 

344 detector_exposure_id=exposureId, 

345 datetime_begin=time, 

346 datetime_end=time, 

347 boresight_airmass=airmass, 

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

349 boresight_rotation_coord="sky", 

350 temperature=lsstTemperature, 

351 pressure=lsstPressure, 

352 relative_humidity=lsstHumidity, 

353 tracking_radec=radec, 

354 altaz_begin=radec.altaz, 

355 observation_type="science", 

356 ) 

357 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

358 return visitInfo 

359 

360 def makeDummyApCorrMap(self): 

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

362 

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

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

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

366 constant. 

367 

368 Returns 

369 ------- 

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

371 Aperture correction map for the exposure. 

372 """ 

373 apCorrMap = afwImage.ApCorrMap() 

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

375 self.bbox, 

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

377 ) 

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

379 self.bbox, 

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

381 ) 

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

383 self.bbox, 

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

385 ) 

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

387 self.bbox, 

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

389 ) 

390 

391 return apCorrMap 

392 

393 def makeTestImage( 

394 self, 

395 expId, 

396 noiseLevel=None, 

397 psfSize=None, 

398 backgroundLevel=None, 

399 detectionSigma=5.0, 

400 badRegionBox=None, 

401 ): 

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

403 

404 Parameters 

405 ---------- 

406 expId : `int` 

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

408 noiseLevel : `float`, optional 

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

410 psfSize : `float`, optional 

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

412 backgroundLevel : `float`, optional 

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

414 detectionSigma : `float`, optional 

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

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

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

418 """ 

419 if backgroundLevel is None: 

420 backgroundLevel = self.backgroundLevel 

421 if noiseLevel is None: 

422 noiseLevel = 5.0 

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

424 

425 if psfSize is None: 

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

427 nSrc = len(self.flux) 

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

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

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

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

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

433 model = plantSources( 

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

435 ) 

436 modelPsfMatched = plantSources( 

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

438 ) 

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

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

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

442 noise -= np.median(noise) 

443 model.image.array += noise 

444 modelPsfMatched.image.array += noise 

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

446 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel 

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

448 

449 if badRegionBox is not None: 

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

451 

452 apCorrMap = self.makeDummyApCorrMap() 

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

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

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

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

457 return exposure, matchedExposure 

458 

459 def makeVisitSummaryTableHandle(self, warpHandle): 

460 schema = afwTable.ExposureTable.makeMinimalSchema() 

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

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

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

464 visitSummaryTable = afwTable.ExposureCatalog(schema) 

465 

466 warp = warpHandle.get() 

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

468 record = visitSummaryTable.addNew() 

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

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

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

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

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

474 

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

476 return handle 

477 

478 @staticmethod 

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

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

481 retrieved using the Gen 3 Butler API. 

482 

483 Parameters 

484 ---------- 

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

486 A mapping of exposure IDs to ExposureF objects. 

487 tract : `int`, optional 

488 Unique identifier for a tract of a skyMap. 

489 patch : `int`, optional 

490 Unique identifier for a subdivision of a tract. 

491 

492 Returns 

493 ------- 

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

495 The data references. 

496 

497 Raises 

498 ------ 

499 ValueError 

500 If an unknown `warpType` is supplied. 

501 """ 

502 dataRefList = [] 

503 for expId in exposures: 

504 exposure = exposures[expId] 

505 dataRef = pipeBase.InMemoryDatasetHandle( 

506 exposure, 

507 storageClass="ExposureF", 

508 copy=True, 

509 dataId=generate_data_id( 

510 tract=tract, 

511 patch=patch, 

512 visit_id=expId, 

513 ), 

514 ) 

515 dataRefList.append(dataRef) 

516 return dataRefList