Coverage for python/lsst/analysis/tools/tasks/associatedSourcesTractAnalysis.py: 26%

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 

22 

23__all__ = ("AssociatedSourcesTractAnalysisConfig", "AssociatedSourcesTractAnalysisTask") 

24 

25import astropy.time 

26import astropy.units as u 

27import lsst.pex.config as pexConfig 

28import numpy as np 

29import pandas as pd 

30from astropy.table import Table, join, vstack 

31from lsst.daf.butler import DatasetProvenance 

32from lsst.drp.tasks.gbdesAstrometricFit import calculate_apparent_motion 

33from lsst.geom import Box2D 

34from lsst.pipe.base import NoWorkFound 

35from lsst.pipe.base import connectionTypes as ct 

36from lsst.skymap import BaseSkyMap 

37from smatch import Matcher 

38 

39from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask 

40 

41 

42class AssociatedSourcesTractAnalysisConnections( 

43 AnalysisBaseConnections, 

44 dimensions=("skymap", "tract", "instrument"), 

45 defaultTemplates={ 

46 "outputName": "isolated_star_presources", 

47 "associatedSourcesInputName": "isolated_star_presources", 

48 "associatedSourceIdsInputName": "isolated_star_presource_associations", 

49 }, 

50): 

51 sourceCatalogs = ct.Input( 

52 doc="Visit based source table to load from the butler", 

53 name="sourceTable_visit", 

54 storageClass="ArrowAstropy", 

55 deferLoad=True, 

56 dimensions=("visit", "band"), 

57 multiple=True, 

58 ) 

59 

60 associatedSources = ct.Input( 

61 doc="Table of associated sources", 

62 name="{associatedSourcesInputName}", 

63 storageClass="ArrowAstropy", 

64 deferLoad=True, 

65 dimensions=("instrument", "skymap", "tract"), 

66 ) 

67 

68 associatedSourceIds = ct.Input( 

69 doc="Table containing unique ids for the associated sources", 

70 name="{associatedSourceIdsInputName}", 

71 storageClass="ArrowAstropy", 

72 deferLoad=True, 

73 dimensions=("instrument", "skymap", "tract"), 

74 ) 

75 

76 skyMap = ct.Input( 

77 doc="Input definition of geometry/bbox and projection/wcs for warped exposures", 

78 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

79 storageClass="SkyMap", 

80 dimensions=("skymap",), 

81 ) 

82 

83 camera = ct.PrerequisiteInput( 

84 doc="Input camera to use for focal plane geometry.", 

85 name="camera", 

86 storageClass="Camera", 

87 dimensions=("instrument",), 

88 isCalibration=True, 

89 ) 

90 

91 astrometricCorrectionCatalogs = ct.Input( 

92 doc="Catalog containing proper motion and parallax.", 

93 name="gbdesAstrometricFit_starCatalog", 

94 storageClass="ArrowNumpyDict", 

95 dimensions=("instrument", "skymap", "tract", "physical_filter"), 

96 multiple=True, 

97 deferLoad=True, 

98 ) 

99 

100 visitTable = ct.Input( 

101 doc="Catalog containing visit information.", 

102 name="visitTable", 

103 storageClass="DataFrame", 

104 dimensions=("instrument",), 

105 ) 

106 

107 def __init__(self, *, config=None): 

108 super().__init__(config=config) 

109 

110 if not config.applyAstrometricCorrections: 

111 self.inputs.remove("astrometricCorrectionCatalogs") 

112 self.inputs.remove("visitTable") 

113 

114 

115class AssociatedSourcesTractAnalysisConfig( 

116 AnalysisBaseConfig, pipelineConnections=AssociatedSourcesTractAnalysisConnections 

117): 

118 applyAstrometricCorrections = pexConfig.Field( 

119 dtype=bool, 

120 default=True, 

121 doc="Apply proper motion and parallax corrections to source positions.", 

122 ) 

123 matchingRadius = pexConfig.Field( 

124 dtype=float, 

125 default=0.2, 

126 doc=( 

127 "Radius in mas with which to match the mean positions of the sources with the positions in the" 

128 " astrometricCorrectionCatalogs." 

129 ), 

130 ) 

131 astrometricCorrectionParameters = pexConfig.DictField( 

132 keytype=str, 

133 itemtype=str, 

134 default={ 

135 "ra": "ra", 

136 "dec": "dec", 

137 "pmRA": "pm_ra", 

138 "pmDec": "pm_dec", 

139 "parallax": "parallax", 

140 }, 

141 doc="Column names for position and motion parameters in the astrometric correction catalogs.", 

142 ) 

143 

144 

145class AssociatedSourcesTractAnalysisTask(AnalysisPipelineTask): 

