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 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
36from lsst.pex.exceptions import InvalidParameterError
37import lsst.pipe.base as pipeBase
40"""Methods for packaging Apdb and Pipelines data into Avro alerts.
41"""
44class PackageAlertsConfig(pexConfig.Config):
45 """Config class for AssociationTask.
46 """
47 schemaFile = pexConfig.Field(
48 dtype=str,
49 doc="Schema definition file for the avro alerts.",
50 default=alertPack.get_path_to_latest_schema()
51 )
52 minCutoutSize = pexConfig.RangeField(
53 dtype=int,
54 min=0,
55 max=1000,
56 default=30,
57 doc="Dimension of the square image cutouts to package in the alert."
58 )
59 alertWriteLocation = pexConfig.Field(
60 dtype=str,
61 doc="Location to write alerts to.",
62 default=os.path.join(os.getcwd(), "alerts"),
63 )
66class PackageAlertsTask(pipeBase.Task):
67 """Tasks for packaging Dia and Pipelines data into Avro alert packages.
68 """
69 ConfigClass = PackageAlertsConfig
70 _DefaultName = "packageAlerts"
72 _scale = (1.0 * geom.arcseconds).asDegrees()
74 def __init__(self, **kwargs):
75 super().__init__(**kwargs)
76 self.alertSchema = alertPack.Schema.from_file(self.config.schemaFile)
77 os.makedirs(self.config.alertWriteLocation, exist_ok=True)
79 @pipeBase.timeMethod
80 def run(self,
81 diaSourceCat,
82 diaObjectCat,
83 diaSrcHistory,
84 diaForcedSources,
85 diffIm,
86 template,
87 ccdExposureIdBits):
88 """Package DiaSources/Object and exposure data into Avro alerts.
90 Writes Avro alerts to a location determined by the
91 ``alertWriteLocation`` configurable.
93 Parameters
94 ----------
95 diaSourceCat : `pandas.DataFrame`
96 New DiaSources to package. DataFrame should be indexed on
97 ``["diaObjectId", "filterName", "diaSourceId"]``
98 diaObjectCat : `pandas.DataFrame`
99 New and updated DiaObjects matched to the new DiaSources. DataFrame
100 is indexed on ``["diaObjectId"]``
101 diaSrcHistory : `pandas.DataFrame`
102 12 month history of DiaSources matched to the DiaObjects. Excludes
103 the newest DiaSource and is indexed on
104 ``["diaObjectId", "filterName", "diaSourceId"]``
105 diaForcedSources : `pandas.DataFrame`
106 12 month history of DiaForcedSources matched to the DiaObjects.
107 ``["diaObjectId"]``
108 diffIm : `lsst.afw.image.ExposureF`
109 Difference image the sources in ``diaSourceCat`` were detected in.
110 template : `lsst.afw.image.ExposureF` or `None`
111 Template image used to create the ``diffIm``.
112 ccdExposureIdBits : `int`
113 Number of bits used in the ccdVisitId.
114 """
115 alerts = []
116 self._patchDiaSources(diaSourceCat)
117 self._patchDiaSources(diaSrcHistory)
118 ccdVisitId = diffIm.getInfo().getVisitInfo().getExposureId()
119 diffImPhotoCalib = diffIm.getPhotoCalib()
120 templatePhotoCalib = template.getPhotoCalib()
121 for srcIndex, diaSource in diaSourceCat.iterrows():
122 # Get all diaSources for the associated diaObject.
123 # TODO: DM-31992 skip DiaSources associated with Solar System
124 # Objects for now.
125 if srcIndex[0] == 0:
126 continue
127 diaObject = diaObjectCat.loc[srcIndex[0]]
128 if diaObject["nDiaSources"] > 1:
129 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
130 else:
131 objSourceHistory = None
132 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]]
133 sphPoint = geom.SpherePoint(diaSource["ra"],
134 diaSource["decl"],
135 geom.degrees)
137 cutoutExtent = self.createDiaSourceExtent(diaSource["bboxSize"])
138 diffImCutout = self.createCcdDataCutout(
139 diffIm,
140 sphPoint,
141 cutoutExtent,
142 diffImPhotoCalib,
143 diaSource["diaSourceId"])
144 templateCutout = self.createCcdDataCutout(
145 template,
146 sphPoint,
147 cutoutExtent,
148 templatePhotoCalib,
149 diaSource["diaSourceId"])
151 # TODO: Create alertIds DM-24858
152 alertId = diaSource["diaSourceId"]
153 alerts.append(
154 self.makeAlertDict(alertId,
155 diaSource,
156 diaObject,
157 objSourceHistory,
158 objDiaForcedSources,
159 diffImCutout,
160 templateCutout))
161 with open(os.path.join(self.config.alertWriteLocation,
162 f"{ccdVisitId}.avro"),
163 "wb") as f:
164 self.alertSchema.store_alerts(f, alerts)
166 def _patchDiaSources(self, diaSources):
167 """Add the ``programId`` column to the data.
169 Parameters
170 ----------
171 diaSources : `pandas.DataFrame`
172 DataFrame of DiaSources to patch.
173 """
174 diaSources["programId"] = 0
176 def createDiaSourceExtent(self, bboxSize):
177 """Create a extent for a box for the cutouts given the size of the
178 square BBox that covers the source footprint.
180 Parameters
181 ----------
182 bboxSize : `int`
183 Size of a side of the square bounding box in pixels.
185 Returns
186 -------
187 extent : `lsst.geom.Extent2I`
188 Geom object representing the size of the bounding box.
189 """
190 if bboxSize < self.config.minCutoutSize:
191 extent = geom.Extent2I(self.config.minCutoutSize,
192 self.config.minCutoutSize)
193 else:
194 extent = geom.Extent2I(bboxSize, bboxSize)
195 return extent
197 def createCcdDataCutout(self, image, skyCenter, extent, photoCalib, srcId):
198 """Grab an image as a cutout and return a calibrated CCDData image.
200 Parameters
201 ----------
202 image : `lsst.afw.image.ExposureF`
203 Image to pull cutout from.
204 skyCenter : `lsst.geom.SpherePoint`
205 Center point of DiaSource on the sky.
206 extent : `lsst.geom.Extent2I`
207 Bounding box to cutout from the image.
208 photoCalib : `lsst.afw.image.PhotoCalib`
209 Calibrate object of the image the cutout is cut from.
210 srcId : `int`
211 Unique id of DiaSource. Used for when an error occurs extracting
212 a cutout.
214 Returns
215 -------
216 ccdData : `astropy.nddata.CCDData` or `None`
217 CCDData object storing the calibrate information from the input
218 difference or template image.
219 """
220 # Catch errors in retrieving the cutout.
221 try:
222 cutout = image.getCutout(skyCenter, extent)
223 except InvalidParameterError:
224 point = image.getWcs().skyToPixel(skyCenter)
225 imBBox = image.getBBox()
226 if not geom.Box2D(image.getBBox()).contains(point):
227 self.log.warning(
228 "DiaSource id=%i centroid lies at pixel (%.2f, %.2f) "
229 "which is outside the Exposure with bounding box "
230 "((%i, %i), (%i, %i)). Returning None for cutout..." %
231 (srcId, point.x, point.y,
232 imBBox.minX, imBBox.maxX, imBBox.minY, imBBox.maxY))
233 else:
234 raise InvalidParameterError(
235 "Failed to retrieve cutout from image for DiaSource with "
236 "id=%i. InvalidParameterError thrown during cutout "
237 "creation. Exiting."
238 % srcId)
239 return None
241 # Find the value of the bottom corner of our cutout's BBox and
242 # subtract 1 so that the CCDData cutout position value will be
243 # [1, 1].
244 cutOutMinX = cutout.getBBox().minX - 1
245 cutOutMinY = cutout.getBBox().minY - 1
246 center = cutout.getWcs().skyToPixel(skyCenter)
247 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
249 cutoutWcs = wcs.WCS(naxis=2)
250 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
251 cutout.getBBox().getWidth())
252 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
253 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
254 skyCenter.getDec().asDegrees()]
255 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
256 center,
257 skyCenter)
259 return CCDData(
260 data=calibCutout.getImage().array,
261 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
262 flags=calibCutout.getMask().array,
263 wcs=cutoutWcs,
264 meta={"cutMinX": cutOutMinX,
265 "cutMinY": cutOutMinY},
266 unit=u.nJy)
268 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
269 """Create a local, linear approximation of the wcs transformation
270 matrix.
272 The approximation is created as if the center is at RA=0, DEC=0. All
273 comparing x,y coordinate are relative to the position of center. Matrix
274 is initially calculated with units arcseconds and then converted to
275 degrees. This yields higher precision results due to quirks in AST.
277 Parameters
278 ----------
279 wcs : `lsst.afw.geom.SkyWcs`
280 Wcs to approximate
281 center : `lsst.geom.Point2D`
282 Point at which to evaluate the LocalWcs.
283 skyCenter : `lsst.geom.SpherePoint`
284 Point on sky to approximate the Wcs.
286 Returns
287 -------
288 localMatrix : `numpy.ndarray`
289 Matrix representation the local wcs approximation with units
290 degrees.
291 """
292 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
293 localGnomonicWcs = afwGeom.makeSkyWcs(
294 center, skyCenter, blankCDMatrix)
295 measurementToLocalGnomonic = wcs.getTransform().then(
296 localGnomonicWcs.getTransform().inverted()
297 )
298 localMatrix = measurementToLocalGnomonic.getJacobian(center)
299 return localMatrix / 3600
301 def makeAlertDict(self,
302 alertId,
303 diaSource,
304 diaObject,
305 objDiaSrcHistory,
306 objDiaForcedSources,
307 diffImCutout,
308 templateCutout):
309 """Convert data and package into a dictionary alert.
311 Parameters
312 ----------
313 diaSource : `pandas.DataFrame`
314 New single DiaSource to package.
315 diaObject : `pandas.DataFrame`
316 DiaObject that ``diaSource`` is matched to.
317 objDiaSrcHistory : `pandas.DataFrame`
318 12 month history of ``diaObject`` excluding the latest DiaSource.
319 objDiaForcedSources : `pandas.DataFrame`
320 12 month history of ``diaObject`` forced measurements.
321 diffImCutout : `astropy.nddata.CCDData` or `None`
322 Cutout of the difference image around the location of ``diaSource``
323 with a min size set by the ``cutoutSize`` configurable.
324 templateCutout : `astropy.nddata.CCDData` or `None`
325 Cutout of the template image around the location of ``diaSource``
326 with a min size set by the ``cutoutSize`` configurable.
327 """
328 alert = dict()
329 alert['alertId'] = alertId
330 alert['diaSource'] = diaSource.to_dict()
332 if objDiaSrcHistory is None:
333 alert['prvDiaSources'] = objDiaSrcHistory
334 else:
335 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
337 if isinstance(objDiaForcedSources, pd.Series):
338 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
339 else:
340 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
341 alert['prvDiaNondetectionLimits'] = None
343 alert['diaObject'] = diaObject.to_dict()
345 alert['ssObject'] = None
347 if diffImCutout is None:
348 alert['cutoutDifference'] = None
349 else:
350 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
352 if templateCutout is None:
353 alert["cutoutTemplate"] = None
354 else:
355 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout)
357 return alert
359 def streamCcdDataToBytes(self, cutout):
360 """Serialize a cutout into bytes.
362 Parameters
363 ----------
364 cutout : `astropy.nddata.CCDData`
365 Cutout to serialize.
367 Returns
368 -------
369 coutputBytes : `bytes`
370 Input cutout serialized into byte data.
371 """
372 with io.BytesIO() as streamer:
373 cutout.write(streamer, format="fits")
374 cutoutBytes = streamer.getvalue()
375 return cutoutBytes