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

183 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-14 04:05 -0800

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 

23import logging 

24from typing import Optional, cast 

25 

26import numpy as np 

27import pandas as pd 

28from astropy import units as u 

29from lsst.pex.config import DictField, Field 

30from lsst.pipe.tasks.configurableActions import ConfigurableActionField, ConfigurableActionStructField 

31 

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

33from .selectors import VectorSelector 

34 

35_LOG = logging.getLogger(__name__) 

36 

37 

38class DownselectVector(VectorAction): 

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

40 shorter Vector. 

41 """ 

42 

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

44 

45 selector = ConfigurableActionField(doc="Action which returns a selection mask", default=VectorSelector) 

46 

47 def getInputSchema(self) -> KeyedDataSchema: 

48 yield (self.vectorKey, Vector) 

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

50 

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

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

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

54 

55 

56class MultiCriteriaDownselectVector(VectorAction): 

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

58 logic, and return the shorter Vector. 

59 """ 

60 

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

62 

63 selectors = ConfigurableActionStructField[VectorAction]( 

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

65 ) 

66 

67 def getInputSchema(self) -> KeyedDataSchema: 

68 yield (self.vectorKey, Vector) 

69 for action in self.selectors: 

70 yield from cast(VectorAction, action).getInputSchema() 

71 

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

73 mask: Optional[Vector] = None 

74 for selector in self.selectors: 

75 subMask = selector(data, **kwargs) 

76 if mask is None: 

77 mask = subMask 

78 else: 

79 mask *= subMask # type: ignore 

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

81 

82 

83class MagColumnNanoJansky(VectorAction): 

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

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

86 

87 def getInputSchema(self) -> KeyedDataSchema: 

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

89 

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

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

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

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

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

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

96 if self.returnMillimags: 

97 mags *= 1000 

98 return mags 

99 

100 

101class FractionalDifference(VectorAction): 

102 """Calculate (A-B)/B""" 

103 

104 actionA = ConfigurableActionField(doc="Action which supplies vector A", dtype=VectorAction) 

105 actionB = ConfigurableActionField(doc="Action which supplies vector B", dtype=VectorAction) 

106 

107 def getInputSchema(self) -> KeyedDataSchema: 

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

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

110 

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

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

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

114 return (vecA - vecB) / vecB 

115 

116 

117class Sn(VectorAction): 

118 """Compute signal-to-noise in the given flux type""" 

119 

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

121 uncertaintySuffix = Field[str]( 

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

123 ) 

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

125 

126 def getInputSchema(self) -> KeyedDataSchema: 

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

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

129 

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

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

132 Parameters 

133 ---------- 

134 df : `Tabular` 

135 Returns 

136 ------- 

137 result : `Vector` 

138 Computed signal-to-noise ratio. 

139 """ 

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

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

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

143 

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

145 

146 

147class ConstantValue(VectorAction): 

148 """Return a constant scalar value""" 

149 

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

151 

152 def getInputSchema(self) -> KeyedDataSchema: 

153 return () 

154 

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

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

157 

158 

159class SubtractVector(VectorAction): 

160 """Calculate (A-B)""" 

161 

162 actionA = ConfigurableActionField(doc="Action which supplies vector A", dtype=VectorAction) 

163 actionB = ConfigurableActionField(doc="Action which supplies vector B", dtype=VectorAction) 

164 

165 def getInputSchema(self) -> KeyedDataSchema: 

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

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

168 

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

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

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

172 return vecA - vecB 

173 

174 

175class DivideVector(VectorAction): 

176 """Calculate (A/B)""" 

177 

178 actionA = ConfigurableActionField(doc="Action which supplies vector A", dtype=VectorAction) 

179 actionB = ConfigurableActionField(doc="Action which supplies vector B", dtype=VectorAction) 

180 

181 def getInputSchema(self) -> KeyedDataSchema: 

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

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

184 

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

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

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

188 return vecA / vecB 

189 

190 

191class LoadVector(VectorAction): 

192 """Load and return a Vector from KeyedData""" 

193 

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

195 

196 def getInputSchema(self) -> KeyedDataSchema: 

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

198 

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

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

201 

202 

203class MagDiff(VectorAction): 

204 """Calculate the difference between two magnitudes; 

205 each magnitude is derived from a flux column. 

206 Parameters 

207 ---------- 

208 TO DO: 

209 Returns 

210 ------- 

211 The magnitude difference in milli mags. 

212 Notes 

