Coverage for tests/assemble_coadd_test_utils.py: 27%

154 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-03-01 15:41 +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 

42from lsst.skymap import Index2D, PatchInfo 

43 

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

45 

46 

47def makeMockSkyInfo(bbox, wcs, patch): 

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

49 

50 Parameters 

51 ---------- 

52 bbox : `lsst.geom.Box` 

53 Bounding box of the patch to be coadded. 

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

55 Coordinate system definition (wcs) for the exposure. 

56 

57 Returns 

58 ------- 

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

60 Patch geometry information. 

61 """ 

62 

63 patchInfo = PatchInfo( 

64 index=Index2D(0, 0), 

65 sequentialIndex=patch, 

66 innerBBox=bbox, 

67 outerBBox=bbox, 

68 tractWcs=wcs, 

69 numCellsPerPatchInner=1, 

70 cellInnerDimensions=(bbox.width, bbox.height), 

71 ) 

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

73 return skyInfo 

74 

75 

76class MockCoaddTestData: 

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

78 are realistic enough to test the image coaddition algorithms. 

79 

80 Notes 

81 ----- 

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

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

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

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

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

87 

88 Parameters 

89 ---------- 

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

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

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

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

94 backgroundLevel : `float`, optional 

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

96 seed : `int`, optional 

97 Seed value to initialize the random number generator. 

98 nSrc : `int`, optional 

99 Number of sources to simulate. 

100 fluxRange : `float`, optional 

101 Range in flux amplitude of the simulated sources. 

102 noiseLevel : `float`, optional 

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

104 sourceSigma : `float`, optional 

105 Average amplitude of the simulated sources, 

106 relative to ``noiseLevel`` 

107 minPsfSize : `float`, optional 

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

109 maxPsfSize : `float`, optional 

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

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

112 The plate scale of the simulated images. 

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

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

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

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

117 ccd : `int`, optional 

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

119 patch : `int`, optional 

120 Unique identifier for a subdivision of a tract. 

121 tract : `int`, optional 

122 Unique identifier for a tract of a skyMap. 

123 

124 Raises 

125 ------ 

126 ValueError 

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

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

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

130 lacking the option to specify the pixel origin. 

131 """ 

132 

133 rotAngle = 0.0 * degrees 

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

135 (`lsst.geom.Angle`). 

136 """ 

137 filterLabel = None 

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

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

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

141 """ 

142 rngData = None 

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

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

145 """ 

146 rngMods = None 

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

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

149 """ 

150 kernelSize = None 

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

152 exposures = {} 

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

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

155 """ 

156 matchedExposures = {} 

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

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

159 """ 

160 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10) 

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

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

163 """ 

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

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

166 the coaddTempExp. 

167 """ 

168 detector = None 

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

170 

171 def __init__( 

172 self, 

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

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

175 backgroundLevel=314.592, 

176 seed=42, 

177 nSrc=37, 

178 fluxRange=2.0, 

179 noiseLevel=5, 

180 sourceSigma=200.0, 

181 minPsfSize=1.5, 

182 maxPsfSize=3.0, 

183 pixelScale=0.2 * arcseconds, 

184 ra=209.0 * degrees, 

185 dec=-20.25 * degrees, 

186 ccd=37, 

187 patch=42, 

188 tract=0, 

189 ): 

190 self.ra = ra 

191 self.dec = dec 

192 self.pixelScale = pixelScale 

193 self.patch = patch 

194 self.tract = tract 

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

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

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

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

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

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

201 self.wcs = self.makeDummyWcs() 

202 

203 # Set up properties of the simulations 

204 nSigmaForKernel = 5 

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

206 

207 bufferSize = self.kernelSize // 2 

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

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

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

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

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

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

214 

215 self.backgroundLevel = backgroundLevel 

216 self.noiseLevel = noiseLevel 

217 self.minPsfSize = minPsfSize 

218 self.maxPsfSize = maxPsfSize 

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

220 

221 def setDummyCoaddInputs(self, exposure, expId): 

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

223 processed using `warpAndPsfMatch`. 

224 

225 Parameters 

226 ---------- 

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

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

229 expId : `int` 

230 A unique identifier for the visit. 

231 """ 

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

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

234 

235 config = CoaddInputRecorderConfig() 

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

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

238 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix) 

239 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix) 

240 

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

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

243 

244 Parameters 

245 ---------- 

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

247 The simulated exposure. 

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

249 VisitInfo containing metadata for the exposure. 

250 expId : `int` 

251 A unique identifier for the visit. 

252 

253 Returns 

254 ------- 

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

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

257 """ 

258 tempExp = rawExposure.clone() 

259 tempExp.setWcs(self.wcs) 

260 

261 tempExp.setFilter(self.filterLabel) 

262 tempExp.setPhotoCalib(self.photoCalib) 

