Coverage for python / lsst / analysis / tools / tasks / associatedSourcesTractAnalysis.py: 24%
120 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__ = ("AssociatedSourcesTractAnalysisConfig", "AssociatedSourcesTractAnalysisTask")
25import astropy.time
26import astropy.units as u
27import lsst.pex.config as pexConfig
28import numpy as np
29from astropy.table import Table, hstack
30from lsst.daf.butler import DatasetProvenance
31from lsst.drp.tasks.gbdesAstrometricFit import calculate_apparent_motion
32from lsst.pipe.base import NoWorkFound
33from lsst.pipe.base import connectionTypes as ct
34from lsst.skymap import BaseSkyMap
35from scipy.spatial import KDTree
37from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask
40class AssociatedSourcesTractAnalysisConnections(
41 AnalysisBaseConnections,
42 dimensions=("skymap", "tract", "instrument"),
43 defaultTemplates={
44 "outputName": "isolated_star_presources",
45 "associatedSourcesInputName": "isolated_star_presources",
46 "associatedSourceIdsInputName": "isolated_star_presource_associations",
47 },
48):
49 sourceCatalogs = ct.Input(
50 doc="Visit based source table to load from the butler",
51 name="sourceTable_visit",
52 storageClass="ArrowAstropy",
53 deferLoad=True,
54 dimensions=("visit", "band"),
55 multiple=True,
56 )
58 associatedSources = ct.Input(
59 doc="Table of associated sources",
60 name="{associatedSourcesInputName}",
61 storageClass="ArrowAstropy",
62 deferLoad=True,
63 dimensions=("instrument", "skymap", "tract"),
64 )
66 associatedSourceIds = ct.Input(
67 doc="Table containing unique ids for the associated sources",
68 name="{associatedSourceIdsInputName}",
69 storageClass="ArrowAstropy",
70 deferLoad=True,
71 dimensions=("instrument", "skymap", "tract"),
72 )
74 skyMap = ct.Input(
75 doc="Input definition of geometry/bbox and projection/wcs for warped exposures",
76 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
77 storageClass="SkyMap",
78 dimensions=("skymap",),
79 )
81 camera = ct.PrerequisiteInput(
82 doc="Input camera to use for focal plane geometry.",
83 name="camera",
84 storageClass="Camera",
85 dimensions=("instrument",),
86 isCalibration=True,
87 )
88 astrometricCorrectionCatalog = ct.Input(
89 doc="Catalog with proper motion and parallax information.",
90 name="isolated_star_stellar_motions",
91 storageClass="ArrowAstropy",
92 deferLoad=True,
93 dimensions=("instrument", "skymap", "tract"),
94 )
96 visitTable = ct.Input(
97 doc="Catalog containing visit information.",
98 name="visitTable",
99 storageClass="DataFrame",
100 dimensions=("instrument",),
101 )
103 def __init__(self, *, config=None):
104 super().__init__(config=config)
106 if not config.applyAstrometricCorrections:
107 self.inputs.remove("astrometricCorrectionCatalog")
108 self.inputs.remove("visitTable")
111class AssociatedSourcesTractAnalysisConfig(
112 AnalysisBaseConfig, pipelineConnections=AssociatedSourcesTractAnalysisConnections
113):
114 applyAstrometricCorrections = pexConfig.Field(
115 dtype=bool,
116 default=True,
117 doc="Apply proper motion and parallax corrections to source positions.",
118 )
119 astrometricCorrectionParameters = pexConfig.DictField(
120 keytype=str,
121 itemtype=str,
122 default={
123 "ra": "ra",
124 "dec": "dec",
125 "pmRA": "raPM",
126 "pmDec": "decPM",
127 "parallax": "parallax",
128 "isolated_star_id": "isolated_star_id",
129 },
130 doc="Column names for position and motion parameters in the astrometric correction catalogs.",
131 )
134class AssociatedSourcesTractAnalysisTask(AnalysisPipelineTask):
135 ConfigClass = AssociatedSourcesTractAnalysisConfig
136 _DefaultName = "associatedSourcesTractAnalysis"
138 @staticmethod
139 def getBoxWcs(skymap, tract):
140 """Get box that defines tract boundaries."""
141 tractInfo = skymap.generateTract(tract)
142 wcs = tractInfo.getWcs()
143 tractBox = tractInfo.getBBox()
144 return tractBox, wcs
146 def callback(self, inputs, dataId):
147 """Callback function to be used with reconstructor."""
148 return self.prepareAssociatedSources(
149 inputs["skyMap"],
150 dataId["tract"],
151 inputs["sourceCatalogs"],
152 inputs["associatedSources"],
153 inputs["associatedSourceIds"],
154 inputs["astrometricCorrectionCatalog"],
155 inputs["visitTable"],
156 )
158 def prepareAssociatedSources(
159 self,
160 skymap,
161 tract,
162 sourceCatalogs,
163 associatedSources,
164 associatedSourceIds,
165 astrometricCorrectionCatalog=None,
166 visitTable=None,
167 ):
168 """Concatenate source catalogs and join on associated source IDs."""
170 # Strip any provenance from tables before merging to prevent
171 # warnings from conflicts being issued by astropy.utils.merge.
172 DatasetProvenance.strip_provenance_from_flat_dict(associatedSources.meta)
173 DatasetProvenance.strip_provenance_from_flat_dict(associatedSourceIds.meta)
175 # associatedSource["obj_index"] refers to the corresponding index (row)
176 # in associatedSourceIds.
177 index = associatedSources["obj_index"]
178 associatedSources["isolated_star_id"] = associatedSourceIds["isolated_star_id"][index]
180 trimmedSourceCatalogs = []
181 fullCatLen = 0
182 # It would be preferable to use astropy's built in functions
183 # but they are too slow so we have this wonderful masterpiece
184 # Which is still not fast but two thirds of the time is the butler get
185 reshapedAssocSources = associatedSources["sourceId"].reshape(len(associatedSources), 1)
186 colsNeeded = list(self.collectInputNames())
187 # Only get the columns needed for the source catalogues.
188 # The isolated_star_id and the obj_index are added later
189 # from other tables so remove these from the list. Also
190 # add the coord_ra and coord_dec as well because this bit
191 # of code needs it even if it isn't requested by a
192 # downstream atool.
193 if "isolated_star_id" in colsNeeded:
194 colsNeeded.remove("isolated_star_id")
195 if "obj_index" in colsNeeded:
196 colsNeeded.remove("obj_index")
197 colsNeeded += ["sourceId", "coord_ra", "coord_dec"]
198 for sourceCatalogRef in sourceCatalogs:
199 sourceCatalog = sourceCatalogRef.get(parameters={"columns": set(colsNeeded)})
200 DatasetProvenance.strip_provenance_from_flat_dict(sourceCatalog.meta)
201 reshapedSourceCat = sourceCatalog["sourceId"].reshape(len(sourceCatalog), 1)
203 tree = KDTree(reshapedSourceCat)
204 _, inds = tree.query(reshapedAssocSources, distance_upper_bound=0.1)
205 ids = inds < len(sourceCatalog)
207 # Keep only the sources in groups that are fully contained within
208 # the tract by matching to the associated sources table
209 trimmedSourceCatalogs.append(hstack([associatedSources[ids], sourceCatalog[inds[ids]]]))
210 fullCatLen += np.sum(ids)
212 columns = trimmedSourceCatalogs[0].columns
213 dtypes = trimmedSourceCatalogs[0].dtype
214 zeros = np.zeros((fullCatLen, len(columns)))
215 fullCat = Table(data=zeros, names=columns, dtype=dtypes)
216 n = 0
217 for trimmedSourceCatalog in trimmedSourceCatalogs:
218 fullCat[n : n + len(trimmedSourceCatalog)] = trimmedSourceCatalog
219 n += len(trimmedSourceCatalog)
221 if astrometricCorrectionCatalog is not None:
222 self.applyAstrometricCorrections(fullCat, astrometricCorrectionCatalog, visitTable)
224 # Keep only finite ras and decs
225 keep = np.isfinite(fullCat["coord_ra"]) & np.isfinite(fullCat["coord_dec"])
226 return fullCat[keep]
228 def applyAstrometricCorrections(self, dataJoined, astrometricCorrectionCatalog, visitTable):
229 """Use proper motion/parallax catalogs to shift positions to median
230 epoch of the visits.
232 Parameters
233 ----------
234 dataJoined : `astropy.table.Table`
235 Table containing source positions, which will be modified in place.
236 astrometricCorrectionCatalog : `astropy.table.Table`
237 Proper motion and parallax catalog.
238 visitTable : `pd.DataFrame`
239 Table containing the MJDs of the visits.
240 """
241 if visitTable.index.name is None:
242 # The expected index may or may not be set, depending on whether
243 # the table was written originally as a DataFrame or something else
244 # Parquet-friendly.
245 visitTable.set_index("visitId", inplace=True)
247 # Get the stellar motion catalog into the right format:
248 for key, value in self.config.astrometricCorrectionParameters.items():
249 astrometricCorrectionCatalog.rename_column(value, key)
250 astrometricCorrectionCatalog["ra"] *= u.degree
251 astrometricCorrectionCatalog["dec"] *= u.degree
252 astrometricCorrectionCatalog["pmRA"] *= u.mas / u.yr
253 astrometricCorrectionCatalog["pmDec"] *= u.mas / u.yr
254 astrometricCorrectionCatalog["parallax"] *= u.mas
256 # Again using astropy join would have been great but this is four
257 # times faster
258 lenAstroCorrCat = len(astrometricCorrectionCatalog)
259 tree = KDTree(astrometricCorrectionCatalog["isolated_star_id"].reshape(lenAstroCorrCat, 1))
260 _, inds = tree.query(
261 dataJoined["isolated_star_id"].reshape(len(dataJoined), 1), distance_upper_bound=0.5
262 )
263 ids = inds < lenAstroCorrCat
265 dataWithPM = hstack([dataJoined[ids], astrometricCorrectionCatalog[inds[ids]]])
267 mjds = visitTable.loc[dataWithPM["visit"]]["expMidptMJD"]
268 times = astropy.time.Time(mjds, format="mjd", scale="tai")
269 dataWithPM["MJD"] = times
270 medianMJD = astropy.time.Time(np.median(mjds), format="mjd", scale="tai")
272 raCorrection, decCorrection = calculate_apparent_motion(dataWithPM, medianMJD)
274 dataJoined["coord_ra"] = dataWithPM["coord_ra"] - raCorrection.value
275 dataJoined["coord_dec"] = dataWithPM["coord_dec"] - decCorrection.value
277 def runQuantum(self, butlerQC, inputRefs, outputRefs):
278 inputs = butlerQC.get(inputRefs)
280 # Load specified columns from source catalogs
281 names = self.collectInputNames()
282 names |= {"sourceId", "coord_ra", "coord_dec"}
283 for item in ["obj_index", "isolated_star_id"]:
284 if item in names:
285 names.remove(item)
287 if self.config.applyAstrometricCorrections:
288 astrometricCorrections = inputs["astrometricCorrectionCatalog"].get(
289 parameters={"columns": self.config.astrometricCorrectionParameters.values()}
290 )
291 inputs["astrometricCorrectionCatalog"] = astrometricCorrections
292 else:
293 inputs["astrometricCorrectionCatalog"] = None
294 inputs["visitTable"] = None
296 dataId = butlerQC.quantum.dataId
297 plotInfo = self.parsePlotInfo(inputs, dataId, connectionName="associatedSources")
299 # TODO: make key used for object index configurable
300 inputs["associatedSources"] = self.loadData(inputs["associatedSources"], ["obj_index", "sourceId"])
301 inputs["associatedSourceIds"] = self.loadData(inputs["associatedSourceIds"], ["isolated_star_id"])
303 if len(inputs["associatedSources"]) == 0:
304 raise NoWorkFound(f"No associated sources in tract {dataId.tract.id}")
306 data = self.callback(inputs, dataId)
308 kwargs = {"data": data, "plotInfo": plotInfo, "skymap": inputs["skyMap"], "camera": inputs["camera"]}
309 outputs = self.run(**kwargs)
310 self.putByBand(butlerQC, outputs, outputRefs)