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", 

46 *[str(x) for x in alertPack.get_latest_schema_version()], 

47 "lsst.alert.avsc") 

48 ) 

49 minCutoutSize = pexConfig.RangeField( 

50 dtype=int, 

51 min=0, 

52 max=1000, 

53 default=30, 

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

55 ) 

56 alertWriteLocation = pexConfig.Field( 

57 dtype=str, 

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

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

60 ) 

61 

62 

63class PackageAlertsTask(pipeBase.Task): 

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

65 """ 

66 ConfigClass = PackageAlertsConfig 

67 _DefaultName = "packageAlerts" 

68 

69 def __init__(self, **kwargs): 

70 super().__init__(**kwargs) 

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

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

73 

74 def run(self, 

75 diaSourceCat, 

76 diaObjectCat, 

77 diaSrcHistory, 

78 diffIm, 

79 template, 

80 ccdExposureIdBits): 

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

82 

83 Writes Avro alerts to a location determined by the 

84 ``alertWriteLocation`` configurable. 

85 

86 Parameters 

87 ---------- 

88 diaSourceCat : `pandas.DataFrame` 

89 New DiaSources to package. DataFrame should be indexed on 

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

91 diaObjectCat : `pandas.DataFrame` 

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

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

94 diaSrcHistory : `pandas.DataFrame` 

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

96 the newest DiaSource and is indexed on 

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

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

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

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

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

102 ccdExposureIdBits : `int` 

103 Number of bits used in the ccdVisitId. 

104 """ 

105 alerts = [] 

106 self._patchDiaSources(diaSourceCat) 

107 self._patchDiaSources(diaSrcHistory) 

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

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

110 # Get all diaSources for the associated diaObject. 

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

112 if diaObject["nDiaSources"] > 1: 

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

114 else: 

115 objSourceHistory = None 

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

117 diaSource["decl"], 

118 geom.degrees) 

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

120 diffImCutout = diffIm.getCutout(sphPoint, cutoutBBox) 

121 templateCutout = None 

122 # TODO: Create alertIds DM-24858 

123 alertId = diaSource["diaSourceId"] 

124 alerts.append( 

125 self.makeAlertDict(alertId, 

126 diaSource, 

127 diaObject, 

128 objSourceHistory, 

129 diffImCutout, 

130 templateCutout)) 

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

132 f"{ccdVisitId}.avro"), 

133 "wb") as f: 

134 self.alertSchema.store_alerts(f, alerts) 

135 

136 def _patchDiaSources(self, diaSources): 

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

138 grouped alert data to ``None``. 

139 

140 Parameters 

141 ---------- 

142 diaSources : `pandas.DataFrame` 

143 DataFrame of DiaSources to patch. 

144 """ 

145 diaSources["programId"] = 0 

146 

147 def createDiaSourceBBox(self, bboxSize): 

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

149 BBox that covers the source footprint. 

150 

151 Parameters 

152 ---------- 

153 bboxSize : `int` 

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

155 

156 Returns 

157 ------- 

158 bbox : `lsst.geom.Extent2I` 

159 Geom object representing the size of the bounding box. 

160 """ 

161 if bboxSize < self.config.minCutoutSize: 

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

163 self.config.minCutoutSize) 

164 else: 

165 bbox = geom.Extent2I(bboxSize, bboxSize) 

166 return bbox 

167 

168 def makeAlertDict(self, 

169 alertId, 

170 diaSource, 

171 diaObject, 

172 objDiaSrcHistory, 

173 diffImCutout, 

174 templateCutout): 

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

176 

177 Parameters 

178 ---------- 

179 diaSource : `pandas.DataFrame` 

180 New single DiaSource to package. 

181 diaObject : `pandas.DataFrame` 

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

183 objDiaSrcHistory : `pandas.DataFrame` 

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

185 diffImCutout : `lsst.afw.image.ExposureF` or `None` 

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

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

188 templateCutout : `lsst.afw.image.ExposureF` or `None` 

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

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

191 """ 

192 alert = dict() 

193 alert['alertId'] = alertId 

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

195 

196 if objDiaSrcHistory is None: 

197 alert['prvDiaSources'] = objDiaSrcHistory 

198 else: 

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

200 

201 alert['prvDiaForcedSources'] = None 

202 alert['prvDiaNondetectionLimits'] = None 

203 

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

205 

206 alert['ssObject'] = None 

207 

208 alert['cutoutDifference'] = self.makeCutoutBytes(diffImCutout) 

209 # TODO: add template cutouts in DM-24327 

210 alert["cutoutTemplate"] = None 

211 

212 return alert 

213 

214 def makeCutoutBytes(self, cutout): 

215 """Serialize a cutout into bytes. 

216 

217 Parameters 

218 ---------- 

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

220 Cutout to serialize. 

221 

222 Returns 

223 ------- 

224 coutputBytes : `bytes` 

225 Input cutout serialized into byte data. 

226 """ 

227 mgr = afwFits.MemFileManager() 

228 cutout.writeFits(mgr) 

229 return mgr.getData()