Coverage for python/lsst/analysis/tools/tasks/sourceObjectTableAnalysis.py: 29%

1# This file is part of analysis_tools.

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.

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, 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.

102

103 TODO: DM-46202, Remove this task once the object epochs are available

104 elsewhere.

105 """

106

107 ConfigClass = ObjectEpochTableConfig

108 _DefaultName = "objectEpochTable"

109

110 def getEpochs(self, cat, epochMapDict):

111 """Get mean epoch of the visits corresponding to object position.

112

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.

120

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"]

141

142 epochTable = Table(allEpochs)

143 return epochTable

144

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

146 inputs = butlerQC.get(inputRefs)

147

148 columns = [f"{band}_{coord}" for band in self.config.bands for coord in ["ra", "dec"]]

149 columns.append("objectId")

150

151 inputs["epochMap"] = {ref.dataId["band"]: ref.get() for ref in inputs["epochMap"]}

152

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])

159

160

161class SourceObjectTableAnalysisConnections(

162 AnalysisBaseConnections,

163 dimensions=("visit",),

164 defaultTemplates={

165 "inputName": "sourceTable_visit",

166 "inputCoaddName": "deep",

167 "associatedSourcesInputName": "isolated_star_presources",

168 "outputName": "sourceObjectTable",

169 },

170):

171 data = ct.Input(

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

173 name="sourceTable_visit",

174 storageClass="ArrowAstropy",

175 dimensions=("visit",),

176 deferLoad=True,

177 )

178

179 associatedSources = ct.Input(

180 doc="Table of associated sources",

181 name="{associatedSourcesInputName}",

182 storageClass="ArrowAstropy",

183 multiple=True,

184 deferLoad=True,

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

186 deferGraphConstraint=True,

187 )

188

189 refCat = ct.Input(

190 doc="Catalog of positions to use as reference.",

191 name="objectTable",

192 storageClass="DataFrame",

193 dimensions=["skymap", "tract", "patch"],

194 multiple=True,

195 deferLoad=True,

196 deferGraphConstraint=True,

197 )

198 astrometricCorrectionCatalog = ct.Input(

199 doc="Catalog containing proper motions and parallaxes.",

200 name="gbdesAstrometricFit_starCatalog",

201 storageClass="ArrowNumpyDict",

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

203 multiple=True,

204 deferLoad=True,

205 )

206 refCatEpochs = ct.Input(

207 doc="Catalog of epochs for refCat objects.",

208 name="object_epoch",

209 storageClass="ArrowAstropy",

210 dimensions=["skymap", "tract", "patch"],

211 multiple=True,

212 deferLoad=True,

213 )

214 visitTable = ct.Input(

215 doc="Catalog containing visit information.",

216 name="visitTable",

217 storageClass="DataFrame",

218 dimensions=("instrument",),

219 )

220

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

222 super().__init__(config=config)

223

224 if not config.applyAstrometricCorrections:

225 self.inputs.remove("astrometricCorrectionCatalog")

226 self.inputs.remove("refCatEpochs")

227 self.inputs.remove("visitTable")

228

229

230class SourceObjectTableAnalysisConfig(

231 AnalysisBaseConfig, pipelineConnections=SourceObjectTableAnalysisConnections

232):

233 ra_column = pexConfig.Field(

234 doc="Name of column in refCat to use for right ascension.",

235 dtype=str,

236 default="r_ra",

237 )

238 dec_column = pexConfig.Field(

239 doc="Name of column in refCat to use for declination.",

240 dtype=str,

241 default="r_dec",

242 )

243 epoch_column = pexConfig.Field(

244 doc=(

245 "Name of column in refCat corresponding to the epoch to which "

246 "sources will be shifted. Should correspond to the positions in "

247 "`ra_column` and `dec_column`."

248 ),

249 dtype=str,

250 default="r_epoch",

251 )

252 refCat_bands = pexConfig.ListField(

253 doc=("Bands in refCat to be combined with `refCat_selectors` to build refCat column names."),

254 dtype=str,

255 default=["u", "g", "r", "i", "z", "y"],

256 )

257 refCat_selectors = pexConfig.ListField(

258 doc=(

259 "Remove objects for which these flags are true. These strings are combined with `refCat_bands`"

260 " to build the full refCat column names"

261 ),

262 dtype=str,

263 default=["pixelFlags_saturated", "pixelFlags_saturatedCenter"],

264 )

265 refCatMatchingRadius = pexConfig.Field(

266 dtype=float,

267 default=1.0,

268 doc=(

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

270 " reference catalog."

271 ),

272 )

273 applyAstrometricCorrections = pexConfig.Field(

274 dtype=bool,

275 default=True,

276 doc="Apply proper motions and parallaxes to source positions.",

277 )

