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

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/>.
22__all__ = ("PackageAlertsConfig", "PackageAlertsTask")
24import os
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
34"""Methods for packaging Apdb and Pipelines data into Avro alerts.
35"""
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 )
61class PackageAlertsTask(pipeBase.Task):
62 """Tasks for packaging Dia and Pipelines data into Avro alert packages.
63 """
64 ConfigClass = PackageAlertsConfig
65 _DefaultName = "packageAlerts"
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)
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.
81 Writes Avro alerts to a location determined by the
82 ``alertWriteLocation`` configurable.
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)
135 def _patchDiaSources(self, diaSources):
136 """Add the ``programId`` column to the data and change currently
137 grouped alert data to ``None``.
139 TODO: The need to change these column values to ``None`` can be removed
140 after DM-24696 is merged.
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
155 def _patchDiaObjects(self, diaObjects):
156 """Change currently grouped alert data to ``None``.
158 TODO: The need to change these column values to ``None`` can be removed
159 after DM-24696 is merged.
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
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.
185 Parameters
186 ----------
187 bboxSize : `int`
188 Size of a side of the square bounding box in pixels.
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
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.
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()
230 if objDiaSrcHistory is None:
231 alert['prvDiaSources'] = objDiaSrcHistory
232 else:
233 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
235 alert['prvDiaForcedSources'] = None
236 alert['prvDiaNondetectionLimits'] = None
238 alert['diaObject'] = diaObject.to_dict()
240 alert['ssObject'] = None
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
250 return alert
252 def makeCutoutBytes(self, cutout):
253 """Serialize a cutout into bytes.
255 Parameters
256 ----------
257 cutout : `lsst.afw.image.ExposureF`
258 Cutout to serialize.
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()