Coverage for python/lsst/ap/association/loadDiaCatalogs.py: 97%
83 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-09-30 11:40 +0000
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 pandas as pd
27import lsst.dax.apdb as daxApdb
28import lsst.geom
29import lsst.pex.config as pexConfig
30import lsst.pipe.base as pipeBase
31import lsst.pipe.base.connectionTypes as connTypes
32import lsst.sphgeom
34from lsst.utils.timer import timeMethod, duration_from_timeMethod
36from lsst.ap.association.utils import getMidpointFromTimespan, paddedRegion, readSchemaFromApdb
37from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema
39__all__ = ("LoadDiaCatalogsTask", "LoadDiaCatalogsConfig", "dropDuplicateRows",
40 "loadDiaObjectsFromApdb", "loadDiaSourcesFromApdb", "loadDiaForcedSourcesFromApdb")
43def dropDuplicateRows(catalog, index, name, log):
44 """Index a catalog loaded from the Apdb and drop any duplicate rows.
46 Where several rows share an index, the first is kept exactly as the Apdb
47 returned it and the others are discarded. The rows stay in the order the
48 Apdb returned them.
50 Parameters
51 ----------
52 catalog : `pandas.DataFrame`
53 Catalog loaded from the Apdb. Left unchanged.
54 index : `str` or `list` [`str`]
55 Column or columns to index the catalog on.
56 name : `str`
57 Name of the catalog, for logging.
58 log : `logging.Logger`
59 Log to report duplicates to.
61 Returns
62 -------
63 catalog : `pandas.DataFrame`
64 A new catalog, indexed by ``index`` and free of duplicates.
65 """
66 catalog = catalog.set_index(index, drop=False)
67 if catalog.index.has_duplicates:
68 log.warning("Duplicate %s loaded from the Apdb. This may cause "
69 "downstream pipeline issues. Dropping duplicated rows.", name)
70 catalog = catalog[~catalog.index.duplicated(keep="first")]
71 return catalog
74def loadDiaObjectsFromApdb(apdb, region, schema, log):
75 """Load DiaObjects from the Apdb based on their HTM location.
77 Parameters
78 ----------
79 apdb : `lsst.dax.apdb.Apdb`
80 Database to load the DiaObjects from.
81 region : `sphgeom.Region`
82 Region of interest, including any padding.
83 schema : `dict` of `lsst.dax.apdb.apdbSchema.ApdbSchema`
84 A dict of the schemas in the apdb.
85 log : `logging.Logger`
86 Log to report the loaded catalog to.
88 Returns
89 -------
90 diaObjects : `pandas.DataFrame`
91 DiaObjects within ``region``, indexed by ``diaObjectId``.
92 """
93 diaObjects = apdb.getDiaObjects(region)
94 diaObjects = dropDuplicateRows(diaObjects, "diaObjectId", "DiaObjects", log)
95 log.info("Loaded %i DiaObjects", len(diaObjects))
96 return convertDataFrameToSdmSchema(schema, diaObjects, tableName="DiaObject", skipIndex=True)
99def loadDiaSourcesFromApdb(apdb, region, diaObjectIds, dateTime, schema, log):
100 """Load DiaSources from the Apdb based on their diaObjectId or location.
102 Parameters
103 ----------
104 apdb : `lsst.dax.apdb.Apdb`
105 Database to load the DiaSources from.
106 region : `sphgeom.Region`
107 Region of interest, including any padding.
108 diaObjectIds : `pandas.Series`
109 Ids of the DiaObjects to load the history for.
110 dateTime : `astropy.time.Time`
111 Time of the current visit.
112 schema : `dict` of `lsst.dax.apdb.apdbSchema.ApdbSchema`
113 A dict of the schemas in the apdb.
114 log : `logging.Logger`
115 Log to report the loaded catalog to.
117 Returns
118 -------
119 diaSources : `pandas.DataFrame`
120 DiaSource history, indexed by ``diaObjectId``, ``band``, and
121 ``diaSourceId``.
122 """
123 diaSources = apdb.getDiaSources(region, diaObjectIds, dateTime)
124 diaSources = dropDuplicateRows(diaSources, ["diaObjectId", "band", "diaSourceId"], "DiaSources", log)
125 log.info("Loaded %i DiaSources", len(diaSources))
126 return convertDataFrameToSdmSchema(schema, diaSources, tableName="DiaSource", skipIndex=True)
129def loadDiaForcedSourcesFromApdb(apdb, region, diaObjectIds, dateTime, schema, log):
130 """Load DiaForcedSources from the Apdb based on their diaObjectId.
132 Parameters
133 ----------
134 apdb : `lsst.dax.apdb.Apdb`
135 Database to load the DiaForcedSources from.
136 region : `sphgeom.Region`
137 Region of interest, including any padding.
138 diaObjectIds : `pandas.Series`
139 Ids of the DiaObjects to load the history for.
140 dateTime : `astropy.time.Time`
141 Time of the current visit.
142 schema : `dict` of `lsst.dax.apdb.apdbSchema.ApdbSchema`
143 A dict of the schemas in the apdb.
144 log : `logging.Logger`
145 Log to report the loaded catalog to.
147 Returns
148 -------
149 diaForcedSources : `pandas.DataFrame`
150 DiaForcedSource history, indexed by ``diaObjectId`` and
151 ``diaForcedSourceId``.
152 """
153 if len(diaObjectIds) == 0:
154 # If no diaObjects are available return an empty DataFrame with
155 # the minimal set of columns.
156 diaForcedSources = pd.DataFrame(columns=["diaObjectId", "diaForcedSourceId"])
157 else:
158 diaForcedSources = apdb.getDiaForcedSources(region, diaObjectIds, dateTime)
159 diaForcedSources = dropDuplicateRows(diaForcedSources, ["diaObjectId", "diaForcedSourceId"],
160 "DiaForcedSources", log)
161 nVisits = 0 if diaForcedSources.empty else len(set(diaForcedSources["visit"]))
162 log.info("Loaded %i DiaForcedSources from %i visits", len(diaForcedSources), nVisits)
163 return convertDataFrameToSdmSchema(schema, diaForcedSources, tableName="DiaForcedSource",
164 skipIndex=True)
167class LoadDiaCatalogsConnections(pipeBase.PipelineTaskConnections,
168 dimensions=("instrument", "group", "detector")):
169 regionTime = connTypes.Input(
170 doc="The predicted exposure region and time",
171 name="regionTimeInfo",
172 storageClass="RegionTimeInfo",
173 dimensions=("instrument", "group", "detector"),
174 )
175 diaObjects = connTypes.Output(
176 doc="DiaObjects preloaded from the APDB.",
177 name="preloaded_diaObjects",
178 storageClass="ArrowAstropy",
179 dimensions=("instrument", "group", "detector"),
180 )
181 diaSources = connTypes.Output(
182 doc="DiaSources preloaded from the APDB.",
183 name="preloaded_diaSources",
184 storageClass="ArrowAstropy",
185 dimensions=("instrument", "group", "detector"),
186 )
187 diaForcedSources = connTypes.Output(
188 doc="DiaForcedSources preloaded from the APDB.",
189 name="preloaded_diaForcedSources",
190 storageClass="ArrowAstropy",
191 dimensions=("instrument", "group", "detector"),
192 )
195class LoadDiaCatalogsConfig(pipeBase.PipelineTaskConfig,
196 pipelineConnections=LoadDiaCatalogsConnections):
197 """Config class for LoadDiaCatalogsConfig.
198 """
199 apdb_config_url = pexConfig.Field(
200 dtype=str,
201 default=None,
202 optional=False,
203 doc="A config file specifying the APDB and its connection parameters, "
204 "typically written by the apdb-cli command-line utility. "
205 "The database must already be initialized.",
206 )
208 pixelMargin = pexConfig.RangeField(
209 doc="Padding to add to 4 all edges of the bounding box (pixels)",
210 dtype=int,
211 default=250,
212 min=0,
213 deprecated="This config has been replaced by `angleMargin`"
214 "Will be removed after v28.",
215 )
216 angleMargin = pexConfig.RangeField(
217 doc="Padding to add to the radius of the bounding circle (arcseconds)",
218 dtype=float,
219 default=20,
220 min=0,
221 )
222 doLoadForcedSources = pexConfig.Field(
223 dtype=bool,
224 default=True,
225 deprecated="Added to allow disabling forced sources for performance "
226 "reasons during the ops rehearsal. "
227 "It is expected to be removed.",
228 doc="Load forced DiaSource history from the APDB? "
229 "This should only be turned off for debugging purposes.",
230 )
233class LoadDiaCatalogsTask(pipeBase.PipelineTask):
234 """Retrieve DiaObjects and associated DiaSources from the Apdb given an
235 input exposure.
236 """
237 ConfigClass = LoadDiaCatalogsConfig
238 _DefaultName = "loadDiaCatalogs"
240 def __init__(self, **kwargs):
241 super().__init__(**kwargs)
242 self.apdb = daxApdb.Apdb.from_uri(self.config.apdb_config_url)
244 @timeMethod
245 def run(self, regionTime):
246 """Preload all DiaObjects and DiaSources from the Apdb given the
247 current exposure.
249 Parameters
250 ----------
251 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
252 A serializable container for a sky region and timespan.
254 Returns
255 -------
256 result : `lsst.pipe.base.Struct`
257 Results struct with components.
259 - ``diaObjects`` : Complete set of DiaObjects covering the input
260 exposure padded by ``angleMargin``. DataFrame is indexed by
261 the ``diaObjectId`` column. (`pandas.DataFrame`)
262 - ``diaSources`` : Complete set of DiaSources covering the input
263 exposure padded by ``angleMargin``. DataFrame is indexed by
264 ``diaObjectId``, ``band``, ``diaSourceId`` columns.
265 (`pandas.DataFrame`)
266 - ``diaForcedSources`` : Complete set of forced photometered
267 fluxes on the past 12 months of difference images at DiaObject
268 locations, indexed by ``diaObjectId`` and
269 ``diaForcedSourceId``. (`pandas.DataFrame`)
271 Raises
272 ------
273 RuntimeError
274 Raised if the Database query failed to load DiaObjects.
275 """
276 region = paddedRegion(regionTime.region,
277 lsst.sphgeom.Angle.fromDegrees(self.config.angleMargin/3600.))
278 schema = readSchemaFromApdb(self.apdb)
280 try:
281 # This is the first database query.
282 try:
283 diaObjects = self.loadDiaObjects(region, schema)
284 finally:
285 self.metadata["loadDiaObjectsDuration"] = duration_from_timeMethod(
286 self.metadata, "loadDiaObjects", clock="Utc")
287 self.log.verbose("DiaObjects: Took %.4f seconds", self.metadata["loadDiaObjectsDuration"])
289 # Load diaSources and forced sources up to the time of the exposure
290 # The timespan may include significant padding, so use the midpoint to
291 # avoid missing valid recent diaSources.
292 visitTime = getMidpointFromTimespan(regionTime.timespan)
294 try:
295 diaSources = self.loadDiaSources(diaObjects, region, visitTime, schema)
296 finally:
297 self.metadata["loadDiaSourcesDuration"] = duration_from_timeMethod(
298 self.metadata, "loadDiaSources", clock="Utc")
299 self.log.verbose("DiaSources: Took %.4f seconds", self.metadata["loadDiaSourcesDuration"])
301 if self.config.doLoadForcedSources: 301 ↛ 310line 301 didn't jump to line 310 because the condition on line 301 was always true
302 try:
303 diaForcedSources = self.loadDiaForcedSources(diaObjects, region, visitTime, schema)
304 finally:
305 self.metadata["loadDiaForcedSourcesDuration"] = duration_from_timeMethod(
306 self.metadata, "loadDiaForcedSources", clock="Utc")
307 self.log.verbose("DiaForcedSources: Took %.4f seconds",
308 self.metadata["loadDiaForcedSourcesDuration"])
309 else:
310 diaForcedSources = pd.DataFrame(columns=["diaObjectId", "diaForcedSourceId"])
311 self.metadata["loadDiaForcedSourcesDuration"] = -1
312 finally:
313 # Loki can add up the three individual times, but a combined log puts less load on the server.
314 self.log.verbose("All catalogs: Took %.4f seconds",
315 self.metadata.get("loadDiaObjectsDuration", 0)
316 + self.metadata.get("loadDiaSourcesDuration", 0)
317 + max(0, self.metadata.get("loadDiaForcedSourcesDuration", 0))
318 )
320 return pipeBase.Struct(
321 diaObjects=diaObjects,
322 diaSources=diaSources,
323 diaForcedSources=diaForcedSources)
325 @timeMethod
326 def loadDiaObjects(self, region, schema):
327 """Load DiaObjects from the Apdb based on their HTM location.
329 Parameters
330 ----------
331 region : `sphgeom.Region`
332 Region of interest.
333 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
334 A dict of the schemas in the apdb.
336 Returns
337 -------
338 diaObjects : `pandas.DataFrame`
339 DiaObjects loaded from the Apdb that are within ``region``,
340 indexed by ``diaObjectId``.
341 """
342 return loadDiaObjectsFromApdb(self.apdb, region, schema, self.log)
344 @timeMethod
345 def loadDiaSources(self, diaObjects, region, dateTime, schema):
346 """Load DiaSources from the Apdb based on their diaObjectId or
347 location.
349 Parameters
350 ----------
351 diaObjects : `pandas.DataFrame`
352 DiaObjects to load the history for, indexed by ``diaObjectId``.
353 region : `sphgeom.Region`
354 Region of interest.
355 dateTime : `astropy.time.Time`
356 Time of the current visit
357 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
358 A dict of the schemas in the apdb.
360 Returns
361 -------
362 diaSources : `pandas.DataFrame`
363 DiaSources loaded from the Apdb that are within ``region`` and
364 associated with ``diaObjects``, indexed by ``diaObjectId``,
365 ``band``, and ``diaSourceId``.
366 """
367 return loadDiaSourcesFromApdb(self.apdb, region, diaObjects.loc[:, "diaObjectId"], dateTime,
368 schema, self.log)
370 @timeMethod
371 def loadDiaForcedSources(self, diaObjects, region, dateTime, schema):
372 """Load DiaForcedSources from the Apdb based on their diaObjectId.
374 Parameters
375 ----------
376 diaObjects : `pandas.DataFrame`
377 DiaObjects to load the history for, indexed by ``diaObjectId``.
378 region : `sphgeom.Region`
379 Region of interest.
380 dateTime : `astropy.time.Time`
381 Time of the current visit
382 schema : 'dict' of `lsst.dax.apdb.apdbSchema.ApdbSchema`
383 A dict of the schemas in the apdb.
385 Returns
386 -------
387 diaForcedSources : `pandas.DataFrame`
388 DiaForcedSources loaded from the Apdb that are associated with
389 ``diaObjects``, indexed by ``diaObjectId`` and
390 ``diaForcedSourceId``.
391 """
392 return loadDiaForcedSourcesFromApdb(self.apdb, region, diaObjects.loc[:, "diaObjectId"],
393 dateTime, schema, self.log)