Coverage for python/lsst/analysis/tools/actions/vector/vectorActions.py: 48%
190 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-04 11:06 +0000
« prev ^ index » next coverage.py v7.5.0, created at 2024-05-04 11:06 +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 "ColorDiff",
31 "ColorError",
32 "MagDiff",
33 "ExtinctionCorrectedMagDiff",
34 "PerGroupStatistic",
35 "ResidualWithPerGroupStatistic",
36 "RAcosDec",
37 "AngularSeparation",
38)
40import logging
41from typing import Optional, cast
43import numpy as np
44import pandas as pd
45from astropy import units as u
46from astropy.coordinates import SkyCoord
47from lsst.pex.config import DictField, Field
48from lsst.pex.config.configurableActions import ConfigurableActionField, ConfigurableActionStructField
50from ...interfaces import KeyedData, KeyedDataSchema, Vector, VectorAction
51from ...math import divide, fluxToMag, log10
52from .selectors import VectorSelector
54_LOG = logging.getLogger(__name__)
56# Basic vectorActions
59class LoadVector(VectorAction):
60 """Load and return a Vector from KeyedData."""
62 vectorKey = Field[str](doc="Key of vector which should be loaded")
64 def getInputSchema(self) -> KeyedDataSchema:
65 return ((self.vectorKey, Vector),)
67 def __call__(self, data: KeyedData, **kwargs) -> Vector:
68 return np.array(cast(Vector, data[self.vectorKey.format(**kwargs)]))
71class DownselectVector(VectorAction):
72 """Get a vector from KeyedData, apply specified selector, return the
73 shorter Vector.
74 """
76 vectorKey = Field[str](doc="column key to load from KeyedData")
78 selector = ConfigurableActionField[VectorAction](
79 doc="Action which returns a selection mask", default=VectorSelector
80 )
82 def getInputSchema(self) -> KeyedDataSchema:
83 yield (self.vectorKey, Vector)
84 yield from cast(VectorAction, self.selector).getInputSchema()
86 def __call__(self, data: KeyedData, **kwargs) -> Vector:
87 mask = cast(VectorAction, self.selector)(data, **kwargs)
88 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask]
91class MultiCriteriaDownselectVector(VectorAction):
92 """Get a vector from KeyedData, apply specified set of selectors with AND
93 logic, and return the shorter Vector.
94 """
96 vectorKey = Field[str](doc="column key to load from KeyedData")
98 selectors = ConfigurableActionStructField[VectorAction](
99 doc="Selectors for selecting rows, will be AND together",
100 )
102 def getInputSchema(self) -> KeyedDataSchema:
103 yield (self.vectorKey, Vector)
104 for action in self.selectors:
105 yield from action.getInputSchema()
107 def __call__(self, data: KeyedData, **kwargs) -> Vector:
108 mask: Optional[Vector] = None
109 for selector in self.selectors:
110 subMask = selector(data, **kwargs)
111 if mask is None:
112 mask = subMask
113 else:
114 mask *= subMask # type: ignore
115 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask]
118# Astronomical vectorActions
121class CalcSn(VectorAction):
122 """Calculate the signal-to-noise ratio from a single flux vector."""
124 fluxType = Field[str](doc="Flux type (vector key) to calculate the S/N.", default="{band}_psfFlux")
125 uncertaintySuffix = Field[str](
126 doc="Suffix to add to fluxType to specify the uncertainty column", default="Err"
127 )
129 def getInputSchema(self) -> KeyedDataSchema:
130 yield self.fluxType, Vector
131 yield f"{self.fluxType}{self.uncertaintySuffix}", Vector
133 def __call__(self, data: KeyedData, **kwargs) -> Vector:
134 signal = np.array(data[self.fluxType.format(**kwargs)])
135 noise = np.array(data[f"{self.fluxType}{self.uncertaintySuffix}".format(**kwargs)])
136 return divide(signal, noise)
139class ColorDiff(VectorAction):
140 """Calculate the difference between two colors from flux actions."""
142 color1_flux1 = ConfigurableActionField[VectorAction](doc="Action providing first color's first flux")
143 color1_flux2 = ConfigurableActionField[VectorAction](doc="Action providing first color's second flux")
144 color2_flux1 = ConfigurableActionField[VectorAction](doc="Action providing second color's first flux")
145 color2_flux2 = ConfigurableActionField[VectorAction](doc="Action providing second color's second flux")
146 returnMillimags = Field[bool](doc="Whether to return color_diff in millimags (mags if not)", default=True)
148 def getInputSchema(self) -> KeyedDataSchema:
149 yield from self.color1_flux1.getInputSchema()
150 yield from self.color1_flux2.getInputSchema()
151 yield from self.color2_flux1.getInputSchema()
152 yield from self.color2_flux2.getInputSchema()
154 def __call__(self, data: KeyedData, **kwargs) -> Vector:
155 color_diff = -2.5 * log10(
156 divide(
157 self.color1_flux1(data, **kwargs) * self.color2_flux2(data, **kwargs),
158 self.color1_flux2(data, **kwargs) * self.color2_flux1(data, **kwargs),
159 )
160 )
162 if self.returnMillimags:
163 color_diff *= 1000
165 return color_diff
168class ColorError(VectorAction):
169 """Calculate the error in a color from two different flux error columns."""
171 flux_err1 = ConfigurableActionField[VectorAction](doc="Action providing error for first flux")
172 flux_err2 = ConfigurableActionField[VectorAction](doc="Action providing error for second flux")
173 returnMillimags = Field[bool](doc="Whether to return color_err in millimags (mags if not)", default=True)
175 def getInputSchema(self) -> KeyedDataSchema:
176 yield from self.flux_err1.getInputSchema()
177 yield from self.flux_err2.getInputSchema()
179 def __call__(self, data: KeyedData, **kwargs) -> Vector:
180 flux_err1 = self.flux_err1(data, **kwargs)
181 flux_err2 = self.flux_err2(data, **kwargs)
182 color_err = (2.5 / np.log(10)) * np.hypot(flux_err1, flux_err2)
184 if self.returnMillimags:
185 color_err *= 1000
187 return color_err
190class ConvertFluxToMag(VectorAction):
191 """Turn nano janskies into magnitudes."""
193 vectorKey = Field[str](doc="Key of flux vector to convert to mags")
194 fluxUnit = Field[str](doc="Astropy unit of flux vector", default="nJy")
195 returnMillimags = Field[bool](doc="Use millimags or not?", default=False)
197 def getInputSchema(self) -> KeyedDataSchema:
198 return ((self.vectorKey, Vector),)
200 def __call__(self, data: KeyedData, **kwargs) -> Vector:
201 return fluxToMag(
202 cast(Vector, data[self.vectorKey.format(**kwargs)]),
203 flux_unit=self.fluxUnit,
204 return_millimags=self.returnMillimags,
205 )
208class ConvertUnits(VectorAction):
209 """Convert the units of a vector."""
211 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector)
212 inUnit = Field[str](doc="input Astropy unit")
213 outUnit = Field[str](doc="output Astropy unit")
215 def getInputSchema(self) -> KeyedDataSchema:
216 return tuple(self.buildAction.getInputSchema())
218 def __call__(self, data: KeyedData, **kwargs) -> Vector:
219 dataWithUnit = self.buildAction(data, **kwargs) * u.Unit(self.inUnit)
220 return dataWithUnit.to(self.outUnit).value
223class MagDiff(VectorAction):
224 """Calculate the difference between two magnitudes;
225 each magnitude is derived from a flux column.
227 Notes
228 -----
229 The flux columns need to be in units (specifiable in
230 the fluxUnits1 and 2 config options) that can be converted
231 to janskies. This action doesn't have any calibration
232 information and assumes that the fluxes are already
233 calibrated.
234 """
236 col1 = Field[str](doc="Column to subtract from")
237 fluxUnits1 = Field[str](doc="Units for col1", default="nanojansky")
238 col2 = Field[str](doc="Column to subtract")
239 fluxUnits2 = Field[str](doc="Units for col2", default="nanojansky")
240 returnMillimags = Field[bool](doc="Use millimags or not?", default=True)
242 def getInputSchema(self) -> KeyedDataSchema:
243 return ((self.col1, Vector), (self.col2, Vector))
245 def __call__(self, data: KeyedData, **kwargs) -> Vector:
246 mag1 = fluxToMag(data[self.col1.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits1))
247 mag2 = fluxToMag(data[self.col2.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits2))
248 magDiff = mag1 - mag2
249 if self.returnMillimags:
250 magDiff *= 1000.0
251 return magDiff
254class ExtinctionCorrectedMagDiff(VectorAction):
255 """Compute the difference between two magnitudes and correct for extinction
256 By default bands are derived from the <band>_ prefix on flux columns,
257 per the naming convention in the Object Table:
258 e.g. the band of 'g_psfFlux' is 'g'. If column names follow another
259 convention, bands can alternatively be supplied via the band1 or band2
260 config parameters.
261 If band1 and band2 are supplied, the flux column names are ignored.
262 """
264 magDiff = ConfigurableActionField[VectorAction](
265 doc="Action that returns a difference in magnitudes", default=MagDiff
266 )
267 ebvCol = Field[str](doc="E(B-V) Column Name", default="ebv")
268 band1 = Field[str](
269 doc="Optional band for magDiff.col1. Supercedes column name prefix",
270 optional=True,
271 default=None,
272 )
273 band2 = Field[str](
274 doc="Optional band for magDiff.col2. Supercedes column name prefix",
275 optional=True,
276 default=None,
277 )
278 extinctionCoeffs = DictField[str, float](
279 doc="Dictionary of extinction coefficients for conversion from E(B-V) to extinction, A_band."
280 "Key must be the band",
281 optional=True,
282 default=None,
283 )
285 def getInputSchema(self) -> KeyedDataSchema:
286 return self.magDiff.getInputSchema() + ((self.ebvCol, Vector),)
288 def __call__(self, data: KeyedData, **kwargs) -> Vector:
289 diff = self.magDiff(data, **kwargs)
290 if not self.extinctionCoeffs:
291 _LOG.debug("No extinction Coefficients. Not applying extinction correction")
292 return diff
294 col1Band = self.band1 if self.band1 else self.magDiff.col1.split("_")[0]
295 col2Band = self.band2 if self.band2 else self.magDiff.col2.split("_")[0]
297 # Return plain MagDiff with warning if either coeff not found
298 for band in (col1Band, col2Band):
299 if band not in self.extinctionCoeffs:
300 _LOG.warning(
301 "%s band not found in coefficients dictionary: %s" " Not applying extinction correction",
302 band,
303 self.extinctionCoeffs,
304 )
305 return diff
307 av1: float = self.extinctionCoeffs[col1Band]
308 av2: float = self.extinctionCoeffs[col2Band]
310 ebv = data[self.ebvCol]
311 # Ignore type until a more complete Vector protocol
312 correction = np.array((av1 - av2) * ebv) * u.mag # type: ignore
314 if self.magDiff.returnMillimags:
315 correction = correction.to(u.mmag)
317 return np.array(diff - correction.value)
320class RAcosDec(VectorAction):
321 """Construct a vector of RA*cos(Dec) in order to have commensurate values
322 between RA and Dec."""
324 raKey = Field[str](doc="RA coordinate", default="coord_ra")
325 decKey = Field[str](doc="Dec coordinate", default="coord_dec")
327 def getInputSchema(self) -> KeyedDataSchema:
328 return ((self.decKey, Vector), (self.raKey, Vector))
330 def __call__(self, data: KeyedData, **kwargs) -> Vector:
331 ra = np.array(data[self.raKey])
332 dec = np.array(data[self.decKey])
333 return ra * np.cos((dec * u.degree).to(u.radian).value)
336class AngularSeparation(VectorAction):
337 """Calculate the angular separation between two coordinate positions."""
339 raKey_A = Field[str](doc="RA coordinate for position A", default="coord_ra")
340 decKey_A = Field[str](doc="Dec coordinate for position A", default="coord_dec")
341 raKey_B = Field[str](doc="RA coordinate for position B", default="coord_ra")
342 decKey_B = Field[str](doc="Dec coordinate for position B", default="coord_dec")
343 outputUnit = Field[str](doc="Output astropy unit", default="milliarcsecond")
345 def getInputSchema(self) -> KeyedDataSchema:
346 return (
347 (self.decKey_A, Vector),
348 (self.raKey_A, Vector),
349 (self.decKey_B, Vector),
350 (self.raKey_B, Vector),
351 )
353 def __call__(self, data: KeyedData, **kwargs) -> Vector:
354 ra_A = np.array(data[self.raKey_A])
355 dec_A = np.array(data[self.decKey_A])
356 ra_B = np.array(data[self.raKey_B])
357 dec_B = np.array(data[self.decKey_B])
358 coord_A = SkyCoord(ra_A * u.degree, dec_A * u.degree)
359 coord_B = SkyCoord(ra_B * u.degree, dec_B * u.degree)
360 return coord_A.separation(coord_B).to(u.Unit(self.outputUnit)).value
363# Statistical vectorActions
366class PerGroupStatistic(VectorAction):
367 """Compute per-group statistic values and return result as a vector with
368 one element per group. The computed statistic can be any function accepted
369 by pandas DataFrameGroupBy.aggregate passed in as a string function name.
370 """
372 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index")
373 buildAction = ConfigurableActionField[VectorAction](doc="Action to build vector", default=LoadVector)
374 func = Field[str](doc="Name of function to be applied per group")
376 def getInputSchema(self) -> KeyedDataSchema:
377 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
379 def __call__(self, data: KeyedData, **kwargs) -> Vector:
380 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": self.buildAction(data, **kwargs)})
381 result = df.groupby("groupKey")["value"].aggregate(self.func)
382 return np.array(result)
385class ResidualWithPerGroupStatistic(VectorAction):
386 """Compute residual between individual elements of group and the per-group
387 statistic."""
389 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index")
390 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector)
391 func = Field[str](doc="Name of function to be applied per group", default="mean")
393 def getInputSchema(self) -> KeyedDataSchema:
394 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
396 def __call__(self, data: KeyedData, **kwargs) -> Vector:
397 values = self.buildAction(data, **kwargs)
398 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": values})
399 result = df.groupby("groupKey")["value"].aggregate(self.func)
401 joinedDf = df.join(result, on="groupKey", validate="m:1", lsuffix="_individual", rsuffix="_group")
403 result = joinedDf["value_individual"] - joinedDf["value_group"]
404 return np.array(result)