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

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 pandas as pd
26import lsst.geom as geom
27import lsst.pex.config as pexConfig
28import lsst.pipe.base as pipeBase
29import lsst.sphgeom as sphgeom
31__all__ = ("LoadDiaCatalogsTask", "LoadDiaCatalogsConfig")
34class LoadDiaCatalogsConfig(pexConfig.Config):
35 """Config class for LoadDiaCatalogsConfig.
36 """
37 htmLevel = pexConfig.RangeField(
38 dtype=int,
39 doc="Level of the HTM pixelization.",
40 default=20,
41 min=1,
42 )
43 htmMaxRanges = pexConfig.RangeField(
44 dtype=int,
45 doc="Maximum number of HTM (min, max) ranges to return.",
46 default=128,
47 min=2,
48 )
49 pixelMargin = pexConfig.RangeField(
50 doc="Padding to add to 4 all edges of the bounding box (pixels)",
51 dtype=int,
52 default=300,
53 min=0,
54 )
55 loadDiaSourcesByPixelId = pexConfig.Field(
56 doc="Load DiaSources by their HTM pixelId instead of by their "
57 "associated diaObjectId",
58 dtype=bool,
59 default=False,
60 )
63class LoadDiaCatalogsTask(pipeBase.Task):
64 """Retrieve DiaObjects and associated DiaSources from the Apdb given an
65 input exposure.
66 """
67 ConfigClass = LoadDiaCatalogsConfig
68 _DefaultName = "loadDiaCatalogs"
70 def __init__(self, **kwargs):
71 pipeBase.Task.__init__(self, **kwargs)
72 self.pixelator = sphgeom.HtmPixelization(self.config.htmLevel)
74 @pipeBase.timeMethod
75 def run(self, exposure, apdb):
76 """Preload all DiaObjects and DiaSources from the Apdb given the
77 current exposure.
79 Parameters
80 ----------
81 exposure : `lsst.afw.image.Exposure`
82 An exposure with a bounding box.
83 apdb : `lsst.dax.apdb.Apdb`
84 AP database connection object.
86 Returns
87 -------
88 result : `lsst.pipe.base.Struct`
89 Results struct with components.
91 - ``diaObjects`` : Complete set of DiaObjects covering the input
92 exposure padded by ``pixelMargin``. DataFrame is indexed by
93 the ``diaObjectId`` column. (`pandas.DataFrame`)
94 - ``diaSources`` : Complete set of DiaSources covering the input
95 exposure padded by ``pixelMargin``. DataFrame is indexed by
96 ``diaObjectId``, ``filterName``, ``diaSourceId`` columns.
97 (`pandas.DataFrame`)
98 """
99 pixelRanges = self._getPixelRanges(exposure)
101 diaObjects = self.loadDiaObjects(pixelRanges, apdb)
103 dateTime = exposure.getInfo().getVisitInfo().getDate().toPython()
105 diaSources = self.loadDiaSources(diaObjects,
106 dateTime,
107 pixelRanges,
108 apdb)
110 return pipeBase.Struct(
111 diaObjects=diaObjects,
112 diaSources=diaSources)
114 @pipeBase.timeMethod
115 def loadDiaObjects(self, pixelRanges, apdb):
116 """Load DiaObjects from the Apdb based on their HTM location.
118 Parameters
119 ----------
120 pixelRanges : `tuple` [`int`]
121 Ranges of pixel values that cover region of interest.
122 apdb : `lsst.dax.apdb.Apdb`
123 Database connection object to load from.
125 Returns
126 -------
127 diaObjects : `pandas.DataFrame`
128 DiaObjects loaded from the Apdb that are within the area defined
129 by ``pixelRanges``.
130 """
131 if len(pixelRanges) == 0:
132 # If no area is specified return an empty DataFrame with the
133 # the column used for indexing later in AssociationTask.
134 diaObjects = pd.DataFrame(columns=["diaObjectId"])
135 else:
136 diaObjects = apdb.getDiaObjects(pixelRanges, return_pandas=True)
137 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
138 return diaObjects
140 @pipeBase.timeMethod
141 def loadDiaSources(self, diaObjects, pixelRanges, dateTime, apdb):
142 """Load DiaSources from the Apdb based on their diaObjectId or
143 pixelId location.
145 Variable used to load sources is set in config.
147 Parameters
148 ----------
149 diaObjects : `pandas.DataFrame`
150 DiaObjects loaded from the Apdb that are within the area defined
151 by ``pixelRanges``.
152 pixelRanges : `list` of `tuples`
153 Ranges of pixelIds that cover region of interest.
154 dataTime : `datetime.datetime`
155 Time of the current visit
156 apdb : `lsst.dax.apdb.Apdb`
157 Database connection object to load from.
159 Returns
160 -------
161 DiaSources : `pandas.DataFrame`
162 DiaSources loaded from the Apdb that are within the area defined
163 by ``pixelRange`` and associated with ``diaObjects``.
164 """
165 if self.config.loadDiaSourcesByPixelId:
166 if len(pixelRanges) == 0:
167 # If no area is specified return an empty DataFrame with the
168 # the column used for indexing later in AssociationTask.
169 diaSources = pd.DataFrame(columns=["diaObjectId",
170 "filterName",
171 "diaSourceId"])
172 else:
173 diaSources = apdb.getDiaSourcesInRegion(pixelRanges,
174 dateTime,
175 return_pandas=True)
176 else:
177 if len(diaObjects) == 0:
178 # If no diaObjects are available return an empty DataFrame with
179 # the the column used for indexing later in AssociationTask.
180 diaSources = pd.DataFrame(columns=["diaObjectId",
181 "filterName",
182 "diaSourceId"])
183 else:
184 diaSources = apdb.getDiaSources(
185 diaObjects.loc[:, "diaObjectId"],
186 dateTime,
187 return_pandas=True)
189 diaSources.set_index(["diaObjectId", "filterName", "diaSourceId"],
190 drop=False,
191 inplace=True)
192 return diaSources
194 @pipeBase.timeMethod
195 def _getPixelRanges(self, exposure):
196 """Calculate covering HTM pixels for the current exposure.
198 Parameters
199 ----------
200 exposure : `lsst.afw.image.Exposure`
201 Exposure object with calibrated WCS.
203 Returns
204 -------
205 htmRanges : `list` of `tuples`
206 A list of tuples containing `int` values.
207 """
208 bbox = geom.Box2D(exposure.getBBox())
209 bbox.grow(self.config.pixelMargin)
210 wcs = exposure.getWcs()
212 region = sphgeom.ConvexPolygon([wcs.pixelToSky(pp).getVector()
213 for pp in bbox.getCorners()])
215 indices = self.pixelator.envelope(region, self.config.htmMaxRanges)
217 return indices.ranges()