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