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

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): 

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

91 

92 Writes Avro alerts to a location determined by the 

93 ``alertWriteLocation`` configurable. 

94 

95 Parameters 

96 ---------- 

97 diaSourceCat : `pandas.DataFrame` 

98 New DiaSources to package. DataFrame should be indexed on 

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

100 diaObjectCat : `pandas.DataFrame` 

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

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

103 diaSrcHistory : `pandas.DataFrame` 

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

105 the newest DiaSource and is indexed on 

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

107 diaForcedSources : `pandas.DataFrame` 

108 12 month history of DiaForcedSources matched to the DiaObjects. 

109 ``["diaObjectId"]`` 

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

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

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

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

114 ccdExposureIdBits : `int`, optional 

115 Unused. Deprecated and will be removed after v27. 

116 """ 

117 if ccdExposureIdBits is not None: 

118 warnings.warn( 

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

120 category=FutureWarning, 

121 ) 

122 alerts = [] 

123 self._patchDiaSources(diaSourceCat) 

124 self._patchDiaSources(diaSrcHistory) 

125 ccdVisitId = diffIm.info.id 

126 diffImPhotoCalib = diffIm.getPhotoCalib() 

127 templatePhotoCalib = template.getPhotoCalib() 

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

129 # Get all diaSources for the associated diaObject. 

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

131 # Objects for now. 

132 if srcIndex[0] == 0: 

133 continue 

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

135 if diaObject["nDiaSources"] > 1: 

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

137 else: 

138 objSourceHistory = None 

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

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

141 diaSource["dec"], 

142 geom.degrees) 

143 

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

145 diffImCutout = self.createCcdDataCutout( 

146 diffIm, 

147 sphPoint, 

148 cutoutExtent, 

149 diffImPhotoCalib, 

150 diaSource["diaSourceId"]) 

151 templateCutout = self.createCcdDataCutout( 

152 template, 

153 sphPoint, 

154 cutoutExtent, 

155 templatePhotoCalib, 

156 diaSource["diaSourceId"]) 

157 

158 # TODO: Create alertIds DM-24858 

159 alertId = diaSource["diaSourceId"] 

160 alerts.append( 

161 self.makeAlertDict(alertId, 

162 diaSource, 

163 diaObject, 

164 objSourceHistory, 

165 objDiaForcedSources, 

166 diffImCutout, 

167 templateCutout)) 

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

169 f"{ccdVisitId}.avro"), 

170 "wb") as f: 

171 self.alertSchema.store_alerts(f, alerts) 

172 

173 def _patchDiaSources(self, diaSources): 

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

175 

176 Parameters 

177 ---------- 

178 diaSources : `pandas.DataFrame` 

179 DataFrame of DiaSources to patch. 

180 """ 

181 diaSources["programId"] = 0 

182 

183 def createDiaSourceExtent(self, bboxSize): 

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

185 square BBox that covers the source footprint. 

186 

187 Parameters 

188 ---------- 

189 bboxSize : `int` 

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

191 

192 Returns 

193 ------- 

194 extent : `lsst.geom.Extent2I` 

195 Geom object representing the size of the bounding box. 

196 """ 

197 if bboxSize < self.config.minCutoutSize: 

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

199 self.config.minCutoutSize) 

200 else: 

201 extent = geom.Extent2I(bboxSize, bboxSize) 

202 return extent 

203 

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

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

206 

207 Parameters 

208 ---------- 

209 image : `lsst.afw.image.ExposureF` 

210 Image to pull cutout from. 

211 skyCenter : `lsst.geom.SpherePoint` 

212 Center point of DiaSource on the sky. 

213 extent : `lsst.geom.Extent2I` 

214 Bounding box to cutout from the image. 

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

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

217 srcId : `int` 

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

219 a cutout. 

220 

221 Returns 

222 ------- 

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

224 CCDData object storing the calibrate information from the input 

225 difference or template image. 

226 """ 

