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
« 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
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)
37import logging
38from typing import Optional, cast
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
46from ...interfaces import KeyedData, KeyedDataSchema, Vector, VectorAction
47from ...math import divide, fluxToMag
48from .selectors import VectorSelector
50_LOG = logging.getLogger(__name__)
52# Basic vectorActions
55class LoadVector(VectorAction):
56 """Load and return a Vector from KeyedData."""
58 vectorKey = Field[str](doc="Key of vector which should be loaded")
60 def getInputSchema(self) -> KeyedDataSchema:
61 return ((self.vectorKey, Vector),)
63 def __call__(self, data: KeyedData, **kwargs) -> Vector:
64 return np.array(cast(Vector, data[self.vectorKey.format(**kwargs)]))
67class DownselectVector(VectorAction):
68 """Get a vector from KeyedData, apply specified selector, return the
69 shorter Vector.
70 """
72 vectorKey = Field[str](doc="column key to load from KeyedData")
74 selector = ConfigurableActionField[VectorAction](
75 doc="Action which returns a selection mask", default=VectorSelector
76 )
78 def getInputSchema(self) -> KeyedDataSchema:
79 yield (self.vectorKey, Vector)
80 yield from cast(VectorAction, self.selector).getInputSchema()
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]
87class MultiCriteriaDownselectVector(VectorAction):
88 """Get a vector from KeyedData, apply specified set of selectors with AND
89 logic, and return the shorter Vector.
90 """
92 vectorKey = Field[str](doc="column key to load from KeyedData")
94 selectors = ConfigurableActionStructField[VectorAction](
95 doc="Selectors for selecting rows, will be AND together",
96 )
98 def getInputSchema(self) -> KeyedDataSchema:
99 yield (self.vectorKey, Vector)
100 for action in self.selectors:
101 yield from action.getInputSchema()
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]
114# Astronomical vectorActions
117class CalcSn(VectorAction):
118 """Calculate the signal-to-noise ratio from a single flux vector."""
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 )
125 def getInputSchema(self) -> KeyedDataSchema:
126 yield self.fluxType, Vector
127 yield f"{self.fluxType}{self.uncertaintySuffix}", Vector
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)
135class ConvertFluxToMag(VectorAction):
136 """Turn nano janskies into magnitudes."""
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)
142 def getInputSchema(self) -> KeyedDataSchema:
143 return ((self.vectorKey, Vector),)
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 )
153class ConvertUnits(VectorAction):
154 """Convert the units of a vector."""
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")
160 def getInputSchema(self) -> KeyedDataSchema:
161 return tuple(self.buildAction.getInputSchema())
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
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 """
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)
192 def getInputSchema(self) -> KeyedDataSchema:
193 return ((self.col1, Vector), (self.col2, Vector))
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
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 """
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 )
235 def getInputSchema(self) -> KeyedDataSchema:
236 return self.magDiff.getInputSchema() + ((self.ebvCol, Vector),)
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
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]
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
257 av1: float = self.extinctionCoeffs[col1Band]
258 av2: float = self.extinctionCoeffs[col2Band]
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
264 if self.magDiff.returnMillimags:
265 correction = correction.to(u.mmag)
267 return np.array(diff - correction.value)
270class RAcosDec(VectorAction):
271 """Construct a vector of RA*cos(Dec) in order to have commensurate values
272 between RA and Dec."""
274 raKey = Field[str](doc="RA coordinate", default="coord_ra")
275 decKey = Field[str](doc="Dec coordinate", default="coord_dec")
277 def getInputSchema(self) -> KeyedDataSchema:
278 return ((self.decKey, Vector), (self.raKey, Vector))
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)
286# Statistical vectorActions
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 """
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")
299 def getInputSchema(self) -> KeyedDataSchema:
300 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
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)
308class ResidualWithPerGroupStatistic(VectorAction):
309 """Compute residual between individual elements of group and the per-group
310 statistic."""
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")
316 def getInputSchema(self) -> KeyedDataSchema:
317 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
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)
324 joinedDf = df.join(result, on="groupKey", validate="m:1", lsuffix="_individual", rsuffix="_group")
326 result = joinedDf["value_individual"] - joinedDf["value_group"]
327 return np.array(result)