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