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

143 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-05 14:05 +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 "LoadVector", 

25 "DownselectVector", 

26 "MultiCriteriaDownselectVector", 

27 "ConvertFluxToMag", 

28 "ConvertUnits", 

29 "CalcSn", 

30 "MagDiff", 

31 "ExtinctionCorrectedMagDiff", 

32 "PerGroupStatistic", 

33 "ResidualWithPerGroupStatistic", 

34 "RAcosDec", 

35) 

36 

37import logging 

38from typing import Optional, cast 

39 

40import numpy as np 

41import pandas as pd 

42from astropy import units as u 

43from lsst.pex.config import DictField, Field 

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

45 

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

47from ...math import divide, fluxToMag 

48from .selectors import VectorSelector 

49 

50_LOG = logging.getLogger(__name__) 

51 

52# Basic vectorActions 

53 

54 

55class LoadVector(VectorAction): 

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

57 

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

59 

60 def getInputSchema(self) -> KeyedDataSchema: 

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

62 

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

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

65 

66 

67class DownselectVector(VectorAction): 

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

69 shorter Vector. 

70 """ 

71 

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

73 

74 selector = ConfigurableActionField[VectorAction]( 

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

76 ) 

77 

78 def getInputSchema(self) -> KeyedDataSchema: 

79 yield (self.vectorKey, Vector) 

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

81 

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

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

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

85 

86 

87class MultiCriteriaDownselectVector(VectorAction): 

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

89 logic, and return the shorter Vector. 

90 """ 

91 

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

93 

94 selectors = ConfigurableActionStructField[VectorAction]( 

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

96 ) 

97 

98 def getInputSchema(self) -> KeyedDataSchema: 

99 yield (self.vectorKey, Vector) 

100 for action in self.selectors: 

101 yield from action.getInputSchema() 

102 

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

104 mask: Optional[Vector] = None 

105 for selector in self.selectors: 

106 subMask = selector(data, **kwargs) 

107 if mask is None: 

108 mask = subMask 

109 else: 

110 mask *= subMask # type: ignore 

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

112 

113 

114# Astronomical vectorActions 

115 

116 

117class CalcSn(VectorAction): 

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

119 

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

121 uncertaintySuffix = Field[str]( 

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

123 ) 

124 

125 def getInputSchema(self) -> KeyedDataSchema: 

126 yield self.fluxType, Vector 

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

128 

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

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

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

132 return divide(signal, noise) 

133 

134 

135class ConvertFluxToMag(VectorAction): 

136 """Turn nano janskies into magnitudes.""" 

137 

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

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

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

141 

142 def getInputSchema(self) -> KeyedDataSchema: 

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

144 

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

146 return fluxToMag( 

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

148 flux_unit=self.fluxUnit, 

149 return_millimags=self.returnMillimags, 

150 ) 

151 

152 

153class ConvertUnits(VectorAction): 

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

155 

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

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

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

159 

160 def getInputSchema(self) -> KeyedDataSchema: 

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

162 

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

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

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

166 

167 

168class MagDiff(VectorAction): 

169 """Calculate the difference between two magnitudes; 

170 each magnitude is derived from a flux column. 

171 Parameters 

172 ---------- 

173 TO DO: 

174 Returns 

175 ------- 

176 The magnitude difference in milli mags. 

177 Notes 

178 ----- 

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

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

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

182 information and assumes that the fluxes are already 

183 calibrated. 

184 """ 

185 

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

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

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

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

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

191 

192 def getInputSchema(self) -> KeyedDataSchema: 

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

194 

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

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

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

198 magDiff = mag1 - mag2 

199 if self.returnMillimags: 

200 magDiff *= 1000.0 

201 return magDiff 

202 

203 

204class ExtinctionCorrectedMagDiff(VectorAction): 

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

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

207 per the naming convention in the Object Table: 

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

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

210 config parameters. 

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

212 """ 

213 

214 magDiff = ConfigurableActionField[VectorAction]( 

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

216 ) 

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

218 band1 = Field[str]( 

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

220 optional=True, 

221 default=None, 

222 ) 

223 band2 = Field[str]( 

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

225 optional=True, 

226 default=None, 

227 ) 

228 extinctionCoeffs = DictField[str, float]( 

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

230 "Key must be the band", 

231 optional=True, 

232 default=None, 

233 ) 

234 

235 def getInputSchema(self) -> KeyedDataSchema: 

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

237 

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

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

240 if not self.extinctionCoeffs: 

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

242 return diff 

243 

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

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

246 

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

248 for band in (col1Band, col2Band): 

249 if band not in self.extinctionCoeffs: 

250 _LOG.warning( 

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

252 band, 

253 self.extinctionCoeffs, 

254 ) 

255 return diff 

256 

257 av1: float = self.extinctionCoeffs[col1Band] 

258 av2: float = self.extinctionCoeffs[col2Band] 

259 

260 ebv = data[self.ebvCol] 

261 # Ignore type until a more complete Vector protocol 

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

263 

264 if self.magDiff.returnMillimags: 

265 correction = correction.to(u.mmag) 

266 

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

268 

269 

270class RAcosDec(VectorAction): 

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

272 between RA and Dec.""" 

273 

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

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

276 

277 def getInputSchema(self) -> KeyedDataSchema: 

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

279 

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

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

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

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

284 

285 

286# Statistical vectorActions 

287 

288 

289class PerGroupStatistic(VectorAction): 

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

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

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

293 """ 

294 

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

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

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

298 

299 def getInputSchema(self) -> KeyedDataSchema: 

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

301 

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

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

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

305 return np.array(result) 

306 

307 

308class ResidualWithPerGroupStatistic(VectorAction): 

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

310 statistic.""" 

311 

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

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

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

315 

316 def getInputSchema(self) -> KeyedDataSchema: 

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

318 

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

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

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

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

323 

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

325 

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

327 return np.array(result)