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

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