Coverage for tests/assembleCoaddTestUtils.py: 25%

161 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-04-21 02:15 -0700

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__ = ["MockWarpReference", "makeMockSkyInfo", "MockCoaddTestData"] 

46 

47 

48class MockWarpReference(pipeBase.InMemoryDatasetHandle): 

49 """Very simple object that looks like a Gen 3 data reference to a warped 

50 exposure. 

51 

52 Parameters 

53 ---------- 

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

55 The exposure to be retrieved by the data reference. 

56 """ 

57 def get(self, *, component=None, parameters=None): 

58 """Retrieve the specified dataset using the API of the Gen 3 Butler. 

59 

60 Parameters 

61 ---------- 

62 component : `str`, optional 

63 If supplied, return the named metadata of the exposure. 

64 parameters : `dict`, optional 

65 If supplied, use the parameters to modify the exposure, 

66 typically by taking a subset. 

67 

68 Returns 

69 ------- 

70 `lsst.afw.image.Exposure` or `lsst.afw.image.VisitInfo` 

71 or `lsst.meas.algorithms.SingleGaussianPsf` 

72 Either the exposure or its metadata, depending on the component 

73 requested. 

74 """ 

75 exp = super().get(component=component, parameters=parameters) 

76 if isinstance(exp, afwImage.ExposureF): 

77 exp = exp.clone() 

78 return exp 

79 

80 

81def makeMockSkyInfo(bbox, wcs, patch): 

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

83 

84 Parameters 

85 ---------- 

86 bbox : `lsst.geom.Box` 

87 Bounding box of the patch to be coadded. 

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

89 Coordinate system definition (wcs) for the exposure. 

90 

91 Returns 

92 ------- 

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

94 Patch geometry information. 

95 """ 

96 def getIndex(): 

97 return patch 

98 patchInfo = pipeBase.Struct(getIndex=getIndex) 

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

100 return skyInfo 

101 

102 

103class MockCoaddTestData: 

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

105 are realistic enough to test the image coaddition algorithms. 

106 

107 Notes 

108 ----- 

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

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

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

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

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

114 

115 Parameters 

116 ---------- 

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

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

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

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

121 backgroundLevel : `float`, optional 

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

123 seed : `int`, optional 

124 Seed value to initialize the random number generator. 

125 nSrc : `int`, optional 

126 Number of sources to simulate. 

127 fluxRange : `float`, optional 

128 Range in flux amplitude of the simulated sources. 

129 noiseLevel : `float`, optional 

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

131 sourceSigma : `float`, optional 

132 Average amplitude of the simulated sources, 

133 relative to ``noiseLevel`` 

134 minPsfSize : `float`, optional 

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

136 maxPsfSize : `float`, optional 

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

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

139 The plate scale of the simulated images. 

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

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

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

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

144 ccd : `int`, optional 

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

146 patch : `int`, optional 

147 Unique identifier for a subdivision of a tract. 

148 tract : `int`, optional 

149 Unique identifier for a tract of a skyMap. 

150 

151 Raises 

152 ------ 

153 ValueError 

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

155 This is due to `GaussianPsf` that is used by `lsst.meas.algorithms.testUtils.plantSources` 

156 lacking the option to specify the pixel origin. 

157 """ 

158 rotAngle = 0.*degrees 

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

160 filterLabel = None 

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

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

163 """ 

164 rngData = None 

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

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

167 """ 

168 rngMods = None 

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

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

171 """ 

172 kernelSize = None 

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

174 exposures = {} 

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

176 matchedExposures = {} 

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

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

179 """ 

180 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

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

184 detector = None 

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

186 

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

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

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

190 minPsfSize=1.5, maxPsfSize=3., 

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

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

193 self.ra = ra 

194 self.dec = dec 

195 self.pixelScale = pixelScale 

196 self.patch = patch 

197 self.tract = tract 

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

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

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

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

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

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

204 self.wcs = self.makeDummyWcs() 

205 

206 # Set up properties of the simulations 

207 nSigmaForKernel = 5 

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

209 

210 bufferSize = self.kernelSize//2 

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

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

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

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

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

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

217 

218 self.backgroundLevel = backgroundLevel 

219 self.noiseLevel = noiseLevel 

220 self.minPsfSize = minPsfSize 

221 self.maxPsfSize = maxPsfSize 

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

223 

224 def setDummyCoaddInputs(self, exposure, expId): 

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

226 processed using `warpAndPsfMatch`. 

227 

228 Parameters 

229 ---------- 

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

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

232 expId : `int` 

233 A unique identifier for the visit. 

234 """ 

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

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

237 

238 config = CoaddInputRecorderConfig() 

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

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

241 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

242 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

243 

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

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

246 

247 Parameters 

248 ---------- 

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

250 The simulated exposure. 

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

252 VisitInfo containing metadata for the exposure. 

253 expId : `int` 

254 A unique identifier for the visit. 

255 

256 Returns 

257 ------- 

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

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

260 """ 

