Coverage for python/lsst/ap/association/packageAlerts.py: 25%

112 statements  

« prev     ^ index     » next       coverage.py v7.3.0, created at 2023-08-17 11:17 +0000

1# This file is part of ap_association. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22__all__ = ("PackageAlertsConfig", "PackageAlertsTask") 

23 

24import io 

25import os 

26import warnings 

27 

28from astropy import wcs 

29import astropy.units as u 

30from astropy.nddata import CCDData, VarianceUncertainty 

31import pandas as pd 

32 

33import lsst.alert.packet as alertPack 

34import lsst.afw.geom as afwGeom 

35import lsst.geom as geom 

36import lsst.pex.config as pexConfig 

37from lsst.pex.exceptions import InvalidParameterError 

38import lsst.pipe.base as pipeBase 

39from lsst.utils.timer import timeMethod 

40 

41 

42"""Methods for packaging Apdb and Pipelines data into Avro alerts. 

43""" 

44 

45 

46class PackageAlertsConfig(pexConfig.Config): 

47 """Config class for AssociationTask. 

48 """ 

49 schemaFile = pexConfig.Field( 

50 dtype=str, 

51 doc="Schema definition file for the avro alerts.", 

52 default=alertPack.get_path_to_latest_schema() 

53 ) 

54 minCutoutSize = pexConfig.RangeField( 

55 dtype=int, 

56 min=0, 

57 max=1000, 

58 default=30, 

59 doc="Dimension of the square image cutouts to package in the alert." 

60 ) 

61 alertWriteLocation = pexConfig.Field( 

62 dtype=str, 

63 doc="Location to write alerts to.", 

64 default=os.path.join(os.getcwd(), "alerts"), 

65 ) 

66 

67 

68class PackageAlertsTask(pipeBase.Task): 

69 """Tasks for packaging Dia and Pipelines data into Avro alert packages. 

70 """ 

71 ConfigClass = PackageAlertsConfig 

72 _DefaultName = "packageAlerts" 

73 

74 _scale = (1.0 * geom.arcseconds).asDegrees() 

75 

76 def __init__(self, **kwargs): 

77 super().__init__(**kwargs) 

78 self.alertSchema = alertPack.Schema.from_file(self.config.schemaFile) 

79 os.makedirs(self.config.alertWriteLocation, exist_ok=True) 

80 

81 @timeMethod 

82 def run(self, 

83 diaSourceCat, 

84 diaObjectCat, 

85 diaSrcHistory, 

86 diaForcedSources, 

87 diffIm, 

88 template, 

89 ccdExposureIdBits=None, # TODO: remove (including docs) on DM-38687. 

90 ): 

91 """Package DiaSources/Object and exposure data into Avro alerts. 

92 

93 Writes Avro alerts to a location determined by the 

94 ``alertWriteLocation`` configurable. 

95 

96 Parameters 

97 ---------- 

98 diaSourceCat : `pandas.DataFrame` 

99 New DiaSources to package. DataFrame should be indexed on 

100 ``["diaObjectId", "band", "diaSourceId"]`` 

101 diaObjectCat : `pandas.DataFrame` 

102 New and updated DiaObjects matched to the new DiaSources. DataFrame 

103 is indexed on ``["diaObjectId"]`` 

104 diaSrcHistory : `pandas.DataFrame` 

105 12 month history of DiaSources matched to the DiaObjects. Excludes 

106 the newest DiaSource and is indexed on 

107 ``["diaObjectId", "band", "diaSourceId"]`` 

108 diaForcedSources : `pandas.DataFrame` 

109 12 month history of DiaForcedSources matched to the DiaObjects. 

110 ``["diaObjectId"]`` 

111 diffIm : `lsst.afw.image.ExposureF` 

112 Difference image the sources in ``diaSourceCat`` were detected in. 

113 template : `lsst.afw.image.ExposureF` or `None` 

114 Template image used to create the ``diffIm``. 

115 ccdExposureIdBits : `int`, optional 

116 Unused. Deprecated and will be removed after v26. 

117 """ 

118 if ccdExposureIdBits is not None: 

119 warnings.warn( 

120 "The 'ccdExposureIdBits' argument is deprecated and unused; it will be removed after v26.", 

121 category=FutureWarning, 

122 stacklevel=3, # Caller + timeMethod 

123 ) 

124 alerts = [] 

125 self._patchDiaSources(diaSourceCat) 

126 self._patchDiaSources(diaSrcHistory) 

127 ccdVisitId = diffIm.info.id 

128 diffImPhotoCalib = diffIm.getPhotoCalib() 

129 templatePhotoCalib = template.getPhotoCalib() 

130 for srcIndex, diaSource in diaSourceCat.iterrows(): 

131 # Get all diaSources for the associated diaObject. 

132 # TODO: DM-31992 skip DiaSources associated with Solar System 

