Coverage for tests/assemble_coadd_test_utils.py: 26%

155 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-12 12:13 +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""" 

28from astropy.time import Time 

29from astropy import units as u 

30from astropy.coordinates import SkyCoord, EarthLocation, Angle 

31import numpy as np 

32 

33from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

34import lsst.afw.geom as afwGeom 

35import lsst.afw.image as afwImage 

36import lsst.geom as geom 

37from lsst.geom import arcseconds, degrees 

38from lsst.meas.algorithms.testUtils import plantSources 

39from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

40import lsst.pipe.base as pipeBase 

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

42 

43from astro_metadata_translator import makeObservationInfo 

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 def getIndex(): 

64 return patch 

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 rotAngle = 0.*degrees 

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

128 (`lsst.geom.Angle`). 

129 """ 

130 filterLabel = None 

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

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

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

134 """ 

135 rngData = None 

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

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

138 """ 

139 rngMods = None 

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

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

142 """ 

143 kernelSize = None 

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

145 exposures = {} 

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

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

148 """ 

149 matchedExposures = {} 

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

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

152 """ 

153 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

156 """ 

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

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

159 the coaddTempExp. 

160 """ 

161 detector = None 

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

163 

164 def __init__(self, shape=geom.Extent2I(201, 301), offset=geom.Point2I(-123, -45), 

165 backgroundLevel=314.592, seed=42, nSrc=37, 

166 fluxRange=2., noiseLevel=5, sourceSigma=200., 

167 minPsfSize=1.5, maxPsfSize=3., 

168 pixelScale=0.2*arcseconds, ra=209.*degrees, dec=-20.25*degrees, 

169 ccd=37, patch=42, tract=0): 

170 self.ra = ra 

171 self.dec = dec 

172 self.pixelScale = pixelScale 

173 self.patch = patch 

174 self.tract = tract 

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

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

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

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

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

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

181 self.wcs = self.makeDummyWcs() 

182 

183 # Set up properties of the simulations 

184 nSigmaForKernel = 5 

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

186 

187 bufferSize = self.kernelSize//2 

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

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

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

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

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

193 self.flux = (self.rngData.random(nSrc)*(fluxRange - 1.) + 1.)*sourceSigma*noiseLevel 

194 

195 self.backgroundLevel = backgroundLevel 

196 self.noiseLevel = noiseLevel 

197 self.minPsfSize = minPsfSize 

198 self.maxPsfSize = maxPsfSize 

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

200 

201 def setDummyCoaddInputs(self, exposure, expId): 

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

203 processed using `warpAndPsfMatch`. 

204 

205 Parameters 

206 ---------- 

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

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

209 expId : `int` 

210 A unique identifier for the visit. 

211 """ 

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

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

214 

215 config = CoaddInputRecorderConfig() 

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

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

218 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

219 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

220 

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

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

223 

224 Parameters 

225 ---------- 

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

227 The simulated exposure. 

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

229 VisitInfo containing metadata for the exposure. 

230 expId : `int` 

231 A unique identifier for the visit. 

232 

233 Returns 

234 ------- 

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

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

237 """ 

238 tempExp = rawExposure.clone() 

239 tempExp.setWcs(self.wcs) 

240 

241 tempExp.setFilter(self.filterLabel) 

242 tempExp.setPhotoCalib(self.photoCalib) 

243 tempExp.getInfo().setVisitInfo(visitInfo) 

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

245 self.setDummyCoaddInputs(tempExp, expId) 

246 return tempExp 

247 

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

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

250 

251 Parameters 

252 ---------- 

253 rotAngle : `lsst.geom.Angle` 

254 Rotation of the CD matrix, East from North 

255 pixelScale : `lsst.geom.Angle` 

256 Pixel scale of the projection. 

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

258 Coordinates of the reference pixel of the wcs. 

259 flipX : `bool`, optional 

260 Flip the direction of increasing Right Ascension. 

261 

262 Returns 

263 ------- 

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

265 A wcs that matches the inputs. 

266 """ 

267 if rotAngle is None: 

268 rotAngle = self.rotAngle 

269 if pixelScale is None: 

270 pixelScale = self.pixelScale 

271 if crval is None: 

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

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

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

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

276 return wcs 