261 tempExp = rawExposure.clone() 

262 tempExp.setWcs(self.wcs) 

263 

264 tempExp.setFilter(self.filterLabel) 

265 tempExp.setPhotoCalib(self.photoCalib) 

266 tempExp.getInfo().setVisitInfo(visitInfo) 

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

268 self.setDummyCoaddInputs(tempExp, expId) 

269 return tempExp 

270 

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

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

273 

274 Parameters 

275 ---------- 

276 rotAngle : `lsst.geom.Angle` 

277 Rotation of the CD matrix, East from North 

278 pixelScale : `lsst.geom.Angle` 

279 Pixel scale of the projection. 

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

281 Coordinates of the reference pixel of the wcs. 

282 flipX : `bool`, optional 

283 Flip the direction of increasing Right Ascension. 

284 

285 Returns 

286 ------- 

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

288 A wcs that matches the inputs. 

289 """ 

290 if rotAngle is None: 

291 rotAngle = self.rotAngle 

292 if pixelScale is None: 

293 pixelScale = self.pixelScale 

294 if crval is None: 

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

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

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

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

299 return wcs 

300 

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

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

303 

304 Parameters 

305 ---------- 

306 exposureId : `int`, optional 

307 Unique integer identifier for this observation. 

308 randomizeTime : `bool`, optional 

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

310 

311 Returns 

312 ------- 

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

314 VisitInfo for the exposure. 

315 """ 

316 lsstLat = -30.244639*u.degree 

317 lsstLon = -70.749417*u.degree 

318 lsstAlt = 2663.*u.m 

319 lsstTemperature = 20.*u.Celsius 

320 lsstHumidity = 40. # in percent 

321 lsstPressure = 73892.*u.pascal 

322 loc = EarthLocation(lat=lsstLat, 

323 lon=lsstLon, 

324 height=lsstAlt) 

325 

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

327 if randomizeTime: 

328 # Pick a random time within a 6 hour window 

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

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

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

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

333 obsInfo = makeObservationInfo(location=loc, 

334 detector_exposure_id=exposureId, 

335 datetime_begin=time, 

336 datetime_end=time, 

337 boresight_airmass=airmass, 

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

339 boresight_rotation_coord='sky', 

340 temperature=lsstTemperature, 

341 pressure=lsstPressure, 

342 relative_humidity=lsstHumidity, 

343 tracking_radec=radec, 

344 altaz_begin=radec.altaz, 

345 observation_type='science', 

346 ) 

347 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

348 return visitInfo 

349 

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

351 detectionSigma=5., badRegionBox=None): 

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

353 

354 Parameters 

355 ---------- 

356 expId : `int` 

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

358 noiseLevel : `float`, optional 

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

360 psfSize : `float`, optional 

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

362 backgroundLevel : `float`, optional 

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

364 detectionSigma : `float`, optional 

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

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

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

368 """ 

369 if backgroundLevel is None: 

370 backgroundLevel = self.backgroundLevel 

371 if noiseLevel is None: 

372 noiseLevel = 5. 

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

374 

375 if psfSize is None: 

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

377 nSrc = len(self.flux) 

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

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

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

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

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

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

384 coordList, addPoissonNoise=False) 

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

386 coordListPsfMatched, addPoissonNoise=False) 

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

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

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

390 noise -= np.median(noise) 

391 model.image.array += noise 

392 modelPsfMatched.image.array += noise 

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

394 detectionThreshold = self.backgroundLevel + detectionSigma*noiseLevel 

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

396 

397 if badRegionBox is not None: 

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

399 

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

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

402 return exposure, matchedExposure 

403 

404 @staticmethod 

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

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

407 retrieved using the Gen 3 Butler API. 

408 

409 Parameters 

410 ---------- 

411 warpType : `str` 

412 Either 'direct' or 'psfMatched'. 

413 tract : `int`, optional 

414 Unique identifier for a tract of a skyMap. 

415 patch : `int`, optional 

416 Unique identifier for a subdivision of a tract. 

417 coaddName : `str`, optional 

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

419 

420 Returns 

421 ------- 

422 dataRefList : `list` of `MockWarpReference` 

423 The data references. 

424 

425 Raises 

426 ------ 

427 ValueError 

428 If an unknown `warpType` is supplied. 

429 """ 

430 dataRefList = [] 

431 for expId in exposures: 

432 if warpType == 'direct': 

433 exposure = exposures[expId] 

434 elif warpType == 'psfMatched': 

435 exposure = matchedExposures[expId] 

436 else: 

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

438 dataRef = MockWarpReference(exposure, storageClass="ExposureF", 

439 tract=tract, patch=patch, visit=expId, coaddName=coaddName) 

440 dataRefList.append(dataRef) 

441 return dataRefList