Coverage for python / lsst / analysis / tools / tasks / sourceObjectTableAnalysis.py: 27%
188 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-12 02:09 -0800
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 NoMatchError(AlgorithmError):
61 """Raised if there are no matches between the source and reference
62 catalogs. This can happen if areas of the source or reference image were
63 not processed successfully."""
65 def __init__(self, targetCatalogSize, refCatalogSize) -> None:
66 self._metadata = {"targetCatalogSize": targetCatalogSize, "refCatalogSize": refCatalogSize}
67 super().__init__("No matches were made between the source and reference catalogs.")
69 @property
70 def metadata(self) -> dict:
71 for key, value in self._metadata.items():
72 if not isinstance(value, int | float | str):
73 raise TypeError(f"{key} is of type {type(value)}, but only (int, float, str) are allowed.")
74 return self._metadata
77class ObjectEpochTableConnections(
78 pipeBase.PipelineTaskConnections,
79 dimensions=("tract", "skymap"),
80):
81 objectCat = ct.Input(
82 doc="Catalog of positions in each patch.",
83 name="objectTable",
84 storageClass="ArrowAstropy",
85 dimensions=["skymap", "tract", "patch"],
86 multiple=True,
87 deferLoad=True,
88 deferGraphConstraint=True,
89 )
91 epochMap = ct.Input(
92 doc="Healsparse map of mean epoch of objectCat in each band.",
93 name="deepCoadd_epoch_map_mean",
94 storageClass="HealSparseMap",
95 dimensions=("skymap", "tract", "band"),
96 multiple=True,
97 deferLoad=True,
98 )
100 objectEpochs = ct.Output(
101 doc="Catalog of epochs for objectCat objects.",
102 name="object_epoch",
103 storageClass="ArrowAstropy",
104 dimensions=["skymap", "tract", "patch"],
105 multiple=True,
106 )
109class ObjectEpochTableConfig(pipeBase.PipelineTaskConfig, pipelineConnections=ObjectEpochTableConnections):
110 bands = pexConfig.ListField(
111 doc=("Bands in objectCat to be combined with `objectCat_selectors` to build objectCat column names."),
112 dtype=str,
113 default=["u", "g", "r", "i", "z", "y"],
114 )
117class ObjectEpochTableTask(pipeBase.PipelineTask):
118 """Collect mean epochs for the observations that went into each object.
120 TODO: DM-46202, Remove this task once the object epochs are available
121 elsewhere.
122 """
124 ConfigClass = ObjectEpochTableConfig
125 _DefaultName = "objectEpochTable"
127 def getEpochs(self, cat, epochMapDict):
128 """Get mean epoch of the visits corresponding to object position.
130 Parameters
131 ----------
132 cat : `astropy.table.Table`
133 Catalog containing object positions.
134 epochMapDict: `dict` [`DeferredDatasetHandle`]
135 Dictionary of handles for healsparse maps containing the mean epoch
136 for positions in the reference catalog.
138 Returns
139 -------
140 epochDf = `astropy.table.Table`
141 Catalog with mean epoch of visits at each object position.
142 """
143 allEpochs = {}
144 for band in self.config.bands:
145 epochs = np.ones(len(cat)) * np.nan
146 validPositions = np.isfinite(cat[f"{band}_ra"]) & np.isfinite(cat[f"{band}_dec"])
147 if validPositions.any():
148 bandEpochs = epochMapDict[band].get_values_pos(
149 cat[f"{band}_ra"][validPositions], cat[f"{band}_dec"][validPositions]
150 )
151 epochsValid = epochMapDict[band].get_values_pos(
152 cat[f"{band}_ra"][validPositions], cat[f"{band}_dec"][validPositions], valid_mask=True
153 )
154 bandEpochs[~epochsValid] = np.nan
155 epochs[validPositions] = bandEpochs
156 allEpochs[f"{band}_epoch"] = epochs
157 allEpochs["objectId"] = cat["objectId"]
159 epochTable = Table(allEpochs)
160 return epochTable
162 def runQuantum(self, butlerQC, inputRefs, outputRefs):
163 inputs = butlerQC.get(inputRefs)
165 columns = [f"{band}_{coord}" for band in self.config.bands for coord in ["ra", "dec"]]
166 columns.append("objectId")
168 inputs["epochMap"] = {ref.dataId["band"]: ref.get() for ref in inputs["epochMap"]}
170 outputEpochRefs = {outputRef.dataId["patch"]: outputRef for outputRef in outputRefs.objectEpochs}
171 for objectCatRef in inputs["objectCat"]:
172 patch = objectCatRef.dataId["patch"]
173 objectCat = objectCatRef.get(parameters={"columns": columns})
174 epochs = self.getEpochs(objectCat, inputs["epochMap"])
175 butlerQC.put(epochs, outputEpochRefs[patch])
178class SourceObjectTableAnalysisConnections(
179 AnalysisBaseConnections,
180 dimensions=("visit",),
181 defaultTemplates={
182 "inputName": "sourceTable_visit",
183 "inputCoaddName": "deep",
184 "associatedSourcesInputName": "isolated_star_presources",
185 "associatedSourceIdsInputName": "isolated_star_presource_associations",
186 "outputName": "sourceObjectTable",
187 },
188):
189 data = ct.Input(
190 doc="Visit based source table to load from the butler",
191 name="sourceTable_visit",
192 storageClass="ArrowAstropy",
193 dimensions=("visit",),
194 deferLoad=True,
195 )
197 associatedSources = ct.Input(
198 doc="Table of associated sources",
199 name="{associatedSourcesInputName}",
200 storageClass="ArrowAstropy",
201 multiple=True,
202 deferLoad=True,
203 dimensions=("instrument", "skymap", "tract"),
204 deferGraphConstraint=True,
205 )
207 associatedSourceIds = ct.Input(
208 doc="Table containing unique ids for the associated sources",
209 name="{associatedSourceIdsInputName}",
210 storageClass="ArrowAstropy",
211 deferLoad=True,
212 multiple=True,
213 dimensions=("instrument", "skymap", "tract"),
214 deferGraphConstraint=True,
215 )
217 refCat = ct.Input(
218 doc="Catalog of positions to use as reference.",
219 name="objectTable",
220 storageClass="DataFrame",
221 dimensions=["skymap", "tract", "patch"],
222 multiple=True,
223 deferLoad=True,
224 deferGraphConstraint=True,
225 )
226 astrometricCorrectionCatalog = ct.Input(
227 doc="Catalog containing proper motions and parallaxes.",
228 name="isolated_star_stellar_motions",
229 storageClass="ArrowAstropy",
230 dimensions=("instrument", "skymap", "tract"),
231 multiple=True,
232 deferLoad=True,
233 )
234 refCatEpochs = ct.Input(
235 doc="Catalog of epochs for refCat objects.",
236 name="object_epoch",
237 storageClass="ArrowAstropy",
238 dimensions=["skymap", "tract", "patch"],
239 multiple=True,
240 deferLoad=True,
241 )
242 visitTable = ct.Input(
243 doc="Catalog containing visit information.",
244 name="visitTable",
245 storageClass="DataFrame",
246 dimensions=("instrument",),
247 )
249 def __init__(self, *, config=None):
250 super().__init__(config=config)
252 if not config.applyAstrometricCorrections:
253 self.inputs.remove("astrometricCorrectionCatalog")
254 self.inputs.remove("refCatEpochs")
255 self.inputs.remove("visitTable")
258class SourceObjectTableAnalysisConfig(
259 AnalysisBaseConfig, pipelineConnections=SourceObjectTableAnalysisConnections
260):
261 ra_column = pexConfig.Field(
262 doc="Name of column in refCat to use for right ascension.",
263 dtype=str,
264 default="r_ra",
265 )
266 dec_column = pexConfig.Field(
267 doc="Name of column in refCat to use for declination.",
268 dtype=str,
269 default="r_dec",
270 )
271 epoch_column = pexConfig.Field(
272 doc=(
273 "Name of column in refCat corresponding to the epoch to which "
274 "sources will be shifted. Should correspond to the positions in "
275 "`ra_column` and `dec_column`."
276 ),
277 dtype=str,
278 default="r_epoch",
279 )
280 refCat_bands = pexConfig.ListField(
281 doc=("Bands in refCat to be combined with `refCat_selectors` to build refCat column names."),
282 dtype=str,
283 default=["u", "g", "r", "i", "z", "y"],
284 )
285 refCat_selectors = pexConfig.ListField(
286 doc=(
287 "Remove objects for which these flags are true. These strings are combined with `refCat_bands`"
288 " to build the full refCat column names"
289 ),
290 dtype=str,
291 default=["pixelFlags_saturated", "pixelFlags_saturatedCenter"],
292 )
293 refCatMatchingRadius = pexConfig.Field(
294 dtype=float,
295 default=1.0,
296 doc=(
297 "Radius in mas with which to match the mean positions of the sources with the positions in the"
298 " reference catalog."
299 ),
300 )
301 applyAstrometricCorrections = pexConfig.Field(
302 dtype=bool,
303 default=True,
304 doc="Apply proper motions and parallaxes to source positions.",
305 )
306 correctionsMatchingRadius = pexConfig.Field(
307 dtype=float,
308 default=0.2,
309 doc=(
310 "Radius in mas with which to match the mean positions of the sources with the positions in the"
311 " astrometricCorrectionCatalog."
312 ),
313 )
314 astrometricCorrectionParameters = pexConfig.DictField(
315 keytype=str,
316 itemtype=str,
317 default={
318 "ra": "ra",
319 "dec": "dec",
320 "pmRA": "raPM",
321 "pmDec": "decPM",
322 "parallax": "parallax",
323 "isolated_star_id": "isolated_star_id",
324 },
325 doc="Column names for position and motion parameters in the astrometric correction catalogs.",
326 )
328 def setDefaults(self):
329 super().setDefaults()
330 from ..atools import TargetRefCatDeltaColorMetrics
332 self.atools.astromColorDiffMetrics = TargetRefCatDeltaColorMetrics
335class SourceObjectTableAnalysisTask(AnalysisPipelineTask):
336 ConfigClass = SourceObjectTableAnalysisConfig
337 _DefaultName = "sourceObjectTableAnalysis"
339 def callback(self, inputs, dataId):
340 """Callback function to be used with reconstructor."""
341 return self.prepareAssociatedSources(
342 dataId["visit"],
343 inputs["data"],
344 inputs["associatedSources"],
345 inputs["associatedSourceIds"],
346 inputs["refCat"],
347 inputs["visitTable"],
348 inputs["astrometricCorrectionCatalog"],
349 )
351 def applyAstrometricCorrections(
352 self, isolatedSources, astrometricCorrectionCatalog, visitTable, visit, refEpochs
353 ):
354 """Shift source positions to match the epoch of the reference catalog
355 objects.
357 Parameters
358 ----------
359 isolatedSources : `astropy.table.Table`
360 Catalog of sources which will be modified in place with the
361 astrometric corrections.
362 astrometricCorrectionCatalog : `astropy.table.Table`
363 Catalog with proper motion and parallax information.
364 visitTable : `pd.DataFrame`
365 Catalog containing the epoch for the visit corresponding to the
366 isolatedSources.
367 visit : `int`
368 Identifier of the isolatedSources' visit.
369 """
370 if visitTable.index.name is None:
371 # The expected index may or may not be set, depending on whether
372 # the table was written originally as a DataFrame or something else
373 # Parquet-friendly.
374 visitTable.set_index("visitId", inplace=True)
375 sourceMjd = visitTable.loc[visit]["expMidptMJD"]
377 # Get target date from reference catalog
378 targetEpochs = refEpochs.to_numpy()
379 # There may not be a valid reference epoch on the edge of a given
380 # region. Do not make an astrometric correction for any sources on the
381 # edge.
382 targetEpochs[~np.isfinite(targetEpochs)] = sourceMjd
384 # Get the stellar motion catalog into the right format:
385 for key, value in self.config.astrometricCorrectionParameters.items():
386 astrometricCorrectionCatalog.rename_column(value, key)
387 astrometricCorrectionCatalog["ra"] *= u.degree
388 astrometricCorrectionCatalog["dec"] *= u.degree
389 astrometricCorrectionCatalog["pmRA"] *= u.mas / u.yr
390 astrometricCorrectionCatalog["pmDec"] *= u.mas / u.yr
391 astrometricCorrectionCatalog["parallax"] *= u.mas
393 joinedData = join(
394 isolatedSources[["isolated_star_id"]],
395 astrometricCorrectionCatalog,
396 keys="isolated_star_id",
397 join_type="left",
398 keep_order=True,
399 metadata_conflicts="silent",
400 )
401 joinedData["MJD"] = astropy.time.Time(sourceMjd, format="mjd", scale="tai")
403 raCorrection, decCorrection = calculate_apparent_motion(
404 joinedData, astropy.time.Time(targetEpochs, format="mjd", scale="tai")
405 )
407 isolatedSources["coord_ra"] -= raCorrection.value
408 isolatedSources["coord_dec"] -= decCorrection.value
410 def prepareAssociatedSources(
411 self,
412 visit,
413 data,
414 associatedSourceRefs,
415 associatedSourceIdRefs,
416 refCats,
417 visitTable,
418 astrometricCorrectionCatalog,
419 ):
420 """Match isolated sources with reference objects and shift the sources
421 to the object epochs if `self.config.applyAstrometricCorrections` is
422 True.
424 Parameters
425 ----------
426 visit : `int`
427 Identifier of the visit corresponding to the data.
428 data : `astropy.table.Table`
429 Catalog of sources to be associated.
430 associatedSourceRefs : `list` [`DeferredDatasetHandle`]
431 Handle for the catalogs of isolated sources. There will be multiple
432 if the visit overlaps with multiple tracts.
433 refCats : `list` [`pd.DataFrame`]
434 Catalog of objects with which the sources will be compared.
435 visitTable : `pd.DataFrame`
436 Catalog containing the epoch for the visit corresponding to the
437 isolatedSources.
438 astrometricCorrectionCatalog : `astropy.table.Table`
439 Catalog with proper motion and parallax information.
440 """
441 isolatedSources = []
442 associatedSourceIds = {
443 ref.dataId["tract"]: ref.get(parameters={"columns": ["isolated_star_id"]})
444 for ref in associatedSourceIdRefs
445 }
446 for associatedSourceRef in associatedSourceRefs:
447 tract = associatedSourceRef.dataId["tract"]
448 associatedSources = associatedSourceRef.get(
449 parameters={"columns": ["visit", "sourceId", "obj_index"]}
450 )
451 index = associatedSources["obj_index"]
452 associatedSources["isolated_star_id"] = associatedSourceIds[tract]["isolated_star_id"][index]
454 visit_sources = associatedSources[associatedSources["visit"] == visit]
455 try:
456 visitData = data.loc[visit_sources["sourceId"]]
457 visitData["isolated_star_id"] = visit_sources["isolated_star_id"]
458 isolatedSources.append(visitData)
459 except KeyError:
460 raise IndexMismatchError()
461 isolatedSources = vstack(isolatedSources)
463 if len(isolatedSources) == 0:
464 raise pipeBase.NoWorkFound(f"No isolated sources found for visit {visit}")
466 with Matcher(np.asarray(isolatedSources["coord_ra"]), np.asarray(isolatedSources["coord_dec"])) as m:
467 idx, isolatedMatchIndices, refMatchIndices, dists = m.query_radius(
468 np.asarray(refCats[self.config.ra_column]),
469 np.asarray(refCats[self.config.dec_column]),
470 self.config.refCatMatchingRadius / 3600.0,
471 return_indices=True,
472 )
474 matchIS = isolatedSources[isolatedMatchIndices]
476 if len(matchIS) == 0:
477 raise NoMatchError(len(isolatedSources), len(refCats))
479 # Apply proper motions and parallaxes to visit sources.
480 if self.config.applyAstrometricCorrections:
481 refCatEpochs = refCats[self.config.epoch_column].iloc[refMatchIndices]
482 self.applyAstrometricCorrections(
483 matchIS, astrometricCorrectionCatalog, visitTable, visit, refCatEpochs
484 )
486 matchRef = refCats.iloc[refMatchIndices]
487 matchIS = matchIS.to_pandas()
489 allCat = pd.concat([matchRef.reset_index(), matchIS.reset_index()], axis=1)
490 return allCat
492 def runQuantum(self, butlerQC, inputRefs, outputRefs):
493 inputs = butlerQC.get(inputRefs)
495 dataId = butlerQC.quantum.dataId
496 plotInfo = self.parsePlotInfo(inputs, dataId)
498 # Get isolated sources:
499 visit = inputs["data"].dataId["visit"]
500 band = inputs["data"].dataId["band"]
501 names = self.collectInputNames()
502 names -= {self.config.ra_column, self.config.dec_column}
503 names.add("sourceId")
504 data = inputs["data"].get(parameters={"columns": names})
505 data.add_index("sourceId")
506 inputs["data"] = data
508 if self.config.applyAstrometricCorrections:
509 refCatEpochs = {
510 epochTable.dataId["patch"]: epochTable.get() for epochTable in inputs["refCatEpochs"]
511 }
512 # Get objects:
513 allRefCats = []
514 refCatSelectors = [
515 f"{refCatBand}_{selector}"
516 for refCatBand in self.config.refCat_bands
517 for selector in self.config.refCat_selectors
518 ]
520 for refCatRef in inputs["refCat"]:
521 refCat = refCatRef.get(
522 parameters={
523 "columns": ["detect_isPrimary", self.config.ra_column, self.config.dec_column, "objectId"]
524 + refCatSelectors
525 }
526 )
527 refCat.set_index("objectId")
528 if self.config.applyAstrometricCorrections:
529 refCat = pd.merge(refCat, refCatEpochs[refCatRef.dataId["patch"]].to_pandas(), on="objectId")
530 goodInds = (
531 refCat["detect_isPrimary"]
532 & np.isfinite(refCat[self.config.ra_column])
533 & np.isfinite(refCat[self.config.dec_column])
534 )
535 goodInds &= ~refCat[refCatSelectors].any(axis=1)
536 allRefCats.append(refCat[goodInds])
538 refCat = pd.concat(allRefCats)
539 inputs["refCat"] = refCat
540 if len(refCat) == 0:
541 raise pipeBase.NoWorkFound(f"No reference catalog objects found to associate with visit {visit}")
543 if self.config.applyAstrometricCorrections:
544 pmCats = []
545 for astrometricCorrectionCatalogRef in inputs["astrometricCorrectionCatalog"]:
546 pmCat = astrometricCorrectionCatalogRef.get(
547 parameters={"columns": self.config.astrometricCorrectionParameters.values()}
548 )
549 pmCats.append(pmCat)
550 inputs["astrometricCorrectionCatalog"] = vstack(pmCats, metadata_conflicts="silent")
551 else:
552 inputs["astrometricCorrectionCatalog"] = None
553 inputs["visitTable"] = None
555 try:
556 allCat = self.callback(inputs, dataId)
557 except pipeBase.AlgorithmError as e:
558 error = pipeBase.AnnotatedPartialOutputsError.annotate(e, self, log=self.log)
559 raise error from e
561 outputs = self.run(data=allCat, bands=band, plotInfo=plotInfo)
562 butlerQC.put(outputs, outputRefs)