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