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

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
31import lsst.alert.packet as alertPack
32import lsst.afw.geom as afwGeom
33import lsst.geom as geom
34import lsst.pex.config as pexConfig
35import lsst.pipe.base as pipeBase
36from lsst.utils import getPackageDir
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=os.path.join(getPackageDir("alert_packet"),
50 "schema",
51 *[str(x) for x in alertPack.get_latest_schema_version()],
52 "lsst.alert.avsc")
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 def run(self,
82 diaSourceCat,
83 diaObjectCat,
84 diaSrcHistory,
85 diffIm,
86 template,
87 ccdExposureIdBits):
88 """Package DiaSources/Object and exposure data into Avro alerts.
90 Writes Avro alerts to a location determined by the
91 ``alertWriteLocation`` configurable.
93 Parameters
94 ----------
95 diaSourceCat : `pandas.DataFrame`
96 New DiaSources to package. DataFrame should be indexed on
97 ``["diaObjectId", "filterName", "diaSourceId"]``
98 diaObjectCat : `pandas.DataFrame`
99 New and updated DiaObjects matched to the new DiaSources. DataFrame
100 is indexed on ``["diaObjectId"]``
101 diaSrcHistory : `pandas.DataFrame`
102 12 month history of DiaSources matched to the DiaObjects. Excludes
103 the newest DiaSource and is indexed on
104 ``["diaObjectId", "filterName", "diaSourceId"]``
105 diffIm : `lsst.afw.image.ExposureF`
106 Difference image the sources in ``diaSourceCat`` were detected in.
107 template : `lsst.afw.image.ExposureF` or `None`
108 Template image used to create the ``diffIm``.
109 ccdExposureIdBits : `int`
110 Number of bits used in the ccdVisitId.
111 """
112 alerts = []
113 self._patchDiaSources(diaSourceCat)
114 self._patchDiaSources(diaSrcHistory)
115 ccdVisitId = diffIm.getInfo().getVisitInfo().getExposureId()
116 diffImPhotoCalib = diffIm.getPhotoCalib()
117 for srcIndex, diaSource in diaSourceCat.iterrows():
118 # Get all diaSources for the associated diaObject.
119 diaObject = diaObjectCat.loc[srcIndex[0]]
120 if diaObject["nDiaSources"] > 1:
121 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
122 else:
123 objSourceHistory = None
124 sphPoint = geom.SpherePoint(diaSource["ra"],
125 diaSource["decl"],
126 geom.degrees)
127 cutoutBBox = self.createDiaSourceBBox(diaSource["bboxSize"])
128 diffImCutout = self.createCcdDataCutout(
129 diffIm.getCutout(sphPoint, cutoutBBox),
130 sphPoint,
131 diffImPhotoCalib)
133 templateCutout = None
134 # TODO: Create alertIds DM-24858
135 alertId = diaSource["diaSourceId"]
136 alerts.append(
137 self.makeAlertDict(alertId,
138 diaSource,
139 diaObject,
140 objSourceHistory,
141 diffImCutout,
142 templateCutout))
143 with open(os.path.join(self.config.alertWriteLocation,
144 f"{ccdVisitId}.avro"),
145 "wb") as f:
146 self.alertSchema.store_alerts(f, alerts)
148 def _patchDiaSources(self, diaSources):
149 """Add the ``programId`` column to the data.
151 Parameters
152 ----------
153 diaSources : `pandas.DataFrame`
154 DataFrame of DiaSources to patch.
155 """
156 diaSources["programId"] = 0
158 def createDiaSourceBBox(self, bboxSize):
159 """Create a bounding box for the cutouts given the size of the square
160 BBox that covers the source footprint.
162 Parameters
163 ----------
164 bboxSize : `int`
165 Size of a side of the square bounding box in pixels.
167 Returns
168 -------
169 bbox : `lsst.geom.Extent2I`
170 Geom object representing the size of the bounding box.
171 """
172 if bboxSize < self.config.minCutoutSize:
173 bbox = geom.Extent2I(self.config.minCutoutSize,
174 self.config.minCutoutSize)
175 else:
176 bbox = geom.Extent2I(bboxSize, bboxSize)
177 return bbox
179 def createCcdDataCutout(self, cutout, skyCenter, photoCalib):
180 """Convert a cutout into a calibrate CCDData image.
182 Parameters
183 ----------
184 cutout : `lsst.afw.image.ExposureF`
185 Cutout to convert.
186 skyCenter : `lsst.geom.SpherePoint`
187 Center point of DiaSource on the sky.
188 photoCalib : `lsst.afw.image.PhotoCalib`
189 Calibrate object of the image the cutout is cut from.
191 Returns
192 -------
193 ccdData : `astropy.nddata.CCDData`
194 CCDData object storing the calibrate information from the input
195 difference or template image.
196 """
197 # Find the value of the bottom corner of our cutout's BBox and
198 # subtract 1 so that the CCDData cutout position value will be
199 # [1, 1].
200 cutOutMinX = cutout.getBBox().minX - 1
201 cutOutMinY = cutout.getBBox().minY - 1
202 center = cutout.getWcs().skyToPixel(skyCenter)
203 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
205 cutoutWcs = wcs.WCS(naxis=2)
206 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
207 cutout.getBBox().getWidth())
208 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
209 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
210 skyCenter.getDec().asDegrees()]
211 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
212 center,
213 skyCenter)
215 return CCDData(
216 data=calibCutout.getImage().array,
217 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
218 flags=calibCutout.getMask().array,
219 wcs=cutoutWcs,
220 meta={"cutMinX": cutOutMinX,
221 "cutMinY": cutOutMinY},
222 unit=u.nJy)
224 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
225 """Create a local, linear approximation of the wcs transformation
226 matrix.
228 The approximation is created as if the center is at RA=0, DEC=0. All
229 comparing x,y coordinate are relative to the position of center. Matrix
230 is initially calculated with units arcseconds and then converted to
231 degrees. This yields higher precision results due to quirks in AST.
233 Parameters
234 ----------
235 wcs : `lsst.afw.geom.SkyWcs`
236 Wcs to approximate
237 center : `lsst.geom.Point2D`
238 Point at which to evaluate the LocalWcs.
239 skyCenter : `lsst.geom.SpherePoint`
240 Point on sky to approximate the Wcs.
242 Returns
243 -------
244 localMatrix : `numpy.ndarray`
245 Matrix representation the local wcs approximation with units
246 degrees.
247 """
248 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
249 localGnomonicWcs = afwGeom.makeSkyWcs(
250 center, skyCenter, blankCDMatrix)
251 measurementToLocalGnomonic = wcs.getTransform().then(
252 localGnomonicWcs.getTransform().inverted()
253 )
254 localMatrix = measurementToLocalGnomonic.getJacobian(center)
255 return localMatrix / 3600
257 def makeAlertDict(self,
258 alertId,
259 diaSource,
260 diaObject,
261 objDiaSrcHistory,
262 diffImCutout,
263 templateCutout):
264 """Convert data and package into a dictionary alert.
266 Parameters
267 ----------
268 diaSource : `pandas.DataFrame`
269 New single DiaSource to package.
270 diaObject : `pandas.DataFrame`
271 DiaObject that ``diaSource`` is matched to.
272 objDiaSrcHistory : `pandas.DataFrame`
273 12 month history of ``diaObject`` excluding the latest DiaSource.
274 diffImCutout : `astropy.nddata.CCDData` or `None`
275 Cutout of the difference image around the location of ``diaSource``
276 with a min size set by the ``cutoutSize`` configurable.
277 templateCutout : `astropy.nddata.CCDData` or `None`
278 Cutout of the template image around the location of ``diaSource``
279 with a min size set by the ``cutoutSize`` configurable.
280 """
281 alert = dict()
282 alert['alertId'] = alertId
283 alert['diaSource'] = diaSource.to_dict()
285 if objDiaSrcHistory is None:
286 alert['prvDiaSources'] = objDiaSrcHistory
287 else:
288 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
290 alert['prvDiaForcedSources'] = None
291 alert['prvDiaNondetectionLimits'] = None
293 alert['diaObject'] = diaObject.to_dict()
295 alert['ssObject'] = None
297 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
298 # TODO: add template cutouts in DM-24327
299 alert["cutoutTemplate"] = None
301 return alert
303 def streamCcdDataToBytes(self, cutout):
304 """Serialize a cutout into bytes.
306 Parameters
307 ----------
308 cutout : `astropy.nddata.CCDData`
309 Cutout to serialize.
311 Returns
312 -------
313 coutputBytes : `bytes`
314 Input cutout serialized into byte data.
315 """
316 with io.BytesIO() as streamer:
317 cutout.write(streamer, format="fits")
318 cutoutBytes = streamer.getvalue()
319 return cutoutBytes