227 # Catch errors in retrieving the cutout. 

228 try: 

229 cutout = image.getCutout(skyCenter, extent) 

230 except InvalidParameterError: 

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

232 imBBox = image.getBBox() 

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

234 self.log.warning( 

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

236 "which is outside the Exposure with bounding box " 

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

238 srcId, point.x, point.y, 

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

240 else: 

241 raise InvalidParameterError( 

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

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

244 "creation. Exiting." 

245 % srcId) 

246 return None 

247 

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

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

250 # [1, 1]. 

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

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

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

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

255 

256 cutoutWcs = wcs.WCS(naxis=2) 

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

258 cutout.getBBox().getWidth()) 

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

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

261 skyCenter.getDec().asDegrees()] 

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

263 center, 

264 skyCenter) 

265 

266 return CCDData( 

267 data=calibCutout.getImage().array, 

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

269 flags=calibCutout.getMask().array, 

270 wcs=cutoutWcs, 

271 meta={"cutMinX": cutOutMinX, 

272 "cutMinY": cutOutMinY}, 

273 unit=u.nJy) 

274 

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

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

277 matrix. 

278 

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

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

281 is initially calculated with units arcseconds and then converted to 

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

283 

284 Parameters 

285 ---------- 

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

287 Wcs to approximate 

288 center : `lsst.geom.Point2D` 

289 Point at which to evaluate the LocalWcs. 

290 skyCenter : `lsst.geom.SpherePoint` 

291 Point on sky to approximate the Wcs. 

292 

293 Returns 

294 ------- 

295 localMatrix : `numpy.ndarray` 

296 Matrix representation the local wcs approximation with units 

297 degrees. 

298 """ 

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

300 localGnomonicWcs = afwGeom.makeSkyWcs( 

301 center, skyCenter, blankCDMatrix) 

302 measurementToLocalGnomonic = wcs.getTransform().then( 

303 localGnomonicWcs.getTransform().inverted() 

304 ) 

305 localMatrix = measurementToLocalGnomonic.getJacobian(center) 

306 return localMatrix / 3600 

307 

308 def makeAlertDict(self, 

309 alertId, 

310 diaSource, 

311 diaObject, 

312 objDiaSrcHistory, 

313 objDiaForcedSources, 

314 diffImCutout, 

315 templateCutout): 

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

317 

318 Parameters 

319 ---------- 

320 diaSource : `pandas.DataFrame` 

321 New single DiaSource to package. 

322 diaObject : `pandas.DataFrame` 

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

324 objDiaSrcHistory : `pandas.DataFrame` 

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

326 objDiaForcedSources : `pandas.DataFrame` 

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

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

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

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

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

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

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

334 """ 

335 alert = dict() 

336 alert['alertId'] = alertId 

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

338 

339 if objDiaSrcHistory is None: 

340 alert['prvDiaSources'] = objDiaSrcHistory 

341 else: 

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

343 

344 if isinstance(objDiaForcedSources, pd.Series): 

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

346 else: 

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

348 alert['prvDiaNondetectionLimits'] = None 

349 

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

351 

352 alert['ssObject'] = None 

353 

354 if diffImCutout is None: 

355 alert['cutoutDifference'] = None 

356 else: 

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

358 

359 if templateCutout is None: 

360 alert["cutoutTemplate"] = None 

361 else: 

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

363 

364 return alert 

365 

366 def streamCcdDataToBytes(self, cutout): 

367 """Serialize a cutout into bytes. 

368 

369 Parameters 

370 ---------- 

371 cutout : `astropy.nddata.CCDData` 

372 Cutout to serialize. 

373 

374 Returns 

375 ------- 

376 coutputBytes : `bytes` 

377 Input cutout serialized into byte data. 

378 """ 

379 with io.BytesIO() as streamer: 

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

381 cutoutBytes = streamer.getvalue() 

382 return cutoutBytes