Coverage for python / lsst / analysis / tools / tasks / sourceObjectTableAnalysis.py: 28%
172 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-15 07:23 -0300
1# This file is part of analysis_tools.
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/>.
21from __future__ import annotations
23__all__ = (
24 "SourceObjectTableAnalysisConfig",
25 "SourceObjectTableAnalysisTask",
26 "ObjectEpochTableConfig",
27 "ObjectEpochTableTask",
28)
30import astropy.time
31import astropy.units as u
32import lsst.pex.config as pexConfig
33import lsst.pipe.base as pipeBase
34import numpy as np
35import pandas as pd
36from astropy.table import Table, join, vstack
37from lsst.drp.tasks.gbdesAstrometricFit import calculate_apparent_motion
38from lsst.pipe.base import AlgorithmError
39from lsst.pipe.base import connectionTypes as ct
40from smatch import Matcher
42from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask
45class IndexMismatchError(AlgorithmError):
46 """Raised if the indices in input associatedSources do not match the input
47 data."""
49 def __init__(self) -> None:
50 super().__init__(
51 "Not all sourceIds in the associated sources catalog are available in the input data."
52 )
54 @property
55 def metadata(self) -> dict:
56 """There is no metadata associated with this error."""
57 return {}
60class ObjectEpochTableConnections(
61 pipeBase.PipelineTaskConnections,
62 dimensions=("tract", "skymap"),
63):
64 objectCat = ct.Input(
65 doc="Catalog of positions in each patch.",
66 name="objectTable",
67 storageClass="ArrowAstropy",
68 dimensions=["skymap", "tract", "patch"],
69 multiple=True,
70 deferLoad=True,
71 deferGraphConstraint=True,
72 )
74 epochMap = ct.Input(
75 doc="Healsparse map of mean epoch of objectCat in each band.",
76 name="deepCoadd_epoch_map_mean",
77 storageClass="HealSparseMap",
78 dimensions=("skymap", "tract", "band"),
79 multiple=True,
80 deferLoad=True,
81 )
83 objectEpochs = ct.Output(
84 doc="Catalog of epochs for objectCat objects.",
85 name="object_epoch",
86 storageClass="ArrowAstropy",
87 dimensions=["skymap", "tract", "patch"],
88 multiple=True,
89 )
92class ObjectEpochTableConfig(pipeBase.PipelineTaskConfig, pipelineConnections=ObjectEpochTableConnections):
93 bands = pexConfig.ListField(
94 doc=("Bands in objectCat to be combined with `objectCat_selectors` to build objectCat column names."),
95 dtype=str,
96 default=["u", "g", "r", "i", "z", "y"],
97 )
100class ObjectEpochTableTask(pipeBase.PipelineTask):
101 """Collect mean epochs for the observations that went into each object.
103 TODO: DM-46202, Remove this task once the object epochs are available
104 elsewhere.
105 """
107 ConfigClass = ObjectEpochTableConfig
108 _DefaultName = "objectEpochTable"
110 def getEpochs(self, cat, epochMapDict):
111 """Get mean epoch of the visits corresponding to object position.
113 Parameters
114 ----------
115 cat : `astropy.table.Table`
116 Catalog containing object positions.
117 epochMapDict: `dict` [`DeferredDatasetHandle`]
118 Dictionary of handles for healsparse maps containing the mean epoch
119 for positions in the reference catalog.
121 Returns
122 -------
123 epochDf = `astropy.table.Table`
124 Catalog with mean epoch of visits at each object position.
125 """
126 allEpochs = {}
127 for band in self.config.bands:
128 epochs = np.ones(len(cat)) * np.nan
129 validPositions = np.isfinite(cat[f"{band}_ra"]) & np.isfinite(cat[f"{band}_dec"])
130 if validPositions.any():
131 bandEpochs = epochMapDict[band].get_values_pos(
132 cat[f"{band}_ra"][validPositions], cat[f"{band}_dec"][validPositions]
133 )
134 epochsValid = epochMapDict[band].get_values_pos(
135 cat[f"{band}_ra"][validPositions], cat[f"{band}_dec"][validPositions], valid_mask=True
136 )
137 bandEpochs[~epochsValid] = np.nan
138 epochs[validPositions] = bandEpochs
139 allEpochs[f"{band}_epoch"] = epochs
140 allEpochs["objectId"] = cat["objectId"]
142 epochTable = Table(allEpochs)
143 return epochTable
145 def runQuantum(self, butlerQC, inputRefs, outputRefs):
146 inputs = butlerQC.get(inputRefs)
148 columns = [f"{band}_{coord}" for band in self.config.bands for coord in ["ra", "dec"]]
149 columns.append("objectId")
151 inputs["epochMap"] = {ref.dataId["band"]: ref.get() for ref in inputs["epochMap"]}
153 outputEpochRefs = {outputRef.dataId["patch"]: outputRef for outputRef in outputRefs.objectEpochs}
154 for objectCatRef in inputs["objectCat"]:
155 patch = objectCatRef.dataId["patch"]
156 objectCat = objectCatRef.get(parameters={"columns": columns})
157 epochs = self.getEpochs(objectCat, inputs["epochMap"])
158 butlerQC.put(epochs, outputEpochRefs[patch])
161class SourceObjectTableAnalysisConnections(
162 AnalysisBaseConnections,
163 dimensions=("visit",),
164 defaultTemplates={
165 "inputName": "sourceTable_visit",
166 "inputCoaddName": "deep",
167 "associatedSourcesInputName": "isolated_star_presources",
168 "associatedSourceIdsInputName": "isolated_star_presource_associations",
169 "outputName": "sourceObjectTable",
170 },
171):
172 data = ct.Input(
173 doc="Visit based source table to load from the butler",
174 name="sourceTable_visit",
175 storageClass="ArrowAstropy",
176 dimensions=("visit",),
177 deferLoad=True,
178 )
180 associatedSources = ct.Input(
181 doc="Table of associated sources",
182 name="{associatedSourcesInputName}",
183 storageClass="ArrowAstropy",
184 multiple=True,
185 deferLoad=True,
186 dimensions=("instrument", "skymap", "tract"),
187 deferGraphConstraint=True,
188 )
190 associatedSourceIds = ct.Input(
191 doc="Table containing unique ids for the associated sources",
192 name="{associatedSourceIdsInputName}",
193 storageClass="ArrowAstropy",
194 deferLoad=True,
195 multiple=True,
196 dimensions=("instrument", "skymap", "tract"),
197 deferGraphConstraint=True,
198 )
200 refCat = ct.Input(
201 doc="Catalog of positions to use as reference.",
202 name="objectTable",
203 storageClass="DataFrame",
204 dimensions=["skymap", "tract", "patch"],
205 multiple=True,
206 deferLoad=True,
207 deferGraphConstraint=True,
208 )
209 astrometricCorrectionCatalog = ct.Input(
210 doc="Catalog containing proper motions and parallaxes.",
211 name="isolated_star_stellar_motions",
212 storageClass="ArrowAstropy",
213 dimensions=("instrument", "skymap", "tract"),
214 multiple=True,
215 deferLoad=True,
216 )
217 refCatEpochs = ct.Input(
218 doc="Catalog of epochs for refCat objects.",
219 name="object_epoch",
220 storageClass="ArrowAstropy",
221 dimensions=["skymap", "tract", "patch"],
222 multiple=True,
223 deferLoad=True,
224 )
225 visitTable = ct.Input(
226 doc="Catalog containing visit information.",
227 name="visitTable",
228 storageClass="DataFrame",
229 dimensions=("instrument",),
230 )
232 def __init__(self, *, config=None):
233 super().__init__(config=config)
235 if not config.applyAstrometricCorrections:
236 self.inputs.remove("astrometricCorrectionCatalog")
237 self.inputs.remove("refCatEpochs")
238 self.inputs.remove("visitTable")
241class SourceObjectTableAnalysisConfig(
242 AnalysisBaseConfig, pipelineConnections=SourceObjectTableAnalysisConnections
243):
244 ra_column = pexConfig.Field(
245 doc="Name of column in refCat to use for right ascension.",
246 dtype=str,
247 default="r_ra",
248 )
249 dec_column = pexConfig.Field(
250 doc="Name of column in refCat to use for declination.",
251 dtype=str,
252 default="r_dec",
253 )
254 epoch_column = pexConfig.Field(
255 doc=(
256 "Name of column in refCat corresponding to the epoch to which "
257 "sources will be shifted. Should correspond to the positions in "
258 "`ra_column` and `dec_column`."
259 ),
260 dtype=str,
261 default="r_epoch",
262 )
263 refCat_bands = pexConfig.ListField(
264 doc=("Bands in refCat to be combined with `refCat_selectors` to build refCat column names."),
265 dtype=str,
266 default=["u", "g", "r", "i", "z", "y"],
267 )
268 refCat_selectors = pexConfig.ListField(
269 doc=(
270 "Remove objects for which these flags are true. These strings are combined with `refCat_bands`"
271 " to build the full refCat column names"
272 ),
273 dtype=str,
274 default=["pixelFlags_saturated", "pixelFlags_saturatedCenter"],
275 )
276 refCatMatchingRadius = pexConfig.Field(
277 dtype=float,
278 default=1.0,
279 doc=(
280 "Radius in mas with which to match the mean positions of the sources with the positions in the"
281 " reference catalog."
282 ),
283 )
284 applyAstrometricCorrections = pexConfig.Field(
285 dtype=bool,
286 default=True,
287 doc="Apply proper motions and parallaxes to source positions.",
288 )
289 correctionsMatchingRadius = pexConfig.Field(
290 dtype=float,
291 default=0.2,
292 doc=(
293 "Radius in mas with which to match the mean positions of the sources with the positions in the"
294 " astrometricCorrectionCatalog."
295 ),
296 )
297 astrometricCorrectionParameters = pexConfig.DictField(
298 keytype=str,
299 itemtype=str,
300 default={
301 "ra": "ra",
302 "dec": "dec",
303 "pmRA": "raPM",
304 "pmDec": "decPM",
305 "parallax": "parallax",
306 "isolated_star_id": "isolated_star_id",
307 },
308 doc="Column names for position and motion parameters in the astrometric correction catalogs.",
309 )
311 def setDefaults(self):
312 super().setDefaults()
313 from ..atools import TargetRefCatDeltaColorMetrics
315 self.atools.astromColorDiffMetrics = TargetRefCatDeltaColorMetrics
318class SourceObjectTableAnalysisTask(AnalysisPipelineTask):
319 ConfigClass = SourceObjectTableAnalysisConfig
320 _DefaultName = "sourceObjectTableAnalysis"
322 def callback(self, inputs, dataId):
323 """Callback function to be used with reconstructor."""
324 return self.prepareAssociatedSources(
325 dataId["visit"],
326 inputs["data"],
327 inputs["associatedSources"],
328 inputs["associatedSourceIds"],
329 inputs["refCat"],
330 inputs["visitTable"],
331 inputs["astrometricCorrectionCatalog"],
332 )
334 def applyAstrometricCorrections(
335 self, isolatedSources, astrometricCorrectionCatalog, visitTable, visit, refEpochs
336 ):
337 """Shift source positions to match the epoch of the reference catalog
338 objects.
340 Parameters
341 ----------
342 isolatedSources : `astropy.table.Table`
343 Catalog of sources which will be modified in place with the
344 astrometric corrections.
345 astrometricCorrectionCatalog : `astropy.table.Table`
346 Catalog with proper motion and parallax information.
347 visitTable : `pd.DataFrame`
348 Catalog containing the epoch for the visit corresponding to the
349 isolatedSources.
350 visit : `int`
351 Identifier of the isolatedSources' visit.
352 """
353 if visitTable.index.name is None:
354 # The expected index may or may not be set, depending on whether
355 # the table was written originally as a DataFrame or something else
356 # Parquet-friendly.
357 visitTable.set_index("visitId", inplace=True)
358 sourceMjd = visitTable.loc[visit]["expMidptMJD"]
360 # Get target date from reference catalog
361 targetEpochs = refEpochs.to_numpy()
362 # There may not be a valid reference epoch on the edge of a given
363 # region. Do not make an astrometric correction for any sources on the
364 # edge.
365 targetEpochs[~np.isfinite(targetEpochs)] = sourceMjd
367 # Get the stellar motion catalog into the right format:
368 for key, value in self.config.astrometricCorrectionParameters.items():
369 astrometricCorrectionCatalog.rename_column(value, key)
370 astrometricCorrectionCatalog["ra"] *= u.degree
371 astrometricCorrectionCatalog["dec"] *= u.degree
372 astrometricCorrectionCatalog["pmRA"] *= u.mas / u.yr
373 astrometricCorrectionCatalog["pmDec"] *= u.mas / u.yr
374 astrometricCorrectionCatalog["parallax"] *= u.mas
376 joinedData = join(
377 isolatedSources[["isolated_star_id"]],
378 astrometricCorrectionCatalog,
379 keys="isolated_star_id",
380 join_type="left",
381 keep_order=True,
382 metadata_conflicts="silent",
383 )
384 joinedData["MJD"] = astropy.time.Time(sourceMjd, format="mjd", scale="tai")
386 raCorrection, decCorrection = calculate_apparent_motion(
387 joinedData, astropy.time.Time(targetEpochs, format="mjd", scale="tai")
388 )
390 isolatedSources["coord_ra"] -= raCorrection.value
391 isolatedSources["coord_dec"] -= decCorrection.value
393 def prepareAssociatedSources(
394 self,
395 visit,
396 data,
397 associatedSourceRefs,
398 associatedSourceIdRefs,
399 refCats,
400 visitTable,
401 astrometricCorrectionCatalog,
402 ):
403 """Match isolated sources with reference objects and shift the sources
404 to the object epochs if `self.config.applyAstrometricCorrections` is
405 True.
407 Parameters
408 ----------
409 visit : `int`
410 Identifier of the visit corresponding to the data.
411 data : `astropy.table.Table`
412 Catalog of sources to be associated.
413 associatedSourceRefs : `list` [`DeferredDatasetHandle`]
414 Handle for the catalogs of isolated sources. There will be multiple
415 if the visit overlaps with multiple tracts.
416 refCats : `list` [`pd.DataFrame`]
417 Catalog of objects with which the sources will be compared.
418 visitTable : `pd.DataFrame`
419 Catalog containing the epoch for the visit corresponding to the
420 isolatedSources.
421 astrometricCorrectionCatalog : `astropy.table.Table`
422 Catalog with proper motion and parallax information.
423 """
424 isolatedSources = []
425 associatedSourceIds = {
426 ref.dataId["tract"]: ref.get(parameters={"columns": ["isolated_star_id"]})
427 for ref in associatedSourceIdRefs
428 }
429 for associatedSourceRef in associatedSourceRefs:
430 tract = associatedSourceRef.dataId["tract"]
431 associatedSources = associatedSourceRef.get(
432 parameters={"columns": ["visit", "sourceId", "obj_index"]}
433 )
434 index = associatedSources["obj_index"]
435 associatedSources["isolated_star_id"] = associatedSourceIds[tract]["isolated_star_id"][index]
437 visit_sources = associatedSources[associatedSources["visit"] == visit]
438 try:
439 visitData = data.loc[visit_sources["sourceId"]]
440 visitData["isolated_star_id"] = visit_sources["isolated_star_id"]
441 isolatedSources.append(visitData)
442 except KeyError:
443 raise IndexMismatchError()
444 isolatedSources = vstack(isolatedSources)
446 if len(isolatedSources) == 0:
447 raise pipeBase.NoWorkFound(f"No isolated sources found for visit {visit}")
449 with Matcher(np.asarray(isolatedSources["coord_ra"]), np.asarray(isolatedSources["coord_dec"])) as m:
450 idx, isolatedMatchIndices, refMatchIndices, dists = m.query_radius(
451 np.asarray(refCats[self.config.ra_column]),
452 np.asarray(refCats[self.config.dec_column]),
453 self.config.refCatMatchingRadius / 3600.0,
454 return_indices=True,
455 )
457 matchIS = isolatedSources[isolatedMatchIndices]
458 # Apply proper motions and parallaxes to visit sources.
459 if self.config.applyAstrometricCorrections:
460 refCatEpochs = refCats[self.config.epoch_column].iloc[refMatchIndices]
461 self.applyAstrometricCorrections(
462 matchIS, astrometricCorrectionCatalog, visitTable, visit, refCatEpochs
463 )
465 matchRef = refCats.iloc[refMatchIndices]
466 matchIS = matchIS.to_pandas()
468 allCat = pd.concat([matchRef.reset_index(), matchIS.reset_index()], axis=1)
469 return allCat
471 def runQuantum(self, butlerQC, inputRefs, outputRefs):
472 inputs = butlerQC.get(inputRefs)
474 dataId = butlerQC.quantum.dataId
475 plotInfo = self.parsePlotInfo(inputs, dataId)
477 # Get isolated sources:
478 visit = inputs["data"].dataId["visit"]
479 band = inputs["data"].dataId["band"]
480 names = self.collectInputNames()
481 names -= {self.config.ra_column, self.config.dec_column}
482 names.add("sourceId")
483 data = inputs["data"].get(parameters={"columns": names})
484 data.add_index("sourceId")
485 inputs["data"] = data
487 if self.config.applyAstrometricCorrections:
488 refCatEpochs = {
489 epochTable.dataId["patch"]: epochTable.get() for epochTable in inputs["refCatEpochs"]
490 }
491 # Get objects:
492 allRefCats = []
493 refCatSelectors = [
494 f"{refCatBand}_{selector}"
495 for refCatBand in self.config.refCat_bands
496 for selector in self.config.refCat_selectors
497 ]
499 for refCatRef in inputs["refCat"]:
500 refCat = refCatRef.get(
501 parameters={
502 "columns": ["detect_isPrimary", self.config.ra_column, self.config.dec_column, "objectId"]
503 + refCatSelectors
504 }
505 )
506 refCat.set_index("objectId")
507 if self.config.applyAstrometricCorrections:
508 refCat = pd.merge(refCat, refCatEpochs[refCatRef.dataId["patch"]].to_pandas(), on="objectId")
509 goodInds = (
510 refCat["detect_isPrimary"]
511 & np.isfinite(refCat[self.config.ra_column])
512 & np.isfinite(refCat[self.config.dec_column])
513 )
514 goodInds &= ~refCat[refCatSelectors].any(axis=1)
515 allRefCats.append(refCat[goodInds])
517 refCat = pd.concat(allRefCats)
518 inputs["refCat"] = refCat
519 if len(refCat) == 0:
520 raise pipeBase.NoWorkFound(f"No reference catalog objects found to associate with visit {visit}")
522 if self.config.applyAstrometricCorrections:
523 pmCats = []
524 for astrometricCorrectionCatalogRef in inputs["astrometricCorrectionCatalog"]:
525 pmCat = astrometricCorrectionCatalogRef.get(
526 parameters={"columns": self.config.astrometricCorrectionParameters.values()}
527 )
528 pmCats.append(pmCat)
529 inputs["astrometricCorrectionCatalog"] = vstack(pmCats, metadata_conflicts="silent")
530 else:
531 inputs["astrometricCorrectionCatalog"] = None
532 inputs["visitTable"] = None
534 allCat = self.callback(inputs, dataId)
536 outputs = self.run(data=allCat, bands=band, plotInfo=plotInfo)
537 butlerQC.put(outputs, outputRefs)