Coverage for tests/assembleCoaddTestUtils.py: 24%

178 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-03-08 02:12 -0800

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.daf.butler 

37import lsst.geom as geom 

38from lsst.geom import arcseconds, degrees 

39from lsst.meas.algorithms.testUtils import plantSources 

40from lsst.obs.base import MakeRawVisitInfoViaObsInfo 

41import lsst.pipe.base as pipeBase 

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

43 

44from astro_metadata_translator import makeObservationInfo 

45 

46__all__ = ["MockWarpReference", "makeMockSkyInfo", "MockCoaddTestData"] 

47 

48 

49class MockWarpReference(lsst.daf.butler.DeferredDatasetHandle): 

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

51 exposure. 

52 

53 Parameters 

54 ---------- 

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

56 The exposure to be retrieved by the data reference. 

57 coaddName : `str` 

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

59 patch : `int` 

60 Unique identifier for a subdivision of a tract. 

61 tract : `int` 

62 Unique identifier for a tract of a skyMap 

63 visit : `int` 

64 Unique identifier for an observation, 

65 potentially consisting of multiple ccds. 

66 """ 

67 def __init__(self, exposure, coaddName='deep', patch=42, tract=0, visit=100): 

68 self.coaddName = coaddName 

69 self.exposure = exposure 

70 self.tract = tract 

71 self.patch = patch 

72 self.visit = visit 

73 

74 def get(self, bbox=None, component=None, parameters=None): 

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

76 

77 Parameters 

78 ---------- 

79 bbox : `lsst.geom.box.Box2I`, optional 

80 If supplied, retrieve only a subregion of the exposure. 

81 component : `str`, optional 

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

83 parameters : `dict`, optional 

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

85 typically by taking a subset. 

86 

87 Returns 

88 ------- 

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

90 or `lsst.meas.algorithms.SingleGaussianPsf` 

91 Either the exposure or its metadata, depending on the datasetType. 

92 """ 

93 if component == 'psf': 

94 return self.exposure.getPsf() 

95 elif component == 'visitInfo': 

96 return self.exposure.getInfo().getVisitInfo() 

97 if parameters is not None: 

98 if "bbox" in parameters: 

99 bbox = parameters["bbox"] 

100 exp = self.exposure.clone() 

101 if bbox is not None: 

102 return exp[bbox] 

103 else: 

104 return exp 

105 

106 @property 

107 def dataId(self): 

108 """Generate a valid data identifier. 

109 

110 Returns 

111 ------- 

112 dataId : `lsst.daf.butler.DataCoordinate` 

113 Data identifier dict for the patch. 

114 """ 

115 return lsst.daf.butler.DataCoordinate.standardize( 

116 tract=self.tract, 

117 patch=self.patch, 

118 visit=self.visit, 

119 instrument="DummyCam", 

120 skymap="Skymap", 

121 universe=lsst.daf.butler.DimensionUniverse(), 

122 ) 

123 

124 

125def makeMockSkyInfo(bbox, wcs, patch): 

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

127 

128 Parameters 

129 ---------- 

130 bbox : `lsst.geom.Box` 

131 Bounding box of the patch to be coadded. 

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

133 Coordinate system definition (wcs) for the exposure. 

134 

135 Returns 

136 ------- 

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

138 Patch geometry information. 

139 """ 

140 def getIndex(): 

141 return patch 

142 patchInfo = pipeBase.Struct(getIndex=getIndex) 

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

144 return skyInfo 

145 

146 

147class MockCoaddTestData: 

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

149 are realistic enough to test the image coaddition algorithms. 

150 

151 Notes 

152 ----- 

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

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

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

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

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

158 

159 Parameters 

160 ---------- 

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

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

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

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

165 backgroundLevel : `float`, optional 

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

167 seed : `int`, optional 

168 Seed value to initialize the random number generator. 

169 nSrc : `int`, optional 

170 Number of sources to simulate. 

171 fluxRange : `float`, optional 

172 Range in flux amplitude of the simulated sources. 

173 noiseLevel : `float`, optional 

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

175 sourceSigma : `float`, optional 

176 Average amplitude of the simulated sources, 

177 relative to ``noiseLevel`` 

178 minPsfSize : `float`, optional 

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

180 maxPsfSize : `float`, optional 

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

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

183 The plate scale of the simulated images. 

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

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

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

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

188 ccd : `int`, optional 

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

190 patch : `int`, optional 

191 Unique identifier for a subdivision of a tract. 

192 tract : `int`, optional 

193 Unique identifier for a tract of a skyMap. 

194 

195 Raises 

196 ------ 

197 ValueError 

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

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

200 lacking the option to specify the pixel origin. 

201 """ 

202 rotAngle = 0.*degrees 

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

204 filterLabel = None 

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

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

207 """ 

208 rngData = None 

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

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

211 """ 

212 rngMods = None 

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

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

215 """ 

216 kernelSize = None 

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

218 exposures = {} 

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

220 matchedExposures = {} 

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

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

223 """ 

224 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

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

228 detector = None 

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

230 

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

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

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

234 minPsfSize=1.5, maxPsfSize=3., 

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

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

237 self.ra = ra 

238 self.dec = dec 

239 self.pixelScale = pixelScale 

240 self.patch = patch 

241 self.tract = tract 

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

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

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

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

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

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

248 self.wcs = self.makeDummyWcs() 

249 

250 # Set up properties of the simulations 

251 nSigmaForKernel = 5 

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

253 

254 bufferSize = self.kernelSize//2 

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

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

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

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

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

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

261 