277 

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

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

280 

281 Parameters 

282 ---------- 

283 exposureId : `int`, optional 

284 Unique integer identifier for this observation. 

285 randomizeTime : `bool`, optional 

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

287 

288 Returns 

289 ------- 

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

291 VisitInfo for the exposure. 

292 """ 

293 lsstLat = -30.244639*u.degree 

294 lsstLon = -70.749417*u.degree 

295 lsstAlt = 2663.*u.m 

296 lsstTemperature = 20.*u.Celsius 

297 lsstHumidity = 40. # in percent 

298 lsstPressure = 73892.*u.pascal 

299 loc = EarthLocation(lat=lsstLat, 

300 lon=lsstLon, 

301 height=lsstAlt) 

302 

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

304 if randomizeTime: 

305 # Pick a random time within a 6 hour window 

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

307 radec = SkyCoord(dec=self.dec.asDegrees(), ra=self.ra.asDegrees(), 

308 unit='deg', obstime=time, frame='icrs', location=loc) 

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

310 obsInfo = makeObservationInfo(location=loc, 

311 detector_exposure_id=exposureId, 

312 datetime_begin=time, 

313 datetime_end=time, 

314 boresight_airmass=airmass, 

315 boresight_rotation_angle=Angle(0.*u.degree), 

316 boresight_rotation_coord='sky', 

317 temperature=lsstTemperature, 

318 pressure=lsstPressure, 

319 relative_humidity=lsstHumidity, 

320 tracking_radec=radec, 

321 altaz_begin=radec.altaz, 

322 observation_type='science', 

323 ) 

324 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

325 return visitInfo 

326 

327 def makeTestImage(self, expId, noiseLevel=None, psfSize=None, backgroundLevel=None, 

328 detectionSigma=5., badRegionBox=None): 

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

330 

331 Parameters 

332 ---------- 

333 expId : `int` 

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

335 noiseLevel : `float`, optional 

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

337 psfSize : `float`, optional 

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

339 backgroundLevel : `float`, optional 

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

341 detectionSigma : `float`, optional 

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

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

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

345 """ 

346 if backgroundLevel is None: 

347 backgroundLevel = self.backgroundLevel 

348 if noiseLevel is None: 

349 noiseLevel = 5. 

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

351 

352 if psfSize is None: 

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

354 nSrc = len(self.flux) 

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

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

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

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

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

360 model = plantSources(self.bbox, self.kernelSize, self.backgroundLevel, 

361 coordList, addPoissonNoise=False) 

362 modelPsfMatched = plantSources(self.bbox, self.kernelSize, self.backgroundLevel, 

363 coordListPsfMatched, addPoissonNoise=False) 

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

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

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

367 noise -= np.median(noise) 

368 model.image.array += noise 

369 modelPsfMatched.image.array += noise 

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

371 detectionThreshold = self.backgroundLevel + detectionSigma*noiseLevel 

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

373 

374 if badRegionBox is not None: 

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

376 

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

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

379 return exposure, matchedExposure 

380 

381 @staticmethod 

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

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

384 retrieved using the Gen 3 Butler API. 

385 

386 Parameters 

387 ---------- 

388 warpType : `str` 

389 Either 'direct' or 'psfMatched'. 

390 tract : `int`, optional 

391 Unique identifier for a tract of a skyMap. 

392 patch : `int`, optional 

393 Unique identifier for a subdivision of a tract. 

394 coaddName : `str`, optional 

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

396 

397 Returns 

398 ------- 

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

400 The data references. 

401 

402 Raises 

403 ------ 

404 ValueError 

405 If an unknown `warpType` is supplied. 

406 """ 

407 dataRefList = [] 

408 for expId in exposures: 

409 if warpType == 'direct': 

410 exposure = exposures[expId] 

411 elif warpType == 'psfMatched': 

412 exposure = matchedExposures[expId] 

413 else: 

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

415 dataRef = pipeBase.InMemoryDatasetHandle( 

416 exposure, 

417 storageClass="ExposureF", 

418 copy=True, 

419 tract=tract, 

420 patch=patch, 

421 visit=expId, 

422 coaddName=coaddName 

423 ) 

424 dataRefList.append(dataRef) 

425 return dataRefList