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

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 # TODO: DM-24926 Create dynamic cutout size based on footprint with max
48 # size 30x30
49 cutoutSize = 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 )
63class PackageAlertsTask(pipeBase.Task):
64 """Tasks for packaging Dia and Pipelines data into Avro alert packages.
65 """
66 ConfigClass = PackageAlertsConfig
67 _DefaultName = "packageAlerts"
69 def __init__(self, **kwargs):
70 super().__init__(**kwargs)
71 self.alertSchema = alertPack.Schema.from_file(self.config.schemaFile)
72 self.cutoutBBox = geom.Extent2I(self.config.cutoutSize,
73 self.config.cutoutSize)
74 os.makedirs(self.config.alertWriteLocation, exist_ok=True)
76 def run(self,
77 diaSourceCat,
78 diaObjectCat,
79 diaSrcHistory,
80 diffIm,
81 template,
82 ccdExposureIdBits):
83 """Package DiaSources/Object and exposure data into Avro alerts.
85 Writes Avro alerts to a location determined by the
86 ``alertWriteLocation`` configurable.
88 Parameters
89 ----------
90 diaSourceCat : `pandas.DataFrame`
91 New DiaSources to package. DataFrame should be indexed on
92 ``["diaObjectId", "filterName", "diaSourceId"]``
93 diaObjectCat : `pandas.DataFrame`
94 New and updated DiaObjects matched to the new DiaSources. DataFrame
95 is indexed on ``["diaObjectId"]``
96 diaSrcHistory : `pandas.DataFrame`
97 12 month history of DiaSources matched to the DiaObjects. Excludes
98 the newest DiaSource and is indexed on
99 ``["diaObjectId", "filterName", "diaSourceId"]``
100 diffIm : `lsst.afw.image.ExposureF`
101 Difference image the sources in ``diaSourceCat`` were detected in.
102 template : `lsst.afw.image.ExposureF` or `None`
103 Template image used to create the ``diffIm``.
104 ccdExposureIdBits : `int`
105 Number of bits used in the ccdVisitId.
106 """
107 alerts = []
108 self._patchDiaSources(diaSourceCat)
109 self._patchDiaSources(diaSrcHistory)
110 self._patchDiaObjects(diaObjectCat)
111 ccdVisitId = diffIm.getInfo().getVisitInfo().getExposureId()
112 for srcIndex, diaSource in diaSourceCat.iterrows():
113 # Get all diaSources for the associated diaObject.
114 diaObject = diaObjectCat.loc[srcIndex[0]]
115 if diaObject["nDiaSources"] > 1:
116 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
117 else:
118 objSourceHistory = None
119 sphPoint = geom.SpherePoint(diaSource["ra"],
120 diaSource["decl"],
121 geom.degrees)
122 diffImCutout = diffIm.getCutout(sphPoint, self.cutoutBBox)
123 templateCutout = None
124 # TODO: Create alertIds DM-24858
125 alertId = diaSource["diaSourceId"]
126 alerts.append(
127 self.makeAlertDict(alertId,
128 diaSource,
129 diaObject,
130 objSourceHistory,
131 diffImCutout,
132 templateCutout))
133 with open(os.path.join(self.config.alertWriteLocation,
134 f"{ccdVisitId}.avro"),
135 "wb") as f:
136 self.alertSchema.store_alerts(f, alerts)
138 def _patchDiaSources(self, diaSources):
139 """Add the ``programId`` column to the data and change currently
140 grouped alert data to ``None``.
142 TODO: The need to change these column values to ``None`` can be removed
143 after DM-24696 is merged.
145 Parameters
146 ----------
147 diaSources : `pandas.DataFrame`
148 DataFrame of DiaSources to patch.
149 """
150 diaSources["programId"] = 0
151 diaSources["ra_decl_Cov"] = None
152 diaSources["x_y_Cov"] = None
153 diaSources["ps_Cov"] = None
154 diaSources["trail_Cov"] = None
155 diaSources["dip_Cov"] = None
156 diaSources["i_cov"] = None
158 def _patchDiaObjects(self, diaObjects):
159 """Change currently grouped alert data to ``None``.
161 TODO: The need to change these column values to ``None`` can be removed
162 after DM-24696 is merged.
164 Parameters
165 ----------
166 diaObjects : `pandas.DataFrame`
167 DataFrame of DiaObjects to patch.
168 """
169 diaObjects["ra_decl_Cov"] = None
170 diaObjects["pm_parallax_Cov"] = None
171 diaObjects["uLcPeriodic"] = None
172 diaObjects["gLcPeriodic"] = None
173 diaObjects["rLcPeriodic"] = None
174 diaObjects["iLcPeriodic"] = None
175 diaObjects["zLcPeriodic"] = None
176 diaObjects["yLcPeriodic"] = None
177 diaObjects["uLcNonPeriodic"] = None
178 diaObjects["gLcNonPeriodic"] = None
179 diaObjects["rLcNonPeriodic"] = None
180 diaObjects["iLcNonPeriodic"] = None
181 diaObjects["zLcNonPeriodic"] = None
182 diaObjects["yLcNonPeriodic"] = None
184 def makeAlertDict(self,
185 alertId,
186 diaSource,
187 diaObject,
188 objDiaSrcHistory,
189 diffImCutout,
190 templateCutout):
191 """Convert data and package into a dictionary alert.
193 Parameters
194 ----------
195 diaSource : `pandas.DataFrame`
196 New single DiaSource to package.
197 diaObject : `pandas.DataFrame`
198 DiaObject that ``diaSource`` is matched to.
199 objDiaSrcHistory : `pandas.DataFrame`
200 12 month history of ``diaObject`` excluding the latest DiaSource.
201 diffImCutout : `lsst.afw.image.ExposureF`
202 Cutout of the difference image around the location of ``diaSource``
203 with a size set by the ``cutoutSize`` configurable.
204 templateCutout : `lsst.afw.image.ExposureF`
205 Cutout of the template image around the location of ``diaSource``
206 with a size set by the ``cutoutSize`` configurable.
207 """
208 alert = dict()
209 alert['alertId'] = alertId
210 alert['diaSource'] = diaSource.to_dict()
212 if objDiaSrcHistory is None:
213 alert['prvDiaSources'] = objDiaSrcHistory
214 else:
215 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
217 alert['prvDiaForcedSources'] = None
218 alert['prvDiaNondetectionLimits'] = None
220 alert['diaObject'] = diaObject.to_dict()
222 alert['ssObject'] = None
224 # TODO: fileName to be removed in future Avro schemas. DM-24696
225 alert['cutoutDifference'] = {
226 'fileName': '',
227 'stampData': self.makeCutoutBytes(diffImCutout),
228 }
229 # TODO: add template cutouts in DM-24327
230 alert["cutoutTemplate"] = None
232 return alert
234 def makeCutoutBytes(self, cutout):
235 """Serialize a cutout into bytes.
237 Parameters
238 ----------
239 cutout : `lsst.afw.image.ExposureF`
240 Cutout to serialize.
242 Returns
243 -------
244 coutputBytes : `bytes`
245 Input cutout serialized into byte data.
246 """
247 mgr = afwFits.MemFileManager()
248 cutout.writeFits(mgr)
249 return mgr.getData()