Coverage for tests/assembleCoaddTestUtils.py: 26%

155 statements  

« prev     ^ index     » next       coverage.py v7.3.1, created at 2023-09-07 10:59 +0000

1# This file is part of pipe_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 `lsst.meas.algorithms.testUtils.plantSources` 

123 lacking the option to specify the pixel origin. 

124 """ 

125 rotAngle = 0.*degrees 

126 "Rotation of the pixel grid on the sky, East from North (`lsst.geom.Angle`)." 

127 filterLabel = None 

128 """The filter definition, usually set in the current instruments' obs package. 

129 For these tests, a simple filter is defined without using an obs package (`lsst.afw.image.FilterLabel`). 

130 """ 

131 rngData = None 

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

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

134 """ 

135 rngMods = None 

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

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

138 """ 

139 kernelSize = None 

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

141 exposures = {} 

142 "The simulated test data, with variable PSF sizes (`dict` of `lsst.afw.image.Exposure`)" 

143 matchedExposures = {} 

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

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

146 """ 

147 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

148 "The photometric zero point to use for converting counts to flux units (`lsst.afw.image.PhotoCalib`)." 

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

150 "Mask planes that, if set, the associated pixel should not be included in the coaddTempExp." 

151 detector = None 

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

153 

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

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

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

157 minPsfSize=1.5, maxPsfSize=3., 

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

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

160 self.ra = ra 

161 self.dec = dec 

162 self.pixelScale = pixelScale 

163 self.patch = patch 

164 self.tract = tract 

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

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

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

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

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

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

171 self.wcs = self.makeDummyWcs() 

172 

173 # Set up properties of the simulations 

174 nSigmaForKernel = 5 

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

176 

177 bufferSize = self.kernelSize//2 

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

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

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

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

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

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

184 

185 self.backgroundLevel = backgroundLevel 

186 self.noiseLevel = noiseLevel 

187 self.minPsfSize = minPsfSize 

188 self.maxPsfSize = maxPsfSize 

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

190 

191 def setDummyCoaddInputs(self, exposure, expId): 

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

193 processed using `warpAndPsfMatch`. 

194 

195 Parameters 

196 ---------- 

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

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

199 expId : `int` 

200 A unique identifier for the visit. 

201 """ 

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

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

204 

205 config = CoaddInputRecorderConfig() 

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

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

208 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

209 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

210 

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

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

213 

214 Parameters 

215 ---------- 

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

217 The simulated exposure. 

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

219 VisitInfo containing metadata for the exposure. 

220 expId : `int` 

221 A unique identifier for the visit. 

222 

223 Returns 

224 ------- 

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

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

227 """ 

228 tempExp = rawExposure.clone() 

229 tempExp.setWcs(self.wcs) 

230 

231 tempExp.setFilter(self.filterLabel) 

232 tempExp.setPhotoCalib(self.photoCalib) 

233 tempExp.getInfo().setVisitInfo(visitInfo) 

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

235 self.setDummyCoaddInputs(tempExp, expId) 

236 return tempExp 

237 

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

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

240 

241 Parameters 

242 ---------- 

243 rotAngle : `lsst.geom.Angle` 

244 Rotation of the CD matrix, East from North 

245 pixelScale : `lsst.geom.Angle` 

246 Pixel scale of the projection. 

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

248 Coordinates of the reference pixel of the wcs. 

249 flipX : `bool`, optional 

250 Flip the direction of increasing Right Ascension. 

251 

252 Returns 

253 ------- 

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

255 A wcs that matches the inputs. 

256 """ 

257 if rotAngle is None: 

258 rotAngle = self.rotAngle 

259 if pixelScale is None: 

260 pixelScale = self.pixelScale 

261 if crval is None: 

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

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

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

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

266 return wcs 

267 

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

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

270 

271 Parameters 

272 ---------- 

273 exposureId : `int`, optional 

274 Unique integer identifier for this observation. 

275 randomizeTime : `bool`, optional 

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

277 

278 Returns 

279 ------- 

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

281 VisitInfo for the exposure. 

282 """ 

283 lsstLat = -30.244639*u.degree 

284 lsstLon = -70.749417*u.degree 

285 lsstAlt = 2663.*u.m 

286 lsstTemperature = 20.*u.Celsius 

287 lsstHumidity = 40. # in percent 

288 lsstPressure = 73892.*u.pascal 

289 loc = EarthLocation(lat=lsstLat, 

290 lon=lsstLon, 

291 height=lsstAlt) 

292 

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

294 if randomizeTime: 

295 # Pick a random time within a 6 hour window 

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

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

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

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

300 obsInfo = makeObservationInfo(location=loc, 

301 detector_exposure_id=exposureId, 

302 datetime_begin=time, 

303 datetime_end=time, 

304 boresight_airmass=airmass, 

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

306 boresight_rotation_coord='sky', 

307 temperature=lsstTemperature, 

308 pressure=lsstPressure, 

309 relative_humidity=lsstHumidity, 

310 tracking_radec=radec, 

311 altaz_begin=radec.altaz, 

312 observation_type='science', 

313 ) 

314 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

315 return visitInfo 

316 

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

318 detectionSigma=5., badRegionBox=None): 

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

320 

321 Parameters 

322 ---------- 

323 expId : `int` 

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

325 noiseLevel : `float`, optional 

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

327 psfSize : `float`, optional 

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

329 backgroundLevel : `float`, optional 

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

331 detectionSigma : `float`, optional 

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

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

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

335 """ 

336 if backgroundLevel is None: 

337 backgroundLevel = self.backgroundLevel 

338 if noiseLevel is None: 

339 noiseLevel = 5. 

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

341 

342 if psfSize is None: 

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

344 nSrc = len(self.flux) 

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

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

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

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

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

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

351 coordList, addPoissonNoise=False) 

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

353 coordListPsfMatched, addPoissonNoise=False) 

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

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

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

357 noise -= np.median(noise) 

358 model.image.array += noise 

359 modelPsfMatched.image.array += noise 

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

361 detectionThreshold = self.backgroundLevel + detectionSigma*noiseLevel 

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

363 

364 if badRegionBox is not None: 

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

366 

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

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

369 return exposure, matchedExposure 

370 

371 @staticmethod 

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

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

374 retrieved using the Gen 3 Butler API. 

375 

376 Parameters 

377 ---------- 

378 warpType : `str` 

379 Either 'direct' or 'psfMatched'. 

380 tract : `int`, optional 

381 Unique identifier for a tract of a skyMap. 

382 patch : `int`, optional 

383 Unique identifier for a subdivision of a tract. 

384 coaddName : `str`, optional 

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

386 

387 Returns 

388 ------- 

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

390 The data references. 

391 

392 Raises 

393 ------ 

394 ValueError 

395 If an unknown `warpType` is supplied. 

396 """ 

397 dataRefList = [] 

398 for expId in exposures: 

399 if warpType == 'direct': 

400 exposure = exposures[expId] 

401 elif warpType == 'psfMatched': 

402 exposure = matchedExposures[expId] 

403 else: 

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

405 dataRef = pipeBase.InMemoryDatasetHandle( 

406 exposure, 

407 storageClass="ExposureF", 

408 copy=True, 

409 tract=tract, 

410 patch=patch, 

411 visit=expId, 

412 coaddName=coaddName 

413 ) 

414 dataRefList.append(dataRef) 

415 return dataRefList