Coverage for python/lsst/ap/association/loadDiaCatalogs.py: 32%
94 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 04:18 -0700
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"""
25import math
27import pandas as pd
29import lsst.dax.apdb as daxApdb
30import lsst.geom
31import lsst.pex.config as pexConfig
32import lsst.pipe.base as pipeBase
33import lsst.pipe.base.connectionTypes as connTypes
34import lsst.sphgeom
36from lsst.utils.timer import timeMethod, duration_from_timeMethod
38from lsst.ap.association.utils import convertDataFrameToSdmSchema, readSchemaFromApdb, getMidpointFromTimespan
40__all__ = ("LoadDiaCatalogsTask", "LoadDiaCatalogsConfig")
43class LoadDiaCatalogsConnections(pipeBase.PipelineTaskConnections,
44 dimensions=("instrument", "group", "detector")):
45 regionTime = connTypes.Input(
46 doc="The predicted exposure region and time",
47 name="regionTimeInfo",
48 storageClass="RegionTimeInfo",
49 dimensions=("instrument", "group", "detector"),
50 )
51 diaObjects = connTypes.Output(
52 doc="DiaObjects preloaded from the APDB.",
53 name="preloaded_diaObjects",
54 storageClass="ArrowAstropy",
55 dimensions=("instrument", "group", "detector"),
56 )
57 diaSources = connTypes.Output(
58 doc="DiaSources preloaded from the APDB.",
59 name="preloaded_diaSources",
60 storageClass="ArrowAstropy",
61 dimensions=("instrument", "group", "detector"),
62 )
63 diaForcedSources = connTypes.Output(
64 doc="DiaForcedSources preloaded from the APDB.",
65 name="preloaded_diaForcedSources",
66 storageClass="ArrowAstropy",
67 dimensions=("instrument", "group", "detector"),
68 )
71class LoadDiaCatalogsConfig(pipeBase.PipelineTaskConfig,
72 pipelineConnections=LoadDiaCatalogsConnections):
73 """Config class for LoadDiaCatalogsConfig.
74 """
75 apdb_config_url = pexConfig.Field(
76 dtype=str,
77 default=None,
78 optional=False,
79 doc="A config file specifying the APDB and its connection parameters, "
80 "typically written by the apdb-cli command-line utility. "
81 "The database must already be initialized.",
82 )
84 pixelMargin = pexConfig.RangeField(
85 doc="Padding to add to 4 all edges of the bounding box (pixels)",
86 dtype=int,
87 default=250,
88 min=0,
89 deprecated="This config has been replaced by `angleMargin`"
90 "Will be removed after v28.",
91 )
92 angleMargin = pexConfig.RangeField(
93 doc="Padding to add to the radius of the bounding circle (arcseconds)",
94 dtype=float,
95 default=20,
96 min=0,
97 )
98 doLoadForcedSources = pexConfig.Field(
99 dtype=bool,
100 default=True,
101 deprecated="Added to allow disabling forced sources for performance "
102 "reasons during the ops rehearsal. "
103 "It is expected to be removed.",
104 doc="Load forced DiaSource history from the APDB? "
105 "This should only be turned off for debugging purposes.",
106 )
109class LoadDiaCatalogsTask(pipeBase.PipelineTask):
110 """Retrieve DiaObjects and associated DiaSources from the Apdb given an
111 input exposure.
112 """
113 ConfigClass = LoadDiaCatalogsConfig
114 _DefaultName = "loadDiaCatalogs"
116 def __init__(self, **kwargs):
117 super().__init__(**kwargs)
118 self.apdb = daxApdb.Apdb.from_uri(self.config.apdb_config_url)
120 @timeMethod
121 def run(self, regionTime):
122 """Preload all DiaObjects and DiaSources from the Apdb given the
123 current exposure.
125 Parameters
126 ----------
127 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
128 A serializable container for a sky region and timespan.
130 Returns
131 -------
132 result : `lsst.pipe.base.Struct`
133 Results struct with components.
135 - ``diaObjects`` : Complete set of DiaObjects covering the input
136 exposure padded by ``pixelMargin``. DataFrame is indexed by
137 the ``diaObjectId`` column. (`pandas.DataFrame`)
138 - ``diaSources`` : Complete set of DiaSources covering the input
139 exposure padded by ``pixelMargin``. DataFrame is indexed by
140 ``diaObjectId``, ``band``, ``diaSourceId`` columns.
141 (`pandas.DataFrame`)
142 - ``diaForcedSources`` : Complete set of forced photometered fluxes
143 on the past 12 months of difference images at DiaObject locations.
145 Raises
146 ------
147 RuntimeError
148 Raised if the Database query failed to load DiaObjects.
149 """
150 region = self._paddedRegion(regionTime.region,
151 lsst.sphgeom.Angle.fromDegrees(self.config.angleMargin/3600.))
152 schema = readSchemaFromApdb(self.apdb)
154 try:
155 # This is the first database query.
156 try:
157 diaObjects = self.loadDiaObjects(region, schema)
158 finally:
159 self.metadata["loadDiaObjectsDuration"] = duration_from_timeMethod(
160 self.metadata, "loadDiaObjects", clock="Utc")
161 self.log.verbose("DiaObjects: Took %.4f seconds", self.metadata["loadDiaObjectsDuration"])
163 # Load diaSources and forced sources up to the time of the exposure
164 # The timespan may include significant padding, so use the midpoint to
165 # avoid missing valid recent diaSources.
166 visitTime = getMidpointFromTimespan(regionTime.timespan)
168 try:
169 diaSources = self.loadDiaSources(diaObjects, region, visitTime, schema)
170 finally:
171 self.metadata["loadDiaSourcesDuration"] = duration_from_timeMethod(
172 self.metadata, "loadDiaSources", clock="Utc")
173 self.log.verbose("DiaSources: Took %.4f seconds", self.metadata["loadDiaSourcesDuration"])
175 if self.config.doLoadForcedSources:
176 try:
177 diaForcedSources = self.loadDiaForcedSources(diaObjects, region, visitTime, schema)
178 finally:
179 self.metadata["loadDiaForcedSourcesDuration"] = duration_from_timeMethod(
180 self.metadata, "loadDiaForcedSources", clock="Utc")
181 self.log.verbose("DiaForcedSources: Took %.4f seconds",
182 self.metadata["loadDiaForcedSourcesDuration"])
183 else:
184 diaForcedSources = pd.DataFrame(columns=["diaObjectId", "diaForcedSourceId"])
185 self.metadata["loadDiaForcedSourcesDuration"] = -1
186 finally:
187 # Loki can add up the three individual times, but a combined log puts less load on the server.
188 self.log.verbose("All catalogs: Took %.4f seconds",
189 self.metadata.get("loadDiaObjectsDuration", 0)
190 + self.metadata.get("loadDiaSourcesDuration", 0)
191 + max(0, self.metadata.get("loadDiaForcedSourcesDuration", 0))
192 )
194 return pipeBase.Struct(
195 diaObjects=diaObjects,
196 diaSources=diaSources,
197 diaForcedSources=diaForcedSources)
199 @staticmethod
200 def _paddedRegion(region, margin):
201 """Return a region that has been expanded by a buffer.
203 Parameters
204 ----------
205 region : `lsst.sphgeom.Region`
206 The region to pad.
207 margin : `lsst.sphgeom.Angle`
208 The amount by which to increase the region.
210 Returns
211 -------
212 padded : `lsst.sphgeom.Region`
213 An enlarged copy of ``region``.
214 """
215 # region is almost certainly a (padded) detector bounding box.
216 if isinstance(region, lsst.sphgeom.ConvexPolygon):
217 # This is an ad-hoc, approximate implementation. It should be good
218 # enough for catalog loading, but is not a general-purpose solution.
219 center = lsst.geom.SpherePoint(region.getCentroid())
220 corners = [lsst.geom.SpherePoint(c) for c in region.getVertices()]
221 # Approximate the region as a Euclidian square
222 # geom.Angle(sphgeom.Angle) converter not pybind-wrapped???
223 diagonal_margin = lsst.geom.Angle(margin.asRadians() * math.sqrt(2.0))
224 padded = [c.offset(center.bearingTo(c), diagonal_margin) for c in corners]
225 return lsst.sphgeom.ConvexPolygon.convexHull([c.getVector() for c in padded])
226 elif hasattr(region, "dilatedBy"):
227 return region.dilatedBy(margin)
228 else:
229 return region.getBoundingCircle().dilatedBy(margin)
231 @timeMethod
232 def loadDiaObjects(self, region, schema):
233 """Load DiaObjects from the Apdb based on their HTM location.
235 Parameters
236 ----------
237 region : `sphgeom.Region`
238 Region of interest.
239 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
240 A dict of the schemas in the apdb.
242 Returns
243 -------
244 diaObjects : `pandas.DataFrame`
245 DiaObjects loaded from the Apdb that are within the area defined
246 by ``pixelRanges``.
247 """
248 diaObjects = self.apdb.getDiaObjects(region)
250 diaObjects.set_index("diaObjectId", drop=False, inplace=True)
251 if diaObjects.index.has_duplicates:
252 self.log.warning(
253 "Duplicate DiaObjects loaded from the Apdb. This may cause "
254 "downstream pipeline issues. Dropping duplicated rows")
255 # Drop duplicates via index and keep the first appearance.
256 diaObjects = diaObjects.groupby(diaObjects.index).first()
257 self.log.info("Loaded %i DiaObjects", len(diaObjects))
259 return convertDataFrameToSdmSchema(schema, diaObjects, tableName="DiaObject")
261 @timeMethod
262 def loadDiaSources(self, diaObjects, region, dateTime, schema):
263 """Load DiaSources from the Apdb based on their diaObjectId or
264 location.
266 Variable used to load sources is set in config.
268 Parameters
269 ----------
270 diaObjects : `pandas.DataFrame`
271 DiaObjects loaded from the Apdb that are within the area defined
272 by ``pixelRanges``.
273 region : `sphgeom.Region`
274 Region of interest.
275 dateTime : `astropy.time.Time`
276 Time of the current visit
277 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
278 A dict of the schemas in the apdb.
280 Returns
281 -------
282 DiaSources : `pandas.DataFrame`
283 DiaSources loaded from the Apdb that are within the area defined
284 by ``pixelRange`` and associated with ``diaObjects``.
285 """
286 diaSources = self.apdb.getDiaSources(region, diaObjects.loc[:, "diaObjectId"], dateTime)
288 diaSources.set_index(["diaObjectId", "band", "diaSourceId"],
289 drop=False,
290 inplace=True)
291 if diaSources.index.has_duplicates:
292 self.log.warning(
293 "Duplicate DiaSources loaded from the Apdb. This may cause "
294 "downstream pipeline issues. Dropping duplicated rows")
295 # Drop duplicates via index and keep the first appearance. Reset
296 # due to the index shape being slight different thatn expected.
297 diaSources = diaSources.groupby(diaSources.index).first().reset_index(drop=True)
298 diaSources.set_index(["diaObjectId", "band", "diaSourceId"],
299 drop=False,
300 inplace=True)
301 self.log.info("Loaded %i DiaSources", len(diaSources))
303 return convertDataFrameToSdmSchema(schema, diaSources, tableName="DiaSource")
305 @timeMethod
306 def loadDiaForcedSources(self, diaObjects, region, dateTime, schema):
307 """Load DiaObjects from the Apdb based on their HTM location.
309 Parameters
310 ----------
311 diaObjects : `pandas.DataFrame`
312 DiaObjects loaded from the Apdb.
313 region : `sphgeom.Region`
314 Region of interest.
315 dateTime : `astropy.time.Time`
316 Time of the current visit
317 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
318 A dict of the schemas in the apdb.
320 Returns
321 -------
322 diaObjects : `pandas.DataFrame`
323 DiaObjects loaded from the Apdb that are within the area defined
324 by ``pixelRanges``.
325 """
327 if len(diaObjects) == 0:
328 # If no diaObjects are available return an empty DataFrame with
329 # the the column used for indexing later in AssociationTask.
330 diaForcedSources = pd.DataFrame(columns=["diaObjectId",
331 "diaForcedSourceId"])
332 else:
333 diaForcedSources = self.apdb.getDiaForcedSources(
334 region,
335 diaObjects.loc[:, "diaObjectId"],
336 dateTime)
338 diaForcedSources.set_index(["diaObjectId", "diaForcedSourceId"],
339 drop=False,
340 inplace=True)
341 if diaForcedSources.index.has_duplicates:
342 self.log.warning(
343 "Duplicate DiaForcedSources loaded from the Apdb. This may "
344 "cause downstream pipeline issues. Dropping duplicated rows.")
345 # Drop duplicates via index and keep the first appearance. Reset
346 # due to the index shape being slightly different than expected.
347 diaForcedSources = diaForcedSources.groupby(diaForcedSources.index).first()
348 diaForcedSources.reset_index(drop=True, inplace=True)
349 diaForcedSources.set_index(["diaObjectId", "diaForcedSourceId"],
350 drop=False,
351 inplace=True)
352 nVisits = 0 if diaForcedSources.empty else len(set(diaForcedSources["visit"]))
353 self.log.info("Loaded %i DiaForcedSources from %i visits", len(diaForcedSources), nVisits)
355 return convertDataFrameToSdmSchema(schema, diaForcedSources, tableName="DiaForcedSource")