Coverage for python/lsst/analysis/tools/actions/vector/vectorActions.py: 48%

190 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-27 04:19 -0700

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 "LoadVector", 

25 "DownselectVector", 

26 "MultiCriteriaDownselectVector", 

27 "ConvertFluxToMag", 

28 "ConvertUnits", 

29 "CalcSn", 

30 "ColorDiff", 

31 "ColorError", 

32 "MagDiff", 

33 "ExtinctionCorrectedMagDiff", 

34 "PerGroupStatistic", 

35 "ResidualWithPerGroupStatistic", 

36 "RAcosDec", 

37 "AngularSeparation", 

38) 

39 

40import logging 

41from typing import Optional, cast 

42 

43import numpy as np 

44import pandas as pd 

45from astropy import units as u 

46from astropy.coordinates import SkyCoord 

47from lsst.pex.config import DictField, Field 

48from lsst.pex.config.configurableActions import ConfigurableActionField, ConfigurableActionStructField 

49 

50from ...interfaces import KeyedData, KeyedDataSchema, Vector, VectorAction 

51from ...math import divide, fluxToMag, log10 

52from .selectors import VectorSelector 

53 

54_LOG = logging.getLogger(__name__) 

55 

56# Basic vectorActions 

57 

58 

59class LoadVector(VectorAction): 

60 """Load and return a Vector from KeyedData.""" 

61 

62 vectorKey = Field[str](doc="Key of vector which should be loaded") 

63 

64 def getInputSchema(self) -> KeyedDataSchema: 

65 return ((self.vectorKey, Vector),) 

66 

67 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

68 return np.array(cast(Vector, data[self.vectorKey.format(**kwargs)])) 

69 

70 

71class DownselectVector(VectorAction): 

72 """Get a vector from KeyedData, apply specified selector, return the 

73 shorter Vector. 

74 """ 

75 

76 vectorKey = Field[str](doc="column key to load from KeyedData") 

77 

78 selector = ConfigurableActionField[VectorAction]( 

79 doc="Action which returns a selection mask", default=VectorSelector 

80 ) 

81 

82 def getInputSchema(self) -> KeyedDataSchema: 

83 yield (self.vectorKey, Vector) 

84 yield from cast(VectorAction, self.selector).getInputSchema() 

85 

86 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

87 mask = cast(VectorAction, self.selector)(data, **kwargs) 

88 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask] 

89 

90 

91class MultiCriteriaDownselectVector(VectorAction): 

92 """Get a vector from KeyedData, apply specified set of selectors with AND 

93 logic, and return the shorter Vector. 

94 """ 

95 

96 vectorKey = Field[str](doc="column key to load from KeyedData") 

97 

98 selectors = ConfigurableActionStructField[VectorAction]( 

99 doc="Selectors for selecting rows, will be AND together", 

100 ) 

101 

102 def getInputSchema(self) -> KeyedDataSchema: 

103 yield (self.vectorKey, Vector) 

104 for action in self.selectors: 

105 yield from action.getInputSchema() 

106 

107 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

108 mask: Optional[Vector] = None 

109 for selector in self.selectors: 

110 subMask = selector(data, **kwargs) 

111 if mask is None: 

112 mask = subMask 

113 else: 

114 mask *= subMask # type: ignore 

115 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask] 

116 

117 

118# Astronomical vectorActions 

119 

120 

121class CalcSn(VectorAction): 

122 """Calculate the signal-to-noise ratio from a single flux vector.""" 

123 

124 fluxType = Field[str](doc="Flux type (vector key) to calculate the S/N.", default="{band}_psfFlux") 

125 uncertaintySuffix = Field[str]( 

126 doc="Suffix to add to fluxType to specify the uncertainty column", default="Err" 

127 ) 

128 

129 def getInputSchema(self) -> KeyedDataSchema: 

130 yield self.fluxType, Vector 

