Coverage for python / lsst / pipe / tasks / drpAssociationPipe.py: 21%
168 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-09 01:51 -0800
1# This file is part of pipe_tasks.
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"""Pipeline for running DiaSource association in a DRP context.
23"""
25__all__ = ["DrpAssociationPipeTask",
26 "DrpAssociationPipeConfig",
27 "DrpAssociationPipeConnections"]
29import os
31import astropy.table as tb
32import numpy as np
33import pandas as pd
36from lsst.pipe.tasks.ssoAssociation import SolarSystemAssociationTask
37import lsst.geom as geom
38import lsst.pex.config as pexConfig
39import lsst.pipe.base as pipeBase
40from lsst.meas.base import SkyMapIdGeneratorConfig
41from lsst.skymap import BaseSkyMap
43from .coaddBase import makeSkyInfo
44from .schemaUtils import convertDataFrameToSdmSchema, readSdmSchemaFile
45from .simpleAssociation import SimpleAssociationTask
48class DrpAssociationPipeConnections(pipeBase.PipelineTaskConnections,
49 dimensions=("tract", "patch", "skymap"),
50 defaultTemplates={"coaddName": "deep",
51 "warpTypeSuffix": "",
52 "fakesType": ""}):
53 diaSourceTables = pipeBase.connectionTypes.Input(
54 doc="Set of catalogs of calibrated DiaSources.",
55 name="{fakesType}{coaddName}Diff_diaSrcTable",
56 storageClass="ArrowAstropy",
57 dimensions=("instrument", "visit", "detector"),
58 deferLoad=True,
59 multiple=True
60 )
61 ssObjectTableRefs = pipeBase.connectionTypes.Input(
62 doc="Reference to catalogs of SolarSolarSystem objects expected to be "
63 "observable in each (visit, detector).",
64 name="preloaded_ss_object_visit",
65 storageClass="ArrowAstropy",
66 dimensions=("instrument", "visit"),
67 minimum=0,
68 deferLoad=True,
69 multiple=True
70 )
71 skyMap = pipeBase.connectionTypes.Input(
72 doc="Input definition of geometry/bbox and projection/wcs for coadded "
73 "exposures",
74 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
75 storageClass="SkyMap",
76 dimensions=("skymap", ),
77 )
78 finalVisitSummaryRefs = pipeBase.connectionTypes.Input(
79 doc="Reference to finalVisitSummary of each exposure, containing "
80 "visitInfo, bbox, and wcs.",
81 name="finalVisitSummary",
82 storageClass="ExposureCatalog",
83 dimensions=("instrument", "visit"),
84 deferLoad=True,
85 multiple=True
86 )
87 assocDiaSourceTable = pipeBase.connectionTypes.Output(
88 doc="Catalog of DiaSources covering the patch and associated with a "
89 "DiaObject.",
90 name="{fakesType}{coaddName}Diff_assocDiaSrcTable",
91 storageClass="DataFrame",
92 dimensions=("tract", "patch"),
93 )
94 associatedSsSources = pipeBase.connectionTypes.Output(
95 doc="Optional output storing ssSource data computed during association.",
96 name="{fakesType}{coaddName}Diff_assocSsSrcTable",
97 storageClass="ArrowAstropy",
98 dimensions=("tract", "patch"),
99 )
100 unassociatedSsObjects = pipeBase.connectionTypes.Output(
101 doc="Expected locations of ssObjects with no associated source.",
102 name="{fakesType}{coaddName}Diff_unassocSsObjTable",
103 storageClass="ArrowAstropy",
104 dimensions=("tract", "patch"),
105 )
106 diaObjectTable = pipeBase.connectionTypes.Output(
107 doc="Catalog of DiaObjects created from spatially associating "
108 "DiaSources.",
109 name="{fakesType}{coaddName}Diff_diaObjTable",
110 storageClass="DataFrame",
111 dimensions=("tract", "patch"),
112 )
114 def __init__(self, *, config=None):
115 super().__init__(config=config)
117 if not config.doSolarSystemAssociation:
118 del self.ssObjectTableRefs
119 del self.associatedSsSources
120 del self.unassociatedSsObjects
121 del self.finalVisitSummaryRefs
124class DrpAssociationPipeConfig(
125 pipeBase.PipelineTaskConfig,
126 pipelineConnections=DrpAssociationPipeConnections):
127 associator = pexConfig.ConfigurableField(
128 target=SimpleAssociationTask,
129 doc="Task used to associate DiaSources with DiaObjects.",
130 )
131 solarSystemAssociator = pexConfig.ConfigurableField(
132 target=SolarSystemAssociationTask,
133 doc="Task used to associate DiaSources with SolarSystemObjects.",
134 )
135 doAddDiaObjectCoords = pexConfig.Field(
136 dtype=bool,
137 default=True,
138 doc="Do pull diaObject's average coordinate as coord_ra and coord_dec"
139 "Duplicates information, but needed for bulk ingest into qserv."
140 )
141 doWriteEmptyTables = pexConfig.Field(
142 dtype=bool,
143 default=False,
144 doc="If True, construct and write out empty diaSource and diaObject "
145 "tables. If False, raise NoWorkFound"
146 )
147 doSolarSystemAssociation = pexConfig.Field(
148 dtype=bool,
149 default=True,
150 doc="Process SolarSystem objects through the pipeline.",
151 )
152 doUseSchema = pexConfig.Field(
153 dtype=bool,
154 default=False,
155 doc="Use an existing schema to coerce the data types of the output columns."
156 )
157 schemaDir = pexConfig.Field(
158 dtype=str,
159 doc="Path to the directory containing schema definitions.",
160 default=os.path.join("${SDM_SCHEMAS_DIR}",
161 "yml"),
162 )
163 schemaFile = pexConfig.Field(
164 dtype=str,
165 doc="Yaml file specifying the schema of the output catalog.",
166 default="lsstcam.yaml",
167 )
168 idGenerator = SkyMapIdGeneratorConfig.make_field()
171class DrpAssociationPipeTask(pipeBase.PipelineTask):
172 """Driver pipeline for loading DiaSource catalogs in a patch/tract
173 region and associating them.
174 """
175 ConfigClass = DrpAssociationPipeConfig
176 _DefaultName = "drpAssociation"
178 def __init__(self, **kwargs):
179 super().__init__(**kwargs)
180 self.makeSubtask('associator')
182 if self.config.doUseSchema:
183 schemaFile = os.path.join(self.config.schemaDir, self.config.schemaFile)
184 self.schema = readSdmSchemaFile(schemaFile)
185 else:
186 self.schema = None
187 if self.config.doSolarSystemAssociation:
188 self.makeSubtask("solarSystemAssociator")
190 def runQuantum(self, butlerQC, inputRefs, outputRefs):
191 inputs = butlerQC.get(inputRefs)
193 inputs["tractId"] = butlerQC.quantum.dataId["tract"]
194 inputs["patchId"] = butlerQC.quantum.dataId["patch"]
195 inputs["idGenerator"] = self.config.idGenerator.apply(butlerQC.quantum.dataId)
196 if not self.config.doSolarSystemAssociation:
197 inputs["ssObjectTableRefs"] = []
198 inputs["finalVisitSummaryRefs"] = []
199 outputs = self.run(**inputs)
200 butlerQC.put(outputs, outputRefs)
202 def run(self,
203 diaSourceTables,
204 ssObjectTableRefs,
205 skyMap,
206 finalVisitSummaryRefs,
207 tractId,
208 patchId,
209 idGenerator=None):
210 """Trim DiaSources to the current Patch and run association.
212 Takes in the set of DiaSource catalogs that covers the current patch,
213 trims them to the dimensions of the patch, and [TODO: eventually]
214 runs association on the concatenated DiaSource Catalog.
216 Parameters
217 ----------
218 diaSourceTables : `list` of `lsst.daf.butler.DeferredDatasetHandle`
219 Set of DiaSource catalogs potentially covering this patch/tract.
220 ssObjectTableRefs : `list` of `lsst.daf.butler.DeferredDatasetHandle`
221 Set of known SSO ephemerides potentially covering this patch/tract.
222 skyMap : `lsst.skymap.BaseSkyMap`
223 SkyMap defining the patch/tract
224 finalVisitSummaryRefs : `list` of `lsst.daf.butler.DeferredDatasetHandle`
225 Reference to finalVisitSummary of each exposure potentially
226 covering this patch/tract, which contain visitInfo, bbox, and wcs
227 tractId : `int`
228 Id of current tract being processed.
229 patchId : `int`
230 Id of current patch being processed.
231 idGenerator : `lsst.meas.base.IdGenerator`, optional
232 Object that generates Object IDs and random number generator seeds.
234 Returns
235 -------
236 output : `lsst.pipe.base.Struct`
237 Results struct with attributes:
239 ``assocDiaSourceTable``
240 Table of DiaSources with updated value for diaObjectId.
241 (`pandas.DataFrame`)
242 ``diaObjectTable``
243 Table of DiaObjects from matching DiaSources
244 (`pandas.DataFrame`).
245 """
246 self.log.info("Running DPR Association on patch %i, tract %i...",
247 patchId, tractId)
249 skyInfo = makeSkyInfo(skyMap, tractId, patchId)
251 # Get the patch bounding box.
252 innerPatchBox = geom.Box2D(skyInfo.patchInfo.getInnerBBox())
253 outerPatchBox = geom.Box2D(skyInfo.patchInfo.getOuterBBox())
254 innerTractSkyRegion = skyInfo.tractInfo.getInnerSkyRegion()
256 # Keep track of our diaCats, ssObject cats, and finalVisitSummaries by their (visit, detector) IDs
257 diaIdDict = prepareCatalogDict(diaSourceTables, useVisitDetector=True)
258 ssObjectIdDict = prepareCatalogDict(ssObjectTableRefs, useVisitDetector=False)
259 finalVisitSummaryIdDict = prepareCatalogDict(finalVisitSummaryRefs, useVisitDetector=False)
261 diaSourceHistory, ssSourceHistory, unassociatedSsObjectHistory = [], [], []
262 nSsSrc, nSsObj = 0, 0
263 visits = set([v for v, _ in diaIdDict.keys()])
264 for visit in visits:
265 # visit summaries and Solar System catalogs are per-visit, so only
266 # load them once for all detectors with that visit
267 visitSummary = finalVisitSummaryIdDict[visit].get() if visit in finalVisitSummaryIdDict else None
268 ssCat = ssObjectIdDict[visit].get() if visit in ssObjectIdDict else None
269 detectors = [det for (v, det) in diaIdDict.keys() if v == visit]
270 for detector in detectors:
271 diaCat = diaIdDict[(visit, detector)].get()
272 nDiaSrcIn = len(diaCat)
273 if (ssCat is not None) and (visitSummary is not None):
274 ssoAssocResult = self.runSolarSystemAssociation(diaCat,
275 ssCat.copy(),
276 visitSummary=visitSummary,
277 patchBbox=innerPatchBox,
278 patchWcs=skyInfo.wcs,
279 innerTractSkyRegion=innerTractSkyRegion,
280 detector=detector,
281 visit=visit,
282 )
284 nSsSrc = len(ssoAssocResult.associatedSsSources)
285 nSsObj = len(ssoAssocResult.unassociatedSsObjects)
286 # If diaSources were associated with Solar System objects,
287 # remove them from the catalog so they won't create new
288 # diaObjects or be associated with other diaObjects.
289 diaCat = ssoAssocResult.unassociatedDiaSources
290 else:
291 nSsSrc, nSsObj = 0, 0
293 # Only trim diaSources to the outer bbox of the patch, so that
294 # diaSources near the patch boundary can be associated.
295 # DiaObjects will be trimmed to the inner patch bbox, and any
296 # diaSources associated with dropped diaObjects will also be dropped
297 diaInPatch = self._trimToPatch(diaCat.to_pandas(), outerPatchBox, skyInfo.wcs)
299 nDiaSrc = diaInPatch.sum()
301 self.log.info(
302 "Read DiaSource catalog of length %i from visit %i, "
303 "detector %i. Found %i sources within the patch/tract "
304 "footprint, including %i associated with SSOs.",
305 nDiaSrcIn, visit, detector, nDiaSrc + nSsSrc, nSsSrc)
307 if nDiaSrc > 0:
308 diaSourceHistory.append(diaCat[diaInPatch])
309 if nSsSrc > 0:
310 ssSourceHistory.append(ssoAssocResult.associatedSsSources)
311 if nSsObj > 0:
312 unassociatedSsObjectHistory.append(ssoAssocResult.unassociatedSsObjects)
314 # After looping over all of the detector-level catalogs that overlap the
315 # patch, combine them into patch-level catalogs
316 diaSourceHistoryCat = self._stackCatalogs(diaSourceHistory)
317 if self.config.doSolarSystemAssociation:
318 ssSourceHistoryCat = self._stackCatalogs(ssSourceHistory, remove_columns=['ra', 'dec'])
319 nSsSrcTotal = len(ssSourceHistoryCat) if ssSourceHistoryCat else 0
320 unassociatedSsObjectHistoryCat = self._stackCatalogs(unassociatedSsObjectHistory)
321 nSsObjTotal = len(unassociatedSsObjectHistoryCat) if unassociatedSsObjectHistoryCat else 0
322 self.log.info("Found %i ssSources and %i missing ssObjects in patch %i, tract %i",
323 nSsSrcTotal, nSsObjTotal, patchId, tractId)
324 else:
325 ssSourceHistoryCat = None
326 unassociatedSsObjectHistoryCat = None
328 if (not diaSourceHistory) and (not ssSourceHistory):
329 if not self.config.doWriteEmptyTables:
330 raise pipeBase.NoWorkFound("Found no overlapping DIASources to associate.")
332 self.log.info("Found %i DiaSources overlapping patch %i, tract %i",
333 len(diaSourceHistoryCat), patchId, tractId)
335 diaSourceTable = diaSourceHistoryCat.to_pandas()
336 diaSourceTable.set_index("diaSourceId", drop=False)
337 # Now run diaObject association on the stacked remaining diaSources
338 assocResult = self.associator.run(diaSourceTable, idGenerator=idGenerator)
340 # Drop any diaObjects that were created outside the inner region of the
341 # patch. These will be associated in the overlapping patch instead.
342 objectsInTractPatch = self._trimToPatch(assocResult.diaObjects,
343 innerPatchBox,
344 skyInfo.wcs,
345 innerTractSkyRegion=innerTractSkyRegion)
346 diaObjects = assocResult.diaObjects[objectsInTractPatch]
347 # Instead of dropping diaSources based on their patch, assign them to a
348 # patch based on whether their diaObject was inside. This means that
349 # some diaSources in the patch catalog will actually have coordinates
350 # just outside the patch.
351 assocDiaSources = self.dropDiaSourceByDiaObjectId(assocResult.diaObjects[~objectsInTractPatch].index,
352 assocResult.assocDiaSources)
354 self.log.info("Associated DiaSources into %i DiaObjects", len(diaObjects))
356 if self.config.doAddDiaObjectCoords:
357 assocDiaSources = self._addDiaObjectCoords(diaObjects, assocDiaSources)
358 if self.config.doUseSchema:
359 diaObjects = convertDataFrameToSdmSchema(self.schema, diaObjects, tableName="DiaObject")
360 assocDiaSources = convertDataFrameToSdmSchema(self.schema, assocDiaSources, tableName="DiaSource")
362 return pipeBase.Struct(
363 diaObjectTable=diaObjects,
364 assocDiaSourceTable=assocDiaSources,
365 associatedSsSources=ssSourceHistoryCat,
366 unassociatedSsObjects=unassociatedSsObjectHistoryCat,
367 )
369 def _stackCatalogs(self, catalogs, remove_columns=None, empty=None):
370 """Stack a list of catalogs.
372 Parameters
373 ----------
374 catalogs : `list` of `astropy.table.Table`
375 Input catalogs with the same columns to be combined.
376 remove_columns : `list` of `str` or None, optional
377 List of column names to drop from the tables before stacking.
379 Returns
380 -------
381 `astropy.table.Table`
382 The combined catalog.
383 """
384 if catalogs:
385 sourceHistory = tb.vstack(catalogs)
386 if remove_columns is not None:
387 sourceHistory.remove_columns(remove_columns)
388 return sourceHistory
389 else:
390 return empty
392 def runSolarSystemAssociation(self, diaCat, ssCat,
393 visitSummary,
394 patchBbox,
395 patchWcs,
396 innerTractSkyRegion,
397 detector,
398 visit,
399 ):
400 """Run Solar System object association and filter the results.
402 Parameters
403 ----------
404 diaCat : `pandas.DataFrame`
405 Catalog of detected diaSources on the image difference.
406 ssCat : `astropy.table.Table`
407 Catalog of predicted coordinates of known Solar System objects.
408 visitSummary : `lsst.afw.table.ExposureCatalog`
409 Table of calibration and metadata for all detectors in a visit.
410 patchBbox : `lsst.geom.Box2D`
411 Bounding box of the patch.
412 patchWcs : `lsst.geom.SkyWcs`
413 Wcs of the tract containing the patch.
414 innerTractSkyRegion : `lsst.sphgeom.Box`
415 Region defining the inner non-overlapping part of a tract.
416 detector : `int`
417 Detector number of the science exposure.
418 visit : `int`
419 Visit number of the science exposure.
421 Returns
422 -------
423 ssoAssocResult : `lsst.pipe.base.Struct`
424 Results struct with attributes:
426 ``associatedSsSources``
427 Table of DiaSources associated with Solar System objects.
428 (`astropy.table.Table`)
429 ``associatedSsDiaSources``
430 Table of Solar System objects associated with DiaSources.
431 (`astropy.table.Table`).
432 ``unassociatedSsObjects``
433 Table of Solar System objects in the patch not associated with
434 any DiaSource (`astropy.table.Table`).
435 ``unassociatedDiaSources``
436 Table of DiaSources not associated with any Solar System object
437 (`astropy.table.Table`).
438 """
439 # Get the exposure metadata from the detector's row in the visitSummary table.
440 ssoAssocResult = self.solarSystemAssociator.run(
441 diaCat,
442 ssCat,
443 visitInfo=visitSummary.find(detector).visitInfo,
444 bbox=visitSummary.find(detector).getBBox(),
445 wcs=visitSummary.find(detector).wcs,
446 )
448 ssInTractPatch = self._trimToPatch(ssoAssocResult.associatedSsSources.to_pandas(),
449 patchBbox,
450 patchWcs,
451 innerTractSkyRegion=innerTractSkyRegion)
452 associatedSsSources = ssoAssocResult.associatedSsSources[ssInTractPatch].copy()
453 assocDiaSrcIds = set(associatedSsSources['diaSourceId'])
454 diaSrcMask = [diaId in assocDiaSrcIds for diaId in ssoAssocResult.ssoAssocDiaSources['diaSourceId']]
455 associatedSsDiaSources = ssoAssocResult.ssoAssocDiaSources[np.array(diaSrcMask)]
457 ssObjInTractPatch = self._trimToPatch(ssoAssocResult.unassociatedSsObjects.to_pandas(),
458 patchBbox,
459 patchWcs,
460 innerTractSkyRegion=innerTractSkyRegion)
461 unassociatedSsObjects = ssoAssocResult.unassociatedSsObjects[ssObjInTractPatch].copy()
462 # Update the table of Solar System objects that were not found with the
463 # visit and detector where they were predicted.
464 if len(unassociatedSsObjects) > 0:
465 unassociatedSsObjects['visit'] = visit
466 unassociatedSsObjects['detector'] = detector
468 return pipeBase.Struct(
469 associatedSsSources=associatedSsSources,
470 associatedSsDiaSources=associatedSsDiaSources,
471 unassociatedSsObjects=unassociatedSsObjects,
472 unassociatedDiaSources=ssoAssocResult.unAssocDiaSources
473 )
475 def _addDiaObjectCoords(self, objects, sources):
476 obj = objects[['ra', 'dec']].rename(columns={"ra": "coord_ra", "dec": "coord_dec"})
477 df = pd.merge(sources.reset_index(), obj, left_on='diaObjectId', right_index=True,
478 how='inner').set_index('diaSourceId')
479 return df
481 def _trimToPatch(self, cat, patchBox, wcs, innerTractSkyRegion=None):
482 """Create generator testing if a set of DiaSources are in the
483 patch/tract.
485 Parameters
486 ----------
487 cat : `pandas.DataFrame`
488 Catalog of DiaSources to test within patch/tract.
489 patchBox : `lsst.geom.Box2D`
490 Bounding box of the patch.
491 wcs : `lsst.geom.SkyWcs`
492 Wcs of the tract.
493 innerTractSkyRegion : `lsst.sphgeom.Box`, optional
494 Region defining the inner non-overlapping part of a tract.
496 Returns
497 -------
498 isInPatch : `numpy.ndarray`, (N,)
499 Booleans representing if the DiaSources are contained within the
500 current patch and tract.
501 """
502 isInPatch = np.zeros(len(cat), dtype=bool)
504 for idx, row in cat.reset_index().iterrows():
505 spPoint = geom.SpherePoint(row["ra"], row["dec"], geom.degrees)
506 pxCoord = wcs.skyToPixel(spPoint)
507 ra_rad = np.deg2rad(row["ra"])
508 dec_rad = np.deg2rad(row["dec"])
509 isInPatch[idx] = patchBox.contains(pxCoord)
511 if innerTractSkyRegion is not None:
512 isInPatch[idx] &= innerTractSkyRegion.contains(ra_rad, dec_rad)
514 return isInPatch
516 def dropDiaSourceByDiaObjectId(self, droppedDiaObjectIds, diaSources):
517 """Drop diaSources with diaObject IDs in the supplied list.
519 Parameters
520 ----------
521 droppedDiaObjectIds : `pandas.DataFrame`
522 DiaObjectIds to match and drop from the list of diaSources.
523 diaSources : `pandas.DataFrame`
524 Catalog of diaSources to check and filter.
526 Returns
527 -------
528 filteredDiaSources : `pandas.DataFrame`
529 The input diaSources with any rows matching the listed diaObjectIds
530 removed.
531 """
532 toDrop = diaSources['diaObjectId'].isin(droppedDiaObjectIds)
534 # Keep only rows that do NOT match
535 return diaSources.loc[~toDrop].copy()
538def prepareCatalogDict(dataRefList, useVisitDetector=True):
539 """Prepare lookup tables of the data references.
541 Parameters
542 ----------
543 dataRefList : `list` of `lsst.daf.butler.DeferredDatasetHandle`
544 The data references to make a lookup table for.
545 useVisitDetector : `bool`, optional
546 Use both visit and detector in the dict key? If False, use only visit.
548 Returns
549 -------
550 `dict` of `lsst.daf.butler.DeferredDatasetHandle`
551 Lookup table of the data references by visit (and optionally detector)
552 """
553 dataDict = {}
555 if useVisitDetector:
556 for dataRef in dataRefList:
557 dataDict[(dataRef.dataId["visit"], dataRef.dataId["detector"])] = dataRef
558 else:
559 for dataRef in dataRefList:
560 dataDict[dataRef.dataId["visit"]] = dataRef
561 return dataDict