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

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__ = ( 

24 "AssociatedSourcesTractAnalysisConfig", 

25 "AssociatedSourcesTractAnalysisTask", 

26 "AssociatedSourcesHealpix3AnalysisTask", 

27) 

28 

29import astropy.time 

30import astropy.units as u 

31import numpy as np 

32from astropy.coordinates import SkyCoord 

33from astropy.table import Table, hstack, vstack 

34 

35import lsst.pex.config as pexConfig 

36from lsst.daf.butler import DatasetProvenance 

37from lsst.drp.tasks.gbdesAstrometricFit import calculate_apparent_motion 

38from lsst.pipe.base import NoWorkFound 

39from lsst.pipe.base import connectionTypes as ct 

40from lsst.skymap import BaseSkyMap 

41from lsst.sphgeom import HealpixPixelization 

42 

43from ..interfaces import AnalysisBaseConfig, AnalysisBaseConnections, AnalysisPipelineTask 

44 

45 

46class AssociatedSourcesTractAnalysisConnections( 

47 AnalysisBaseConnections, 

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

49 defaultTemplates={ 

50 "outputName": "isolated_star_presources", 

51 "associatedSourcesInputName": "isolated_star_presources", 

52 "associatedSourceIdsInputName": "isolated_star_presource_associations", 

53 }, 

54): 

55 sourceCatalogs = ct.Input( 

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

57 name="sourceTable_visit", 

58 storageClass="ArrowAstropy", 

59 deferLoad=True, 

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

61 multiple=True, 

62 ) 

63 

64 associatedSources = ct.Input( 

65 doc="Table of associated sources", 

66 name="{associatedSourcesInputName}", 

67 storageClass="ArrowAstropy", 

68 deferLoad=True, 

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

70 ) 

71 

72 associatedSourceIds = ct.Input( 

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

74 name="{associatedSourceIdsInputName}", 

75 storageClass="ArrowAstropy", 

76 deferLoad=True, 

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

78 ) 

79 

80 skyMap = ct.Input( 

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

82 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

83 storageClass="SkyMap", 

84 dimensions=("skymap",), 

85 ) 

86 

87 camera = ct.PrerequisiteInput( 

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

89 name="camera", 

90 storageClass="Camera", 

91 dimensions=("instrument",), 

92 isCalibration=True, 

93 ) 

94 astrometricCorrectionCatalog = ct.Input( 

95 doc="Catalog with proper motion and parallax information.", 

96 name="isolated_star_stellar_motions", 

97 storageClass="ArrowAstropy", 

98 deferLoad=True, 

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

100 ) 

101 

102 visitTable = ct.Input( 

103 doc="Catalog containing visit information.", 

104 name="visitTable", 

105 storageClass="DataFrame", 

106 dimensions=("instrument",), 

107 ) 

108 

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

110 super().__init__(config=config) 

111 

112 if not config.applyAstrometricCorrections: 

113 self.inputs.remove("astrometricCorrectionCatalog") 

114 self.inputs.remove("visitTable") 

115 

116 

117class AssociatedSourcesTractAnalysisConfig( 

118 AnalysisBaseConfig, pipelineConnections=AssociatedSourcesTractAnalysisConnections 

119): 

120 applyAstrometricCorrections = pexConfig.Field( 

121 dtype=bool, 

122 default=True, 

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

124 ) 

125 astrometricCorrectionParameters = pexConfig.DictField( 

126 keytype=str, 

127 itemtype=str, 

128 default={ 

129 "ra": "ra", 

130 "dec": "dec", 

131 "pmRA": "raPM", 

132 "pmDec": "decPM", 

133 "parallax": "parallax", 

134 "isolated_star_id": "isolated_star_id", 

135 }, 

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

137 ) 

138 maxVisitCount = pexConfig.Field( 

139 dtype=int, 

140 default=None, 

141 doc="Maximum number of visits to use in calculating the metrics.", 

142 optional=True, 

143 ) 

144 maxObjects = pexConfig.Field( 

145 dtype=int, 

146 default=None, 

147 doc="Maximum number of associated objects to use in calculating the metrics.", 

148 optional=True, 

149 ) 

150 

151 

152class AssociatedSourcesTractAnalysisTask(AnalysisPipelineTask): 

153 ConfigClass = AssociatedSourcesTractAnalysisConfig 

154 _DefaultName = "associatedSourcesTractAnalysis" 

155 

156 @staticmethod 

157 def getBoxWcs(skymap, tract): 

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

