Coverage for tests/assemble_coadd_test_utils.py: 26%

155 statements  

« prev     ^ index     » next       coverage.py v7.3.3, created at 2023-12-16 15:07 +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 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.geom import arcseconds, degrees 

39from lsst.meas.algorithms.testUtils import plantSources 

40from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

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

42 

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

44 

45 

46def makeMockSkyInfo(bbox, wcs, patch): 

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

48 

49 Parameters 

50 ---------- 

51 bbox : `lsst.geom.Box` 

52 Bounding box of the patch to be coadded. 

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

54 Coordinate system definition (wcs) for the exposure. 

55 

56 Returns 

57 ------- 

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

59 Patch geometry information. 

60 """ 

61 

62 def getIndex(): 

63 return patch 

64 

65 patchInfo = pipeBase.Struct(getIndex=getIndex) 

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

67 return skyInfo 

68 

69 

70class MockCoaddTestData: 

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

72 are realistic enough to test the image coaddition algorithms. 

73 

74 Notes 

75 ----- 

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

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

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

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

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

81 

82 Parameters 

83 ---------- 

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

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

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

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

88 backgroundLevel : `float`, optional 

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

90 seed : `int`, optional 

91 Seed value to initialize the random number generator. 

92 nSrc : `int`, optional 

93 Number of sources to simulate. 

94 fluxRange : `float`, optional 

95 Range in flux amplitude of the simulated sources. 

96 noiseLevel : `float`, optional 

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

98 sourceSigma : `float`, optional 

99 Average amplitude of the simulated sources, 

100 relative to ``noiseLevel`` 

101 minPsfSize : `float`, optional 

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

103 maxPsfSize : `float`, optional 

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

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

106 The plate scale of the simulated images. 

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

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

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

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

111 ccd : `int`, optional 

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

113 patch : `int`, optional 

114 Unique identifier for a subdivision of a tract. 

115 tract : `int`, optional 

116 Unique identifier for a tract of a skyMap. 

117 

118 Raises 

119 ------ 

120 ValueError 

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

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

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

124 lacking the option to specify the pixel origin. 

125 """ 

126 

127 rotAngle = 0.0 * degrees 

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

129 (`lsst.geom.Angle`). 

130 """ 

131 filterLabel = None 

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

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

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

135 """ 

136 rngData = None 

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

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

139 """ 

140 rngMods = None 

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

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

143 """ 

144 kernelSize = None 

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

146 exposures = {} 

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

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

149 """ 

150 matchedExposures = {} 

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

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

153 """ 

154 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

157 """ 

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

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

160 the coaddTempExp. 

161 """ 

162 detector = None 

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

164 

165 def __init__( 

166 self, 

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

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

169 backgroundLevel=314.592, 

170 seed=42, 

171 nSrc=37, 

172 fluxRange=2.0, 

173 noiseLevel=5, 

174 sourceSigma=200.0, 

175 minPsfSize=1.5, 

176 maxPsfSize=3.0, 

177 pixelScale=0.2 * arcseconds, 

178 ra=209.0 * degrees, 

179 dec=-20.25 * degrees, 

180 ccd=37, 

181 patch=42, 

182 tract=0, 

183 ): 

184 self.ra = ra 

185 self.dec = dec 

186 self.pixelScale = pixelScale 

187 self.patch = patch 

188 self.tract = tract 

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

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

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

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

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

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

195 self.wcs = self.makeDummyWcs() 

196 

197 # Set up properties of the simulations 

198 nSigmaForKernel = 5 

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

200 

201 bufferSize = self.kernelSize // 2 

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

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

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

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

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

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

208 

209 self.backgroundLevel = backgroundLevel 

210 self.noiseLevel = noiseLevel 

211 self.minPsfSize = minPsfSize 

212 self.maxPsfSize = maxPsfSize 

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

214 

215 def setDummyCoaddInputs(self, exposure, expId): 

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

217 processed using `warpAndPsfMatch`. 

218 

219 Parameters 

220 ---------- 

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

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

223 expId : `int` 

224 A unique identifier for the visit. 

225 """ 

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

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

228 

229 config = CoaddInputRecorderConfig() 

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

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

232 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

233 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

234 

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

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

237 

238 Parameters 

239 ---------- 

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

241 The simulated exposure. 

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

243 VisitInfo containing metadata for the exposure. 

244 expId : `int` 

245 A unique identifier for the visit. 

246 

247 Returns 

248 ------- 

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

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

251 """ 

252 tempExp = rawExposure.clone() 

253 tempExp.setWcs(self.wcs) 

254 

255 tempExp.setFilter(self.filterLabel) 

256 tempExp.setPhotoCalib(self.photoCalib) 

257 tempExp.getInfo().setVisitInfo(visitInfo) 

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

259 self.setDummyCoaddInputs(tempExp, expId) 

260 return tempExp 

261 

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

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

264 

265 Parameters 

