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

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
27import lsst.geom as geom
28import lsst.pex.config as pexConfig
29import lsst.pipe.base as pipeBase
30import lsst.sphgeom as sphgeom
32__all__ = ("LoadDiaCatalogsTask", "LoadDiaCatalogsConfig")
35class LoadDiaCatalogsConfig(pexConfig.Config):
36 """Config class for LoadDiaCatalogsConfig.
37 """
38 htmLevel = pexConfig.RangeField(
39 dtype=int,
40 doc="Level of the HTM pixelization.",
41 default=20,
42 min=1,
43 )
44 htmMaxRanges = pexConfig.RangeField(
45 dtype=int,
46 doc="Maximum number of HTM (min, max) ranges to return.",
47 default=128,
48 min=2,
49 )
50 pixelMargin = pexConfig.RangeField(
51 doc="Padding to add to 4 all edges of the bounding box (pixels)",
52 dtype=int,
53 default=300,
54 min=0,
55 )
56 loadDiaSourcesByPixelId = pexConfig.Field(
57 doc="Load DiaSources by their HTM pixelId instead of by their "
58 "associated diaObjectId",
59 dtype=bool,
60 default=False,
61 )
64class LoadDiaCatalogsTask(pipeBase.Task):
65 """Retrieve DiaObjects and associated DiaSources from the Apdb given an
66 input exposure.
67 """
68 ConfigClass = LoadDiaCatalogsConfig
69 _DefaultName = "loadDiaCatalogs"
71 def __init__(self, **kwargs):
72 pipeBase.Task.__init__(self, **kwargs)
73 self.pixelator = sphgeom.HtmPixelization(self.config.htmLevel)
75 @pipeBase.timeMethod
76 def run(self, exposure, apdb):
77 """Preload all DiaObjects and DiaSources from the Apdb given the
78 current exposure.
80 Parameters
81 ----------
82 exposure : `lsst.afw.image.Exposure`
83 An exposure with a bounding box.
84 apdb : `lsst.dax.apdb.Apdb`
85 AP database connection object.
87 Returns
88 -------
89 result : `lsst.pipe.base.Struct`
90 Results struct with components.
92 - ``diaObjects`` : Complete set of DiaObjects covering the input
93 exposure padded by ``pixelMargin``. DataFrame is indexed by
94 the ``diaObjectId`` column. (`pandas.DataFrame`)
95 - ``diaSources`` : Complete set of DiaSources covering the input
96 exposure padded by ``pixelMargin``. DataFrame is indexed by
97 ``diaObjectId``, ``filterName``, ``diaSourceId`` columns.
98 (`pandas.DataFrame`)
99 """
100 pixelRanges = self._getPixelRanges(exposure)
102 diaObjects = self.loadDiaObjects(pixelRanges, apdb)
104 dateTime = exposure.getInfo().getVisitInfo().getDate().toPython()
106 diaSources = self.loadDiaSources(diaObjects,
107 dateTime,
108 pixelRanges,
109 apdb)
111 diaForcedSources = self.loadDiaForcedSources(diaObjects,
112 dateTime,
113 apdb)
115 return pipeBase.Struct(
116 diaObjects=diaObjects,
117 diaSources=diaSources,
118 diaForcedSources=diaForcedSources)
120 @pipeBase.timeMethod
121 def loadDiaObjects(self, pixelRanges, apdb):
122 """Load DiaObjects from the Apdb based on their HTM location.
124 Parameters
125 ----------
126 pixelRanges : `tuple` [`int`]
127 Ranges of pixel values that cover region of interest.
128 apdb : `lsst.dax.apdb.Apdb`
129 Database connection object to load from.
131 Returns
132 -------
133 diaObjects : `pandas.DataFrame`
134 DiaObjects loaded from the Apdb that are within the area defined
135 by ``pixelRanges``.
136 """
137 if len(pixelRanges) == 0:
138 # If no area is specified return an empty DataFrame with the
139 # the column used for indexing later in AssociationTask.
140 diaObjects = pd.DataFrame(columns=["diaObjectId"])
141 else:
142 diaObjects = apdb.getDiaObjects(pixelRanges, return_pandas=True)
143 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
144 return diaObjects.replace(to_replace=[None], value=np.nan)
146 @pipeBase.timeMethod
147 def loadDiaSources(self, diaObjects, pixelRanges, dateTime, apdb):
148 """Load DiaSources from the Apdb based on their diaObjectId or
149 pixelId location.
151 Variable used to load sources is set in config.
153 Parameters
154 ----------
155 diaObjects : `pandas.DataFrame`
156 DiaObjects loaded from the Apdb that are within the area defined
157 by ``pixelRanges``.
158 pixelRanges : `list` of `tuples`
159 Ranges of pixelIds that cover region of interest.
160 dateTime : `datetime.datetime`
161 Time of the current visit
162 apdb : `lsst.dax.apdb.Apdb`
163 Database connection object to load from.
165 Returns
166 -------
167 DiaSources : `pandas.DataFrame`
168 DiaSources loaded from the Apdb that are within the area defined
169 by ``pixelRange`` and associated with ``diaObjects``.
170 """
171 if self.config.loadDiaSourcesByPixelId:
172 if len(pixelRanges) == 0:
173 # If no area is specified return an empty DataFrame with the
174 # the column used for indexing later in AssociationTask.
175 diaSources = pd.DataFrame(columns=["diaObjectId",
176 "filterName",
177 "diaSourceId"])
178 else:
179 diaSources = apdb.getDiaSourcesInRegion(pixelRanges,
180 dateTime,
181 return_pandas=True)
182 else:
183 if len(diaObjects) == 0:
184 # If no diaObjects are available return an empty DataFrame with
185 # the the column used for indexing later in AssociationTask.
186 diaSources = pd.DataFrame(columns=["diaObjectId",
187 "filterName",
188 "diaSourceId"])
189 else:
190 diaSources = apdb.getDiaSources(
191 diaObjects.loc[:, "diaObjectId"],
192 dateTime,
193 return_pandas=True)
195 diaSources.set_index(["diaObjectId", "filterName", "diaSourceId"],
196 drop=False,
197 inplace=True)
198 return diaSources.replace(to_replace=[None], value=np.nan)
200 @pipeBase.timeMethod
201 def loadDiaForcedSources(self, diaObjects, dateTime, apdb):
202 """Load DiaObjects from the Apdb based on their HTM location.
204 Parameters
205 ----------
206 diaObjects : `pandas.DataFrame`
207 DiaObjects loaded from the Apdb.
208 dateTime : `datetime.datetime`
209 Time of the current visit
210 apdb : `lsst.dax.apdb.Apdb`
211 Database connection object to load from.
213 Returns
214 -------
215 diaObjects : `pandas.DataFrame`
216 DiaObjects loaded from the Apdb that are within the area defined
217 by ``pixelRanges``.
218 """
219 if len(diaObjects) == 0:
220 # If no diaObjects are available return an empty DataFrame with
221 # the the column used for indexing later in AssociationTask.
222 diaForcedSources = pd.DataFrame(columns=["diaObjectId",
223 "diaForcedSourceId"])
224 else:
225 diaForcedSources = apdb.getDiaForcedSources(
226 diaObjects.loc[:, "diaObjectId"],
227 dateTime,
228 return_pandas=True)
229 diaForcedSources.set_index(["diaObjectId", "diaForcedSourceId"],
230 drop=False,
231 inplace=True)
232 return diaForcedSources.replace(to_replace=[None], value=np.nan)
234 @pipeBase.timeMethod
235 def _getPixelRanges(self, exposure):
236 """Calculate covering HTM pixels for the current exposure.
238 Parameters
239 ----------
240 exposure : `lsst.afw.image.Exposure`
241 Exposure object with calibrated WCS.
243 Returns
244 -------
245 htmRanges : `list` of `tuples`
246 A list of tuples containing `int` values.
247 """
248 bbox = geom.Box2D(exposure.getBBox())
249 bbox.grow(self.config.pixelMargin)
250 wcs = exposure.getWcs()
252 region = sphgeom.ConvexPolygon([wcs.pixelToSky(pp).getVector()
253 for pp in bbox.getCorners()])
255 indices = self.pixelator.envelope(region, self.config.htmMaxRanges)
257 return indices.ranges()