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 for srcIndex, diaSource in diaSourceCat.iterrows():
123 # Get all diaSources for the associated diaObject.
124 diaObject = diaObjectCat.loc[srcIndex[0]]
125 if diaObject["nDiaSources"] > 1:
126 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
127 else:
128 objSourceHistory = None
129 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]]
130 sphPoint = geom.SpherePoint(diaSource["ra"],
131 diaSource["decl"],
132 geom.degrees)
133 cutoutBBox = self.createDiaSourceBBox(diaSource["bboxSize"])
134 diffImCutout = self.createCcdDataCutout(
135 diffIm.getCutout(sphPoint, cutoutBBox),
136 sphPoint,
137 diffImPhotoCalib)
139 templateCutout = None
140 # TODO: Create alertIds DM-24858
141 alertId = diaSource["diaSourceId"]
142 alerts.append(
143 self.makeAlertDict(alertId,
144 diaSource,
145 diaObject,
146 objSourceHistory,
147 objDiaForcedSources,
148 diffImCutout,
149 templateCutout))
150 with open(os.path.join(self.config.alertWriteLocation,
151 f"{ccdVisitId}.avro"),
152 "wb") as f:
153 self.alertSchema.store_alerts(f, alerts)
155 def _patchDiaSources(self, diaSources):
156 """Add the ``programId`` column to the data.
158 Parameters
159 ----------
160 diaSources : `pandas.DataFrame`
161 DataFrame of DiaSources to patch.
162 """
163 diaSources["programId"] = 0
165 def createDiaSourceBBox(self, bboxSize):
166 """Create a bounding box for the cutouts given the size of the square
167 BBox that covers the source footprint.
169 Parameters
170 ----------
171 bboxSize : `int`
172 Size of a side of the square bounding box in pixels.
174 Returns
175 -------
176 bbox : `lsst.geom.Extent2I`
177 Geom object representing the size of the bounding box.
178 """
179 if bboxSize < self.config.minCutoutSize:
180 bbox = geom.Extent2I(self.config.minCutoutSize,
181 self.config.minCutoutSize)
182 else:
183 bbox = geom.Extent2I(bboxSize, bboxSize)
184 return bbox
186 def createCcdDataCutout(self, cutout, skyCenter, photoCalib):
187 """Convert a cutout into a calibrate CCDData image.
189 Parameters
190 ----------
191 cutout : `lsst.afw.image.ExposureF`
192 Cutout to convert.
193 skyCenter : `lsst.geom.SpherePoint`
194 Center point of DiaSource on the sky.
195 photoCalib : `lsst.afw.image.PhotoCalib`
196 Calibrate object of the image the cutout is cut from.
198 Returns
199 -------
200 ccdData : `astropy.nddata.CCDData`
201 CCDData object storing the calibrate information from the input
202 difference or template image.
203 """
204 # Find the value of the bottom corner of our cutout's BBox and
205 # subtract 1 so that the CCDData cutout position value will be
206 # [1, 1].
207 cutOutMinX = cutout.getBBox().minX - 1
208 cutOutMinY = cutout.getBBox().minY - 1
209 center = cutout.getWcs().skyToPixel(skyCenter)
210 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
212 cutoutWcs = wcs.WCS(naxis=2)
213 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
214 cutout.getBBox().getWidth())
215 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
216 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
217 skyCenter.getDec().asDegrees()]
218 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
219 center,
220 skyCenter)
222 return CCDData(
223 data=calibCutout.getImage().array,
224 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
225 flags=calibCutout.getMask().array,
226 wcs=cutoutWcs,
227 meta={"cutMinX": cutOutMinX,
228 "cutMinY": cutOutMinY},
229 unit=u.nJy)
231 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
232 """Create a local, linear approximation of the wcs transformation
233 matrix.
235 The approximation is created as if the center is at RA=0, DEC=0. All
236 comparing x,y coordinate are relative to the position of center. Matrix
237 is initially calculated with units arcseconds and then converted to
238 degrees. This yields higher precision results due to quirks in AST.
240 Parameters
241 ----------
242 wcs : `lsst.afw.geom.SkyWcs`
243 Wcs to approximate
244 center : `lsst.geom.Point2D`
245 Point at which to evaluate the LocalWcs.
246 skyCenter : `lsst.geom.SpherePoint`
247 Point on sky to approximate the Wcs.
249 Returns
250 -------
251 localMatrix : `numpy.ndarray`
252 Matrix representation the local wcs approximation with units
253 degrees.
254 """
255 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
256 localGnomonicWcs = afwGeom.makeSkyWcs(
257 center, skyCenter, blankCDMatrix)
258 measurementToLocalGnomonic = wcs.getTransform().then(
259 localGnomonicWcs.getTransform().inverted()
260 )
261 localMatrix = measurementToLocalGnomonic.getJacobian(center)
262 return localMatrix / 3600
264 def makeAlertDict(self,
265 alertId,
266 diaSource,
267 diaObject,
268 objDiaSrcHistory,
269 objDiaForcedSources,
270 diffImCutout,
271 templateCutout):
272 """Convert data and package into a dictionary alert.
274 Parameters
275 ----------
276 diaSource : `pandas.DataFrame`
277 New single DiaSource to package.
278 diaObject : `pandas.DataFrame`
279 DiaObject that ``diaSource`` is matched to.
280 objDiaSrcHistory : `pandas.DataFrame`
281 12 month history of ``diaObject`` excluding the latest DiaSource.
282 objDiaForcedSources : `pandas.DataFrame`
283 12 month history of ``diaObject`` forced measurements.
284 diffImCutout : `astropy.nddata.CCDData` or `None`
285 Cutout of the difference image around the location of ``diaSource``
286 with a min size set by the ``cutoutSize`` configurable.
287 templateCutout : `astropy.nddata.CCDData` or `None`
288 Cutout of the template image around the location of ``diaSource``
289 with a min size set by the ``cutoutSize`` configurable.
290 """
291 alert = dict()
292 alert['alertId'] = alertId
293 alert['diaSource'] = diaSource.to_dict()
295 if objDiaSrcHistory is None:
296 alert['prvDiaSources'] = objDiaSrcHistory
297 else:
298 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
300 if isinstance(objDiaForcedSources, pd.Series):
301 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
302 else:
303 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
304 alert['prvDiaNondetectionLimits'] = None
306 alert['diaObject'] = diaObject.to_dict()
308 alert['ssObject'] = None
310 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
311 # TODO: add template cutouts in DM-24327
312 alert["cutoutTemplate"] = None
314 return alert
316 def streamCcdDataToBytes(self, cutout):
317 """Serialize a cutout into bytes.
319 Parameters
320 ----------
321 cutout : `astropy.nddata.CCDData`
322 Cutout to serialize.
324 Returns
325 -------
326 coutputBytes : `bytes`
327 Input cutout serialized into byte data.
328 """
329 with io.BytesIO() as streamer:
330 cutout.write(streamer, format="fits")
331 cutoutBytes = streamer.getvalue()
332 return cutoutBytes