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

160 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-04 04:15 -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 "MagDiff", 

31 "ExtinctionCorrectedMagDiff", 

32 "PerGroupStatistic", 

33 "ResidualWithPerGroupStatistic", 

34 "RAcosDec", 

35 "AngularSeparation", 

36) 

37 

38import logging 

39from typing import Optional, cast 

40 

41import numpy as np 

42import pandas as pd 

43from astropy import units as u 

44from astropy.coordinates import SkyCoord 

45from lsst.pex.config import DictField, Field 

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

47 

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

49from ...math import divide, fluxToMag 

50from .selectors import VectorSelector 

51 

52_LOG = logging.getLogger(__name__) 

53 

54# Basic vectorActions 

55 

56 

57class LoadVector(VectorAction): 

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

59 

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

61 

62 def getInputSchema(self) -> KeyedDataSchema: 

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

64 

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

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

67 

68 

69class DownselectVector(VectorAction): 

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

71 shorter Vector. 

72 """ 

73 

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

75 

76 selector = ConfigurableActionField[VectorAction]( 

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

78 ) 

79 

80 def getInputSchema(self) -> KeyedDataSchema: 

81 yield (self.vectorKey, Vector) 

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

83 

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

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

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

87 

88 

89class MultiCriteriaDownselectVector(VectorAction): 

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

91 logic, and return the shorter Vector. 

92 """ 

93 

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

95 

96 selectors = ConfigurableActionStructField[VectorAction]( 

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

98 ) 

99 

100 def getInputSchema(self) -> KeyedDataSchema: 

101 yield (self.vectorKey, Vector) 

102 for action in self.selectors: 

103 yield from action.getInputSchema() 

104 

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

106 mask: Optional[Vector] = None 

107 for selector in self.selectors: 

108 subMask = selector(data, **kwargs) 

109 if mask is None: 

110 mask = subMask 

111 else: 

112 mask *= subMask # type: ignore 

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

114 

115 

116# Astronomical vectorActions 

117 

118 

119class CalcSn(VectorAction): 

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

121 

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

123 uncertaintySuffix = Field[str]( 

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

125 ) 

126 

127 def getInputSchema(self) -> KeyedDataSchema: 

128 yield self.fluxType, Vector 

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

130 

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

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

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

134 return divide(signal, noise) 

135 

136 

137class ConvertFluxToMag(VectorAction): 

138 """Turn nano janskies into magnitudes.""" 

139 

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

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

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

143 

144 def getInputSchema(self) -> KeyedDataSchema: 

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

146 

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

148 return fluxToMag( 

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

150 flux_unit=self.fluxUnit, 

151 return_millimags=self.returnMillimags, 

152 ) 

153 

154 

155class ConvertUnits(VectorAction): 

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

157 

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

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

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

161 

162 def getInputSchema(self) -> KeyedDataSchema: 

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

164 

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

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

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

168 

169 

170class MagDiff(VectorAction): 

171 """Calculate the difference between two magnitudes; 

172 each magnitude is derived from a flux column. 

173 Parameters 

174 ---------- 

175 TO DO: 

176 Returns 

177 ------- 

178 The magnitude difference in milli mags. 

179 Notes 

180 ----- 

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

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

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

184 information and assumes that the fluxes are already 

185 calibrated. 

186 """ 

187 

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

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

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

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

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

193 

194 def getInputSchema(self) -> KeyedDataSchema: 

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

196 

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

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

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

200 magDiff = mag1 - mag2 

201 if self.returnMillimags: 

202 magDiff *= 1000.0 

203 return magDiff 

204 

205 

206class ExtinctionCorrectedMagDiff(VectorAction): 

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

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

209 per the naming convention in the Object Table: 

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

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

212 config parameters. 

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

214 """ 

215 

216 magDiff = ConfigurableActionField[VectorAction]( 

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

218 ) 

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

220 band1 = Field[str]( 

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

222 optional=True, 

223 default=None, 

224 ) 

225 band2 = Field[str]( 

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

227 optional=True, 

228 default=None, 

229 ) 

230 extinctionCoeffs = DictField[str, float]( 

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

232 "Key must be the band", 

233 optional=True, 

234 default=None, 

235 ) 

236 

237 def getInputSchema(self) -> KeyedDataSchema: 

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

239 

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

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

242 if not self.extinctionCoeffs: 

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

244 return diff 

245 

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

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

248 

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

250 for band in (col1Band, col2Band): 

251 if band not in self.extinctionCoeffs: 

252 _LOG.warning( 

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

254 band, 

255 self.extinctionCoeffs, 

256 ) 

257 return diff 

258 

259 av1: float = self.extinctionCoeffs[col1Band] 

260 av2: float = self.extinctionCoeffs[col2Band] 

261 

262 ebv = data[self.ebvCol] 

263 # Ignore type until a more complete Vector protocol 

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

265 

266 if self.magDiff.returnMillimags: 

267 correction = correction.to(u.mmag) 

268 

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

270 

271 

272class RAcosDec(VectorAction): 

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

274 between RA and Dec.""" 

275 

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

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

278 

279 def getInputSchema(self) -> KeyedDataSchema: 

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

281 

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

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

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

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

286 

287 

288class AngularSeparation(VectorAction): 

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

290 

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

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

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

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

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

296 

297 def getInputSchema(self) -> KeyedDataSchema: 

298 return ( 

299 (self.decKey_A, Vector), 

300 (self.raKey_A, Vector), 

301 (self.decKey_B, Vector), 

302 (self.raKey_B, Vector), 

303 ) 

304 

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

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

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

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

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

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

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

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

313 

314 

315# Statistical vectorActions 

316 

317 

318class PerGroupStatistic(VectorAction): 

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

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

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

322 """ 

323 

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

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

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

327 

328 def getInputSchema(self) -> KeyedDataSchema: 

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

330 

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

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

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

334 return np.array(result) 

335 

336 

337class ResidualWithPerGroupStatistic(VectorAction): 

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

339 statistic.""" 

340 

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

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

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

344 

345 def getInputSchema(self) -> KeyedDataSchema: 

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

347 

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

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

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

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

352 

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

354 

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

356 return np.array(result)