Coverage for python/lsst/ap/association/packageAlerts.py: 27%
105 statements
« prev ^ index » next coverage.py v6.4.1, created at 2022-07-11 08:03 +0000
« prev ^ index » next coverage.py v6.4.1, created at 2022-07-11 08:03 +0000
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 diaObject = diaObjectCat.loc[srcIndex[0]]
124 if diaObject["nDiaSources"] > 1:
125 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
126 else:
127 objSourceHistory = None
128 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]]
129 sphPoint = geom.SpherePoint(diaSource["ra"],
130 diaSource["decl"],
131 geom.degrees)
133 cutoutExtent = self.createDiaSourceExtent(diaSource["bboxSize"])
134 diffImCutout = self.createCcdDataCutout(
135 diffIm,
136 sphPoint,
137 cutoutExtent,
138 diffImPhotoCalib,
139 diaSource["diaSourceId"])
140 templateCutout = self.createCcdDataCutout(
141 template,
142 sphPoint,
143 cutoutExtent,
144 templatePhotoCalib,
145 diaSource["diaSourceId"])
147 # TODO: Create alertIds DM-24858
148 alertId = diaSource["diaSourceId"]
149 alerts.append(
150 self.makeAlertDict(alertId,
151 diaSource,
152 diaObject,
153 objSourceHistory,
154 objDiaForcedSources,
155 diffImCutout,
156 templateCutout))
157 with open(os.path.join(self.config.alertWriteLocation,
158 f"{ccdVisitId}.avro"),
159 "wb") as f:
160 self.alertSchema.store_alerts(f, alerts)
162 def _patchDiaSources(self, diaSources):
163 """Add the ``programId`` column to the data.
165 Parameters
166 ----------
167 diaSources : `pandas.DataFrame`
168 DataFrame of DiaSources to patch.
169 """
170 diaSources["programId"] = 0
172 def createDiaSourceExtent(self, bboxSize):
173 """Create a extent for a box for the cutouts given the size of the
174 square BBox that covers the source footprint.
176 Parameters
177 ----------
178 bboxSize : `int`
179 Size of a side of the square bounding box in pixels.
181 Returns
182 -------
183 extent : `lsst.geom.Extent2I`
184 Geom object representing the size of the bounding box.
185 """
186 if bboxSize < self.config.minCutoutSize:
187 extent = geom.Extent2I(self.config.minCutoutSize,
188 self.config.minCutoutSize)
189 else:
190 extent = geom.Extent2I(bboxSize, bboxSize)
191 return extent
193 def createCcdDataCutout(self, image, skyCenter, extent, photoCalib, srcId):
194 """Grab an image as a cutout and return a calibrated CCDData image.
196 Parameters
197 ----------
198 image : `lsst.afw.image.ExposureF`
199 Image to pull cutout from.
200 skyCenter : `lsst.geom.SpherePoint`
201 Center point of DiaSource on the sky.
202 extent : `lsst.geom.Extent2I`
203 Bounding box to cutout from the image.
204 photoCalib : `lsst.afw.image.PhotoCalib`
205 Calibrate object of the image the cutout is cut from.
206 srcId : `int`
207 Unique id of DiaSource. Used for when an error occurs extracting
208 a cutout.
210 Returns
211 -------
212 ccdData : `astropy.nddata.CCDData` or `None`
213 CCDData object storing the calibrate information from the input
214 difference or template image.
215 """
216 # Catch errors in retrieving the cutout.
217 try:
218 cutout = image.getCutout(skyCenter, extent)
219 except InvalidParameterError:
220 point = image.getWcs().skyToPixel(skyCenter)
221 imBBox = image.getBBox()
222 if not geom.Box2D(image.getBBox()).contains(point):
223 self.log.warn(
224 "DiaSource id=%i centroid lies at pixel (%.2f, %.2f) "
225 "which is outside the Exposure with bounding box "
226 "((%i, %i), (%i, %i)). Returning None for cutout..." %
227 (srcId, point.x, point.y,
228 imBBox.minX, imBBox.maxX, imBBox.minY, imBBox.maxY))
229 else:
230 raise InvalidParameterError(
231 "Failed to retrieve cutout from image for DiaSource with "
232 "id=%i. InvalidParameterError thrown during cutout "
233 "creation. Exiting."
234 % srcId)
235 return None
237 # Find the value of the bottom corner of our cutout's BBox and
238 # subtract 1 so that the CCDData cutout position value will be
239 # [1, 1].
240 cutOutMinX = cutout.getBBox().minX - 1
241 cutOutMinY = cutout.getBBox().minY - 1
242 center = cutout.getWcs().skyToPixel(skyCenter)
243 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
245 cutoutWcs = wcs.WCS(naxis=2)
246 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
247 cutout.getBBox().getWidth())
248 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
249 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
250 skyCenter.getDec().asDegrees()]
251 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
252 center,
253 skyCenter)
255 return CCDData(
256 data=calibCutout.getImage().array,
257 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
258 flags=calibCutout.getMask().array,
259 wcs=cutoutWcs,
260 meta={"cutMinX": cutOutMinX,
261 "cutMinY": cutOutMinY},
262 unit=u.nJy)
264 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
265 """Create a local, linear approximation of the wcs transformation
266 matrix.
268 The approximation is created as if the center is at RA=0, DEC=0. All
269 comparing x,y coordinate are relative to the position of center. Matrix
270 is initially calculated with units arcseconds and then converted to
271 degrees. This yields higher precision results due to quirks in AST.
273 Parameters
274 ----------
275 wcs : `lsst.afw.geom.SkyWcs`
276 Wcs to approximate
277 center : `lsst.geom.Point2D`
278 Point at which to evaluate the LocalWcs.
279 skyCenter : `lsst.geom.SpherePoint`
280 Point on sky to approximate the Wcs.
282 Returns
283 -------
284 localMatrix : `numpy.ndarray`
285 Matrix representation the local wcs approximation with units
286 degrees.
287 """
288 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
289 localGnomonicWcs = afwGeom.makeSkyWcs(
290 center, skyCenter, blankCDMatrix)
291 measurementToLocalGnomonic = wcs.getTransform().then(
292 localGnomonicWcs.getTransform().inverted()
293 )
294 localMatrix = measurementToLocalGnomonic.getJacobian(center)
295 return localMatrix / 3600
297 def makeAlertDict(self,
298 alertId,
299 diaSource,
300 diaObject,
301 objDiaSrcHistory,
302 objDiaForcedSources,
303 diffImCutout,
304 templateCutout):
305 """Convert data and package into a dictionary alert.
307 Parameters
308 ----------
309 diaSource : `pandas.DataFrame`
310 New single DiaSource to package.
311 diaObject : `pandas.DataFrame`
312 DiaObject that ``diaSource`` is matched to.
313 objDiaSrcHistory : `pandas.DataFrame`
314 12 month history of ``diaObject`` excluding the latest DiaSource.
315 objDiaForcedSources : `pandas.DataFrame`
316 12 month history of ``diaObject`` forced measurements.
317 diffImCutout : `astropy.nddata.CCDData` or `None`
318 Cutout of the difference image around the location of ``diaSource``
319 with a min size set by the ``cutoutSize`` configurable.
320 templateCutout : `astropy.nddata.CCDData` or `None`
321 Cutout of the template image around the location of ``diaSource``
322 with a min size set by the ``cutoutSize`` configurable.
323 """
324 alert = dict()
325 alert['alertId'] = alertId
326 alert['diaSource'] = diaSource.to_dict()
328 if objDiaSrcHistory is None:
329 alert['prvDiaSources'] = objDiaSrcHistory
330 else:
331 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
333 if isinstance(objDiaForcedSources, pd.Series):
334 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
335 else:
336 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
337 alert['prvDiaNondetectionLimits'] = None
339 alert['diaObject'] = diaObject.to_dict()
341 alert['ssObject'] = None
343 if diffImCutout is None:
344 alert['cutoutDifference'] = None
345 else:
346 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
348 if templateCutout is None:
349 alert["cutoutTemplate"] = None
350 else:
351 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout)
353 return alert
355 def streamCcdDataToBytes(self, cutout):
356 """Serialize a cutout into bytes.
358 Parameters
359 ----------
360 cutout : `astropy.nddata.CCDData`
361 Cutout to serialize.
363 Returns
364 -------
365 coutputBytes : `bytes`
366 Input cutout serialized into byte data.
367 """
368 with io.BytesIO() as streamer:
369 cutout.write(streamer, format="fits")
370 cutoutBytes = streamer.getvalue()
371 return cutoutBytes