133 # Objects for now. 

134 if srcIndex[0] == 0: 

135 continue 

136 diaObject = diaObjectCat.loc[srcIndex[0]] 

137 if diaObject["nDiaSources"] > 1: 

138 objSourceHistory = diaSrcHistory.loc[srcIndex[0]] 

139 else: 

140 objSourceHistory = None 

141 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]] 

142 sphPoint = geom.SpherePoint(diaSource["ra"], 

143 diaSource["dec"], 

144 geom.degrees) 

145 

146 cutoutExtent = self.createDiaSourceExtent(diaSource["bboxSize"]) 

147 diffImCutout = self.createCcdDataCutout( 

148 diffIm, 

149 sphPoint, 

150 cutoutExtent, 

151 diffImPhotoCalib, 

152 diaSource["diaSourceId"]) 

153 templateCutout = self.createCcdDataCutout( 

154 template, 

155 sphPoint, 

156 cutoutExtent, 

157 templatePhotoCalib, 

158 diaSource["diaSourceId"]) 

159 

160 # TODO: Create alertIds DM-24858 

161 alertId = diaSource["diaSourceId"] 

162 alerts.append( 

163 self.makeAlertDict(alertId, 

164 diaSource, 

165 diaObject, 

166 objSourceHistory, 

167 objDiaForcedSources, 

168 diffImCutout, 

169 templateCutout)) 

170 with open(os.path.join(self.config.alertWriteLocation, 

171 f"{ccdVisitId}.avro"), 

172 "wb") as f: 

173 self.alertSchema.store_alerts(f, alerts) 

174 

175 def _patchDiaSources(self, diaSources): 

176 """Add the ``programId`` column to the data. 

177 

178 Parameters 

179 ---------- 

180 diaSources : `pandas.DataFrame` 

181 DataFrame of DiaSources to patch. 

182 """ 

183 diaSources["programId"] = 0 

184 

185 def createDiaSourceExtent(self, bboxSize): 

186 """Create a extent for a box for the cutouts given the size of the 

187 square BBox that covers the source footprint. 

188 

189 Parameters 

190 ---------- 

191 bboxSize : `int` 

192 Size of a side of the square bounding box in pixels. 

193 

194 Returns 

195 ------- 

196 extent : `lsst.geom.Extent2I` 

197 Geom object representing the size of the bounding box. 

198 """ 

199 if bboxSize < self.config.minCutoutSize: 

200 extent = geom.Extent2I(self.config.minCutoutSize, 

201 self.config.minCutoutSize) 

202 else: 

203 extent = geom.Extent2I(bboxSize, bboxSize) 

204 return extent 

205 

206 def createCcdDataCutout(self, image, skyCenter, extent, photoCalib, srcId): 

207 """Grab an image as a cutout and return a calibrated CCDData image. 

208 

209 Parameters 

210 ---------- 

211 image : `lsst.afw.image.ExposureF` 

212 Image to pull cutout from. 

213 skyCenter : `lsst.geom.SpherePoint` 

214 Center point of DiaSource on the sky. 

215 extent : `lsst.geom.Extent2I` 

216 Bounding box to cutout from the image. 

217 photoCalib : `lsst.afw.image.PhotoCalib` 

218 Calibrate object of the image the cutout is cut from. 

219 srcId : `int` 

220 Unique id of DiaSource. Used for when an error occurs extracting 

221 a cutout. 

222 

223 Returns 

224 ------- 

225 ccdData : `astropy.nddata.CCDData` or `None` 

226 CCDData object storing the calibrate information from the input 

227 difference or template image. 

228 """ 

229 # Catch errors in retrieving the cutout. 

230 try: 

231 cutout = image.getCutout(skyCenter, extent) 

232 except InvalidParameterError: 

233 point = image.getWcs().skyToPixel(skyCenter) 

234 imBBox = image.getBBox() 

235 if not geom.Box2D(image.getBBox()).contains(point): 

236 self.log.warning( 

237 "DiaSource id=%i centroid lies at pixel (%.2f, %.2f) " 

238 "which is outside the Exposure with bounding box " 

239 "((%i, %i), (%i, %i)). Returning None for cutout...", 

240 srcId, point.x, point.y, 

241 imBBox.minX, imBBox.maxX, imBBox.minY, imBBox.maxY) 

242 else: 

243 raise InvalidParameterError( 

244 "Failed to retrieve cutout from image for DiaSource with " 

245 "id=%i. InvalidParameterError thrown during cutout " 

246 "creation. Exiting." 

247 % srcId) 

248 return None 

249 

250 # Find the value of the bottom corner of our cutout's BBox and 

251 # subtract 1 so that the CCDData cutout position value will be 

252 # [1, 1]. 

253 cutOutMinX = cutout.getBBox().minX - 1 

