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

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