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

199 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-10 10:36 +0000

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

25 "MultiCriteriaDownselectVector", 

26 "MagColumnNanoJansky", 

27 "FractionalDifference", 

28 "Sn", 

29 "ConstantValue", 

30 "SubtractVector", 

31 "DivideVector", 

32 "LoadVector", 

33 "MagDiff", 

34 "SNCalculator", 

35 "ExtinctionCorrectedMagDiff", 

36 "PerGroupStatistic", 

37 "ResidualWithPerGroupStatistic", 

38 "RAcosDec", 

39 "ConvertUnits", 

40) 

41 

42import logging 

43from typing import Optional, cast 

44 

45import numpy as np 

46import pandas as pd 

47from astropy import units as u 

48from lsst.pex.config import DictField, Field 

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

50 

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

52from .selectors import VectorSelector 

53 

54_LOG = logging.getLogger(__name__) 

55 

56 

57class DownselectVector(VectorAction): 

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

59 shorter Vector. 

60 """ 

61 

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

63 

64 selector = ConfigurableActionField[VectorAction]( 

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

66 ) 

67 

68 def getInputSchema(self) -> KeyedDataSchema: 

69 yield (self.vectorKey, Vector) 

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

71 

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

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

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

75 

76 

77class MultiCriteriaDownselectVector(VectorAction): 

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

79 logic, and return the shorter Vector. 

80 """ 

81 

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

83 

84 selectors = ConfigurableActionStructField[VectorAction]( 

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

86 ) 

87 

88 def getInputSchema(self) -> KeyedDataSchema: 

89 yield (self.vectorKey, Vector) 

90 for action in self.selectors: 

91 yield from action.getInputSchema() 

92 

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

94 mask: Optional[Vector] = None 

95 for selector in self.selectors: 

96 subMask = selector(data, **kwargs) 

97 if mask is None: 

98 mask = subMask 

99 else: 

100 mask *= subMask # type: ignore 

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

102 

103 

104class MagColumnNanoJansky(VectorAction): 

105 """Turn nano janskies into magnitudes.""" 

106 

107 vectorKey = Field[str](doc="column key to use for this transformation") 

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

109 

110 def getInputSchema(self) -> KeyedDataSchema: 

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

112 

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

114 with np.warnings.catch_warnings(): # type: ignore 

115 np.warnings.filterwarnings("ignore", r"invalid value encountered") # type: ignore 

116 np.warnings.filterwarnings("ignore", r"divide by zero") # type: ignore 

117 vec = cast(Vector, data[self.vectorKey.format(**kwargs)]) 

118 mags = (np.array(vec) * u.nJy).to(u.ABmag).value # type: ignore 

119 if self.returnMillimags: 

120 mags *= 1000 

121 return mags 

122 

123 

124class FractionalDifference(VectorAction): 

125 """Calculate (A-B)/B.""" 

126 

127 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A") 

128 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B") 

129 

130 def getInputSchema(self) -> KeyedDataSchema: 

131 yield from self.actionA.getInputSchema() # type: ignore 

132 yield from self.actionB.getInputSchema() # type: ignore 

133 

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

135 vecA = self.actionA(data, **kwargs) # type: ignore 

136 vecB = self.actionB(data, **kwargs) # type: ignore 

137 return (vecA - vecB) / vecB 

138 

139 

140class Sn(VectorAction): 

141 """Compute signal-to-noise in the given flux type.""" 

142 

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

144 uncertaintySuffix = Field[str]( 

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

146 ) 

147 band = Field[str](doc="Band to calculate the S/N in.", default="i") 

148 

149 def getInputSchema(self) -> KeyedDataSchema: 

150 yield (fluxCol := self.fluxType), Vector 

151 yield f"{fluxCol}{self.uncertaintySuffix}", Vector 

152 

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

154 """Computes S/N in self.fluxType 

155 

156 Parameters 

157 ---------- 

158 df : `Tabular` 

159 

160 Returns 

161 ------- 

162 result : `Vector` 

163 Computed signal-to-noise ratio. 

164 """ 

165 fluxCol = self.fluxType.format(**(kwargs | dict(band=self.band))) 

166 errCol = f"{fluxCol}{self.uncertaintySuffix.format(**kwargs)}" 

167 result = cast(Vector, data[fluxCol]) / data[errCol] # type: ignore 

168 

169 return np.array(cast(Vector, result)) 

170 

171 

172class ConstantValue(VectorAction): 

173 """Return a constant scalar value.""" 

174 

175 value = Field[float](doc="A single constant value", optional=False) 

176 

177 def getInputSchema(self) -> KeyedDataSchema: 

178 return () 

179 

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

181 return np.array([self.value]) 

182 

183 

184class SubtractVector(VectorAction): 

185 """Calculate (A-B).""" 

186 

187 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A") 

188 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B") 

189 

190 def getInputSchema(self) -> KeyedDataSchema: 

191 yield from self.actionA.getInputSchema() # type: ignore 

