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

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 return pipeBase.Struct(
112 diaObjects=diaObjects,
113 diaSources=diaSources)
115 @pipeBase.timeMethod
116 def loadDiaObjects(self, pixelRanges, apdb):
117 """Load DiaObjects from the Apdb based on their HTM location.
119 Parameters
120 ----------
121 pixelRanges : `tuple` [`int`]
122 Ranges of pixel values that cover region of interest.
123 apdb : `lsst.dax.apdb.Apdb`
124 Database connection object to load from.
126 Returns
127 -------
128 diaObjects : `pandas.DataFrame`
129 DiaObjects loaded from the Apdb that are within the area defined
130 by ``pixelRanges``.
131 """
132 if len(pixelRanges) == 0:
133 # If no area is specified return an empty DataFrame with the
134 # the column used for indexing later in AssociationTask.
135 diaObjects = pd.DataFrame(columns=["diaObjectId"])
136 else:
137 diaObjects = apdb.getDiaObjects(pixelRanges, return_pandas=True)
138 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
139 return diaObjects.replace(to_replace=[None], value=np.nan)
141 @pipeBase.timeMethod
142 def loadDiaSources(self, diaObjects, pixelRanges, dateTime, apdb):
143 """Load DiaSources from the Apdb based on their diaObjectId or
144 pixelId location.
146 Variable used to load sources is set in config.
148 Parameters
149 ----------
150 diaObjects : `pandas.DataFrame`
151 DiaObjects loaded from the Apdb that are within the area defined
152 by ``pixelRanges``.
153 pixelRanges : `list` of `tuples`
154 Ranges of pixelIds that cover region of interest.
155 dataTime : `datetime.datetime`
156 Time of the current visit
157 apdb : `lsst.dax.apdb.Apdb`
158 Database connection object to load from.
160 Returns
161 -------
162 DiaSources : `pandas.DataFrame`
163 DiaSources loaded from the Apdb that are within the area defined
164 by ``pixelRange`` and associated with ``diaObjects``.
165 """
166 if self.config.loadDiaSourcesByPixelId:
167 if len(pixelRanges) == 0:
168 # If no area is specified return an empty DataFrame with the
169 # the column used for indexing later in AssociationTask.
170 diaSources = pd.DataFrame(columns=["diaObjectId",
171 "filterName",
172 "diaSourceId"])
173 else:
174 diaSources = apdb.getDiaSourcesInRegion(pixelRanges,
175 dateTime,
176 return_pandas=True)
177 else:
178 if len(diaObjects) == 0:
179 # If no diaObjects are available return an empty DataFrame with
180 # the the column used for indexing later in AssociationTask.
181 diaSources = pd.DataFrame(columns=["diaObjectId",
182 "filterName",
183 "diaSourceId"])
184 else:
185 diaSources = apdb.getDiaSources(
186 diaObjects.loc[:, "diaObjectId"],
187 dateTime,
188 return_pandas=True)
190 diaSources.set_index(["diaObjectId", "filterName", "diaSourceId"],
191 drop=False,
192 inplace=True)
193 return diaSources.replace(to_replace=[None], value=np.nan)
195 @pipeBase.timeMethod
196 def _getPixelRanges(self, exposure):
197 """Calculate covering HTM pixels for the current exposure.
199 Parameters
200 ----------
201 exposure : `lsst.afw.image.Exposure`
202 Exposure object with calibrated WCS.
204 Returns
205 -------
206 htmRanges : `list` of `tuples`
207 A list of tuples containing `int` values.
208 """
209 bbox = geom.Box2D(exposure.getBBox())
210 bbox.grow(self.config.pixelMargin)
211 wcs = exposure.getWcs()
213 region = sphgeom.ConvexPolygon([wcs.pixelToSky(pp).getVector()
214 for pp in bbox.getCorners()])
216 indices = self.pixelator.envelope(region, self.config.htmMaxRanges)
218 return indices.ranges()