131 yield f"{self.fluxType}{self.uncertaintySuffix}", Vector 

132 

133 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

134 signal = np.array(data[self.fluxType.format(**kwargs)]) 

135 noise = np.array(data[f"{self.fluxType}{self.uncertaintySuffix}".format(**kwargs)]) 

136 return divide(signal, noise) 

137 

138 

139class ColorDiff(VectorAction): 

140 """Calculate the difference between two colors from flux actions.""" 

141 

142 color1_flux1 = ConfigurableActionField[VectorAction](doc="Action providing first color's first flux") 

143 color1_flux2 = ConfigurableActionField[VectorAction](doc="Action providing first color's second flux") 

144 color2_flux1 = ConfigurableActionField[VectorAction](doc="Action providing second color's first flux") 

145 color2_flux2 = ConfigurableActionField[VectorAction](doc="Action providing second color's second flux") 

146 returnMillimags = Field[bool](doc="Whether to return color_diff in millimags (mags if not)", default=True) 

147 

148 def getInputSchema(self) -> KeyedDataSchema: 

149 yield from self.color1_flux1.getInputSchema() 

150 yield from self.color1_flux2.getInputSchema() 

151 yield from self.color2_flux1.getInputSchema() 

152 yield from self.color2_flux2.getInputSchema() 

153 

154 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

155 color_diff = -2.5 * log10( 

156 divide( 

157 self.color1_flux1(data, **kwargs) * self.color2_flux2(data, **kwargs), 

158 self.color1_flux2(data, **kwargs) * self.color2_flux1(data, **kwargs), 

159 ) 

160 ) 

161 

162 if self.returnMillimags: 

163 color_diff *= 1000 

164 

165 return color_diff 

166 

167 

168class ColorError(VectorAction): 

169 """Calculate the error in a color from two different flux error columns.""" 

170 

171 flux_err1 = ConfigurableActionField[VectorAction](doc="Action providing error for first flux") 

172 flux_err2 = ConfigurableActionField[VectorAction](doc="Action providing error for second flux") 

173 returnMillimags = Field[bool](doc="Whether to return color_err in millimags (mags if not)", default=True) 

174 

175 def getInputSchema(self) -> KeyedDataSchema: 

176 yield from self.flux_err1.getInputSchema() 

177 yield from self.flux_err2.getInputSchema() 

178 

179 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

180 flux_err1 = self.flux_err1(data, **kwargs) 

181 flux_err2 = self.flux_err2(data, **kwargs) 

182 color_err = (2.5 / np.log(10)) * np.hypot(flux_err1, flux_err2) 

183 

184 if self.returnMillimags: 

185 color_err *= 1000 

186 

187 return color_err 

188 

189 

190class ConvertFluxToMag(VectorAction): 

191 """Turn nano janskies into magnitudes.""" 

192 

193 vectorKey = Field[str](doc="Key of flux vector to convert to mags") 

194 fluxUnit = Field[str](doc="Astropy unit of flux vector", default="nJy") 

195 returnMillimags = Field[bool](doc="Use millimags or not?", default=False) 

196 

197 def getInputSchema(self) -> KeyedDataSchema: 

198 return ((self.vectorKey, Vector),) 

199 

200 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

201 return fluxToMag( 

202 cast(Vector, data[self.vectorKey.format(**kwargs)]), 

203 flux_unit=self.fluxUnit, 

204 return_millimags=self.returnMillimags, 

205 ) 

206 

207 

208class ConvertUnits(VectorAction): 

209 """Convert the units of a vector.""" 

210 

211 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector) 

212 inUnit = Field[str](doc="input Astropy unit") 

213 outUnit = Field[str](doc="output Astropy unit") 

214 

215 def getInputSchema(self) -> KeyedDataSchema: 

216 return tuple(self.buildAction.getInputSchema()) 

217 

218 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

219 dataWithUnit = self.buildAction(data, **kwargs) * u.Unit(self.inUnit) 