278 correctionsMatchingRadius = pexConfig.Field(

279 dtype=float,

280 default=0.2,

281 doc=(

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

283 " astrometricCorrectionCatalog."

284 ),

285 )

286 astrometricCorrectionParameters = pexConfig.DictField(

287 keytype=str,

288 itemtype=str,

289 default={

290 "ra": "ra",

291 "dec": "dec",

292 "pmRA": "pm_ra",

293 "pmDec": "pm_dec",

294 "parallax": "parallax",

295 },

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

297 )

298

299 def setDefaults(self):

300 super().setDefaults()

301 from ..atools import TargetRefCatDeltaColorMetrics

302

303 self.atools.astromColorDiffMetrics = TargetRefCatDeltaColorMetrics

304

305

306class SourceObjectTableAnalysisTask(AnalysisPipelineTask):

307 ConfigClass = SourceObjectTableAnalysisConfig

308 _DefaultName = "sourceObjectTableAnalysis"

309

310 def callback(self, inputs, dataId):

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

312 return self.prepareAssociatedSources(

313 dataId["visit"],

314 inputs["data"],

315 inputs["associatedSources"],

316 inputs["refCat"],

317 inputs["visitTable"],

318 inputs["astrometricCorrectionCatalog"],

319 )

320

321 def applyAstrometricCorrections(

322 self, isolatedSources, astrometricCorrectionCatalog, visitTable, visit, refEpochs

323 ):

324 """Shift source positions to match the epoch of the reference catalog

325 objects.

326

327 Parameters

328 ----------

329 isolatedSources : `astropy.table.Table`

330 Catalog of sources which will be modified in place with the

331 astrometric corrections.

332 astrometricCorrectionCatalog : `pd.DataFrame`

333 Catalog with proper motion and parallax information.

334 visitTable : `pd.DataFrame`

335 Catalog containing the epoch for the visit corresponding to the

336 isolatedSources.

337 visit : `int`

338 Identifier of the isolatedSources' visit.

339 """

340 if visitTable.index.name is None:

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

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

343 # Parquet-friendly.

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

345 sourceMjd = visitTable.loc[visit]["expMidptMJD"]

346

347 # Get target date from reference catalog

348 targetEpochs = refEpochs.to_numpy()

349 # There may not be a valid reference epoch on the edge of a given

350 # region. Do not make an astrometric correction for any sources on the

351 # edge.

352 targetEpochs[~np.isfinite(targetEpochs)] = sourceMjd

353

354 with Matcher(isolatedSources["coord_ra"], isolatedSources["coord_dec"]) as m:

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

356 astrometricCorrectionCatalog[self.config.astrometricCorrectionParameters["ra"]].to_numpy(),

357 astrometricCorrectionCatalog[self.config.astrometricCorrectionParameters["dec"]].to_numpy(),

358 self.config.correctionsMatchingRadius / 3600,

359 return_indices=True,

360 )

361

362 params = self.config.astrometricCorrectionParameters

363 pmDf = Table(

364 {

365 "ra": astrometricCorrectionCatalog[params["ra"]].to_numpy() * u.degree,

366 "dec": astrometricCorrectionCatalog[params["dec"]].to_numpy() * u.degree,

367 "pmRA": astrometricCorrectionCatalog[params["pmRA"]].to_numpy() * u.mas / u.yr,

368 "pmDec": astrometricCorrectionCatalog[params["pmDec"]].to_numpy() * u.mas / u.yr,

369 "parallax": astrometricCorrectionCatalog[params["parallax"]].to_numpy() * u.mas,

370 }

371 )

372

373 pmDf["MJD"] = astropy.time.Time(sourceMjd, format="mjd", scale="tai")

374 raCorrection, decCorrection = calculate_apparent_motion(

375 pmDf[i2], astropy.time.Time(targetEpochs[i1], format="mjd", scale="tai")

376 )

377

378 isolatedSources["coord_ra"][i1] -= raCorrection.value

379 isolatedSources["coord_dec"][i1] -= decCorrection.value

380

381 def prepareAssociatedSources(

382 self, visit, data, associatedSourceRefs, refCats, visitTable, astrometricCorrectionCatalog

383 ):

384 """Match isolated sources with reference objects and shift the sources

385 to the object epochs if `self.config.applyAstrometricCorrections` is

386 True.

387

388 Parameters

389 ----------

390 visit : `int`

391 Identifier of the visit corresponding to the data.

392 data : `astropy.table.Table`

393 Catalog of sources to be associated.

394 associatedSourceRefs : `list` [`DeferredDatasetHandle`]

395 Handle for the catalogs of isolated sources. There will be multiple

396 if the visit overlaps with multiple tracts.

397 refCats : `list` [`pd.DataFrame`]

398 Catalog of objects with which the sources will be compared.

399 visitTable : `pd.DataFrame`

400 Catalog containing the epoch for the visit corresponding to the

401 isolatedSources.

402 astrometricCorrectionCatalog : `pd.DataFrame`

403 Catalog with proper motion and parallax information.

404 """

