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

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 io
25import os
27from astropy import wcs
28import astropy.units as u
29from astropy.nddata import CCDData, VarianceUncertainty
30import pandas as pd
32import lsst.alert.packet as alertPack
33import lsst.afw.geom as afwGeom
34import lsst.geom as geom
35import lsst.pex.config as pexConfig
36import lsst.pipe.base as pipeBase
39"""Methods for packaging Apdb and Pipelines data into Avro alerts.
40"""
43class PackageAlertsConfig(pexConfig.Config):
44 """Config class for AssociationTask.
45 """
46 schemaFile = pexConfig.Field(
47 dtype=str,
48 doc="Schema definition file for the avro alerts.",
49 default=alertPack.get_path_to_latest_schema()
50 )
51 minCutoutSize = pexConfig.RangeField(
52 dtype=int,
53 min=0,
54 max=1000,
55 default=30,
56 doc="Dimension of the square image cutouts to package in the alert."
57 )
58 alertWriteLocation = pexConfig.Field(
59 dtype=str,
60 doc="Location to write alerts to.",
61 default=os.path.join(os.getcwd(), "alerts"),
62 )
65class PackageAlertsTask(pipeBase.Task):
66 """Tasks for packaging Dia and Pipelines data into Avro alert packages.
67 """
68 ConfigClass = PackageAlertsConfig
69 _DefaultName = "packageAlerts"
71 _scale = (1.0 * geom.arcseconds).asDegrees()
73 def __init__(self, **kwargs):
74 super().__init__(**kwargs)
75 self.alertSchema = alertPack.Schema.from_file(self.config.schemaFile)
76 os.makedirs(self.config.alertWriteLocation, exist_ok=True)
78 def run(self,
79 diaSourceCat,
80 diaObjectCat,
81 diaSrcHistory,
82 diaForcedSources,
83 diffIm,
84 template,
85 ccdExposureIdBits):
86 """Package DiaSources/Object and exposure data into Avro alerts.
88 Writes Avro alerts to a location determined by the
89 ``alertWriteLocation`` configurable.
91 Parameters
92 ----------
93 diaSourceCat : `pandas.DataFrame`
94 New DiaSources to package. DataFrame should be indexed on
95 ``["diaObjectId", "filterName", "diaSourceId"]``
96 diaObjectCat : `pandas.DataFrame`
97 New and updated DiaObjects matched to the new DiaSources. DataFrame
98 is indexed on ``["diaObjectId"]``
99 diaSrcHistory : `pandas.DataFrame`
100 12 month history of DiaSources matched to the DiaObjects. Excludes
101 the newest DiaSource and is indexed on
102 ``["diaObjectId", "filterName", "diaSourceId"]``
103 diaForcedSources : `pandas.DataFrame`
104 12 month history of DiaForcedSources matched to the DiaObjects.
105 ``["diaObjectId"]``
106 diffIm : `lsst.afw.image.ExposureF`
107 Difference image the sources in ``diaSourceCat`` were detected in.
108 template : `lsst.afw.image.ExposureF` or `None`
109 Template image used to create the ``diffIm``.
110 ccdExposureIdBits : `int`
111 Number of bits used in the ccdVisitId.
112 """
113 alerts = []
114 self._patchDiaSources(diaSourceCat)
115 self._patchDiaSources(diaSrcHistory)
116 ccdVisitId = diffIm.getInfo().getVisitInfo().getExposureId()
117 diffImPhotoCalib = diffIm.getPhotoCalib()
118 templatePhotoCalib = template.getPhotoCalib()
119 for srcIndex, diaSource in diaSourceCat.iterrows():
120 # Get all diaSources for the associated diaObject.
121 diaObject = diaObjectCat.loc[srcIndex[0]]
122 if diaObject["nDiaSources"] > 1:
123 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
124 else:
125 objSourceHistory = None
126 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]]
127 sphPoint = geom.SpherePoint(diaSource["ra"],
128 diaSource["decl"],
129 geom.degrees)
130 cutoutBBox = self.createDiaSourceBBox(diaSource["bboxSize"])
131 diffImCutout = self.createCcdDataCutout(
132 diffIm.getCutout(sphPoint, cutoutBBox),
133 sphPoint,
134 diffImPhotoCalib)
136 templateBBox = self.createDiaSourceBBox(diaSource["bboxSize"])
137 templateCutout = self.createCcdDataCutout(
138 template.getCutout(sphPoint, templateBBox),
139 sphPoint,
140 templatePhotoCalib)
142 # TODO: Create alertIds DM-24858
143 alertId = diaSource["diaSourceId"]
144 alerts.append(
145 self.makeAlertDict(alertId,
146 diaSource,
147 diaObject,
148 objSourceHistory,
149 objDiaForcedSources,
150 diffImCutout,
151 templateCutout))
152 with open(os.path.join(self.config.alertWriteLocation,
153 f"{ccdVisitId}.avro"),
154 "wb") as f:
155 self.alertSchema.store_alerts(f, alerts)
157 def _patchDiaSources(self, diaSources):
158 """Add the ``programId`` column to the data.
160 Parameters
161 ----------
162 diaSources : `pandas.DataFrame`
163 DataFrame of DiaSources to patch.
164 """
165 diaSources["programId"] = 0
167 def createDiaSourceBBox(self, bboxSize):
168 """Create a bounding box for the cutouts given the size of the square
169 BBox that covers the source footprint.
171 Parameters
172 ----------
173 bboxSize : `int`
174 Size of a side of the square bounding box in pixels.
176 Returns
177 -------
178 bbox : `lsst.geom.Extent2I`
179 Geom object representing the size of the bounding box.
180 """
181 if bboxSize < self.config.minCutoutSize:
182 bbox = geom.Extent2I(self.config.minCutoutSize,
183 self.config.minCutoutSize)
184 else:
185 bbox = geom.Extent2I(bboxSize, bboxSize)
186 return bbox
188 def createCcdDataCutout(self, cutout, skyCenter, photoCalib):
189 """Convert a cutout into a calibrate CCDData image.
191 Parameters
192 ----------
193 cutout : `lsst.afw.image.ExposureF`
194 Cutout to convert.
195 skyCenter : `lsst.geom.SpherePoint`
196 Center point of DiaSource on the sky.
197 photoCalib : `lsst.afw.image.PhotoCalib`
198 Calibrate object of the image the cutout is cut from.
200 Returns
201 -------
202 ccdData : `astropy.nddata.CCDData`
203 CCDData object storing the calibrate information from the input
204 difference or template image.
205 """
206 # Find the value of the bottom corner of our cutout's BBox and
207 # subtract 1 so that the CCDData cutout position value will be
208 # [1, 1].
209 cutOutMinX = cutout.getBBox().minX - 1
210 cutOutMinY = cutout.getBBox().minY - 1
211 center = cutout.getWcs().skyToPixel(skyCenter)
212 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
214 cutoutWcs = wcs.WCS(naxis=2)
215 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
216 cutout.getBBox().getWidth())
217 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
218 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
219 skyCenter.getDec().asDegrees()]
220 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
221 center,
222 skyCenter)
224 return CCDData(
225 data=calibCutout.getImage().array,
226 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
227 flags=calibCutout.getMask().array,
228 wcs=cutoutWcs,
229 meta={"cutMinX": cutOutMinX,
230 "cutMinY": cutOutMinY},
231 unit=u.nJy)
233 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
234 """Create a local, linear approximation of the wcs transformation
235 matrix.
237 The approximation is created as if the center is at RA=0, DEC=0. All
238 comparing x,y coordinate are relative to the position of center. Matrix
239 is initially calculated with units arcseconds and then converted to
240 degrees. This yields higher precision results due to quirks in AST.
242 Parameters
243 ----------
244 wcs : `lsst.afw.geom.SkyWcs`
245 Wcs to approximate
246 center : `lsst.geom.Point2D`
247 Point at which to evaluate the LocalWcs.
248 skyCenter : `lsst.geom.SpherePoint`
249 Point on sky to approximate the Wcs.
251 Returns
252 -------
253 localMatrix : `numpy.ndarray`
254 Matrix representation the local wcs approximation with units
255 degrees.
256 """
257 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
258 localGnomonicWcs = afwGeom.makeSkyWcs(
259 center, skyCenter, blankCDMatrix)
260 measurementToLocalGnomonic = wcs.getTransform().then(
261 localGnomonicWcs.getTransform().inverted()
262 )
263 localMatrix = measurementToLocalGnomonic.getJacobian(center)
264 return localMatrix / 3600
266 def makeAlertDict(self,
267 alertId,
268 diaSource,
269 diaObject,
270 objDiaSrcHistory,
271 objDiaForcedSources,
272 diffImCutout,
273 templateCutout):
274 """Convert data and package into a dictionary alert.
276 Parameters
277 ----------
278 diaSource : `pandas.DataFrame`
279 New single DiaSource to package.
280 diaObject : `pandas.DataFrame`
281 DiaObject that ``diaSource`` is matched to.
282 objDiaSrcHistory : `pandas.DataFrame`
283 12 month history of ``diaObject`` excluding the latest DiaSource.
284 objDiaForcedSources : `pandas.DataFrame`
285 12 month history of ``diaObject`` forced measurements.
286 diffImCutout : `astropy.nddata.CCDData` or `None`
287 Cutout of the difference image around the location of ``diaSource``
288 with a min size set by the ``cutoutSize`` configurable.
289 templateCutout : `astropy.nddata.CCDData` or `None`
290 Cutout of the template image around the location of ``diaSource``
291 with a min size set by the ``cutoutSize`` configurable.
292 """
293 alert = dict()
294 alert['alertId'] = alertId
295 alert['diaSource'] = diaSource.to_dict()
297 if objDiaSrcHistory is None:
298 alert['prvDiaSources'] = objDiaSrcHistory
299 else:
300 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
302 if isinstance(objDiaForcedSources, pd.Series):
303 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
304 else:
305 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
306 alert['prvDiaNondetectionLimits'] = None
308 alert['diaObject'] = diaObject.to_dict()
310 alert['ssObject'] = None
312 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
313 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout)
315 return alert
317 def streamCcdDataToBytes(self, cutout):
318 """Serialize a cutout into bytes.
320 Parameters
321 ----------
322 cutout : `astropy.nddata.CCDData`
323 Cutout to serialize.
325 Returns
326 -------
327 coutputBytes : `bytes`
328 Input cutout serialized into byte data.
329 """
330 with io.BytesIO() as streamer:
331 cutout.write(streamer, format="fits")
332 cutoutBytes = streamer.getvalue()
333 return cutoutBytes