220 return dataWithUnit.to(self.outUnit).value 

221 

222 

223class MagDiff(VectorAction): 

224 """Calculate the difference between two magnitudes; 

225 each magnitude is derived from a flux column. 

226 

227 Notes 

228 ----- 

229 The flux columns need to be in units (specifiable in 

230 the fluxUnits1 and 2 config options) that can be converted 

231 to janskies. This action doesn't have any calibration 

232 information and assumes that the fluxes are already 

233 calibrated. 

234 """ 

235 

236 col1 = Field[str](doc="Column to subtract from") 

237 fluxUnits1 = Field[str](doc="Units for col1", default="nanojansky") 

238 col2 = Field[str](doc="Column to subtract") 

239 fluxUnits2 = Field[str](doc="Units for col2", default="nanojansky") 

240 returnMillimags = Field[bool](doc="Use millimags or not?", default=True) 

241 

242 def getInputSchema(self) -> KeyedDataSchema: 

243 return ((self.col1, Vector), (self.col2, Vector)) 

244 

245 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

246 mag1 = fluxToMag(data[self.col1.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits1)) 

247 mag2 = fluxToMag(data[self.col2.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits2)) 

248 magDiff = mag1 - mag2 

249 if self.returnMillimags: 

250 magDiff *= 1000.0 

251 return magDiff 

252 

253 

254class ExtinctionCorrectedMagDiff(VectorAction): 

255 """Compute the difference between two magnitudes and correct for extinction 

256 By default bands are derived from the <band>_ prefix on flux columns, 

257 per the naming convention in the Object Table: 

258 e.g. the band of 'g_psfFlux' is 'g'. If column names follow another 

259 convention, bands can alternatively be supplied via the band1 or band2 

260 config parameters. 

261 If band1 and band2 are supplied, the flux column names are ignored. 

262 """ 

263 

264 magDiff = ConfigurableActionField[VectorAction]( 

265 doc="Action that returns a difference in magnitudes", default=MagDiff 

266 ) 

267 ebvCol = Field[str](doc="E(B-V) Column Name", default="ebv") 

268 band1 = Field[str]( 

269 doc="Optional band for magDiff.col1. Supercedes column name prefix", 

270 optional=True, 

271 default=None, 

272 ) 

273 band2 = Field[str]( 

274 doc="Optional band for magDiff.col2. Supercedes column name prefix", 

275 optional=True, 

276 default=None, 

277 ) 

278 extinctionCoeffs = DictField[str, float]( 

279 doc="Dictionary of extinction coefficients for conversion from E(B-V) to extinction, A_band." 

280 "Key must be the band", 

281 optional=True, 

282 default=None, 

283 ) 

284 

285 def getInputSchema(self) -> KeyedDataSchema: 

286 return self.magDiff.getInputSchema() + ((self.ebvCol, Vector),) 

287 

288 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

289 diff = self.magDiff(data, **kwargs) 

290 if not self.extinctionCoeffs: 

291 _LOG.debug("No extinction Coefficients. Not applying extinction correction") 

292 return diff 

293 

294 col1Band = self.band1 if self.band1 else self.magDiff.col1.split("_")[0] 

295 col2Band = self.band2 if self.band2 else self.magDiff.col2.split("_")[0] 

296 

297 # Return plain MagDiff with warning if either coeff not found 

298 for band in (col1Band, col2Band): 

299 if band not in self.extinctionCoeffs: 

300 _LOG.warning( 

301 "%s band not found in coefficients dictionary: %s" " Not applying extinction correction", 

302 band, 

303 self.extinctionCoeffs, 

304 ) 

305 return diff 

306 

307 av1: float = self.extinctionCoeffs[col1Band] 

308 av2: float = self.extinctionCoeffs[col2Band] 

309 

310 ebv = data[self.ebvCol] 

311 # Ignore type until a more complete Vector protocol 

312 correction = np.array((av1 - av2) * ebv) * u.mag # type: ignore 