213 ----- 

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

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

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

217 information and assumes that the fluxes are already 

218 calibrated. 

219 """ 

220 

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

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

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

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

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

226 

227 def getInputSchema(self) -> KeyedDataSchema: 

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

229 

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

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

232 mag1 = flux1.to(u.ABmag) 

233 

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

235 mag2 = flux2.to(u.ABmag) 

236 

237 magDiff = mag1 - mag2 

238 

239 if self.returnMillimags: 

240 magDiff = magDiff.to(u.mmag) 

241 

242 return np.array(magDiff.value) 

243 

244 

245class SNCalculator(VectorAction): 

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

247 

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

249 uncertaintySuffix = Field[str]( 

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

251 ) 

252 

253 def getInputSchema(self) -> KeyedDataSchema: 

254 yield self.fluxType, Vector 

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

256 

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

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

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

260 sn = signal / noise 

261 

262 return np.array(sn) 

263 

264 

265class ExtinctionCorrectedMagDiff(VectorAction): 

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

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

268 per the naming convention in the Object Table: 

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

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

271 config parameters. 

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

273 """ 

274 

275 magDiff = ConfigurableActionField( 

276 doc="Action that returns a difference in magnitudes", default=MagDiff, dtype=VectorAction 

277 ) 

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

279 band1 = Field[str]( 

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

281 optional=True, 

282 default=None, 

283 ) 

284 band2 = Field[str]( 

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

286 optional=True, 

287 default=None, 

288 ) 

289 extinctionCoeffs = DictField[str, float]( 

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

291 "Key must be the band", 

292 optional=True, 

293 default=None, 

294 ) 

295 

296 def getInputSchema(self) -> KeyedDataSchema: 

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

298 

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

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

301 if not self.extinctionCoeffs: 

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

303 return diff 

304 

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

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

307 

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

309 for band in (col1Band, col2Band): 

310 if band not in self.extinctionCoeffs: 

311 _LOG.warning( 

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

313 band, 

314 self.extinctionCoeffs, 

315 ) 

316 return diff 

317 

318 av1: float = self.extinctionCoeffs[col1Band] 

319 av2: float = self.extinctionCoeffs[col2Band] 

320 

321 ebv = data[self.ebvCol] 

322 # Ignore type until a more complete Vector protocol 

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

324 

325 if self.magDiff.returnMillimags: 

326 correction = correction.to(u.mmag) 

327 

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

329 

330 

331class AstromDiff(VectorAction): 

332 """Calculate the difference between two columns, assuming their units 

333 are degrees, and convert the difference to arcseconds. 

334 Parameters 

335 ---------- 

336 df : `pandas.core.frame.DataFrame` 

337 The catalog to calculate the position difference from. 

338 Returns 

339 ------- 

340 angleDiffValue : `np.ndarray` 

341 The difference between two columns, either in the input units or in 

342 milliarcseconds. 

343 Notes 

344 ----- 

345 The columns need to be in units (specifiable in the radecUnits1 and 2 

346 config options) that can be converted to arcseconds. This action doesn't 

347 have any calibration information and assumes that the positions are already 

348 calibrated. 

349 """ 

350 

351 col1 = Field[str](doc="Column to subtract from", dtype=str) 

352 radecUnits1 = Field[str](doc="Units for col1", dtype=str, default="degree") 

353 col2 = Field[str](doc="Column to subtract", dtype=str) 

354 radecUnits2 = Field[str](doc="Units for col2", dtype=str, default="degree") 

355 returnMilliArcsecs = Field[bool](doc="Use marcseconds or not?", dtype=bool, default=True) 

356 

357 def getInputSchema(self) -> KeyedDataSchema: 

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

359 

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

361 angle1 = np.array(data[self.col1.format(**kwargs)]) * u.Unit(self.radecUnits1) 

362 

363 angle2 = np.array(data[self.col2.format(**kwargs)]) * u.Unit(self.radecUnits2) 

364 

365 angleDiff = angle1 - angle2 

366 

367 if self.returnMilliArcsecs: 

368 angleDiffValue = angleDiff.to(u.arcsec).value * 1000 

369 else: 

370 angleDiffValue = angleDiff.value 

371 return angleDiffValue 

372 

373 

374class PerGroupStatistic(VectorAction): 

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

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

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

378 """ 

379 

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

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

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

383 

384 def getInputSchema(self) -> KeyedDataSchema: 

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

386 

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

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

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

390 return np.array(result)