159 tractInfo = skymap.generateTract(tract) 

160 wcs = tractInfo.getWcs() 

161 tractBox = tractInfo.getBBox() 

162 return tractBox, wcs 

163 

164 def callback(self, inputs, dataId): 

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

166 return self.prepareAssociatedSources( 

167 inputs["sourceCatalogs"], 

168 inputs["associatedSources"], 

169 inputs["associatedSourceIds"], 

170 inputs["astrometricCorrectionCatalog"], 

171 inputs["visitTable"], 

172 ) 

173 

174 def prepareAssociatedSources( 

175 self, 

176 sourceCatalogs, 

177 associatedSources, 

178 associatedSourceIds, 

179 astrometricCorrectionCatalog=None, 

180 visitTable=None, 

181 ): 

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

183 rng = np.random.default_rng() 

184 

185 # Strip any provenance from tables before merging to prevent 

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

187 DatasetProvenance.strip_provenance_from_flat_dict(associatedSources.meta) 

188 DatasetProvenance.strip_provenance_from_flat_dict(associatedSourceIds.meta) 

189 

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

191 # in associatedSourceIds. 

192 index = associatedSources["obj_index"] 

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

194 

195 if self.config.maxObjects: 

196 objectChoice = rng.permutation(associatedSourceIds["isolated_star_id"])[: self.config.maxObjects] 

197 objectChoice.sort() 

198 sub1 = np.clip( 

199 np.searchsorted(objectChoice, associatedSources["isolated_star_id"]), 

200 0, 

201 len(objectChoice) - 1, 

202 ) 

203 matched = objectChoice[sub1] == associatedSources["isolated_star_id"] 

204 associatedSources = associatedSources[matched] 

205 

206 trimmedSourceCatalogs = [] 

207 fullCatLen = 0 

208 # It would be preferable to use astropy's built in functions 

209 # but they are too slow so instead we use a numpy searchsorted 

210 # maneuver. 

211 sortedAssocSources = associatedSources["sourceId"].copy() 

212 assocSourcesSort = associatedSources["sourceId"].argsort() 

213 sortedAssocSources.sort() 

214 nAssocSources = len(sortedAssocSources) 

215 colsNeeded = list(self.collectInputNames()) 

216 # Only get the columns needed for the source catalogues. 

217 # The isolated_star_id and the obj_index are added later 

218 # from other tables so remove these from the list. Also 

219 # add the coord_ra and coord_dec as well because this bit 

220 # of code needs it even if it isn't requested by a 

221 # downstream atool. 

222 if "isolated_star_id" in colsNeeded: 

223 colsNeeded.remove("isolated_star_id") 

224 if "obj_index" in colsNeeded: 

225 colsNeeded.remove("obj_index") 

226 colsNeeded += ["sourceId", "coord_ra", "coord_dec"] 

227 

228 if self.config.maxVisitCount: 

229 sourceCatalogs = rng.permutation(sourceCatalogs)[: self.config.maxVisitCount] 

230 

231 for sourceCatalogRef in sourceCatalogs: 

232 sourceCatalog = sourceCatalogRef.get(parameters={"columns": set(colsNeeded)}) 

233 DatasetProvenance.strip_provenance_from_flat_dict(sourceCatalog.meta) 

234 

235 sub = np.clip( 

236 np.searchsorted(sortedAssocSources, sourceCatalog["sourceId"]), 0, nAssocSources - 1 

237 ) 

238 sourceCatalogInds = sortedAssocSources[sub] == sourceCatalog["sourceId"] 

239 assocCatalogInds = sub[sourceCatalogInds] 

240 

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

242 # the tract by matching to the associated sources table 

243 trimmedSourceCatalogs.append( 

244 hstack( 

245 [associatedSources[assocSourcesSort][assocCatalogInds], sourceCatalog[sourceCatalogInds]] 

246 ) 

247 ) 

248 fullCatLen += np.sum(sourceCatalogInds) 

249 

250 columns = trimmedSourceCatalogs[0].columns 

251 dtypes = trimmedSourceCatalogs[0].dtype 

252 zeros = np.zeros((fullCatLen, len(columns))) 

253 fullCat = Table(data=zeros, names=columns, dtype=dtypes) 

254 n = 0 

255 for trimmedSourceCatalog in trimmedSourceCatalogs: 

256 fullCat[n : n + len(trimmedSourceCatalog)] = trimmedSourceCatalog 

257 n += len(trimmedSourceCatalog) 

258 