313 

314 if self.magDiff.returnMillimags: 

315 correction = correction.to(u.mmag) 

316 

317 return np.array(diff - correction.value) 

318 

319 

320class RAcosDec(VectorAction): 

321 """Construct a vector of RA*cos(Dec) in order to have commensurate values 

322 between RA and Dec.""" 

323 

324 raKey = Field[str](doc="RA coordinate", default="coord_ra") 

325 decKey = Field[str](doc="Dec coordinate", default="coord_dec") 

326 

327 def getInputSchema(self) -> KeyedDataSchema: 

328 return ((self.decKey, Vector), (self.raKey, Vector)) 

329 

330 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

331 ra = np.array(data[self.raKey]) 

332 dec = np.array(data[self.decKey]) 

333 return ra * np.cos((dec * u.degree).to(u.radian).value) 

334 

335 

336class AngularSeparation(VectorAction): 

337 """Calculate the angular separation between two coordinate positions.""" 

338 

339 raKey_A = Field[str](doc="RA coordinate for position A", default="coord_ra") 

340 decKey_A = Field[str](doc="Dec coordinate for position A", default="coord_dec") 

341 raKey_B = Field[str](doc="RA coordinate for position B", default="coord_ra") 

342 decKey_B = Field[str](doc="Dec coordinate for position B", default="coord_dec") 

343 outputUnit = Field[str](doc="Output astropy unit", default="milliarcsecond") 

344 

345 def getInputSchema(self) -> KeyedDataSchema: 

346 return ( 

347 (self.decKey_A, Vector), 

348 (self.raKey_A, Vector), 

349 (self.decKey_B, Vector), 

350 (self.raKey_B, Vector), 

351 ) 

352 

353 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

354 ra_A = np.array(data[self.raKey_A]) 

355 dec_A = np.array(data[self.decKey_A]) 

356 ra_B = np.array(data[self.raKey_B]) 

357 dec_B = np.array(data[self.decKey_B]) 

358 coord_A = SkyCoord(ra_A * u.degree, dec_A * u.degree) 

359 coord_B = SkyCoord(ra_B * u.degree, dec_B * u.degree) 

360 return coord_A.separation(coord_B).to(u.Unit(self.outputUnit)).value 

361 

362 

363# Statistical vectorActions 

364 

365 

366class PerGroupStatistic(VectorAction): 

367 """Compute per-group statistic values and return result as a vector with 

368 one element per group. The computed statistic can be any function accepted 

369 by pandas DataFrameGroupBy.aggregate passed in as a string function name. 

370 """ 

371 

372 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index") 

373 buildAction = ConfigurableActionField[VectorAction](doc="Action to build vector", default=LoadVector) 

374 func = Field[str](doc="Name of function to be applied per group") 

375 

376 def getInputSchema(self) -> KeyedDataSchema: 

377 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),) 

378 

379 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

380 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": self.buildAction(data, **kwargs)}) 

381 result = df.groupby("groupKey")["value"].aggregate(self.func) 

382 return np.array(result) 

383 

384 

385class ResidualWithPerGroupStatistic(VectorAction): 

386 """Compute residual between individual elements of group and the per-group 

387 statistic.""" 

388 

389 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index") 

390 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector) 

391 func = Field[str](doc="Name of function to be applied per group", default="mean") 

392 

393 def getInputSchema(self) -> KeyedDataSchema: 

394 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),) 

395 

396 def __call__(self, data: KeyedData, **kwargs) -> Vector: 

397 values = self.buildAction(data, **kwargs) 

398 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": values}) 

399 result = df.groupby("groupKey")["value"].aggregate(self.func) 

400 

401 joinedDf = df.join(result, on="groupKey", validate="m:1", lsuffix="_individual", rsuffix="_group") 

402 

403 result = joinedDf["value_individual"] - joinedDf["value_group"] 

404 return np.array(result)