266 ---------- 

267 rotAngle : `lsst.geom.Angle` 

268 Rotation of the CD matrix, East from North 

269 pixelScale : `lsst.geom.Angle` 

270 Pixel scale of the projection. 

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

272 Coordinates of the reference pixel of the wcs. 

273 flipX : `bool`, optional 

274 Flip the direction of increasing Right Ascension. 

275 

276 Returns 

277 ------- 

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

279 A wcs that matches the inputs. 

280 """ 

281 if rotAngle is None: 

282 rotAngle = self.rotAngle 

283 if pixelScale is None: 

284 pixelScale = self.pixelScale 

285 if crval is None: 

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

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

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

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

290 return wcs 

291 

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

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

294 

295 Parameters 

296 ---------- 

297 exposureId : `int`, optional 

298 Unique integer identifier for this observation. 

299 randomizeTime : `bool`, optional 

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

301 

302 Returns 

303 ------- 

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

305 VisitInfo for the exposure. 

306 """ 

307 lsstLat = -30.244639 * u.degree 

308 lsstLon = -70.749417 * u.degree 

309 lsstAlt = 2663.0 * u.m 

310 lsstTemperature = 20.0 * u.Celsius 

311 lsstHumidity = 40.0 # in percent 

312 lsstPressure = 73892.0 * u.pascal 

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

314 

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

316 if randomizeTime: 

317 # Pick a random time within a 6 hour window 

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

319 radec = SkyCoord( 

320 dec=self.dec.asDegrees(), 

321 ra=self.ra.asDegrees(), 

322 unit="deg", 

323 obstime=time, 

324 frame="icrs", 

325 location=loc, 

326 ) 

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

328 obsInfo = makeObservationInfo( 

329 location=loc, 

330 detector_exposure_id=exposureId, 

331 datetime_begin=time, 

332 datetime_end=time, 

333 boresight_airmass=airmass, 

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

335 boresight_rotation_coord="sky", 

336 temperature=lsstTemperature, 

337 pressure=lsstPressure, 

338 relative_humidity=lsstHumidity, 

339 tracking_radec=radec, 

340 altaz_begin=radec.altaz, 

341 observation_type="science", 

342 ) 

343 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

344 return visitInfo 

345 

346 def makeTestImage( 

347 self, 

348 expId, 

349 noiseLevel=None, 

350 psfSize=None, 

351 backgroundLevel=None, 

352 detectionSigma=5.0, 

353 badRegionBox=None, 

354 ): 

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

356 

357 Parameters 

358 ---------- 

359 expId : `int` 

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

361 noiseLevel : `float`, optional 

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

363 psfSize : `float`, optional 

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

365 backgroundLevel : `float`, optional 

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

367 detectionSigma : `float`, optional 

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

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

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

371 """ 

372 if backgroundLevel is None: 

373 backgroundLevel = self.backgroundLevel 

374 if noiseLevel is None: 

375 noiseLevel = 5.0 

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

377 

378 if psfSize is None: 

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

380 nSrc = len(self.flux) 

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

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

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

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

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

386 model = plantSources( 

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

388 ) 

389 modelPsfMatched = plantSources( 

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

391 ) 

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

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

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

395 noise -= np.median(noise) 

396 model.image.array += noise 

397 modelPsfMatched.image.array += noise 

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

399 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel 

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

401 

402 if badRegionBox is not None: 

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

404 

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

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

407 return exposure, matchedExposure 

408 

409 @staticmethod 

410 def makeDataRefList(exposures, matchedExposures, warpType, tract=0, patch=42, coaddName="deep"): 

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

412 retrieved using the Gen 3 Butler API. 

413 

414 Parameters 

415 ---------- 

416 warpType : `str` 

417 Either 'direct' or 'psfMatched'. 

418 tract : `int`, optional 

419 Unique identifier for a tract of a skyMap. 

420 patch : `int`, optional 

421 Unique identifier for a subdivision of a tract. 

422 coaddName : `str`, optional 

423 The type of coadd being produced. Typically 'deep'. 

424 

425 Returns 

426 ------- 

427 dataRefList : `list` of `~lsst.pipe.base.InMemoryDatasetHandle` 

428 The data references. 

429 

430 Raises 

431 ------ 

432 ValueError 

433 If an unknown `warpType` is supplied. 

434 """ 

435 dataRefList = [] 

436 for expId in exposures: 

437 if warpType == "direct": 

438 exposure = exposures[expId] 

439 elif warpType == "psfMatched": 

440 exposure = matchedExposures[expId] 

441 else: 

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

443 dataRef = pipeBase.InMemoryDatasetHandle( 

444 exposure, 

445 storageClass="ExposureF", 

446 copy=True, 

447 tract=tract, 

448 patch=patch, 

449 visit=expId, 

450 coaddName=coaddName, 

451 ) 

452 dataRefList.append(dataRef) 

453 return dataRefList