Coverage for tests/assemble_coadd_test_utils.py: 27%

155 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-26 04:19 -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""" 

28import lsst.afw.geom as afwGeom 

29import lsst.afw.image as afwImage 

30import lsst.geom as geom 

31import lsst.pipe.base as pipeBase 

32import numpy as np 

33from astro_metadata_translator import makeObservationInfo 

34from astropy import units as u 

35from astropy.coordinates import Angle, EarthLocation, SkyCoord 

36from astropy.time import Time 

37from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

38from lsst.cell_coadds.test_utils import generate_data_id 

39from lsst.geom import arcseconds, degrees 

40from lsst.meas.algorithms.testUtils import plantSources 

41from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

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

43from lsst.skymap import Index2D, PatchInfo 

44 

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

46 

47 

48def makeMockSkyInfo(bbox, wcs, patch): 

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

50 

51 Parameters 

52 ---------- 

53 bbox : `lsst.geom.Box` 

54 Bounding box of the patch to be coadded. 

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

56 Coordinate system definition (wcs) for the exposure. 

57 

58 Returns 

59 ------- 

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

61 Patch geometry information. 

62 """ 

63 

64 patchInfo = PatchInfo( 

65 index=Index2D(0, 0), 

66 sequentialIndex=patch, 

67 innerBBox=bbox, 

68 outerBBox=bbox, 

69 tractWcs=wcs, 

70 numCellsPerPatchInner=1, 

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

72 ) 

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

74 return skyInfo 

75 

76 

77class MockCoaddTestData: 

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

79 are realistic enough to test the image coaddition algorithms. 

80 

81 Notes 

82 ----- 

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

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

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

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

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

88 

89 Parameters 

90 ---------- 

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

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

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

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

95 backgroundLevel : `float`, optional 

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

97 seed : `int`, optional 

98 Seed value to initialize the random number generator. 

99 nSrc : `int`, optional 

100 Number of sources to simulate. 

101 fluxRange : `float`, optional 

102 Range in flux amplitude of the simulated sources. 

103 noiseLevel : `float`, optional 

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

105 sourceSigma : `float`, optional 

106 Average amplitude of the simulated sources, 

107 relative to ``noiseLevel`` 

108 minPsfSize : `float`, optional 

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

110 maxPsfSize : `float`, optional 

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

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

113 The plate scale of the simulated images. 

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

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

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

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

118 ccd : `int`, optional 

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

120 patch : `int`, optional 

121 Unique identifier for a subdivision of a tract. 

122 tract : `int`, optional 

123 Unique identifier for a tract of a skyMap. 

124 

125 Raises 

126 ------ 

127 ValueError 

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

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

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

131 lacking the option to specify the pixel origin. 

132 """ 

133 

134 rotAngle = 0.0 * degrees 

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

136 (`lsst.geom.Angle`). 

137 """ 

138 filterLabel = None 

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

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

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

142 """ 

143 rngData = None 

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

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

146 """ 

147 rngMods = None 

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

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

150 """ 

151 kernelSize = None 

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

153 exposures = {} 

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

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

156 """ 

157 matchedExposures = {} 

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

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

160 """ 

161 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

164 """ 

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

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

167 the coaddTempExp. 

168 """ 

169 detector = None 

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

171 

172 def __init__( 

173 self, 

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

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

176 backgroundLevel=314.592, 

177 seed=42, 

178 nSrc=37, 

179 fluxRange=2.0, 

180 noiseLevel=5, 

181 sourceSigma=200.0, 

182 minPsfSize=1.5, 

183 maxPsfSize=3.0, 

184 pixelScale=0.2 * arcseconds, 

185 ra=209.0 * degrees, 

186 dec=-20.25 * degrees, 

187 ccd=37, 

188 patch=42, 

189 tract=0, 

190 ): 

191 self.ra = ra 

192 self.dec = dec 

193 self.pixelScale = pixelScale 

194 self.patch = patch 

195 self.tract = tract 

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

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

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

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

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

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

202 self.wcs = self.makeDummyWcs() 

203 

204 # Set up properties of the simulations 

205 nSigmaForKernel = 5 

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

207 

208 bufferSize = self.kernelSize // 2 

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

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

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

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

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

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

215 

216 self.backgroundLevel = backgroundLevel 

217 self.noiseLevel = noiseLevel 

218 self.minPsfSize = minPsfSize 

219 self.maxPsfSize = maxPsfSize 

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

221 

222 def setDummyCoaddInputs(self, exposure, expId): 

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

224 processed using `warpAndPsfMatch`. 

225 

226 Parameters 

227 ---------- 

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

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

230 expId : `int` 

231 A unique identifier for the visit. 

232 """ 

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

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

235 

236 config = CoaddInputRecorderConfig() 

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

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

239 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

240 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

241 

242 def makeCoaddTempExp(self, rawExposure, visitInfo, expId): 

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

244 

245 Parameters 

246 ---------- 

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

248 The simulated exposure. 

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

250 VisitInfo containing metadata for the exposure. 

251 expId : `int` 

252 A unique identifier for the visit. 

253 

254 Returns 

255 ------- 

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

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

