Coverage for python/lsst/ap/association/packageAlerts.py: 17%
280 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 04:18 -0700
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 sys
27import time
29from astropy import wcs
30import astropy.units as u
31from astropy.nddata import CCDData, VarianceUncertainty
32import pandas as pd
33import struct
34import fastavro
35from confluent_kafka.cimpl import KafkaError
37# confluent_kafka is not in the standard Rubin environment as it is a third
38# party package and is only needed when producing alerts.
39try:
40 import confluent_kafka
41 from confluent_kafka import KafkaException
42 from confluent_kafka.admin import AdminClient
43except ImportError:
44 confluent_kafka = None
46import lsst.alert.packet as alertPack
47from lsst.afw.detection import InvalidPsfError
48import lsst.afw.geom as afwGeom
49import lsst.geom as geom
50import lsst.pex.config as pexConfig
51from lsst.pex.exceptions import InvalidParameterError
52import lsst.pipe.base as pipeBase
53from lsst.pipe.tasks.associationUtils import ss_object_id_to_obj_id
54import lsst.utils.logging
55from lsst.utils.timer import timeMethod
58class PackageAlertsConfig(pexConfig.Config):
59 """Config class for AssociationTask.
60 """
61 schemaFile = pexConfig.Field(
62 dtype=str,
63 doc="Schema definition file URI for the avro alerts.",
64 default=str(alertPack.get_uri_to_latest_schema())
65 )
66 minCutoutSize = pexConfig.RangeField(
67 dtype=int,
68 min=0,
69 max=1000,
70 default=30,
71 doc="Dimension of the square image cutouts to package in the alert."
72 )
73 maxCutoutSize = pexConfig.RangeField(
74 dtype=int,
75 min=0,
76 max=1000,
77 default=102,
78 doc="Dimension of the square image cutouts to package in the alert. The"
79 "default size comes from the max trail length in arcseconds "
80 "(10 deg/day) for a 30 second observation divided by the"
81 "arcseconds per pixel (0.2), which is 62.5 pixels. The effective"
82 "size of the psf (40 pixels) is then added for a total of 102 pixels. "
83 )
84 alertWriteLocation = pexConfig.Field(
85 dtype=str,
86 doc="Location to write alerts to.",
87 default=os.path.join(os.getcwd(), "alerts"),
88 )
90 doProduceAlerts = pexConfig.Field(
91 dtype=bool,
92 doc="Turn on alert production to kafka if true and if confluent_kafka is in the environment.",
93 default=False,
94 )
96 doWriteAlerts = pexConfig.Field(
97 dtype=bool,
98 doc="Write alerts to disk if true.",
99 default=False,
100 )
102 doWriteFailedAlerts = pexConfig.Field(
103 dtype=bool,
104 doc="If an alert cannot be sent when doProduceAlerts is set, "
105 "write it to disk for debugging purposes.",
106 default=False,
107 )
109 maxTimeout = pexConfig.Field(
110 dtype=float,
111 doc="Sets the maximum time in seconds to wait for the alert stream "
112 "broker to respond to a query before timing out.",
113 default=15.0,
114 )
116 deliveryTimeout = pexConfig.Field(
117 dtype=float,
118 doc="Sets the time to wait for the producer to wait to deliver an "
119 "alert in milliseconds.",
120 default=1200.0,
121 )
123 useAveragePsf = pexConfig.Field(
124 dtype=bool,
125 doc="Use the average PSF for the image, instead of the PSF for each cutout. "
126 "This option is much less accurate, but much faster.",
127 default=False,
128 )
131class PackageAlertsTask(pipeBase.Task):
132 """Tasks for packaging Dia and Pipelines data into Avro alert packages.
133 """
134 ConfigClass = PackageAlertsConfig
135 _DefaultName = "packageAlerts"
137 _scale = (1.0 * geom.arcseconds).asDegrees()
139 def __init__(self, **kwargs):
140 super().__init__(**kwargs)
141 self.alertSchema = alertPack.Schema.from_uri(self.config.schemaFile)
142 os.makedirs(self.config.alertWriteLocation, exist_ok=True)
144 if self.config.doProduceAlerts:
145 if confluent_kafka is not None:
146 self.password = os.getenv("AP_KAFKA_PRODUCER_PASSWORD")
147 if not self.password:
148 raise ValueError("Kafka password environment variable was not set.")
149 self.username = os.getenv("AP_KAFKA_PRODUCER_USERNAME")
150 if not self.username:
151 raise ValueError("Kafka username environment variable was not set.")
152 self.server = os.getenv("AP_KAFKA_SERVER")
153 if not self.server:
154 raise ValueError("Kafka server environment variable was not set.")
155 self.kafkaTopic = os.getenv("AP_KAFKA_TOPIC")
156 if not self.kafkaTopic:
157 raise ValueError("Kafka topic environment variable was not set.")
159 # confluent_kafka configures all of its classes with dictionaries. This one
160 # sets up the bare minimum that is needed.
161 self.kafkaConfig = {
162 # This is the URL to use to connect to the Kafka cluster.
163 "bootstrap.servers": self.server,
164 # These next two properties tell the Kafka client about the specific
165 # authentication and authorization protocols that should be used when
166 # connecting.
167 "security.protocol": "SASL_PLAINTEXT",
168 "sasl.mechanisms": "SCRAM-SHA-512",
169 # The sasl.username and sasl.password are passed through over
170 # SCRAM-SHA-512 auth to connect to the cluster. The username is not
171 # sensitive, but the password is (of course) a secret value which
172 # should never be committed to source code.
173 "sasl.username": self.username,
174 "sasl.password": self.password,
175 # Batch size limits the largest size of a kafka alert that can be sent.
176 # We set the batch size to 2 Mb.
177 "batch.size": 2097152,
178 "linger.ms": 5,
179 "delivery.timeout.ms": self.config.deliveryTimeout,
180 # Compression options are snappy, lz4, zstd, and gzip.
181 "compression.type": 'snappy'
182 }
183 self.kafkaAdminConfig = {
184 # This is the URL to use to connect to the Kafka cluster.
185 "bootstrap.servers": self.server,
186 # These next two properties tell the Kafka client about the specific
187 # authentication and authorization protocols that should be used when
188 # connecting.
189 "security.protocol": "SASL_PLAINTEXT",
190 "sasl.mechanisms": "SCRAM-SHA-512",
191 # The sasl.username and sasl.password are passed through over
192 # SCRAM-SHA-512 auth to connect to the cluster. The username is not
193 # sensitive, but the password is (of course) a secret value which
194 # should never be committed to source code.
195 "sasl.username": self.username,
196 "sasl.password": self.password,
197 }
199 self._server_check()
200 self.producer = confluent_kafka.Producer(**self.kafkaConfig)
202 else:
203 raise RuntimeError("Produce alerts is set but confluent_kafka is not present in "
204 "the environment. Alerts will not be sent to the alert stream.")
206 @timeMethod
207 def run(self,
208 diaSourceCat,
209 diaObjectCat,
210 diaSrcHistory,
211 diaForcedSources,
212 diffIm,
213 calexp,
214 template,
215 ssSrc=None,
216 doRunForcedMeasurement=True,
217 forcedSourceHistoryThreshold=0,
218 ):
219 """Package DiaSources/Object and exposure data into Avro alerts.
221 Alerts can be sent to the alert stream if ``doProduceAlerts`` is set
222 and written to disk if ``doWriteAlerts`` is set. Both can be set at the
223 same time, and are independent of one another.
225 Writes Avro alerts to a location determined by the
226 ``alertWriteLocation`` configurable.
228 Parameters
229 ----------
230 diaSourceCat : `pandas.DataFrame`
231 New DiaSources to package. DataFrame should be indexed on
232 ``["diaObjectId", "band", "diaSourceId"]``
233 diaObjectCat : `pandas.DataFrame`
234 New and updated DiaObjects matched to the new DiaSources. DataFrame
235 is indexed on ``["diaObjectId"]``
236 diaSrcHistory : `pandas.DataFrame`
237 12 month history of DiaSources matched to the DiaObjects. Excludes
238 the newest DiaSource and is indexed on
239 ``["diaObjectId", "band", "diaSourceId"]``
240 diaForcedSources : `pandas.DataFrame`
241 12 month history of DiaForcedSources matched to the DiaObjects.
242 ``["diaObjectId"]``
243 diffIm : `lsst.afw.image.ExposureF`
244 Difference image the sources in ``diaSourceCat`` were detected in.
245 calexp : `lsst.afw.image.ExposureF`
246 Calexp used to create the ``diffIm``.
247 template : `lsst.afw.image.ExposureF` or `None`
248 Template image used to create the ``diffIm``.
249 ssSrc : `astropy.table.Table`, optional
250 Solar system specific information for diaSources associated to ssObjects.
251 doRunForcedMeasurement : `bool`, optional
252 Flag to indicate whether forced measurement was run.
253 This should only be turned off for debugging purposes.
254 Added to allow disabling forced sources for performance
255 reasons during the ops rehearsal.
256 forcedSourceHistoryThreshold : `int`, optional
257 Minimum number of detections of a diaObject required
258 to run forced photometry. Set to 1 to include all diaObjects.
259 """
261 alerts = []
262 self._patchDiaSources(diaSourceCat)
263 self._patchDiaSources(diaSrcHistory)
264 detector = diffIm.detector.getId()
265 visit = diffIm.visitInfo.id
266 midpoint_unix = diffIm.visitInfo.date.toAstropy().tai.unix
267 exposure_time = diffIm.visitInfo.exposureTime
268 diffImPhotoCalib = diffIm.getPhotoCalib()
269 calexpPhotoCalib = calexp.getPhotoCalib()
270 templatePhotoCalib = template.getPhotoCalib()
271 diffImPsf = self._computePsf(diffIm, diffIm.psf.getAveragePosition())
272 sciencePsf = self._computePsf(calexp, calexp.psf.getAveragePosition())
273 templatePsf = self._computePsf(template, template.psf.getAveragePosition())
275 # cast to object; works around pandas coercing Series to nullable
276 # Float64 type when selecting single rows below
277 diaSourceCat = diaSourceCat.astype('object')
278 diaObjectCat = diaObjectCat.astype('object')
279 diaSrcHistory = diaSrcHistory.astype('object')
280 diaForcedSources = diaForcedSources.astype('object')
282 if ssSrc is not None:
283 ssSrc = ssSrc.set_index('diaSourceId')
284 ssSrc = ssSrc.astype('object')
286 n_sources = len(diaSourceCat)
287 self.log.info("Packaging alerts for %d DiaSources.", n_sources)
288 # Log every 10 seconds as proof of liveness.
289 loop_logger = lsst.utils.logging.PeriodicLogger(self.log, interval=10.0)
291 for srcIndex, diaSource in diaSourceCat.iterrows():
292 loop_logger.log("%s/%s sources have been packaged.", len(alerts), n_sources)
294 alertId = diaSource["diaSourceId"]
295 if srcIndex[0] == 0:
296 diaObject = None
297 objSourceHistory = None
298 objDiaForcedSources = None
299 ssSource = ssSrc.loc[[alertId]].to_dict("records")[0]
300 else:
301 ssSource = None
302 diaObject = diaObjectCat.loc[srcIndex[0]]
303 if diaObject["nDiaSources"] > 1:
304 objSourceHistory = diaSrcHistory.loc[srcIndex[0]]
305 else:
306 objSourceHistory = None
307 if doRunForcedMeasurement and diaObject["nDiaSources"] >= forcedSourceHistoryThreshold:
308 objDiaForcedSources = diaForcedSources.loc[srcIndex[0]]
309 else:
310 # Send empty table with correct columns
311 objDiaForcedSources = diaForcedSources.loc[[]]
313 sphPoint = geom.SpherePoint(diaSource["ra"],
314 diaSource["dec"],
315 geom.degrees)
316 pixelPoint = geom.Point2D(diaSource["x"],
317 diaSource["y"])
319 cutoutExtent = self.createDiaSourceExtent(diaSource["bboxSize"])
320 diffImCutout = self.createCcdDataCutout(
321 diffIm,
322 sphPoint,
323 pixelPoint,
324 cutoutExtent,
325 diffImPhotoCalib,
326 diaSource["diaSourceId"],
327 averagePsf=diffImPsf,
328 rotPa=diffIm.visitInfo.boresightRotAngle.asDegrees())
329 calexpCutout = self.createCcdDataCutout(
330 calexp,
331 sphPoint,
332 pixelPoint,
333 cutoutExtent,
334 calexpPhotoCalib,
335 diaSource["diaSourceId"],
336 averagePsf=sciencePsf,
337 rotPa=calexp.visitInfo.boresightRotAngle.asDegrees())
338 templateCutout = self.createCcdDataCutout(
339 template,
340 sphPoint,
341 pixelPoint,
342 cutoutExtent,
343 templatePhotoCalib,
344 diaSource["diaSourceId"],
345 averagePsf=templatePsf,
346 rotPa=calexp.visitInfo.boresightRotAngle.asDegrees())
348 alerts.append(
349 self.makeAlertDict(diaSource["diaSourceId"],
350 diaSource,
351 diaObject,
352 objSourceHistory,
353 objDiaForcedSources,
354 diffImCutout,
355 calexpCutout,
356 templateCutout,
357 ssSource))
359 if self.config.doProduceAlerts:
360 self.log.info("Producing alerts to %s.", self.kafkaTopic)
361 self.produceAlerts(alerts, visit, detector, midpoint_unix, exposure_time)
362 else:
363 # Fill values for the metadata so that downstream metrics don't crash
364 self.metadata['visit_midpoint'] = midpoint_unix
365 self.metadata['produce_end_timestamp'] = -1
366 self.metadata['produce_start_timestamp'] = -1
367 self.metadata['alert_timing_since_shutter_close'] = -1
368 self.metadata['total_alerts'] = -1
370 if self.config.doWriteAlerts:
371 avro_path = os.path.join(self.config.alertWriteLocation, f"{visit}_{detector}.avro")
372 self.log.info("Writing alerts to %s.", avro_path)
373 with open(avro_path, "wb") as f:
374 self.alertSchema.store_alerts(f, alerts)
376 def _patchDiaSources(self, diaSources):
377 """Add the ``programId`` column to the data.
379 Parameters
380 ----------
381 diaSources : `pandas.DataFrame`
382 DataFrame of DiaSources to patch.
383 """
384 diaSources["programId"] = 0
386 def createDiaSourceExtent(self, bboxSize):
387 """Create an extent for a box for the cutouts given the size of the
388 square BBox that covers the source footprint.
390 Parameters
391 ----------
392 bboxSize : `int`
393 Size of a side of the square bounding box in pixels.
395 Returns
396 -------
397 extent : `lsst.geom.Extent2I`
398 Geom object representing the size of the bounding box.
399 """
400 if bboxSize < self.config.minCutoutSize:
401 extent = geom.Extent2I(self.config.minCutoutSize,
402 self.config.minCutoutSize)
403 elif bboxSize > self.config.maxCutoutSize:
404 extent = geom.Extent2I(self.config.maxCutoutSize,
405 self.config.maxCutoutSize)
406 else:
407 extent = geom.Extent2I(bboxSize, bboxSize)
408 return extent
410 @timeMethod
411 def produceAlerts(self, alerts, visit, detector, midpoint_unix, exposure_time):
412 """Serialize alerts and send them to the alert stream using
413 confluent_kafka's producer.
415 Parameters
416 ----------
417 alerts : `dict`
418 Dictionary of alerts to be sent to the alert stream.
419 visit, detector : `int`
420 Visit and detector ids of these alerts. Used to write out alerts
421 which fail to be sent to the alert stream.
422 """
423 schema_id = self.alertSchema.get_schema_id()
424 # Serialize and send alerts with the producer timestamp.
425 total_alerts = 0
426 produce_start_timestamp = time.time()
427 try:
428 for alert in alerts:
429 alertBytes = self._serializeAlert(alert,
430 schema=self.alertSchema.definition, schema_id=schema_id)
431 try:
432 timestamp = time.time()*1000 # Current time in milliseconds
433 headers = [("producer_timestamp", str(timestamp).encode('utf-8'))]
434 self.producer.produce(self.kafkaTopic, alertBytes, callback=self._delivery_callback,
435 headers=headers)
436 self.producer.flush()
437 total_alerts += 1
439 except KafkaException as e:
440 self.log.warning('Kafka error: %s, message was %s bytes',
441 e, sys.getsizeof(alertBytes))
443 if self.config.doWriteFailedAlerts:
444 with open(os.path.join(self.config.alertWriteLocation,
445 f"{visit}_{detector}_{alert['diaSourceId']}.avro"), "wb") as f:
446 f.write(alertBytes)
448 self.producer.flush()
449 finally:
450 produce_end_timestamp = time.time() # Current time in seconds
451 total_time = produce_end_timestamp - (midpoint_unix + exposure_time/2.0)
452 self.metadata['visit_midpoint'] = midpoint_unix
453 self.metadata['produce_end_timestamp'] = produce_end_timestamp
454 self.metadata['produce_start_timestamp'] = produce_start_timestamp
455 self.metadata['alert_timing_since_shutter_close'] = total_time
456 self.metadata['total_alerts'] = total_alerts
457 # A single log message is easier for Loki to parse than timeMethod's start+end pairs.
458 self.log.verbose("Producing alerts: Took %.4f seconds",
459 produce_end_timestamp - produce_start_timestamp)
461 self.log.info(f"Total time since shutter close to produce alerts for"
462 f" visit {visit} detector {detector}: {total_time} seconds")
464 def createCcdDataCutout(self, image, skyCenter, pixelCenter, extent, photoCalib, srcId, averagePsf=None,
465 rotPa=None):
466 """Grab an image as a cutout and return a calibrated CCDData image.
468 Parameters
469 ----------
470 image : `lsst.afw.image.ExposureF`
471 Image to pull cutout from.
472 skyCenter : `lsst.geom.SpherePoint`
473 Center point of DiaSource on the sky.
474 pixelCenter : `lsst.geom.Point2D`
475 Pixel center of DiaSource on the sky.
476 extent : `lsst.geom.Extent2I`
477 Bounding box to cutout from the image.
478 photoCalib : `lsst.afw.image.PhotoCalib`
479 Calibrate object of the image the cutout is cut from.
480 srcId : `int`
481 Unique id of DiaSource. Used for when an error occurs extracting
482 a cutout.
483 averagePsf : `numpy.array`, optional
484 Average PSF to attach to the cutout.
485 Used if ``self.config.useAveragePsf`` is set.
486 cutoutType : str, optional
487 Type of cutout being created ('difference', 'template', or 'science')
489 Returns
490 -------
491 ccdData : `astropy.nddata.CCDData` or `None`
492 CCDData object storing the calibrate information from the input
493 difference or template image.
494 """
495 imBBox = image.getBBox()
496 if not geom.Box2D(image.getBBox()).contains(pixelCenter):
497 self.log.warning(
498 "DiaSource id=%i centroid lies at pixel (%.2f, %.2f) "
499 "which is outside the Exposure with bounding box "
500 "((%i, %i), (%i, %i)). Returning None for cutout...",
501 srcId, pixelCenter.x, pixelCenter.y,
502 imBBox.minX, imBBox.maxX, imBBox.minY, imBBox.maxY)
503 return None
504 # Catch errors in retrieving the cutout.
505 try:
506 cutout = image.getCutout(pixelCenter, extent)
507 except InvalidParameterError:
508 self.log.warning(
509 "Failed to retrieve cutout from image for DiaSource with "
510 "id=%i. InvalidParameterError thrown during cutout "
511 "creation. Returning None for cutout..."
512 % srcId)
513 if self.config.useAveragePsf:
514 if averagePsf is None:
515 self.log.info("Using source id=%i to compute the average PSF.", srcId)
516 averagePsf = self._computePsf(image, pixelCenter, srcId=srcId)
517 cutoutPsf = averagePsf
518 else:
519 cutoutPsf = self._computePsf(image, pixelCenter, srcId=srcId)
521 # Find the value of the bottom corner of our cutout's BBox and
522 # subtract 1 so that the CCDData cutout position value will be
523 # [1, 1].
524 cutOutMinX = cutout.getBBox().minX - 1
525 cutOutMinY = cutout.getBBox().minY - 1
526 center = cutout.getWcs().skyToPixel(skyCenter)
527 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
529 cutoutWcs = wcs.WCS(naxis=2)
530 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
531 cutout.getBBox().getWidth())
532 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
533 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
534 skyCenter.getDec().asDegrees()]
535 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
536 center,
537 skyCenter)
538 ccdData = CCDData(
539 data=calibCutout.getImage().array,
540 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
541 flags=calibCutout.getMask().array,
542 wcs=cutoutWcs,
543 psf=cutoutPsf,
544 meta={"cutMinX": cutOutMinX,
545 "cutMinY": cutOutMinY},
546 unit=u.nJy)
548 ccdData.header['ROTPA'] = (rotPa, 'Pos angle in deg of focal plane +Y wrt North')
550 return ccdData
552 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
553 """Create a local, linear approximation of the wcs transformation
554 matrix.
556 The approximation is created as if the center is at RA=0, DEC=0. All
557 comparing x,y coordinate are relative to the position of center. Matrix
558 is initially calculated with units arcseconds and then converted to
559 degrees. This yields higher precision results due to quirks in AST.
561 Parameters
562 ----------
563 wcs : `lsst.afw.geom.SkyWcs`
564 Wcs to approximate
565 center : `lsst.geom.Point2D`
566 Point at which to evaluate the LocalWcs.
567 skyCenter : `lsst.geom.SpherePoint`
568 Point on sky to approximate the Wcs.
570 Returns
571 -------
572 localMatrix : `numpy.ndarray`
573 Matrix representation the local wcs approximation with units
574 degrees.
575 """
576 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
577 localGnomonicWcs = afwGeom.makeSkyWcs(
578 center, skyCenter, blankCDMatrix)
579 measurementToLocalGnomonic = wcs.getTransform().then(
580 localGnomonicWcs.getTransform().inverted()
581 )
582 localMatrix = measurementToLocalGnomonic.getJacobian(center)
583 return localMatrix / 3600
585 def makeAlertDict(self,
586 diaSourceId,
587 diaSource,
588 diaObject,
589 objDiaSrcHistory,
590 objDiaForcedSources,
591 diffImCutout,
592 calexpCutout,
593 templateCutout,
594 ssSource=None):
595 """Convert data and package into a dictionary alert.
597 Parameters
598 ----------
599 diaSourceId : `int`
600 Unique identifier of the triggering diaSource
601 diaSource : `pandas.DataFrame`
602 New single DiaSource to package.
603 diaObject : `pandas.DataFrame`
604 DiaObject that ``diaSource`` is matched to.
605 objDiaSrcHistory : `pandas.DataFrame`
606 12 month history of ``diaObject`` excluding the latest DiaSource.
607 objDiaForcedSources : `pandas.DataFrame`
608 12 month history of ``diaObject`` forced measurements.
609 diffImCutout : `astropy.nddata.CCDData` or `None`
610 Cutout of the difference image around the location of ``diaSource``
611 with a min size set by the ``cutoutSize`` configurable.
612 calexpCutout : `astropy.nddata.CCDData` or `None`
613 Cutout of the calexp around the location of ``diaSource``
614 with a min size set by the ``cutoutSize`` configurable.
615 templateCutout : `astropy.nddata.CCDData` or `None`
616 Cutout of the template image around the location of ``diaSource``
617 with a min size set by the ``cutoutSize`` configurable.
618 """
619 alert = dict()
620 alert['diaSourceId'] = diaSourceId
621 alert['diaSource'] = diaSource.to_dict()
623 if objDiaSrcHistory is None:
624 alert['prvDiaSources'] = None
625 else:
626 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
628 if isinstance(objDiaForcedSources, pd.Series):
629 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
630 elif objDiaForcedSources is None:
631 alert['prvDiaForcedSources'] = None
632 else:
633 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
635 alert['prvDiaNondetectionLimits'] = None
637 if diaObject is None:
638 alert['diaObject'] = None
639 else:
640 alert['diaObject'] = diaObject.to_dict()
642 if ssSource is None:
643 alert['ssSource'] = None
644 alert['MPCORB'] = None
645 else:
646 mpcorbColumns = [col for col in ssSource if col[:7] == 'MPCORB_']
647 mpcOrbit = {key[7:]: ssSource[key] for key in mpcorbColumns}
648 mpcOrbit['ssObjectId'] = ssSource['ssObjectId']
649 mpcOrbit['mpcDesignation'] = ss_object_id_to_obj_id(ssSource['ssObjectId'])[0]
650 ssSource = {key: ssSource[key] for key in ssSource if key not in mpcorbColumns}
651 ssSource['diaSourceId'] = diaSourceId
652 alert['ssSource'] = ssSource
653 alert['MPCORB'] = mpcOrbit
655 if diffImCutout is None:
656 alert['cutoutDifference'] = None
657 else:
658 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
660 if calexpCutout is None:
661 alert['cutoutScience'] = None
662 else:
663 alert['cutoutScience'] = self.streamCcdDataToBytes(calexpCutout)
665 if templateCutout is None:
666 alert["cutoutTemplate"] = None
667 else:
668 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout)
670 return alert
672 def streamCcdDataToBytes(self, cutout):
673 """Serialize a cutout into bytes.
675 Parameters
676 ----------
677 cutout : `astropy.nddata.CCDData`
678 Cutout to serialize.
680 Returns
681 -------
682 coutputBytes : `bytes`
683 Input cutout serialized into byte data.
684 """
685 with io.BytesIO() as streamer:
686 cutout.write(streamer, format="fits")
687 cutoutBytes = streamer.getvalue()
688 return cutoutBytes
690 def _serializeAlert(self, alert, schema=None, schema_id=0):
691 """Serialize an alert to a byte sequence for sending to Kafka.
693 Parameters
694 ----------
695 alert : `dict`
696 An alert payload to be serialized.
697 schema : `dict`, optional
698 An Avro schema definition describing how to encode `alert`. By default,
699 the schema is None, which sets it to the latest schema available.
700 schema_id : `int`, optional
701 The Confluent Schema Registry ID of the schema. By default, 0 (an
702 invalid ID) is used, indicating that the schema is not registered.
704 Returns
705 -------
706 serialized : `bytes`
707 The byte sequence describing the alert, including the Confluent Wire
708 Format prefix.
709 """
710 if schema is None:
711 schema = self.alertSchema.definition
713 buf = io.BytesIO()
714 buf.write(self._serializeConfluentWireHeader(schema_id))
715 fastavro.schemaless_writer(buf, schema, alert)
716 return buf.getvalue()
718 @staticmethod
719 def _serializeConfluentWireHeader(schema_version):
720 """Returns the byte prefix for Confluent Wire Format-style Kafka messages.
722 Parameters
723 ----------
724 schema_version : `int`
725 A version number which indicates the Confluent Schema Registry ID
726 number of the Avro schema used to encode the message that follows this
727 header.
729 Returns
730 -------
731 header : `bytes`
732 The 5-byte encoded message prefix.
734 Notes
735 -----
736 The Confluent Wire Format is described more fully here:
737 https://docs.confluent.io/current/schema-registry/serdes-develop/index.html#wire-format
738 """
739 ConfluentWireFormatHeader = struct.Struct(">bi")
740 return ConfluentWireFormatHeader.pack(0, schema_version)
742 def _delivery_callback(self, err, msg):
743 if err:
744 if err.code() == KafkaError.UNKNOWN_TOPIC_OR_PART:
745 topic = msg.topic() if msg and hasattr(msg,
746 'topic') else 'unknown'
747 self.log.warning(
748 'Message failed delivery. Topic %s unknown: %s',
749 topic, err)
750 else:
751 self.log.debug('Message delivered to %s [%d] @ %d', msg.topic(), msg.partition(), msg.offset())
753 def _server_check(self):
754 """Checks if the alert stream credentials are still valid and the
755 server is contactable.
757 Raises
758 -------
759 KafkaException
760 Raised if the server us not contactable.
761 RuntimeError
762 Raised if the server is contactable but there are no topics
763 present.
764 """
765 admin_client = AdminClient(self.kafkaAdminConfig)
766 topics = admin_client.list_topics(timeout=self.config.maxTimeout).topics
768 if not topics:
769 raise RuntimeError()
770 if self.kafkaTopic not in topics:
771 raise RuntimeError(
772 f"Topic '{self.kafkaTopic}' not found on the Kafka server. "
773 f"Check that the correct topic has been provided and exists "
774 f"on the Kafka server.")
776 def _computePsf(self, exposure, pixelCenter, srcId=None):
777 """Compute the PSF at a location and catch errors.
779 Parameters
780 ----------
781 exposure : `lsst.afw.image.Exposure`
782 The image to compute the PSF for.
783 pixelCenter : `lsst.geom.Point2D`
784 The location on the image to compute the PSF.
785 srcId : `int`, optional
786 Unique id of DiaSource. Used for when an error occurs extracting
787 a cutout.
789 Returns
790 -------
791 cutoutPsf : `numpy.array`
792 Array of the PSF values.
793 """
794 try:
795 # use exposure.psf.computeKernelImage to provide PSF centered in the array
796 cutoutPsf = exposure.psf.computeKernelImage(pixelCenter).array
797 except InvalidParameterError:
798 if srcId is not None:
799 msg = "Could not calculate PSF for DiaSource with "\
800 "id=%i. InvalidParameterError encountered. Exiting."\
801 % srcId
802 else:
803 msg = "Could not calculate average PSF for the image"
804 self.log.warning(msg)
805 cutoutPsf = None
806 except InvalidPsfError:
807 if srcId is not None:
808 msg = "Could not calculate PSF for DiaSource with "\
809 "id=%i. InvalidPsfError encountered. Exiting."\
810 % srcId
811 else:
812 msg = "Could not calculate average PSF for the image"
813 self.log.warning(msg)
814 cutoutPsf = None
815 return cutoutPsf