192 yield from self.actionB.getInputSchema() # type: ignore 

193 

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

195 vecA = self.actionA(data, **kwargs) # type: ignore 

196 vecB = self.actionB(data, **kwargs) # type: ignore 

197 return vecA - vecB 

198 

199 

200class DivideVector(VectorAction): 

201 """Calculate (A/B)""" 

202 

203 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A") 

204 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B") 

205 

206 def getInputSchema(self) -> KeyedDataSchema: 

207 yield from self.actionA.getInputSchema() # type: ignore 

208 yield from self.actionB.getInputSchema() # type: ignore 

209 

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

211 vecA = self.actionA(data, **kwargs) # type: ignore 

212 vecB = self.actionB(data, **kwargs) # type: ignore 

213 return vecA / vecB 

214 

215 

216class LoadVector(VectorAction): 

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

218 

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

220 

221 def getInputSchema(self) -> KeyedDataSchema: 

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

223 

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

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

226 

227 

228class MagDiff(VectorAction): 

229 """Calculate the difference between two magnitudes; 

230 each magnitude is derived from a flux column. 

231 Parameters 

232 ---------- 

233 TO DO: 

234 Returns 

235 ------- 

236 The magnitude difference in milli mags. 

237 Notes 

238 ----- 

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

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

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

242 information and assumes that the fluxes are already 

243 calibrated. 

244 """ 

245 

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

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

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

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

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

251 

252 def getInputSchema(self) -> KeyedDataSchema: 

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

254 

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

256 flux1 = np.array(data[self.col1.format(**kwargs)]) * u.Unit(self.fluxUnits1) 

257 mag1 = flux1.to(u.ABmag) 

258 

259 flux2 = np.array(data[self.col2.format(**kwargs)]) * u.Unit(self.fluxUnits2) 

260 mag2 = flux2.to(u.ABmag) 

261 

262 magDiff = mag1 - mag2 

263 

264 if self.returnMillimags: 

265 magDiff = magDiff.to(u.mmag) 

266 

267 return np.array(magDiff.value) 

268 

269 

270class SNCalculator(VectorAction): 

271 """Calculate the signal-to-noise.""" 

272 

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

274 uncertaintySuffix = Field[str]( 

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

276 ) 

277 

278 def getInputSchema(self) -> KeyedDataSchema: 

279 yield self.fluxType, Vector 

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

281 

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

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

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

285 sn = signal / noise 

286 

287 return np.array(sn) 

288 

289 

290class ExtinctionCorrectedMagDiff(VectorAction): 

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

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

293 per the naming convention in the Object Table: 

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

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

296 config parameters. 

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

298 """ 

299 

300 magDiff = ConfigurableActionField[VectorAction]( 

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

302 ) 

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

304 band1 = Field[str]( 

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

306 optional=True, 

307 default=None, 

308 ) 

309 band2 = Field[str]( 

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

311 optional=True, 

312 default=None, 

313 ) 

314 extinctionCoeffs = DictField[str, float]( 

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

316 "Key must be the band", 

317 optional=True, 

318 default=None, 

319 ) 

320 

321 def getInputSchema(self) -> KeyedDataSchema: 

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

323 

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

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

326 if not self.extinctionCoeffs: 

327 _LOG.warning("No extinction Coefficients. Not applying extinction correction") 

328 return diff 

329 

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

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

332 

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

334 for band in (col1Band, col2Band): 

335 if band not in self.extinctionCoeffs: 

336 _LOG.warning( 

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

338 band, 

339 self.extinctionCoeffs, 

340 ) 

341 return diff 

342 

343 av1: float = self.extinctionCoeffs[col1Band] 

344 av2: float = self.extinctionCoeffs[col2Band] 

345 

346 ebv = data[self.ebvCol] 

347 # Ignore type until a more complete Vector protocol 

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

349 

350 if self.magDiff.returnMillimags: 

351 correction = correction.to(u.mmag) 

352 

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

354 

355 

356class PerGroupStatistic(VectorAction): 

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

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

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

360 """ 

361 

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

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

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

365 

366 def getInputSchema(self) -> KeyedDataSchema: 

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

368 

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

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

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

372 return np.array(result) 

373 

374 

375class ResidualWithPerGroupStatistic(VectorAction): 

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

377 statistic.""" 

378 

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

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

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

382 

383 def getInputSchema(self) -> KeyedDataSchema: 

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

385 

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

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

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

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

390 

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

392 

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

394 return np.array(result) 

395 

396 

397class RAcosDec(VectorAction): 

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

399 between RA and Dec.""" 

400 

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

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

403 

404 def getInputSchema(self) -> KeyedDataSchema: 

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

406 

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

408 ra = data[self.raKey] 

409 dec = data[self.decKey] 

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

411 

412 

413class ConvertUnits(VectorAction): 

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

415 

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

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

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

419 

420 def getInputSchema(self) -> KeyedDataSchema: 

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

422 

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

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

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