258 """ 

259 tempExp = rawExposure.clone() 

260 tempExp.setWcs(self.wcs) 

261 

262 tempExp.setFilter(self.filterLabel) 

263 tempExp.setPhotoCalib(self.photoCalib) 

264 tempExp.getInfo().setVisitInfo(visitInfo) 

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

266 self.setDummyCoaddInputs(tempExp, expId) 

267 return tempExp 

268 

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

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

271 

272 Parameters 

273 ---------- 

274 rotAngle : `lsst.geom.Angle` 

275 Rotation of the CD matrix, East from North 

276 pixelScale : `lsst.geom.Angle` 

277 Pixel scale of the projection. 

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

279 Coordinates of the reference pixel of the wcs. 

280 flipX : `bool`, optional 

281 Flip the direction of increasing Right Ascension. 

282 

283 Returns 

284 ------- 

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

286 A wcs that matches the inputs. 

287 """ 

288 if rotAngle is None: 

289 rotAngle = self.rotAngle 

290 if pixelScale is None: 

291 pixelScale = self.pixelScale 

292 if crval is None: 

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

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

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

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

297 return wcs 

298 

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

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

301 

302 Parameters 

303 ---------- 

304 exposureId : `int`, optional 

305 Unique integer identifier for this observation. 

306 randomizeTime : `bool`, optional 

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

308 

309 Returns 

310 ------- 

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

312 VisitInfo for the exposure. 

313 """ 

314 lsstLat = -30.244639 * u.degree 

315 lsstLon = -70.749417 * u.degree 

316 lsstAlt = 2663.0 * u.m 

317 lsstTemperature = 20.0 * u.Celsius 

318 lsstHumidity = 40.0 # in percent 

319 lsstPressure = 73892.0 * u.pascal 

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

321 

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

323 if randomizeTime: 

324 # Pick a random time within a 6 hour window 

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

326 radec = SkyCoord( 

327 dec=self.dec.asDegrees(), 

328 ra=self.ra.asDegrees(), 

329 unit="deg", 

330 obstime=time, 

331 frame="icrs", 

332 location=loc, 

333 ) 

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

335 obsInfo = makeObservationInfo( 

336 location=loc, 

337 detector_exposure_id=exposureId, 

338 datetime_begin=time, 

339 datetime_end=time, 

340 boresight_airmass=airmass, 

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

342 boresight_rotation_coord="sky", 

343 temperature=lsstTemperature, 

344 pressure=lsstPressure, 

345 relative_humidity=lsstHumidity, 

346 tracking_radec=radec, 

347 altaz_begin=radec.altaz, 

348 observation_type="science", 

349 ) 

350 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

351 return visitInfo 

352 

353 def makeTestImage( 

354 self, 

355 expId, 

356 noiseLevel=None, 

357 psfSize=None, 

358 backgroundLevel=None, 

359 detectionSigma=5.0, 

360 badRegionBox=None, 

361 ): 

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

363 

364 Parameters 

365 ---------- 

366 expId : `int` 

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

368 noiseLevel : `float`, optional 

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

370 psfSize : `float`, optional 

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

372 backgroundLevel : `float`, optional 

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

374 detectionSigma : `float`, optional 

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

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

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

378 """ 

379 if backgroundLevel is None: 

380 backgroundLevel = self.backgroundLevel 

381 if noiseLevel is None: 

382 noiseLevel = 5.0 

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

384 

385 if psfSize is None: 

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

387 nSrc = len(self.flux) 

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

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

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

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

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

393 model = plantSources( 

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

395 ) 

396 modelPsfMatched = plantSources( 

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

398 ) 

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

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

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

402 noise -= np.median(noise) 

403 model.image.array += noise 

404 modelPsfMatched.image.array += noise 

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

406 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel 

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

408 

409 if badRegionBox is not None: 

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

411 

412 exposure = self.makeCoaddTempExp(model, visitInfo, expId) 

413 matchedExposure = self.makeCoaddTempExp(modelPsfMatched, visitInfo, expId) 

414 return exposure, matchedExposure 

415 

416 @staticmethod 

417 def makeDataRefList(exposures, matchedExposures, warpType, tract=0, patch=42): 

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

419 retrieved using the Gen 3 Butler API. 

420 

421 Parameters 

422 ---------- 

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

424 A mapping of exposure IDs to ExposureF objects that correspond to 

425 directWarp datasets. 

426 matchedExposures : `Mapping` [`Any`, `~lsst.afw.image.ExposureF`] 

427 A mapping of exposure IDs to ExposureF objects that correspond to 

428 psfMatchedWarp datasets. 

429 warpType : `str` 

430 Either 'direct' or 'psfMatched'. 

431 tract : `int`, optional 

432 Unique identifier for a tract of a skyMap. 

433 patch : `int`, optional 

434 Unique identifier for a subdivision of a tract. 

435 

436 Returns 

437 ------- 

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

439 The data references. 

440 

441 Raises 

442 ------ 

443 ValueError 

444 If an unknown `warpType` is supplied. 

445 """ 

446 dataRefList = [] 

447 for expId in exposures: 

448 if warpType == "direct": 

449 exposure = exposures[expId] 

450 elif warpType == "psfMatched": 

451 exposure = matchedExposures[expId] 

452 else: 

453 raise ValueError("warpType must be one of 'direct' or 'psfMatched'") 

454 dataRef = pipeBase.InMemoryDatasetHandle( 

455 exposure, 

456 storageClass="ExposureF", 

457 copy=True, 

458 dataId=generate_data_id( 

459 tract=tract, 

460 patch=patch, 

461 visit_id=expId, 

462 ), 

463 ) 

464 dataRefList.append(dataRef) 

465 return dataRefList