146 ConfigClass = AssociatedSourcesTractAnalysisConfig 

147 _DefaultName = "associatedSourcesTractAnalysis" 

148 

149 @staticmethod 

150 def getBoxWcs(skymap, tract): 

151 """Get box that defines tract boundaries.""" 

152 tractInfo = skymap.generateTract(tract) 

153 wcs = tractInfo.getWcs() 

154 tractBox = tractInfo.getBBox() 

155 return tractBox, wcs 

156 

157 def callback(self, inputs, dataId): 

158 """Callback function to be used with reconstructor.""" 

159 return self.prepareAssociatedSources( 

160 inputs["skyMap"], 

161 dataId["tract"], 

162 inputs["sourceCatalogs"], 

163 inputs["associatedSources"], 

164 inputs["associatedSourceIds"], 

165 inputs["astrometricCorrectionCatalogs"], 

166 inputs["visitTable"], 

167 ) 

168 

169 def prepareAssociatedSources( 

170 self, 

171 skymap, 

172 tract, 

173 sourceCatalogs, 

174 associatedSources, 

175 associatedSourceIds, 

176 astrometricCorrectionCatalogs=None, 

177 visitTable=None, 

178 ): 

179 """Concatenate source catalogs and join on associated source IDs.""" 

180 

181 # Strip any provenance from tables before merging to prevent 

182 # warnings from conflicts being issued by astropy.utils.merge. 

183 for srcCat in sourceCatalogs: 

184 DatasetProvenance.strip_provenance_from_flat_dict(srcCat.meta) 

185 DatasetProvenance.strip_provenance_from_flat_dict(associatedSources.meta) 

186 DatasetProvenance.strip_provenance_from_flat_dict(associatedSourceIds.meta) 

187 

188 # associatedSource["obj_index"] refers to the corresponding index (row) 

189 # in associatedSourceIds. 

190 index = associatedSources["obj_index"] 

191 associatedSources["isolated_star_id"] = associatedSourceIds["isolated_star_id"][index] 

192 

193 # Keep only sources with associations 

194 sourceCatalogStack = vstack(sourceCatalogs, join_type="exact") 

195 dataJoined = join(sourceCatalogStack, associatedSources, keys="sourceId", join_type="inner") 

196 

197 if astrometricCorrectionCatalogs is not None: 

198 self.applyAstrometricCorrections(dataJoined, astrometricCorrectionCatalogs, visitTable) 

199 

200 # Determine which sources are contained in tract 

201 ra = np.radians(dataJoined["coord_ra"]) 

202 dec = np.radians(dataJoined["coord_dec"]) 

203 box, wcs = self.getBoxWcs(skymap, tract) 

204 box = Box2D(box) 

205 x, y = wcs.skyToPixelArray(ra, dec) 

206 boxSelection = box.contains(x, y) 

207 

208 # Keep only the sources in groups that are fully contained within the 

209 # tract 

210 dataFiltered = dataJoined[boxSelection] 

211 

212 return dataFiltered 

213 

214 def applyAstrometricCorrections(self, dataJoined, astrometricCorrectionCatalogs, visitTable): 

215 """Use proper motion/parallax catalogs to shift positions to median 

216 epoch of the visits. 

217 

218 Parameters 

219 ---------- 

220 dataJoined : `astropy.table.Table` 

221 Table containing source positions, which will be modified in place. 

222 astrometricCorrectionCatalogs: `dict` [`pd.DataFrame`] 

223 Dictionary keyed by band with proper motion and parallax catalogs. 

224 visitTable : `pd.DataFrame` 

225 Table containing the MJDs of the visits. 

226 """ 

227 if visitTable.index.name is None: 

228 # The expected index may or may not be set, depending on whether 

229 # the table was written originally as a DataFrame or something else 

230 # Parquet-friendly. 

231 visitTable.set_index("visitId", inplace=True) 

232 for band in np.unique(dataJoined["band"]): 

233 bandInd = dataJoined["band"] == band 

234 bandSources = dataJoined[bandInd] 

235 # Add key for sorting below. 

236 bandSources["__index__"] = np.arange(len(bandSources)) 

237 bandSourcesDf = bandSources.to_pandas() 

238 meanRAs = bandSourcesDf.groupby("obj_index")["coord_ra"].aggregate("mean") 

239 meanDecs = bandSourcesDf.groupby("obj_index")["coord_dec"].aggregate("mean") 

240 

241 bandPMs = astrometricCorrectionCatalogs[band] 

242 with Matcher(meanRAs, meanDecs) as m: 

243 idx, i1, i2, d = m.query_radius( 

244 bandPMs[self.config.astrometricCorrectionParameters["ra"]], 

245 bandPMs[self.config.astrometricCorrectionParameters["dec"]], 

246 (self.config.matchingRadius * u.mas).to(u.degree), 

247 return_indices=True, 

248 ) 