405 isolatedSources = []

406 for associatedSourceRef in associatedSourceRefs:

407 associatedSources = associatedSourceRef.get(

408 parameters={"columns": ["visit", "source_row", "sourceId"]}

409 )

410 visit_sources = associatedSources[associatedSources["visit"] == visit]

411 try:

412 isolatedSources.append(data.loc[visit_sources["sourceId"]])

413 except KeyError:

414 raise IndexMismatchError()

415 isolatedSources = vstack(isolatedSources)

416

417 if len(isolatedSources) == 0:

418 raise pipeBase.NoWorkFound(f"No isolated sources found for visit {visit}")

419

420 with Matcher(isolatedSources["coord_ra"], isolatedSources["coord_dec"]) as m:

421 idx, isolatedMatchIndices, refMatchIndices, dists = m.query_radius(

422 refCats[self.config.ra_column].values,

423 refCats[self.config.dec_column].values,

424 self.config.refCatMatchingRadius / 3600.0,

425 return_indices=True,

426 )

427

428 matchIS = isolatedSources[isolatedMatchIndices]

429 # Apply proper motions and parallaxes to visit sources.

430 if self.config.applyAstrometricCorrections:

431 refCatEpochs = refCats[self.config.epoch_column].iloc[refMatchIndices]

432 self.applyAstrometricCorrections(

433 matchIS, astrometricCorrectionCatalog, visitTable, visit, refCatEpochs

434 )

435

436 matchRef = refCats.iloc[refMatchIndices]

437 matchIS = matchIS.to_pandas()

438

439 allCat = pd.concat([matchRef.reset_index(), matchIS.reset_index()], axis=1)

440 return allCat

441

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

443 inputs = butlerQC.get(inputRefs)

444

445 dataId = butlerQC.quantum.dataId

446 plotInfo = self.parsePlotInfo(inputs, dataId)

447

448 # Get isolated sources:

449 visit = inputs["data"].dataId["visit"]

450 band = inputs["data"].dataId["band"]

451 names = self.collectInputNames()

452 names -= {self.config.ra_column, self.config.dec_column}

453 names.add("sourceId")

454 data = inputs["data"].get(parameters={"columns": names})

455 data.add_index("sourceId")

456 inputs["data"] = data

457

458 if self.config.applyAstrometricCorrections:

459 refCatEpochs = {

460 epochTable.dataId["patch"]: epochTable.get() for epochTable in inputs["refCatEpochs"]

461 }

462 # Get objects:

463 allRefCats = []

464 refCatSelectors = [

465 f"{refCatBand}_{selector}"

466 for refCatBand in self.config.refCat_bands

467 for selector in self.config.refCat_selectors

468 ]

469

470 for refCatRef in inputs["refCat"]:

471 refCat = refCatRef.get(

472 parameters={

473 "columns": ["detect_isPrimary", self.config.ra_column, self.config.dec_column, "objectId"]

474 + refCatSelectors

475 }

476 )

477 refCat.set_index("objectId")

478 if self.config.applyAstrometricCorrections:

479 refCat = pd.merge(refCat, refCatEpochs[refCatRef.dataId["patch"]].to_pandas(), on="objectId")

480 goodInds = (

481 refCat["detect_isPrimary"]

482 & np.isfinite(refCat[self.config.ra_column])

483 & np.isfinite(refCat[self.config.dec_column])

484 )

485 goodInds &= ~refCat[refCatSelectors].any(axis=1)

486 allRefCats.append(refCat[goodInds])

487

488 refCat = pd.concat(allRefCats)

489 inputs["refCat"] = refCat

490 if len(refCat) == 0:

491 raise pipeBase.NoWorkFound(f"No reference catalog objects found to associate with visit {visit}")

492

493 if self.config.applyAstrometricCorrections:

494 pmCats = []

495 for astrometricCorrectionCatalogRef in inputs["astrometricCorrectionCatalog"]:

496 pmCat = astrometricCorrectionCatalogRef.get(

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

498 )

499 pmCats.append(pd.DataFrame(pmCat))

500 inputs["astrometricCorrectionCatalog"] = pd.concat(pmCats)

501 else:

502 inputs["astrometricCorrectionCatalog"] = None

503 inputs["visitTable"] = None

504

505 allCat = self.callback(inputs, dataId)

506

507 outputs = self.run(data=allCat, bands=band, plotInfo=plotInfo)

508 butlerQC.put(outputs, outputRefs)