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-12 03:08 -0700
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-12 03:08 -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
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)
42import logging
43from typing import Optional, cast
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
51from ...interfaces import KeyedData, KeyedDataSchema, Vector, VectorAction
52from .selectors import VectorSelector
54_LOG = logging.getLogger(__name__)
57class DownselectVector(VectorAction):
58 """Get a vector from KeyedData, apply specified selector, return the
59 shorter Vector.
60 """
62 vectorKey = Field[str](doc="column key to load from KeyedData")
64 selector = ConfigurableActionField[VectorAction](
65 doc="Action which returns a selection mask", default=VectorSelector
66 )
68 def getInputSchema(self) -> KeyedDataSchema:
69 yield (self.vectorKey, Vector)
70 yield from cast(VectorAction, self.selector).getInputSchema()
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]
77class MultiCriteriaDownselectVector(VectorAction):
78 """Get a vector from KeyedData, apply specified set of selectors with AND
79 logic, and return the shorter Vector.
80 """
82 vectorKey = Field[str](doc="column key to load from KeyedData")
84 selectors = ConfigurableActionStructField[VectorAction](
85 doc="Selectors for selecting rows, will be AND together",
86 )
88 def getInputSchema(self) -> KeyedDataSchema:
89 yield (self.vectorKey, Vector)
90 for action in self.selectors:
91 yield from action.getInputSchema()
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]
104class MagColumnNanoJansky(VectorAction):
105 """Turn nano janskies into magnitudes."""
107 vectorKey = Field[str](doc="column key to use for this transformation")
108 returnMillimags = Field[bool](doc="Use millimags or not?", default=False)
110 def getInputSchema(self) -> KeyedDataSchema:
111 return ((self.vectorKey, Vector),)
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
124class FractionalDifference(VectorAction):
125 """Calculate (A-B)/B."""
127 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A")
128 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B")
130 def getInputSchema(self) -> KeyedDataSchema:
131 yield from self.actionA.getInputSchema() # type: ignore
132 yield from self.actionB.getInputSchema() # type: ignore
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
140class Sn(VectorAction):
141 """Compute signal-to-noise in the given flux type."""
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")
149 def getInputSchema(self) -> KeyedDataSchema:
150 yield (fluxCol := self.fluxType), Vector
151 yield f"{fluxCol}{self.uncertaintySuffix}", Vector
153 def __call__(self, data: KeyedData, **kwargs) -> Vector:
154 """Computes S/N in self.fluxType
156 Parameters
157 ----------
158 df : `Tabular`
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
169 return np.array(cast(Vector, result))
172class ConstantValue(VectorAction):
173 """Return a constant scalar value."""
175 value = Field[float](doc="A single constant value", optional=False)
177 def getInputSchema(self) -> KeyedDataSchema:
178 return ()
180 def __call__(self, data: KeyedData, **kwargs) -> Vector:
181 return np.array([self.value])
184class SubtractVector(VectorAction):
185 """Calculate (A-B)."""
187 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A")
188 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B")
190 def getInputSchema(self) -> KeyedDataSchema:
191 yield from self.actionA.getInputSchema() # type: ignore
192 yield from self.actionB.getInputSchema() # type: ignore
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
200class DivideVector(VectorAction):
201 """Calculate (A/B)"""
203 actionA = ConfigurableActionField[VectorAction](doc="Action which supplies vector A")
204 actionB = ConfigurableActionField[VectorAction](doc="Action which supplies vector B")
206 def getInputSchema(self) -> KeyedDataSchema:
207 yield from self.actionA.getInputSchema() # type: ignore
208 yield from self.actionB.getInputSchema() # type: ignore
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
216class LoadVector(VectorAction):
217 """Load and return a Vector from KeyedData."""
219 vectorKey = Field[str](doc="Key of vector which should be loaded")
221 def getInputSchema(self) -> KeyedDataSchema:
222 return ((self.vectorKey, Vector),)
224 def __call__(self, data: KeyedData, **kwargs) -> Vector:
225 return np.array(cast(Vector, data[self.vectorKey.format(**kwargs)]))
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 """
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)
252 def getInputSchema(self) -> KeyedDataSchema:
253 return ((self.col1, Vector), (self.col2, Vector))
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)
259 flux2 = np.array(data[self.col2.format(**kwargs)]) * u.Unit(self.fluxUnits2)
260 mag2 = flux2.to(u.ABmag)
262 magDiff = mag1 - mag2
264 if self.returnMillimags:
265 magDiff = magDiff.to(u.mmag)
267 return np.array(magDiff.value)
270class SNCalculator(VectorAction):
271 """Calculate the signal-to-noise."""
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 )
278 def getInputSchema(self) -> KeyedDataSchema:
279 yield self.fluxType, Vector
280 yield f"{self.fluxType}{self.uncertaintySuffix}", Vector
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
287 return np.array(sn)
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 """
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 )
321 def getInputSchema(self) -> KeyedDataSchema:
322 return self.magDiff.getInputSchema() + ((self.ebvCol, Vector),)
324 def __call__(self, data: KeyedData, **kwargs) -> Vector:
325 diff = self.magDiff(data, **kwargs)
326 if not self.extinctionCoeffs:
327 _LOG.debug("No extinction Coefficients. Not applying extinction correction")
328 return diff
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]
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
343 av1: float = self.extinctionCoeffs[col1Band]
344 av2: float = self.extinctionCoeffs[col2Band]
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
350 if self.magDiff.returnMillimags:
351 correction = correction.to(u.mmag)
353 return np.array(diff - correction.value)
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 """
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")
366 def getInputSchema(self) -> KeyedDataSchema:
367 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
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)
375class ResidualWithPerGroupStatistic(VectorAction):
376 """Compute residual between individual elements of group and the per-group
377 statistic."""
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")
383 def getInputSchema(self) -> KeyedDataSchema:
384 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
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)
391 joinedDf = df.join(result, on="groupKey", validate="m:1", lsuffix="_individual", rsuffix="_group")
393 result = joinedDf["value_individual"] - joinedDf["value_group"]
394 return np.array(result)
397class RAcosDec(VectorAction):
398 """Construct a vector of RA*cos(Dec) in order to have commensurate values
399 between RA and Dec."""
401 raKey = Field[str](doc="RA coordinate", default="coord_ra")
402 decKey = Field[str](doc="Dec coordinate", default="coord_dec")
404 def getInputSchema(self) -> KeyedDataSchema:
405 return ((self.decKey, Vector), (self.raKey, Vector))
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)
413class ConvertUnits(VectorAction):
414 """Convert the units of a vector."""
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")
420 def getInputSchema(self) -> KeyedDataSchema:
421 return tuple(self.buildAction.getInputSchema())
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