Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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 os 

25 

26import lsst.afw.fits as afwFits 

27import lsst.alert.packet as alertPack 

28import lsst.geom as geom 

29import lsst.pex.config as pexConfig 

30import lsst.pipe.base as pipeBase 

31from lsst.utils import getPackageDir 

32 

33 

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

35""" 

36 

37 

38class PackageAlertsConfig(pexConfig.Config): 

39 """Config class for AssociationTask. 

40 """ 

41 schemaFile = pexConfig.Field( 

42 dtype=str, 

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

44 default=os.path.join(getPackageDir("alert_packet"), 

45 "schema/latest/lsst.alert.avsc"), 

46 ) 

47 minCutoutSize = pexConfig.RangeField( 

48 dtype=int, 

49 min=0, 

50 max=1000, 

51 default=30, 

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

53 ) 

54 alertWriteLocation = pexConfig.Field( 

55 dtype=str, 

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

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

58 ) 

59 

60 

61class PackageAlertsTask(pipeBase.Task): 

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

63 """ 

64 ConfigClass = PackageAlertsConfig 

65 _DefaultName = "packageAlerts" 

66 

67 def __init__(self, **kwargs): 

68 super().__init__(**kwargs) 

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

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

71 

72 def run(self, 

73 diaSourceCat, 

74 diaObjectCat, 

75 diaSrcHistory, 

76 diffIm, 

77 template, 

78 ccdExposureIdBits): 

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

80 

81 Writes Avro alerts to a location determined by the 

82 ``alertWriteLocation`` configurable. 

83 

84 Parameters 

85 ---------- 

86 diaSourceCat : `pandas.DataFrame` 

87 New DiaSources to package. DataFrame should be indexed on 

88 ``["diaObjectId", "filterName", "diaSourceId"]`` 

89 diaObjectCat : `pandas.DataFrame` 

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

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

92 diaSrcHistory : `pandas.DataFrame` 

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

94 the newest DiaSource and is indexed on 

95 ``["diaObjectId", "filterName", "diaSourceId"]`` 

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

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

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

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

100 ccdExposureIdBits : `int` 

101 Number of bits used in the ccdVisitId. 

102 """ 

103 alerts = [] 

104 self._patchDiaSources(diaSourceCat) 

105 self._patchDiaSources(diaSrcHistory) 

106 self._patchDiaObjects(diaObjectCat) 

107 ccdVisitId = diffIm.getInfo().getVisitInfo().getExposureId() 

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

109 # Get all diaSources for the associated diaObject. 

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

111 if diaObject["nDiaSources"] > 1: 

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

113 else: 

114 objSourceHistory = None 

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

116 diaSource["decl"], 

117 geom.degrees) 

118 cutoutBBox = self.createDiaSourceBBox(diaSource["bboxSize"]) 

119 diffImCutout = diffIm.getCutout(sphPoint, cutoutBBox) 

120 templateCutout = None 

121 # TODO: Create alertIds DM-24858 

122 alertId = diaSource["diaSourceId"] 

123 alerts.append( 

124 self.makeAlertDict(alertId, 

125 diaSource, 

126 diaObject, 

127 objSourceHistory, 

128 diffImCutout, 

129 templateCutout)) 

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

131 f"{ccdVisitId}.avro"), 

132 "wb") as f: 

133 self.alertSchema.store_alerts(f, alerts) 

134 

135 def _patchDiaSources(self, diaSources): 

136 """Add the ``programId`` column to the data and change currently 

137 grouped alert data to ``None``. 

138 

139 TODO: The need to change these column values to ``None`` can be removed 

140 after DM-24696 is merged. 

141 

142 Parameters 

143 ---------- 

144 diaSources : `pandas.DataFrame` 

145 DataFrame of DiaSources to patch. 

146 """ 

147 diaSources["programId"] = 0 

148 diaSources["ra_decl_Cov"] = None 

149 diaSources["x_y_Cov"] = None 

150 diaSources["ps_Cov"] = None 

151 diaSources["trail_Cov"] = None 

152 diaSources["dip_Cov"] = None 

153 diaSources["i_cov"] = None 

154 

155 def _patchDiaObjects(self, diaObjects): 

156 """Change currently grouped alert data to ``None``. 