249 

250 catRAs = np.zeros_like(meanRAs) 

251 catDecs = np.zeros_like(meanRAs) 

252 pmRAs = np.zeros_like(meanRAs) 

253 pmDecs = np.zeros_like(meanRAs) 

254 parallaxes = np.zeros(len(meanRAs)) 

255 catRAs[i1] = bandPMs[self.config.astrometricCorrectionParameters["ra"]][i2] 

256 catDecs[i1] = bandPMs[self.config.astrometricCorrectionParameters["dec"]][i2] 

257 pmRAs[i1] = bandPMs[self.config.astrometricCorrectionParameters["pmRA"]][i2] 

258 pmDecs[i1] = bandPMs[self.config.astrometricCorrectionParameters["pmDec"]][i2] 

259 parallaxes[i1] = bandPMs[self.config.astrometricCorrectionParameters["parallax"]][i2] 

260 

261 pmDf = Table( 

262 { 

263 "ra": catRAs * u.degree, 

264 "dec": catDecs * u.degree, 

265 "pmRA": pmRAs * u.mas / u.yr, 

266 "pmDec": pmDecs * u.mas / u.yr, 

267 "parallax": parallaxes * u.mas, 

268 "obj_index": meanRAs.index, 

269 } 

270 ) 

271 

272 dataWithPM = join(bandSources, pmDf, keys="obj_index", join_type="left") 

273 

274 visits = bandSourcesDf["visit"].unique() 

275 mjds = [visitTable.loc[visit]["expMidptMJD"] for visit in visits] 

276 mjdTable = Table( 

277 [astropy.time.Time(mjds, format="mjd", scale="tai"), visits], names=["MJD", "visit"] 

278 ) 

279 dataWithMJD = join(dataWithPM, mjdTable, keys="visit", join_type="left") 

280 # After astropy 7.0, it should be possible to use "keep_order=True" 

281 # in the join and avoid sorting. 

282 dataWithMJD.sort("__index__") 

283 medianMJD = astropy.time.Time(np.median(mjds), format="mjd", scale="tai") 

284 

285 raCorrection, decCorrection = calculate_apparent_motion(dataWithMJD, medianMJD) 

286 

287 dataJoined["coord_ra"][bandInd] = dataWithMJD["coord_ra"] - raCorrection.value 

288 dataJoined["coord_dec"][bandInd] = dataWithMJD["coord_dec"] - decCorrection.value 

289 

290 def runQuantum(self, butlerQC, inputRefs, outputRefs): 

291 inputs = butlerQC.get(inputRefs) 

292 

293 # Load specified columns from source catalogs 

294 names = self.collectInputNames() 

295 names |= {"sourceId", "coord_ra", "coord_dec"} 

296 for item in ["obj_index", "isolated_star_id"]: 

297 if item in names: 

298 names.remove(item) 

299 

300 sourceCatalogs = [] 

301 for handle in inputs["sourceCatalogs"]: 

302 sourceCatalogs.append(self.loadData(handle, names)) 

303 inputs["sourceCatalogs"] = sourceCatalogs 

304 

305 if self.config.applyAstrometricCorrections: 

306 astrometricCorrections = {} 

307 for pmCatRef in inputs["astrometricCorrectionCatalogs"]: 

308 pmCat = pmCatRef.get( 

309 parameters={"columns": self.config.astrometricCorrectionParameters.values()} 

310 ) 

311 astrometricCorrections[pmCatRef.dataId["band"]] = pd.DataFrame(pmCat) 

312 inputs["astrometricCorrectionCatalogs"] = astrometricCorrections 

313 else: 

314 inputs["astrometricCorrectionCatalogs"] = None 

315 inputs["visitTable"] = None 

316 

317 dataId = butlerQC.quantum.dataId 

318 plotInfo = self.parsePlotInfo(inputs, dataId, connectionName="associatedSources") 

319 

320 # TODO: make key used for object index configurable 

321 inputs["associatedSources"] = self.loadData(inputs["associatedSources"], ["obj_index", "sourceId"]) 

322 inputs["associatedSourceIds"] = self.loadData(inputs["associatedSourceIds"], ["isolated_star_id"]) 

323 

324 if len(inputs["associatedSources"]) == 0: 

325 raise NoWorkFound(f"No associated sources in tract {dataId.tract.id}") 

326 

327 data = self.callback(inputs, dataId) 

328 

329 kwargs = {"data": data, "plotInfo": plotInfo, "skymap": inputs["skyMap"], "camera": inputs["camera"]} 

330 outputs = self.run(**kwargs) 

331 self.putByBand(butlerQC, outputs, outputRefs)