259 if (astrometricCorrectionCatalog is not None) and (len(fullCat) != 0): 

260 self.applyAstrometricCorrections(fullCat, astrometricCorrectionCatalog, visitTable) 

261 

262 # Keep only finite ras and decs 

263 keep = np.isfinite(fullCat["coord_ra"]) & np.isfinite(fullCat["coord_dec"]) 

264 return fullCat[keep] 

265 

266 def applyAstrometricCorrections(self, dataJoined, astrometricCorrectionCatalog, visitTable): 

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

268 epoch of the visits. 

269 

270 Parameters 

271 ---------- 

272 dataJoined : `astropy.table.Table` 

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

274 astrometricCorrectionCatalog : `astropy.table.Table` 

275 Proper motion and parallax catalog. 

276 visitTable : `pd.DataFrame` 

277 Table containing the MJDs of the visits. 

278 """ 

279 if visitTable.index.name is None: 

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

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

282 # Parquet-friendly. 

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

284 

285 # Get the stellar motion catalog into the right format: 

286 for key, value in self.config.astrometricCorrectionParameters.items(): 

287 astrometricCorrectionCatalog.rename_column(value, key) 

288 astrometricCorrectionCatalog["ra"] *= u.degree 

289 astrometricCorrectionCatalog["dec"] *= u.degree 

290 astrometricCorrectionCatalog["pmRA"] *= u.mas / u.yr 

291 astrometricCorrectionCatalog["pmDec"] *= u.mas / u.yr 

292 astrometricCorrectionCatalog["parallax"] *= u.mas 

293 

294 # Join the dataJoined catalog with the astrometricCorrectionCatalog. 

295 sourceIds = dataJoined["isolated_star_id"].copy() 

296 sortOrder = np.argsort(sourceIds) 

297 sub = np.clip( 

298 np.searchsorted(sourceIds, astrometricCorrectionCatalog["isolated_star_id"], sorter=sortOrder), 

299 0, 

300 len(sourceIds) - 1, 

301 ) 

302 correctionsCatalogInds = sourceIds[sortOrder[sub]] == astrometricCorrectionCatalog["isolated_star_id"] 

303 sourceIdsInds = sortOrder[sub[correctionsCatalogInds]] 

304 

305 dataWithPM = hstack([dataJoined[sourceIdsInds], astrometricCorrectionCatalog[correctionsCatalogInds]]) 

306 

307 mjds = visitTable.loc[dataWithPM["visit"]]["expMidptMJD"] 

308 times = astropy.time.Time(mjds, format="mjd", scale="tai") 

309 dataWithPM["MJD"] = times 

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

311 

312 raCorrection, decCorrection = calculate_apparent_motion(dataWithPM, medianMJD) 

313 dataJoined["coord_ra"][sourceIdsInds] = dataWithPM["coord_ra"] - raCorrection.value 

314 dataJoined["coord_dec"][sourceIdsInds] = dataWithPM["coord_dec"] - decCorrection.value 

315 

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

317 inputs = butlerQC.get(inputRefs) 

318 

319 # Load specified columns from source catalogs 

320 names = self.collectInputNames() 

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

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

323 if item in names: 

324 names.remove(item) 

325 

326 if self.config.applyAstrometricCorrections: 

327 astrometricCorrections = inputs["astrometricCorrectionCatalog"].get( 

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

329 ) 

330 inputs["astrometricCorrectionCatalog"] = astrometricCorrections 

331 else: 

332 inputs["astrometricCorrectionCatalog"] = None 

333 inputs["visitTable"] = None 

334 

335 dataId = butlerQC.quantum.dataId 

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

337 

338 # TODO: make key used for object index configurable 

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

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

341 

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

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

344 

345 data = self.callback(inputs, dataId) 

346 

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

348 outputs = self.run(**kwargs) 

349 self.putByBand(butlerQC, outputs, outputRefs) 

350 

351 

352class AssociatedSourcesHealpix3AnalysisConnections( 

353 AssociatedSourcesTractAnalysisConnections, 

354 dimensions=("healpix3", "instrument"), 

355): 

356 associatedSources = ct.Input( 

357 doc="Table of associated sources", 

358 name="{associatedSourcesInputName}", 

359 storageClass="ArrowAstropy", 

360 deferLoad=True, 

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

362 multiple=True, 

363 ) 

364 

365 associatedSourceIds = ct.Input( 

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

367 name="{associatedSourceIdsInputName}", 

368 storageClass="ArrowAstropy", 

369 deferLoad=True, 

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

371 multiple=True, 

372 ) 

373 astrometricCorrectionCatalog = ct.Input( 

374 doc="Catalog with proper motion and parallax information.", 

375 name="isolated_star_stellar_motions", 

376 storageClass="ArrowAstropy", 

377 deferLoad=True, 

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

379 multiple=True, 

380 ) 

381 

382 

383class AssociatedSourcesHealpix3AnalysisConfig( 

384 AssociatedSourcesTractAnalysisConfig, pipelineConnections=AssociatedSourcesHealpix3AnalysisConnections 

385): 

386 pass 

387 

388 

389class AssociatedSourcesHealpix3AnalysisTask(AssociatedSourcesTractAnalysisTask): 

390 ConfigClass = AssociatedSourcesHealpix3AnalysisConfig 

391 _DefaultName = "associatedSourcesHealpix3Analysis" 

392 

393 def getHealpixOverlap(self, sources, sourceIds, pixelId, astrometricCorrections=None): 

394 

395 pixelization = HealpixPixelization(3) 

396 pixelRegion = pixelization.pixel(pixelId) 

397 

398 sourceCoords = SkyCoord(sourceIds["ra"] * u.degree, sourceIds["dec"] * u.degree).cartesian.xyz 

399 

400 inPixel = pixelRegion.contains(*sourceCoords.value) 

401 if not inPixel.any(): 

402 return sources[:0] 

403 pixelIds = sourceIds[inPixel]["isolated_star_id"] 

404 

405 sub1 = np.clip(np.searchsorted(pixelIds, sources["isolated_star_id"]), 0, len(pixelIds) - 1) 

406 matched = pixelIds[sub1] == sources["isolated_star_id"] 

407 

408 return sources[matched] 

409 

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

411 inputs = butlerQC.get(inputRefs) 

412 

413 # Load specified columns from source catalogs 

414 names = self.collectInputNames() 

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

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

417 if item in names: 

418 names.remove(item) 

419 

420 dataId = butlerQC.quantum.dataId 

421 plotInfo = self.parsePlotInfo( 

422 {"associatedSources": inputs["associatedSources"][0]}, dataId, connectionName="associatedSources" 

423 ) 

424 

425 # Loop over tract inputs, keeping only objects that in this healpix, 

426 # then stack in one big table. 

427 pixelId = dataId["healpix3"] 

428 associatedSourceRefs = { 

429 assocRef.dataId["tract"]: assocRef for assocRef in inputs["associatedSources"] 

430 } 

431 associatedSourceIdRefs = { 

432 assocRef.dataId["tract"]: assocRef for assocRef in inputs["associatedSourceIds"] 

433 } 

434 astrometricCorrectionRefs = { 

435 assocRef.dataId["tract"]: assocRef for assocRef in inputs["astrometricCorrectionCatalog"] 

436 } 

437 data = [] 

438 for tract in associatedSourceRefs: 

439 tractAssociatedSources = self.loadData(associatedSourceRefs[tract], ["obj_index", "sourceId"]) 

440 tractAssociatedSourceIds = self.loadData( 

441 associatedSourceIdRefs[tract], ["isolated_star_id", "ra", "dec"] 

442 ) 

443 if self.config.applyAstrometricCorrections: 

444 astromCorrections = astrometricCorrectionRefs[tract].get( 

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

446 ) 

447 else: 

448 astromCorrections = None 

449 tractInput = { 

450 "associatedSources": tractAssociatedSources, 

451 "associatedSourceIds": tractAssociatedSourceIds, 

452 "astrometricCorrectionCatalog": astromCorrections, 

453 "sourceCatalogs": inputs["sourceCatalogs"], 

454 "visitTable": inputs["visitTable"], 

455 } 

456 tractData = self.callback(tractInput, dataId) 

457 if len(tractData) == 0: 

458 continue 

459 trimmedData = self.getHealpixOverlap(tractData, tractAssociatedSourceIds, pixelId) 

460 trimmedData["tract"] = tract 

461 data.append(trimmedData) 

462 data = vstack(data) 

463 

464 if len(data["associatedSources"]) == 0: 

465 raise NoWorkFound(f"No associated sources in healpix {dataId.healpix3.id}") 

466 

467 kwargs = { 

468 "data": data, 

469 "plotInfo": plotInfo, 

470 "camera": inputs["camera"], 

471 } 

472 outputs = self.run(**kwargs) 

473 self.putByBand(butlerQC, outputs, outputRefs)