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