263 tempExp.getInfo().setVisitInfo(visitInfo) 

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

265 self.setDummyCoaddInputs(tempExp, expId) 

266 return tempExp 

267 

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

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

270 

271 Parameters 

272 ---------- 

273 rotAngle : `lsst.geom.Angle` 

274 Rotation of the CD matrix, East from North 

275 pixelScale : `lsst.geom.Angle` 

276 Pixel scale of the projection. 

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

278 Coordinates of the reference pixel of the wcs. 

279 flipX : `bool`, optional 

280 Flip the direction of increasing Right Ascension. 

281 

282 Returns 

283 ------- 

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

285 A wcs that matches the inputs. 

286 """ 

287 if rotAngle is None: 

288 rotAngle = self.rotAngle 

289 if pixelScale is None: 

290 pixelScale = self.pixelScale 

291 if crval is None: 

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

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

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

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

296 return wcs 

297 

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

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

300 

301 Parameters 

302 ---------- 

303 exposureId : `int`, optional 

304 Unique integer identifier for this observation. 

305 randomizeTime : `bool`, optional 

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

307 

308 Returns 

309 ------- 

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

311 VisitInfo for the exposure. 

312 """ 

313 lsstLat = -30.244639 * u.degree 

314 lsstLon = -70.749417 * u.degree 

315 lsstAlt = 2663.0 * u.m 

316 lsstTemperature = 20.0 * u.Celsius 

317 lsstHumidity = 40.0 # in percent 

318 lsstPressure = 73892.0 * u.pascal 

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

320 

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

322 if randomizeTime: 

323 # Pick a random time within a 6 hour window 

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

325 radec = SkyCoord( 

326 dec=self.dec.asDegrees(), 

327 ra=self.ra.asDegrees(), 

328 unit="deg", 

329 obstime=time, 

330 frame="icrs", 

331 location=loc, 

332 ) 

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

334 obsInfo = makeObservationInfo( 

335 location=loc, 

336 detector_exposure_id=exposureId, 

337 datetime_begin=time, 

338 datetime_end=time, 

339 boresight_airmass=airmass, 

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

341 boresight_rotation_coord="sky", 

342 temperature=lsstTemperature, 

343 pressure=lsstPressure, 

344 relative_humidity=lsstHumidity, 

345 tracking_radec=radec, 

346 altaz_begin=radec.altaz, 

347 observation_type="science", 

348 ) 

349 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo) 

350 return visitInfo 

351 

352 def makeTestImage( 

353 self, 

354 expId, 

355 noiseLevel=None, 

356 psfSize=None, 

357 backgroundLevel=None, 

358 detectionSigma=5.0, 

359 badRegionBox=None, 

360 ): 

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

362 

363 Parameters 

364 ---------- 

365 expId : `int` 

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

367 noiseLevel : `float`, optional 

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

369 psfSize : `float`, optional 

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

371 backgroundLevel : `float`, optional 

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

373 detectionSigma : `float`, optional 

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

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

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

377 """ 

378 if backgroundLevel is None: 

379 backgroundLevel = self.backgroundLevel 

380 if noiseLevel is None: 

381 noiseLevel = 5.0 

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

383 

384 if psfSize is None: 

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

386 nSrc = len(self.flux) 

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

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

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

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

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

392 model = plantSources( 

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

394 ) 

395 modelPsfMatched = plantSources( 

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

397 ) 

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

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

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

401 noise -= np.median(noise) 

402 model.image.array += noise 

403 modelPsfMatched.image.array += noise 

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

405 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel 

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

407 

408 if badRegionBox is not None: 

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

410 

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

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

413 return exposure, matchedExposure 

414 

415 @staticmethod 

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

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

418 retrieved using the Gen 3 Butler API. 

419 

420 Parameters 

421 ---------- 

422 warpType : `str` 

423 Either 'direct' or 'psfMatched'. 

424 tract : `int`, optional 

425 Unique identifier for a tract of a skyMap. 

426 patch : `int`, optional 

427 Unique identifier for a subdivision of a tract. 

428 coaddName : `str`, optional 

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

430 

431 Returns 

432 ------- 

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

434 The data references. 

435 

436 Raises 

437 ------ 

438 ValueError 

439 If an unknown `warpType` is supplied. 

440 """ 

441 dataRefList = [] 

442 for expId in exposures: 

443 if warpType == "direct": 

444 exposure = exposures[expId] 

445 elif warpType == "psfMatched": 

446 exposure = matchedExposures[expId] 

447 else: 

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

449 dataRef = pipeBase.InMemoryDatasetHandle( 

450 exposure, 

451 storageClass="ExposureF", 

452 copy=True, 

453 tract=tract, 

454 patch=patch, 

455 visit=expId, 

456 coaddName=coaddName, 

457 ) 

458 dataRefList.append(dataRef) 

459 return dataRefList