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

112 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-25 11:11 +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 ) 

123 alerts = [] 

124 self._patchDiaSources(diaSourceCat) 

125 self._patchDiaSources(diaSrcHistory) 

126 ccdVisitId = diffIm.info.id 

127 diffImPhotoCalib = diffIm.getPhotoCalib() 

128 templatePhotoCalib = template.getPhotoCalib() 

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

130 # Get all diaSources for the associated diaObject. 

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

132 # Objects for now. 

133 if srcIndex[0] == 0: 

134 continue 

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

136 if diaObject["nDiaSources"] > 1: 

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

138 else: 

139 objSourceHistory = None 

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

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

142 diaSource["dec"], 

143 geom.degrees) 

144 

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

146 diffImCutout = self.createCcdDataCutout( 

147 diffIm, 

148 sphPoint, 

149 cutoutExtent, 

150 diffImPhotoCalib, 

151 diaSource["diaSourceId"]) 

152 templateCutout = self.createCcdDataCutout( 

153 template, 

154 sphPoint, 

155 cutoutExtent, 

156 templatePhotoCalib, 

157 diaSource["diaSourceId"]) 

158 

159 # TODO: Create alertIds DM-24858 

160 alertId = diaSource["diaSourceId"] 

161 alerts.append( 

162 self.makeAlertDict(alertId, 

163 diaSource, 

164 diaObject, 

165 objSourceHistory, 

166 objDiaForcedSources, 

167 diffImCutout, 

168 templateCutout)) 

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

170 f"{ccdVisitId}.avro"), 

171 "wb") as f: 

172 self.alertSchema.store_alerts(f, alerts) 

173 

174 def _patchDiaSources(self, diaSources): 

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

176 

177 Parameters 

178 ---------- 

179 diaSources : `pandas.DataFrame` 

180 DataFrame of DiaSources to patch. 

181 """ 

182 diaSources["programId"] = 0 

183 

184 def createDiaSourceExtent(self, bboxSize): 

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

186 square BBox that covers the source footprint. 

187 

188 Parameters 

189 ---------- 

190 bboxSize : `int` 

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

192 

193 Returns 

194 ------- 

195 extent : `lsst.geom.Extent2I` 

196 Geom object representing the size of the bounding box. 

197 """ 

198 if bboxSize < self.config.minCutoutSize: 

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

200 self.config.minCutoutSize) 

201 else: 

202 extent = geom.Extent2I(bboxSize, bboxSize) 

203 return extent 

204 

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

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

207 

208 Parameters 

209 ---------- 

210 image : `lsst.afw.image.ExposureF` 

211 Image to pull cutout from. 

212 skyCenter : `lsst.geom.SpherePoint` 

213 Center point of DiaSource on the sky. 

214 extent : `lsst.geom.Extent2I` 

215 Bounding box to cutout from the image. 

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

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

218 srcId : `int` 

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

220 a cutout. 

221 

222 Returns 

223 ------- 

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

225 CCDData object storing the calibrate information from the input 

226 difference or template image. 

227 """ 

228 # Catch errors in retrieving the cutout. 

229 try: 

230 cutout = image.getCutout(skyCenter, extent) 

231 except InvalidParameterError: 

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

233 imBBox = image.getBBox() 

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

235 self.log.warning( 

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

237 "which is outside the Exposure with bounding box " 

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

239 srcId, point.x, point.y, 

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

241 else: 

242 raise InvalidParameterError( 

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

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

245 "creation. Exiting." 

246 % srcId) 

247 return None 

248 

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

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

251 # [1, 1]. 

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

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

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

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

256 

257 cutoutWcs = wcs.WCS(naxis=2) 

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

259 cutout.getBBox().getWidth()) 

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

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

262 skyCenter.getDec().asDegrees()] 

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

264 center, 

265 skyCenter) 

266 

267 return CCDData( 

268 data=calibCutout.getImage().array, 

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

270 flags=calibCutout.getMask().array, 

271 wcs=cutoutWcs, 

272 meta={"cutMinX": cutOutMinX, 

273 "cutMinY": cutOutMinY}, 

274 unit=u.nJy) 

275 

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

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

278 matrix. 

279 

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

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

282 is initially calculated with units arcseconds and then converted to 

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

284 

285 Parameters 

286 ---------- 

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

288 Wcs to approximate 

289 center : `lsst.geom.Point2D` 

290 Point at which to evaluate the LocalWcs. 

291 skyCenter : `lsst.geom.SpherePoint` 

292 Point on sky to approximate the Wcs. 

293 

294 Returns 

295 ------- 

296 localMatrix : `numpy.ndarray` 

297 Matrix representation the local wcs approximation with units 

298 degrees. 

299 """ 

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

301 localGnomonicWcs = afwGeom.makeSkyWcs( 

302 center, skyCenter, blankCDMatrix) 

303 measurementToLocalGnomonic = wcs.getTransform().then( 

304 localGnomonicWcs.getTransform().inverted() 

305 ) 

306 localMatrix = measurementToLocalGnomonic.getJacobian(center) 

307 return localMatrix / 3600 

308 

309 def makeAlertDict(self, 

310 alertId, 

311 diaSource, 

312 diaObject, 

313 objDiaSrcHistory, 

314 objDiaForcedSources, 

315 diffImCutout, 

316 templateCutout): 

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

318 

319 Parameters 

320 ---------- 

321 diaSource : `pandas.DataFrame` 

322 New single DiaSource to package. 

323 diaObject : `pandas.DataFrame` 

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

325 objDiaSrcHistory : `pandas.DataFrame` 

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

327 objDiaForcedSources : `pandas.DataFrame` 

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

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

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

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

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

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

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

335 """ 

336 alert = dict() 

337 alert['alertId'] = alertId 

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

339 

340 if objDiaSrcHistory is None: 

341 alert['prvDiaSources'] = objDiaSrcHistory 

342 else: 

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

344 

345 if isinstance(objDiaForcedSources, pd.Series): 

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

347 else: 

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

349 alert['prvDiaNondetectionLimits'] = None 

350 

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

352 

353 alert['ssObject'] = None 

354 

355 if diffImCutout is None: 

356 alert['cutoutDifference'] = None 

357 else: 

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

359 

360 if templateCutout is None: 

361 alert["cutoutTemplate"] = None 

362 else: 

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

364 

365 return alert 

366 

367 def streamCcdDataToBytes(self, cutout): 

368 """Serialize a cutout into bytes. 

369 

370 Parameters 

371 ---------- 

372 cutout : `astropy.nddata.CCDData` 

373 Cutout to serialize. 

374 

375 Returns 

376 ------- 

377 coutputBytes : `bytes` 

378 Input cutout serialized into byte data. 

379 """ 

380 with io.BytesIO() as streamer: 

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

382 cutoutBytes = streamer.getvalue() 

383 return cutoutBytes