262 self.backgroundLevel = backgroundLevel 

263 self.noiseLevel = noiseLevel 

264 self.minPsfSize = minPsfSize 

265 self.maxPsfSize = maxPsfSize 

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

267 

268 def setDummyCoaddInputs(self, exposure, expId): 

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

270 processed using `warpAndPsfMatch`. 

271 

272 Parameters 

273 ---------- 

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

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

276 expId : `int` 

277 A unique identifier for the visit. 

278 """ 

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

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

281 

282 config = CoaddInputRecorderConfig() 

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

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

285 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

286 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

287 

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

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

290 

291 Parameters 

292 ---------- 

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

294 The simulated exposure. 

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

296 VisitInfo containing metadata for the exposure. 

297 expId : `int` 

298 A unique identifier for the visit. 

299 

300 Returns 

301 ------- 

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

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

304 """ 

305 tempExp = rawExposure.clone() 

306 tempExp.setWcs(self.wcs) 

307 

308 tempExp.setFilter(self.filterLabel) 

309 tempExp.setPhotoCalib(self.photoCalib) 

310 tempExp.getInfo().setVisitInfo(visitInfo) 

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

312 self.setDummyCoaddInputs(tempExp, expId) 

313 return tempExp 

314 

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

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

317 

318 Parameters 

319 ---------- 

320 rotAngle : `lsst.geom.Angle` 

321 Rotation of the CD matrix, East from North 

322 pixelScale : `lsst.geom.Angle` 

323 Pixel scale of the projection. 

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

325 Coordinates of the reference pixel of the wcs. 

326 flipX : `bool`, optional 

327 Flip the direction of increasing Right Ascension. 

328 

329 Returns 

330 ------- 

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

332 A wcs that matches the inputs. 

333 """ 

334 if rotAngle is None: 

335 rotAngle = self.rotAngle 

336 if pixelScale is None: 

337 pixelScale = self.pixelScale 

338 if crval is None: 

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

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

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

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

343 return wcs 

344 

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

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

347 

348 Parameters 

349 ---------- 

350 exposureId : `int`, optional 

351 Unique integer identifier for this observation. 

352 randomizeTime : `bool`, optional 

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

354 

355 Returns 

356 ------- 

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

358 VisitInfo for the exposure. 

359 """ 

360 lsstLat = -30.244639*u.degree 

361 lsstLon = -70.749417*u.degree 

362 lsstAlt = 2663.*u.m 

363 lsstTemperature = 20.*u.Celsius 

364 lsstHumidity = 40. # in percent 

365 lsstPressure = 73892.*u.pascal 

366 loc = EarthLocation(lat=lsstLat, 

367 lon=lsstLon, 

368 height=lsstAlt) 

369 

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

371 if randomizeTime: 

372 # Pick a random time within a 6 hour window 

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

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

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

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

377 obsInfo = makeObservationInfo(location=loc, 

378 detector_exposure_id=exposureId, 

379 datetime_begin=time, 

380 datetime_end=time, 

381 boresight_airmass=airmass, 

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

383 boresight_rotation_coord='sky', 

384 temperature=lsstTemperature, 

385 pressure=lsstPressure, 

386 relative_humidity=lsstHumidity, 

387 tracking_radec=radec, 

388 altaz_begin=radec.altaz, 

389 observation_type='science', 

390 ) 

391 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

392 return visitInfo 

393 

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

395 detectionSigma=5., badRegionBox=None): 

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

397 

398 Parameters 

399 ---------- 

400 expId : `int` 

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

402 noiseLevel : `float`, optional 

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

404 psfSize : `float`, optional 

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

406 backgroundLevel : `float`, optional 

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

408 detectionSigma : `float`, optional 

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

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

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

412 """ 

413 if backgroundLevel is None: 

414 backgroundLevel = self.backgroundLevel 

415 if noiseLevel is None: 

416 noiseLevel = 5. 

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

418 

419 if psfSize is None: 

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

421 nSrc = len(self.flux) 

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

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

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

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

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

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

428 coordList, addPoissonNoise=False) 

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

430 coordListPsfMatched, addPoissonNoise=False) 

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

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

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

434 noise -= np.median(noise) 

435 model.image.array += noise 

436 modelPsfMatched.image.array += noise 

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

438 detectionThreshold = self.backgroundLevel + detectionSigma*noiseLevel 

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

440 

441 if badRegionBox is not None: 

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

443 

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

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

446 return exposure, matchedExposure 

447 

448 @staticmethod 

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

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

451 retrieved using the Gen 3 Butler API. 

452 

453 Parameters 

454 ---------- 

455 warpType : `str` 

456 Either 'direct' or 'psfMatched'. 

457 tract : `int`, optional 

458 Unique identifier for a tract of a skyMap. 

459 patch : `int`, optional 

460 Unique identifier for a subdivision of a tract. 

461 coaddName : `str`, optional 

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

463 

464 Returns 

465 ------- 

466 dataRefList : `list` of `MockWarpReference` 

467 The data references. 

468 

469 Raises 

470 ------ 

471 ValueError 

472 If an unknown `warpType` is supplied. 

473 """ 

474 dataRefList = [] 

475 for expId in exposures: 

476 if warpType == 'direct': 

477 exposure = exposures[expId] 

478 elif warpType == 'psfMatched': 

479 exposure = matchedExposures[expId] 

480 else: 

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

482 dataRef = MockWarpReference(exposure, coaddName=coaddName, 

483 tract=tract, patch=patch, visit=expId) 

484 dataRefList.append(dataRef) 

485 return dataRefList