Coverage for python/lsst/ap/association/packageAlerts.py: 78%
305 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-10 01:47 -0800
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 findMissingDiaSourceColumns(self, diaSources):
396 """Check that a DiaSource catalog has the correct schema for alerts.
398 Parameters
399 ----------
400 diaSources : `pandas.DataFrame`
401 DiaSources to include in alerts.
403 Returns
404 -------
405 missing : `list` [`str`]
406 Names of the DiaSource fields in the alert schema that have no
407 default, do not accept null, and are not columns of
408 ``diaSources``, in schema order.
409 """
410 diaSourceSchema = next(field["type"] for field in self.alertSchema.definition["fields"]
411 if field["name"] == "diaSource")
412 return [field["name"] for field in diaSourceSchema["fields"]
413 if "default" not in field
414 and not (isinstance(field["type"], list) and "null" in field["type"])
415 and field["name"] not in diaSources.columns]
417 def createDiaSourceExtent(self, bboxSize):
418 """Create an extent for a box for the cutouts given the size of the
419 square BBox that covers the source footprint.
421 Parameters
422 ----------
423 bboxSize : `int`
424 Size of a side of the square bounding box in pixels.
426 Returns
427 -------
428 extent : `lsst.geom.Extent2I`
429 Geom object representing the size of the bounding box.
430 """
431 if bboxSize < self.config.minCutoutSize:
432 extent = geom.Extent2I(self.config.minCutoutSize,
433 self.config.minCutoutSize)
434 elif bboxSize > self.config.maxCutoutSize:
435 extent = geom.Extent2I(self.config.maxCutoutSize,
436 self.config.maxCutoutSize)
437 else:
438 extent = geom.Extent2I(bboxSize, bboxSize)
439 return extent
441 @timeMethod
442 def produceAlerts(self, alerts, visit, detector, midpoint_unix, exposure_time):
443 """Serialize alerts and send them to the alert stream using
444 confluent_kafka's producer.
446 Parameters
447 ----------
448 alerts : `dict`
449 Dictionary of alerts to be sent to the alert stream.
450 visit, detector : `int`
451 Visit and detector ids of these alerts. Used to write out alerts
452 which fail to be sent to the alert stream.
453 """
454 schema_id = self.alertSchema.get_schema_id()
455 # Serialize and send alerts with the producer timestamp.
456 total_alerts = 0
457 produce_start_timestamp = time.time()
458 try:
459 for alert in alerts:
460 alertBytes = self._serializeAlert(alert,
461 schema=self.alertSchema.definition, schema_id=schema_id)
462 try:
463 timestamp = time.time()*1000 # Current time in milliseconds
464 headers = [("producer_timestamp", str(timestamp).encode('utf-8'))]
465 self.producer.produce(self.kafkaTopic, alertBytes, callback=self._delivery_callback,
466 headers=headers)
467 self.producer.flush()
468 total_alerts += 1
470 except KafkaException as e:
471 self.log.warning('Kafka error: %s, message was %s bytes',
472 e, sys.getsizeof(alertBytes))
474 if self.config.doWriteFailedAlerts: 474 ↛ 459line 474 didn't jump to line 459 because the condition on line 474 was always true
475 with open(os.path.join(self.config.alertWriteLocation,
476 f"{visit}_{detector}_{alert['diaSourceId']}.avro"), "wb") as f:
477 f.write(alertBytes)
479 self.producer.flush()
480 finally:
481 produce_end_timestamp = time.time() # Current time in seconds
482 total_time = produce_end_timestamp - (midpoint_unix + exposure_time/2.0)
483 self.metadata['visit_midpoint'] = midpoint_unix
484 self.metadata['produce_end_timestamp'] = produce_end_timestamp
485 self.metadata['produce_start_timestamp'] = produce_start_timestamp
486 self.metadata['alert_timing_since_shutter_close'] = total_time
487 self.metadata['total_alerts'] = total_alerts
488 # A single log message is easier for Loki to parse than timeMethod's start+end pairs.
489 self.log.verbose("Producing alerts: Took %.4f seconds",
490 produce_end_timestamp - produce_start_timestamp)
492 self.log.info(f"Total time since shutter close to produce alerts for"
493 f" visit {visit} detector {detector}: {total_time} seconds")
495 def createCcdDataCutout(self, image, skyCenter, pixelCenter, extent, photoCalib, srcId, averagePsf=None,
496 rotPa=None):
497 """Grab an image as a cutout and return a calibrated CCDData image.
499 Parameters
500 ----------
501 image : `lsst.afw.image.ExposureF`
502 Image to pull cutout from.
503 skyCenter : `lsst.geom.SpherePoint`
504 Center point of DiaSource on the sky.
505 pixelCenter : `lsst.geom.Point2D`
506 Pixel center of DiaSource on the sky.
507 extent : `lsst.geom.Extent2I`
508 Bounding box to cutout from the image.
509 photoCalib : `lsst.afw.image.PhotoCalib`
510 Calibrate object of the image the cutout is cut from.
511 srcId : `int`
512 Unique id of DiaSource. Used for when an error occurs extracting
513 a cutout.
514 averagePsf : `numpy.array`, optional
515 Average PSF to attach to the cutout.
516 Used if ``self.config.useAveragePsf`` is set.
517 cutoutType : str, optional
518 Type of cutout being created ('difference', 'template', or 'science')
520 Returns
521 -------
522 ccdData : `astropy.nddata.CCDData` or `None`
523 CCDData object storing the calibrate information from the input
524 difference or template image.
525 """
526 imBBox = image.getBBox()
527 if not geom.Box2D(image.getBBox()).contains(pixelCenter):
528 self.log.warning(
529 "DiaSource id=%i centroid lies at pixel (%.2f, %.2f) "
530 "which is outside the Exposure with bounding box "
531 "((%i, %i), (%i, %i)). Returning None for cutout...",
532 srcId, pixelCenter.x, pixelCenter.y,
533 imBBox.minX, imBBox.maxX, imBBox.minY, imBBox.maxY)
534 return None
535 # Catch errors in retrieving the cutout.
536 try:
537 cutout = image.getCutout(pixelCenter, extent)
538 except InvalidParameterError:
539 self.log.warning(
540 "Failed to retrieve cutout from image for DiaSource with "
541 "id=%i. InvalidParameterError thrown during cutout "
542 "creation. Returning None for cutout..."
543 % srcId)
544 if self.config.useAveragePsf:
545 if averagePsf is None:
546 self.log.info("Using source id=%i to compute the average PSF.", srcId)
547 averagePsf = self._computePsf(image, pixelCenter, srcId=srcId)
548 cutoutPsf = averagePsf
549 else:
550 cutoutPsf = self._computePsf(image, pixelCenter, srcId=srcId)
552 # Find the value of the bottom corner of our cutout's BBox and
553 # subtract 1 so that the CCDData cutout position value will be
554 # [1, 1].
555 cutOutMinX = cutout.getBBox().minX - 1
556 cutOutMinY = cutout.getBBox().minY - 1
557 center = cutout.getWcs().skyToPixel(skyCenter)
558 calibCutout = photoCalib.calibrateImage(cutout.getMaskedImage())
560 cutoutWcs = wcs.WCS(naxis=2)
561 cutoutWcs.array_shape = (cutout.getBBox().getWidth(),
562 cutout.getBBox().getWidth())
563 cutoutWcs.wcs.crpix = [center.x - cutOutMinX, center.y - cutOutMinY]
564 cutoutWcs.wcs.crval = [skyCenter.getRa().asDegrees(),
565 skyCenter.getDec().asDegrees()]
566 cutoutWcs.wcs.cd = self.makeLocalTransformMatrix(cutout.getWcs(),
567 center,
568 skyCenter)
569 ccdData = CCDData(
570 data=calibCutout.getImage().array,
571 uncertainty=VarianceUncertainty(calibCutout.getVariance().array),
572 flags=calibCutout.getMask().array,
573 wcs=cutoutWcs,
574 psf=cutoutPsf,
575 meta={"cutMinX": cutOutMinX,
576 "cutMinY": cutOutMinY},
577 unit=u.nJy)
579 ccdData.header['ROTPA'] = (rotPa, 'Pos angle in deg of focal plane +Y wrt North')
581 return ccdData
583 def makeLocalTransformMatrix(self, wcs, center, skyCenter):
584 """Create a local, linear approximation of the wcs transformation
585 matrix.
587 The approximation is created as if the center is at RA=0, DEC=0. All
588 comparing x,y coordinate are relative to the position of center. Matrix
589 is initially calculated with units arcseconds and then converted to
590 degrees. This yields higher precision results due to quirks in AST.
592 Parameters
593 ----------
594 wcs : `lsst.afw.geom.SkyWcs`
595 Wcs to approximate
596 center : `lsst.geom.Point2D`
597 Point at which to evaluate the LocalWcs.
598 skyCenter : `lsst.geom.SpherePoint`
599 Point on sky to approximate the Wcs.
601 Returns
602 -------
603 localMatrix : `numpy.ndarray`
604 Matrix representation the local wcs approximation with units
605 degrees.
606 """
607 blankCDMatrix = [[self._scale, 0], [0, self._scale]]
608 localGnomonicWcs = afwGeom.makeSkyWcs(
609 center, skyCenter, blankCDMatrix)
610 measurementToLocalGnomonic = wcs.getTransform().then(
611 localGnomonicWcs.getTransform().inverted()
612 )
613 localMatrix = measurementToLocalGnomonic.getJacobian(center)
614 return localMatrix / 3600
616 def makeAlertDict(self,
617 diaSourceId,
618 observationReason,
619 targetName,
620 diaSource,
621 diaObject,
622 objDiaSrcHistory,
623 objDiaForcedSources,
624 diffImCutout,
625 calexpCutout,
626 templateCutout,
627 ssSource=None):
628 """Convert data and package into a dictionary alert.
630 Parameters
631 ----------
632 diaSourceId : `int`
633 Unique identifier of the triggering diaSource
634 diaSource : `pandas.DataFrame`
635 New single DiaSource to package.
636 observationReason : `str`
637 Scheduler reason for the image containing this diaSource.
638 targetName : `str`
639 Scheduler target for the image containing this diaSource.
640 diaObject : `pandas.DataFrame`
641 DiaObject that ``diaSource`` is matched to.
642 objDiaSrcHistory : `pandas.DataFrame`
643 12 month history of ``diaObject`` excluding the latest DiaSource.
644 objDiaForcedSources : `pandas.DataFrame`
645 12 month history of ``diaObject`` forced measurements.
646 diffImCutout : `astropy.nddata.CCDData` or `None`
647 Cutout of the difference image around the location of ``diaSource``
648 with a min size set by the ``cutoutSize`` configurable.
649 calexpCutout : `astropy.nddata.CCDData` or `None`
650 Cutout of the calexp around the location of ``diaSource``
651 with a min size set by the ``cutoutSize`` configurable.
652 templateCutout : `astropy.nddata.CCDData` or `None`
653 Cutout of the template image around the location of ``diaSource``
654 with a min size set by the ``cutoutSize`` configurable.
655 """
656 alert = dict()
657 alert['diaSourceId'] = diaSourceId
658 # alerts should have NULL if scheduler fields are not populated
659 if not len(observationReason):
660 observationReason = None
661 if not len(targetName):
662 targetName = None
663 alert['observation_reason'] = observationReason
664 alert['target_name'] = targetName
665 alert['diaSource'] = diaSource.to_dict()
667 if objDiaSrcHistory is None: 667 ↛ 668line 667 didn't jump to line 668 because the condition on line 667 was never true
668 alert['prvDiaSources'] = None
669 else:
670 alert['prvDiaSources'] = objDiaSrcHistory.to_dict("records")
672 if isinstance(objDiaForcedSources, pd.Series): 672 ↛ 673line 672 didn't jump to line 673 because the condition on line 672 was never true
673 alert['prvDiaForcedSources'] = [objDiaForcedSources.to_dict()]
674 elif objDiaForcedSources is None: 674 ↛ 675line 674 didn't jump to line 675 because the condition on line 674 was never true
675 alert['prvDiaForcedSources'] = None
676 else:
677 alert['prvDiaForcedSources'] = objDiaForcedSources.to_dict("records")
679 alert['prvDiaNondetectionLimits'] = None
681 if diaObject is None: 681 ↛ 682line 681 didn't jump to line 682 because the condition on line 681 was never true
682 alert['diaObject'] = None
683 else:
684 alert['diaObject'] = diaObject.to_dict()
686 if ssSource is None: 686 ↛ 690line 686 didn't jump to line 690 because the condition on line 686 was always true
687 alert['ssSource'] = None
688 alert['mpc_orbits'] = None
689 else:
690 mpcorbColumns = [col for col in ssSource if col[:7] == 'MPCORB_']
691 mpcOrbit = {key[7:]: ssSource[key] for key in mpcorbColumns}
693 fields_to_cast = [
694 "arc_length_total", "arc_length_sel", "a", "mean_anomaly", "period", "mean_motion", "a_unc",
695 "mean_anomaly_unc", "period_unc", "mean_motion_unc", "yarkovsky", "srp", "a1", "a2", "a3",
696 "dt", "yarkovsky_unc", "srp_unc", "a1_unc", "a2_unc", "a3_unc", "dt_unc",
697 "not_normalized_rms", "earth_moid"
698 ]
699 for key in fields_to_cast:
700 if key in mpcOrbit and mpcOrbit[key] is not None:
701 mpcOrbit[key] = float(mpcOrbit[key])
703 if 'packed_primary_provisional_designation' not in mpcOrbit:
704 unpacked_desig = ssSource['designation']
705 packed_desig = ss_object_id_to_obj_id(ssSource['ssObjectId'], packed=True)
706 mpcOrbit['designation'] = unpacked_desig
707 mpcOrbit['packed_primary_provisional_designation'] = packed_desig
708 mpcOrbit['unpacked_primary_provisional_designation'] = unpacked_desig
710 mpcOrbit['id'] = 0
711 ssSource = {key: ssSource[key] for key in ssSource if key not in mpcorbColumns}
712 ssSource['diaSourceId'] = diaSourceId
713 alert['ssSource'] = ssSource
714 alert['mpc_orbits'] = mpcOrbit
716 if diffImCutout is None:
717 alert['cutoutDifference'] = None
718 else:
719 alert['cutoutDifference'] = self.streamCcdDataToBytes(diffImCutout)
721 if calexpCutout is None:
722 alert['cutoutScience'] = None
723 else:
724 alert['cutoutScience'] = self.streamCcdDataToBytes(calexpCutout)
726 if templateCutout is None:
727 alert["cutoutTemplate"] = None
728 else:
729 alert["cutoutTemplate"] = self.streamCcdDataToBytes(templateCutout)
731 return alert
733 def streamCcdDataToBytes(self, cutout):
734 """Serialize a cutout into bytes.
736 Parameters
737 ----------
738 cutout : `astropy.nddata.CCDData`
739 Cutout to serialize.
741 Returns
742 -------
743 coutputBytes : `bytes`
744 Input cutout serialized into byte data.
745 """
746 with io.BytesIO() as streamer:
747 cutout.write(streamer, format="fits", hdu_flags='MASKPLANE')
748 cutoutBytes = streamer.getvalue()
749 return cutoutBytes
751 def _serializeAlert(self, alert, schema=None, schema_id=0):
752 """Serialize an alert to a byte sequence for sending to Kafka.
754 Parameters
755 ----------
756 alert : `dict`
757 An alert payload to be serialized.
758 schema : `dict`, optional
759 An Avro schema definition describing how to encode `alert`. By default,
760 the schema is None, which sets it to the latest schema available.
761 schema_id : `int`, optional
762 The Confluent Schema Registry ID of the schema. By default, 0 (an
763 invalid ID) is used, indicating that the schema is not registered.
765 Returns
766 -------
767 serialized : `bytes`
768 The byte sequence describing the alert, including the Confluent Wire
769 Format prefix.
770 """
771 if schema is None:
772 schema = self.alertSchema.definition
774 buf = io.BytesIO()
775 buf.write(self._serializeConfluentWireHeader(schema_id))
776 fastavro.schemaless_writer(buf, schema, alert)
777 return buf.getvalue()
779 @staticmethod
780 def _serializeConfluentWireHeader(schema_version):
781 """Returns the byte prefix for Confluent Wire Format-style Kafka messages.
783 Parameters
784 ----------
785 schema_version : `int`
786 A version number which indicates the Confluent Schema Registry ID
787 number of the Avro schema used to encode the message that follows this
788 header.
790 Returns
791 -------
792 header : `bytes`
793 The 5-byte encoded message prefix.
795 Notes
796 -----
797 The Confluent Wire Format is described more fully here:
798 https://docs.confluent.io/current/schema-registry/serdes-develop/index.html#wire-format
799 """
800 ConfluentWireFormatHeader = struct.Struct(">bi")
801 return ConfluentWireFormatHeader.pack(0, schema_version)
803 def _delivery_callback(self, err, msg):
804 if err:
805 if err.code() == KafkaError.UNKNOWN_TOPIC_OR_PART:
806 topic = msg.topic() if msg and hasattr(msg,
807 'topic') else 'unknown'
808 self.log.warning(
809 'Message failed delivery. Topic %s unknown: %s',
810 topic, err)
811 else:
812 self.log.debug('Message delivered to %s [%d] @ %d', msg.topic(), msg.partition(), msg.offset())
814 def _server_check(self):
815 """Checks if the alert stream credentials are still valid and the
816 server is contactable.
818 Raises
819 -------
820 KafkaException
821 Raised if the server us not contactable.
822 RuntimeError
823 Raised if the server is contactable but there are no topics
824 present.
825 """
826 admin_client = AdminClient(self.kafkaAdminConfig)
827 topics = admin_client.list_topics(timeout=self.config.maxTimeout).topics
829 if not topics:
830 raise RuntimeError()
831 if self.kafkaTopic not in topics:
832 raise RuntimeError(
833 f"Topic '{self.kafkaTopic}' not found on the Kafka server. "
834 f"Check that the correct topic has been provided and exists "
835 f"on the Kafka server.")
837 def _computePsf(self, exposure, pixelCenter, srcId=None):
838 """Compute the PSF at a location and catch errors.
840 Parameters
841 ----------
842 exposure : `lsst.afw.image.Exposure`
843 The image to compute the PSF for.
844 pixelCenter : `lsst.geom.Point2D`
845 The location on the image to compute the PSF.
846 srcId : `int`, optional
847 Unique id of DiaSource. Used for when an error occurs extracting
848 a cutout.
850 Returns
851 -------
852 cutoutPsf : `numpy.array`
853 Array of the PSF values.
854 """
855 try:
856 # use exposure.psf.computeKernelImage to provide PSF centered in the array
857 cutoutPsf = exposure.psf.computeKernelImage(pixelCenter).array
858 except InvalidParameterError:
859 if srcId is not None:
860 msg = "Could not calculate PSF for DiaSource with "\
861 "id=%i. InvalidParameterError encountered. Exiting."\
862 % srcId
863 else:
864 msg = "Could not calculate average PSF for the image"
865 self.log.warning(msg)
866 cutoutPsf = None
867 except InvalidPsfError:
868 if srcId is not None:
869 msg = "Could not calculate PSF for DiaSource with "\
870 "id=%i. InvalidPsfError encountered. Exiting."\
871 % srcId
872 else:
873 msg = "Could not calculate average PSF for the image"
874 self.log.warning(msg)
875 cutoutPsf = None
876 # Cast the PSF to float32 to reduce the size of the alert.
877 if cutoutPsf is not None: 877 ↛ 880line 877 didn't jump to line 880 because the condition on line 877 was always true
878 cutoutPsf = cutoutPsf.astype(np.float32)
880 return cutoutPsf