Coverage for python/lsst/ap/association/loadDiaCatalogs.py : 30%

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 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"""Task for pre-loading DiaSources and DiaObjects within ap_pipe.
23"""
24import numpy as np
25import pandas as pd
26from sqlalchemy.exc import OperationalError, ProgrammingError
28import lsst.geom as geom
29import lsst.pex.config as pexConfig
30import lsst.pipe.base as pipeBase
31import lsst.sphgeom as sphgeom
33__all__ = ("LoadDiaCatalogsTask", "LoadDiaCatalogsConfig")
36class LoadDiaCatalogsConfig(pexConfig.Config):
37 """Config class for LoadDiaCatalogsConfig.
38 """
39 pixelMargin = pexConfig.RangeField(
40 doc="Padding to add to 4 all edges of the bounding box (pixels)",
41 dtype=int,
42 default=250,
43 min=0,
44 )
47class LoadDiaCatalogsTask(pipeBase.Task):
48 """Retrieve DiaObjects and associated DiaSources from the Apdb given an
49 input exposure.
50 """
51 ConfigClass = LoadDiaCatalogsConfig
52 _DefaultName = "loadDiaCatalogs"
54 def __init__(self, **kwargs):
55 pipeBase.Task.__init__(self, **kwargs)
57 @pipeBase.timeMethod
58 def run(self, exposure, apdb):
59 """Preload all DiaObjects and DiaSources from the Apdb given the
60 current exposure.
62 Parameters
63 ----------
64 exposure : `lsst.afw.image.Exposure`
65 An exposure with a bounding box.
66 apdb : `lsst.dax.apdb.Apdb`
67 AP database connection object.
69 Returns
70 -------
71 result : `lsst.pipe.base.Struct`
72 Results struct with components.
74 - ``diaObjects`` : Complete set of DiaObjects covering the input
75 exposure padded by ``pixelMargin``. DataFrame is indexed by
76 the ``diaObjectId`` column. (`pandas.DataFrame`)
77 - ``diaSources`` : Complete set of DiaSources covering the input
78 exposure padded by ``pixelMargin``. DataFrame is indexed by
79 ``diaObjectId``, ``filterName``, ``diaSourceId`` columns.
80 (`pandas.DataFrame`)
81 """
82 visiInfo = exposure.getInfo().getVisitInfo()
83 region = self._getRegion(exposure)
85 # This is the first database query
86 try:
87 diaObjects = self.loadDiaObjects(region, apdb)
88 except (OperationalError, ProgrammingError) as e:
89 raise RuntimeError(
90 "Database query failed to load DiaObjects; did you call "
91 "make_apdb.py first? If you did, some other error occurred "
92 "during database access of the DiaObject table.") from e
94 dateTime = visiInfo.getDate()
96 diaSources = self.loadDiaSources(diaObjects,
97 region,
98 dateTime,
99 apdb)
101 diaForcedSources = self.loadDiaForcedSources(diaObjects,
102 region,
103 dateTime,
104 apdb)
106 return pipeBase.Struct(
107 diaObjects=diaObjects,
108 diaSources=diaSources,
109 diaForcedSources=diaForcedSources)
111 @pipeBase.timeMethod
112 def loadDiaObjects(self, region, apdb):
113 """Load DiaObjects from the Apdb based on their HTM location.
115 Parameters
116 ----------
117 region : `sphgeom.Region`
118 Region of interest.
119 apdb : `lsst.dax.apdb.Apdb`
120 Database connection object to load from.
122 Returns
123 -------
124 diaObjects : `pandas.DataFrame`
125 DiaObjects loaded from the Apdb that are within the area defined
126 by ``pixelRanges``.
127 """
128 if region is None:
129 # If no area is specified return an empty DataFrame with the
130 # the column used for indexing later in AssociationTask.
131 diaObjects = pd.DataFrame(columns=["diaObjectId"])
132 else:
133 diaObjects = apdb.getDiaObjects(region)
135 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
136 if diaObjects.index.has_duplicates:
137 self.log.warn(
138 "Duplicate DiaObjects loaded from the Apdb. This may cause "
139 "downstream pipeline issues. Dropping duplicated rows")
140 # Drop duplicates via index and keep the first appearance.
141 diaObjects = diaObjects.groupby(diaObjects.index).first()
143 return diaObjects.replace(to_replace=[None], value=np.nan)
145 @pipeBase.timeMethod
146 def loadDiaSources(self, diaObjects, region, dateTime, apdb):
147 """Load DiaSources from the Apdb based on their diaObjectId or
148 pixelId location.
150 Variable used to load sources is set in config.
152 Parameters
153 ----------
154 diaObjects : `pandas.DataFrame`
155 DiaObjects loaded from the Apdb that are within the area defined
156 by ``pixelRanges``.
157 region : `sphgeom.Region`
158 Region of interest.
159 dateTime : `lsst.daf.base.DateTime`
160 Time of the current visit
161 apdb : `lsst.dax.apdb.Apdb`
162 Database connection object to load from.
164 Returns
165 -------
166 DiaSources : `pandas.DataFrame`
167 DiaSources loaded from the Apdb that are within the area defined
168 by ``pixelRange`` and associated with ``diaObjects``.
169 """
170 if region is None:
171 # If no area is specified return an empty DataFrame with the
172 # the column used for indexing later in AssociationTask.
173 diaSources = pd.DataFrame(columns=["diaObjectId",
174 "filterName",
175 "diaSourceId"])
176 else:
177 diaSources = apdb.getDiaSources(region, diaObjects.loc[:, "diaObjectId"], dateTime)
179 diaSources.set_index(["diaObjectId", "filterName", "diaSourceId"],
180 drop=False,
181 inplace=True)
182 if diaSources.index.has_duplicates:
183 self.log.warn(
184 "Duplicate DiaSources loaded from the Apdb. This may cause "
185 "downstream pipeline issues. Dropping duplicated rows")
186 # Drop duplicates via index and keep the first appearance. Reset
187 # due to the index shape being slight different thatn expected.
188 diaSources = diaSources.groupby(diaSources.index).first().reset_index(drop=True)
189 diaSources.set_index(["diaObjectId", "filterName", "diaSourceId"],
190 drop=False,
191 inplace=True)
193 return diaSources.replace(to_replace=[None], value=np.nan)
195 @pipeBase.timeMethod
196 def loadDiaForcedSources(self, diaObjects, region, dateTime, apdb):
197 """Load DiaObjects from the Apdb based on their HTM location.
199 Parameters
200 ----------
201 diaObjects : `pandas.DataFrame`
202 DiaObjects loaded from the Apdb.
203 region : `sphgeom.Region`
204 Region of interest.
205 dateTime : `lsst.daf.base.DateTime`
206 Time of the current visit
207 apdb : `lsst.dax.apdb.Apdb`
208 Database connection object to load from.
210 Returns
211 -------
212 diaObjects : `pandas.DataFrame`
213 DiaObjects loaded from the Apdb that are within the area defined
214 by ``pixelRanges``.
215 """
216 if len(diaObjects) == 0:
217 # If no diaObjects are available return an empty DataFrame with
218 # the the column used for indexing later in AssociationTask.
219 diaForcedSources = pd.DataFrame(columns=["diaObjectId",
220 "diaForcedSourceId"])
221 else:
222 diaForcedSources = apdb.getDiaForcedSources(
223 region,
224 diaObjects.loc[:, "diaObjectId"],
225 dateTime)
227 diaForcedSources.set_index(["diaObjectId", "diaForcedSourceId"],
228 drop=False,
229 inplace=True)
230 if diaForcedSources.index.has_duplicates:
231 self.log.warn(
232 "Duplicate DiaForcedSources loaded from the Apdb. This may "
233 "cause downstream pipeline issues. Dropping duplicated rows.")
234 # Drop duplicates via index and keep the first appearance. Reset
235 # due to the index shape being slight different thatn expected.
236 diaForcedSources = diaForcedSources.groupby(diaForcedSources.index).first()
237 diaForcedSources.reset_index(drop=True, inplace=True)
238 diaForcedSources.set_index(["diaObjectId", "diaForcedSourceId"],
239 drop=False,
240 inplace=True)
242 return diaForcedSources.replace(to_replace=[None], value=np.nan)
244 @pipeBase.timeMethod
245 def _getRegion(self, exposure):
246 """Calculate an enveloping region for an exposure.
248 Parameters
249 ----------
250 exposure : `lsst.afw.image.Exposure`
251 Exposure object with calibrated WCS.
253 Returns
254 -------
255 region : `sphgeom.Region`
256 Region enveloping an exposure.
257 """
258 bbox = geom.Box2D(exposure.getBBox())
259 bbox.grow(self.config.pixelMargin)
260 wcs = exposure.getWcs()
262 region = sphgeom.ConvexPolygon([wcs.pixelToSky(pp).getVector()
263 for pp in bbox.getCorners()])
265 return region