157 

158 TODO: The need to change these column values to ``None`` can be removed 

159 after DM-24696 is merged. 

160 

161 Parameters 

162 ---------- 

163 diaObjects : `pandas.DataFrame` 

164 DataFrame of DiaObjects to patch. 

165 """ 

166 diaObjects["ra_decl_Cov"] = None 

167 diaObjects["pm_parallax_Cov"] = None 

168 diaObjects["uLcPeriodic"] = None 

169 diaObjects["gLcPeriodic"] = None 

170 diaObjects["rLcPeriodic"] = None 

171 diaObjects["iLcPeriodic"] = None 

172 diaObjects["zLcPeriodic"] = None 

173 diaObjects["yLcPeriodic"] = None 

174 diaObjects["uLcNonPeriodic"] = None 

175 diaObjects["gLcNonPeriodic"] = None 

176 diaObjects["rLcNonPeriodic"] = None 

177 diaObjects["iLcNonPeriodic"] = None 

178 diaObjects["zLcNonPeriodic"] = None 

179 diaObjects["yLcNonPeriodic"] = None 

180 

181 def createDiaSourceBBox(self, bboxSize): 

182 """Create a bounding box for the cutouts given the size of the square 

183 BBox that covers the source footprint. 

184 

185 Parameters 

186 ---------- 

187 bboxSize : `int` 

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

189 

190 Returns 

191 ------- 

192 bbox : `lsst.geom.Extent2I` 

193 Geom object representing the size of the bounding box. 

194 """ 

195 if bboxSize < self.config.minCutoutSize: 

196 bbox = geom.Extent2I(self.config.minCutoutSize, 

197 self.config.minCutoutSize) 

198 else: 

199 bbox = geom.Extent2I(bboxSize, bboxSize) 

200 return bbox 

201 

202 def makeAlertDict(self, 

203 alertId, 

204 diaSource, 

205 diaObject, 

206 objDiaSrcHistory, 

207 diffImCutout, 

208 templateCutout): 

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

210 

211 Parameters 

212 ---------- 

213 diaSource : `pandas.DataFrame` 

214 New single DiaSource to package. 

215 diaObject : `pandas.DataFrame` 

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

217 objDiaSrcHistory : `pandas.DataFrame` 

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

219 diffImCutout : `lsst.afw.image.ExposureF` 

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

221 with a size set by the ``cutoutSize`` configurable. 

222 templateCutout : `lsst.afw.image.ExposureF` 

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

224 with a size set by the ``cutoutSize`` configurable. 

225 """ 

226 alert = dict() 

227 alert['alertId'] = alertId 

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

229 

230 if objDiaSrcHistory is None: 

231 alert['prvDiaSources'] = objDiaSrcHistory 

232 else: 

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

234 

235 alert['prvDiaForcedSources'] = None 

236 alert['prvDiaNondetectionLimits'] = None 

237 

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

239 

240 alert['ssObject'] = None 

241 

242 # TODO: fileName to be removed in future Avro schemas. DM-24696 

243 alert['cutoutDifference'] = { 

244 'fileName': '', 

245 'stampData': self.makeCutoutBytes(diffImCutout), 

246 } 

247 # TODO: add template cutouts in DM-24327 

248 alert["cutoutTemplate"] = None 

249 

250 return alert 

251 

252 def makeCutoutBytes(self, cutout): 

253 """Serialize a cutout into bytes. 

254 

255 Parameters 

256 ---------- 

257 cutout : `lsst.afw.image.ExposureF` 

258 Cutout to serialize. 

259 

260 Returns 

261 ------- 

262 coutputBytes : `bytes` 

263 Input cutout serialized into byte data. 

264 """ 

265 mgr = afwFits.MemFileManager() 

266 cutout.writeFits(mgr) 

267 return mgr.getData()