Coverage for python/lsst/pipe/tasks/drpAssociationPipe.py : 34%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of pipe_tasks.
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.
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.
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.
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/>.
21#
23"""Pipeline for running DiaSource association in a DRP context.
24"""
26import numpy as np
27import pandas as pd
29import lsst.geom as geom
30import lsst.pex.config as pexConfig
31import lsst.pipe.base as pipeBase
32from lsst.skymap import BaseSkyMap
34from .coaddBase import makeSkyInfo
35from .simpleAssociation import SimpleAssociationTask
37__all__ = ["DrpAssociationPipeTask",
38 "DrpAssociationPipeConfig",
39 "DrpAssociationPipeConnections"]
42class DrpAssociationPipeConnections(pipeBase.PipelineTaskConnections,
43 dimensions=("tract", "patch", "skymap"),
44 defaultTemplates={"coaddName": "deep",
45 "warpTypeSuffix": "",
46 "fakesType": ""}):
47 diaSourceTables = pipeBase.connectionTypes.Input(
48 doc="Set of catalogs of calibrated DiaSources.",
49 name="{fakesType}{coaddName}Diff_diaSrcTable",
50 storageClass="DataFrame",
51 dimensions=("instrument", "visit", "detector"),
52 deferLoad=True,
53 multiple=True
54 )
55 skyMap = pipeBase.connectionTypes.Input(
56 doc="Input definition of geometry/bbox and projection/wcs for coadded "
57 "exposures",
58 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
59 storageClass="SkyMap",
60 dimensions=("skymap", ),
61 )
62 assocDiaSourceTable = pipeBase.connectionTypes.Output(
63 doc="Catalog of DiaSources covering the patch and associated with a "
64 "DiaObject.",
65 name="{fakesType}{coaddName}Diff_assocDiaSrcTable",
66 storageClass="DataFrame",
67 dimensions=("tract", "patch"),
68 )
69 diaObjectTable = pipeBase.connectionTypes.Output(
70 doc="Catalog of DiaObjects created from spatially associating "
71 "DiaSources.",
72 name="{fakesType}{coaddName}Diff_diaObjTable",
73 storageClass="DataFrame",
74 dimensions=("tract", "patch"),
75 )
78class DrpAssociationPipeConfig(
79 pipeBase.PipelineTaskConfig,
80 pipelineConnections=DrpAssociationPipeConnections):
81 associator = pexConfig.ConfigurableField(
82 target=SimpleAssociationTask,
83 doc="Task used to associate DiaSources with DiaObjects.",
84 )
85 doAddDiaObjectCoords = pexConfig.Field(
86 dtype=bool,
87 default=True,
88 doc="Do pull diaObject's average coordinate as coord_ra and coord_dec"
89 "Duplicates information, but needed for bulk ingest into qserv."
90 )
91 doWriteEmptyTables = pexConfig.Field(
92 dtype=bool,
93 default=False,
94 doc="If True, construct and write out empty diaSource and diaObject "
95 "tables. If False, raise NoWorkFound"
96 )
99class DrpAssociationPipeTask(pipeBase.PipelineTask):
100 """Driver pipeline for loading DiaSource catalogs in a patch/tract
101 region and associating them.
102 """
103 ConfigClass = DrpAssociationPipeConfig
104 _DefaultName = "drpAssociation"
106 def __init__(self, **kwargs):
107 super().__init__(**kwargs)
108 self.makeSubtask('associator')
110 def runQuantum(self, butlerQC, inputRefs, outputRefs):
111 inputs = butlerQC.get(inputRefs)
113 inputs["tractId"] = butlerQC.quantum.dataId["tract"]
114 inputs["patchId"] = butlerQC.quantum.dataId["patch"]
115 tractPatchId, skymapBits = butlerQC.quantum.dataId.pack(
116 "tract_patch",
117 returnMaxBits=True)
118 inputs["tractPatchId"] = tractPatchId
119 inputs["skymapBits"] = skymapBits
121 outputs = self.run(**inputs)
122 butlerQC.put(outputs, outputRefs)
124 def run(self,
125 diaSourceTables,
126 skyMap,
127 tractId,
128 patchId,
129 tractPatchId,
130 skymapBits):
131 """Trim DiaSources to the current Patch and run association.
133 Takes in the set of DiaSource catalogs that covers the current patch,
134 trims them to the dimensions of the patch, and [TODO: eventually]
135 runs association on the concatenated DiaSource Catalog.
137 Parameters
138 ----------
139 diaSourceTables : `list` of `lst.daf.butler.DeferredDatasetHandle`
140 Set of DiaSource catalogs potentially covering this patch/tract.
141 skyMap : `lsst.skymap.BaseSkyMap`
142 SkyMap defining the patch/tract
143 tractId : `int`
144 Id of current tract being processed.
145 patchId : `int`
146 Id of current patch being processed
148 Returns
149 -------
150 output : `lsst.pipe.base.Struct`
151 Results struct with attributes:
153 ``assocDiaSourceTable``
154 Table of DiaSources with updated value for diaObjectId.
155 (`pandas.DataFrame`)
156 ``diaObjectTable``
157 Table of DiaObjects from matching DiaSources
158 (`pandas.DataFrame`).
159 """
160 self.log.info("Running DPR Association on patch %i, tract %i...",
161 patchId, tractId)
163 skyInfo = makeSkyInfo(skyMap, tractId, patchId)
165 # Get the patch bounding box.
166 innerPatchBox = geom.Box2D(skyInfo.patchInfo.getInnerBBox())
168 diaSourceHistory = []
169 for catRef in diaSourceTables:
170 cat = catRef.get(
171 datasetType=self.config.connections.diaSourceTables,
172 immediate=True)
174 isInTractPatch = self._trimToPatch(cat,
175 innerPatchBox,
176 skyInfo.wcs)
178 nDiaSrc = isInTractPatch.sum()
179 self.log.info(
180 "Read DiaSource catalog of length %i from visit %i, "
181 "detector %i. Found %i sources within the patch/tract "
182 "footprint.",
183 len(cat), catRef.dataId["visit"],
184 catRef.dataId["detector"], nDiaSrc)
186 if nDiaSrc <= 0:
187 continue
189 cutCat = cat[isInTractPatch]
190 diaSourceHistory.append(cutCat)
192 if diaSourceHistory:
193 diaSourceHistoryCat = pd.concat(diaSourceHistory)
194 else:
195 # No rows to associate
196 if self.config.doWriteEmptyTables:
197 self.log.info("Constructing empty table")
198 # Construct empty table using last table and dropping all the rows
199 diaSourceHistoryCat = cat.drop(cat.index)
200 else:
201 raise pipeBase.NoWorkFound("Found no overlapping DIASources to associate.")
203 self.log.info("Found %i DiaSources overlapping patch %i, tract %i",
204 len(diaSourceHistoryCat), patchId, tractId)
206 assocResult = self.associator.run(diaSourceHistoryCat,
207 tractPatchId,
208 skymapBits)
210 self.log.info("Associated DiaSources into %i DiaObjects",
211 len(assocResult.diaObjects))
213 if self.config.doAddDiaObjectCoords:
214 assocResult.assocDiaSources = self._addDiaObjectCoords(assocResult.diaObjects,
215 assocResult.assocDiaSources)
217 return pipeBase.Struct(
218 diaObjectTable=assocResult.diaObjects,
219 assocDiaSourceTable=assocResult.assocDiaSources)
221 def _addDiaObjectCoords(self, objects, sources):
222 obj = objects[['ra', 'decl']].rename(columns={"ra": "coord_ra", "decl": "coord_dec"})
223 df = pd.merge(sources.reset_index(), obj, left_on='diaObjectId', right_index=True,
224 how='inner').set_index('diaSourceId')
225 return df
227 def _trimToPatch(self, cat, innerPatchBox, wcs):
228 """Create generator testing if a set of DiaSources are in the
229 patch/tract.
231 Parameters
232 ----------
233 cat : `pandas.DataFrame`
234 Catalog of DiaSources to test within patch/tract.
235 innerPatchBox : `lsst.geom.Box2D`
236 Bounding box of the patch.
237 wcs : `lsst.geom.SkyWcs`
238 Wcs of the tract.
240 Returns
241 ------
242 isInPatch : `numpy.ndarray`, (N,)
243 Booleans representing if the DiaSources are contained within the
244 current patch and tract.
245 """
246 isInPatch = np.array([
247 innerPatchBox.contains(
248 wcs.skyToPixel(
249 geom.SpherePoint(row["ra"], row["decl"], geom.degrees)))
250 for idx, row in cat.iterrows()])
251 return isInPatch