254 cutOutMinY = cutout.getBBox().minY - 1 

255 center = cutout.getWcs().skyToPixel(skyCenter) 

256 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage()) 

257 

258 cutoutWcs = wcs.WCS(naxis=2) 

259 cutoutWcs.array_shape = (cutout.getBBox().getWidth(), 

260 cutout.getBBox().getWidth()) 

261 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY] 

262 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(), 

263 skyCenter.getDec().asDegrees()] 

264 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(), 

265 center, 

266 skyCenter) 

267 

268 return CCDData( 

269 data=calibCutout.getImage().array, 

270 uncertainty=VarianceUncertainty(calibCutout.getVariance().array), 

271 flags=calibCutout.getMask().array, 

272 wcs=cutoutWcs, 

273 meta={"cutMinX": cutOutMinX, 

274 "cutMinY": cutOutMinY}, 

275 unit=u.nJy) 

276 

277 def makeLocalTransformMatrix(self, wcs, center, skyCenter): 

278 """Create a local, linear approximation of the wcs transformation 

279 matrix. 

280 

281 The approximation is created as if the center is at RA=0, DEC=0. All 

282 comparing x,y coordinate are relative to the position of center. Matrix 

283 is initially calculated with units arcseconds and then converted to 

284 degrees. This yields higher precision results due to quirks in AST. 

285 

286 Parameters 

287 ---------- 

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

289 Wcs to approximate 

290 center : `lsst.geom.Point2D` 

291 Point at which to evaluate the LocalWcs. 

292 skyCenter : `lsst.geom.SpherePoint` 

293 Point on sky to approximate the Wcs. 

294 

295 Returns 

296 ------- 

297 localMatrix : `numpy.ndarray` 

298 Matrix representation the local wcs approximation with units 

299 degrees. 

300 """ 

301 blankCDMatrix = [[self._scale, 0], [0, self._scale]] 

302 localGnomonicWcs = afwGeom.makeSkyWcs( 

303 center, skyCenter, blankCDMatrix) 

304 measurementToLocalGnomonic = wcs.getTransform().then( 

305 localGnomonicWcs.getTransform().inverted() 

306 ) 

307 localMatrix = measurementToLocalGnomonic.getJacobian(center) 

308 return localMatrix / 3600 

309 

310 def makeAlertDict(self, 

311 alertId, 

312 diaSource, 

313 diaObject, 

314 objDiaSrcHistory, 

315 objDiaForcedSources, 

316 diffImCutout, 

317 templateCutout): 

318 """Convert data and package into a dictionary alert. 

319 

320 Parameters 

321 ---------- 

322 diaSource : `pandas.DataFrame` 

323 New single DiaSource to package. 

324 diaObject : `pandas.DataFrame` 

325 DiaObject that ``diaSource`` is matched to. 

326 objDiaSrcHistory : `pandas.DataFrame` 

327 12 month history of ``diaObject`` excluding the latest DiaSource. 

328 objDiaForcedSources : `pandas.DataFrame` 

329 12 month history of ``diaObject`` forced measurements. 

330 diffImCutout : `astropy.nddata.CCDData` or `None` 

331 Cutout of the difference image around the location of ``diaSource`` 

332 with a min size set by the ``cutoutSize`` configurable. 

333 templateCutout : `astropy.nddata.CCDData` or `None` 

334 Cutout of the template image around the location of ``diaSource`` 

335 with a min size set by the ``cutoutSize`` configurable. 

336 """ 

337 alert = dict() 

338 alert['alertId'] = alertId 

339 alert['diaSource'] = diaSource.to_dict() 

340 

341 if objDiaSrcHistory is None: 

342 alert['prvDiaSources'] = objDiaSrcHistory 

343 else: 

344 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records") 

345 

346 if isinstance(objDiaForcedSources, pd.Series): 

347 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()] 

348 else: 

349 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records") 

350 alert['prvDiaNondetectionLimits'] = None 

351 

352 alert['diaObject'] = diaObject.to_dict() 

353 

354 alert['ssObject'] = None 

355 

356 if diffImCutout is None: 

357 alert['cutoutDifference'] = None 

358 else: 

359 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout) 

360 

361 if templateCutout is None: 

362 alert["cutoutTemplate"] = None 

363 else: 

364 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout) 

365 

366 return alert 

367 

368 def streamCcdDataToBytes(self, cutout): 

369 """Serialize a cutout into bytes. 

370 

371 Parameters 

372 ---------- 

373 cutout : `astropy.nddata.CCDData` 

374 Cutout to serialize. 

375 

376 Returns 

377 ------- 

378 coutputBytes : `bytes` 

379 Input cutout serialized into byte data. 

380 """ 

381 with io.BytesIO() as streamer: 

382 cutout.write(streamer, format="fits") 

383 cutoutBytes = streamer.getvalue() 

384 return cutoutBytes