Coverage for python/lsst/analysis/tools/actions/vector/vectorActions.py: 49%
160 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-23 04:51 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-23 04:51 -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 "LoadVector",
25 "DownselectVector",
26 "MultiCriteriaDownselectVector",
27 "ConvertFluxToMag",
28 "ConvertUnits",
29 "CalcSn",
30 "MagDiff",
31 "ExtinctionCorrectedMagDiff",
32 "PerGroupStatistic",
33 "ResidualWithPerGroupStatistic",
34 "RAcosDec",
35 "AngularSeparation",
36)
38import logging
39from typing import Optional, cast
41import numpy as np
42import pandas as pd
43from astropy import units as u
44from astropy.coordinates import SkyCoord
45from lsst.pex.config import DictField, Field
46from lsst.pex.config.configurableActions import ConfigurableActionField, ConfigurableActionStructField
48from ...interfaces import KeyedData, KeyedDataSchema, Vector, VectorAction
49from ...math import divide, fluxToMag
50from .selectors import VectorSelector
52_LOG = logging.getLogger(__name__)
54# Basic vectorActions
57class LoadVector(VectorAction):
58 """Load and return a Vector from KeyedData."""
60 vectorKey = Field[str](doc="Key of vector which should be loaded")
62 def getInputSchema(self) -> KeyedDataSchema:
63 return ((self.vectorKey, Vector),)
65 def __call__(self, data: KeyedData, **kwargs) -> Vector:
66 return np.array(cast(Vector, data[self.vectorKey.format(**kwargs)]))
69class DownselectVector(VectorAction):
70 """Get a vector from KeyedData, apply specified selector, return the
71 shorter Vector.
72 """
74 vectorKey = Field[str](doc="column key to load from KeyedData")
76 selector = ConfigurableActionField[VectorAction](
77 doc="Action which returns a selection mask", default=VectorSelector
78 )
80 def getInputSchema(self) -> KeyedDataSchema:
81 yield (self.vectorKey, Vector)
82 yield from cast(VectorAction, self.selector).getInputSchema()
84 def __call__(self, data: KeyedData, **kwargs) -> Vector:
85 mask = cast(VectorAction, self.selector)(data, **kwargs)
86 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask]
89class MultiCriteriaDownselectVector(VectorAction):
90 """Get a vector from KeyedData, apply specified set of selectors with AND
91 logic, and return the shorter Vector.
92 """
94 vectorKey = Field[str](doc="column key to load from KeyedData")
96 selectors = ConfigurableActionStructField[VectorAction](
97 doc="Selectors for selecting rows, will be AND together",
98 )
100 def getInputSchema(self) -> KeyedDataSchema:
101 yield (self.vectorKey, Vector)
102 for action in self.selectors:
103 yield from action.getInputSchema()
105 def __call__(self, data: KeyedData, **kwargs) -> Vector:
106 mask: Optional[Vector] = None
107 for selector in self.selectors:
108 subMask = selector(data, **kwargs)
109 if mask is None:
110 mask = subMask
111 else:
112 mask *= subMask # type: ignore
113 return cast(Vector, data[self.vectorKey.format(**kwargs)])[mask]
116# Astronomical vectorActions
119class CalcSn(VectorAction):
120 """Calculate the signal-to-noise ratio from a single flux vector."""
122 fluxType = Field[str](doc="Flux type (vector key) to calculate the S/N.", default="{band}_psfFlux")
123 uncertaintySuffix = Field[str](
124 doc="Suffix to add to fluxType to specify the uncertainty column", default="Err"
125 )
127 def getInputSchema(self) -> KeyedDataSchema:
128 yield self.fluxType, Vector
129 yield f"{self.fluxType}{self.uncertaintySuffix}", Vector
131 def __call__(self, data: KeyedData, **kwargs) -> Vector:
132 signal = np.array(data[self.fluxType.format(**kwargs)])
133 noise = np.array(data[f"{self.fluxType}{self.uncertaintySuffix}".format(**kwargs)])
134 return divide(signal, noise)
137class ConvertFluxToMag(VectorAction):
138 """Turn nano janskies into magnitudes."""
140 vectorKey = Field[str](doc="Key of flux vector to convert to mags")
141 fluxUnit = Field[str](doc="Astropy unit of flux vector", default="nJy")
142 returnMillimags = Field[bool](doc="Use millimags or not?", default=False)
144 def getInputSchema(self) -> KeyedDataSchema:
145 return ((self.vectorKey, Vector),)
147 def __call__(self, data: KeyedData, **kwargs) -> Vector:
148 return fluxToMag(
149 cast(Vector, data[self.vectorKey.format(**kwargs)]),
150 flux_unit=self.fluxUnit,
151 return_millimags=self.returnMillimags,
152 )
155class ConvertUnits(VectorAction):
156 """Convert the units of a vector."""
158 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector)
159 inUnit = Field[str](doc="input Astropy unit")
160 outUnit = Field[str](doc="output Astropy unit")
162 def getInputSchema(self) -> KeyedDataSchema:
163 return tuple(self.buildAction.getInputSchema())
165 def __call__(self, data: KeyedData, **kwargs) -> Vector:
166 dataWithUnit = self.buildAction(data, **kwargs) * u.Unit(self.inUnit)
167 return dataWithUnit.to(self.outUnit).value
170class MagDiff(VectorAction):
171 """Calculate the difference between two magnitudes;
172 each magnitude is derived from a flux column.
173 Parameters
174 ----------
175 TO DO:
176 Returns
177 -------
178 The magnitude difference in milli mags.
179 Notes
180 -----
181 The flux columns need to be in units (specifiable in
182 the fluxUnits1 and 2 config options) that can be converted
183 to janskies. This action doesn't have any calibration
184 information and assumes that the fluxes are already
185 calibrated.
186 """
188 col1 = Field[str](doc="Column to subtract from")
189 fluxUnits1 = Field[str](doc="Units for col1", default="nanojansky")
190 col2 = Field[str](doc="Column to subtract")
191 fluxUnits2 = Field[str](doc="Units for col2", default="nanojansky")
192 returnMillimags = Field[bool](doc="Use millimags or not?", default=True)
194 def getInputSchema(self) -> KeyedDataSchema:
195 return ((self.col1, Vector), (self.col2, Vector))
197 def __call__(self, data: KeyedData, **kwargs) -> Vector:
198 mag1 = fluxToMag(data[self.col1.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits1))
199 mag2 = fluxToMag(data[self.col2.format(**kwargs)], flux_unit=u.Unit(self.fluxUnits2))
200 magDiff = mag1 - mag2
201 if self.returnMillimags:
202 magDiff *= 1000.0
203 return magDiff
206class ExtinctionCorrectedMagDiff(VectorAction):
207 """Compute the difference between two magnitudes and correct for extinction
208 By default bands are derived from the <band>_ prefix on flux columns,
209 per the naming convention in the Object Table:
210 e.g. the band of 'g_psfFlux' is 'g'. If column names follow another
211 convention, bands can alternatively be supplied via the band1 or band2
212 config parameters.
213 If band1 and band2 are supplied, the flux column names are ignored.
214 """
216 magDiff = ConfigurableActionField[VectorAction](
217 doc="Action that returns a difference in magnitudes", default=MagDiff
218 )
219 ebvCol = Field[str](doc="E(B-V) Column Name", default="ebv")
220 band1 = Field[str](
221 doc="Optional band for magDiff.col1. Supercedes column name prefix",
222 optional=True,
223 default=None,
224 )
225 band2 = Field[str](
226 doc="Optional band for magDiff.col2. Supercedes column name prefix",
227 optional=True,
228 default=None,
229 )
230 extinctionCoeffs = DictField[str, float](
231 doc="Dictionary of extinction coefficients for conversion from E(B-V) to extinction, A_band."
232 "Key must be the band",
233 optional=True,
234 default=None,
235 )
237 def getInputSchema(self) -> KeyedDataSchema:
238 return self.magDiff.getInputSchema() + ((self.ebvCol, Vector),)
240 def __call__(self, data: KeyedData, **kwargs) -> Vector:
241 diff = self.magDiff(data, **kwargs)
242 if not self.extinctionCoeffs:
243 _LOG.debug("No extinction Coefficients. Not applying extinction correction")
244 return diff
246 col1Band = self.band1 if self.band1 else self.magDiff.col1.split("_")[0]
247 col2Band = self.band2 if self.band2 else self.magDiff.col2.split("_")[0]
249 # Return plain MagDiff with warning if either coeff not found
250 for band in (col1Band, col2Band):
251 if band not in self.extinctionCoeffs:
252 _LOG.warning(
253 "%s band not found in coefficients dictionary: %s" " Not applying extinction correction",
254 band,
255 self.extinctionCoeffs,
256 )
257 return diff
259 av1: float = self.extinctionCoeffs[col1Band]
260 av2: float = self.extinctionCoeffs[col2Band]
262 ebv = data[self.ebvCol]
263 # Ignore type until a more complete Vector protocol
264 correction = np.array((av1 - av2) * ebv) * u.mag # type: ignore
266 if self.magDiff.returnMillimags:
267 correction = correction.to(u.mmag)
269 return np.array(diff - correction.value)
272class RAcosDec(VectorAction):
273 """Construct a vector of RA*cos(Dec) in order to have commensurate values
274 between RA and Dec."""
276 raKey = Field[str](doc="RA coordinate", default="coord_ra")
277 decKey = Field[str](doc="Dec coordinate", default="coord_dec")
279 def getInputSchema(self) -> KeyedDataSchema:
280 return ((self.decKey, Vector), (self.raKey, Vector))
282 def __call__(self, data: KeyedData, **kwargs) -> Vector:
283 ra = np.array(data[self.raKey])
284 dec = np.array(data[self.decKey])
285 return ra * np.cos((dec * u.degree).to(u.radian).value)
288class AngularSeparation(VectorAction):
289 """Calculate the angular separation between two coordinate positions."""
291 raKey_A = Field[str](doc="RA coordinate for position A", default="coord_ra")
292 decKey_A = Field[str](doc="Dec coordinate for position A", default="coord_dec")
293 raKey_B = Field[str](doc="RA coordinate for position B", default="coord_ra")
294 decKey_B = Field[str](doc="Dec coordinate for position B", default="coord_dec")
295 outputUnit = Field[str](doc="Output astropy unit", default="milliarcsecond")
297 def getInputSchema(self) -> KeyedDataSchema:
298 return (
299 (self.decKey_A, Vector),
300 (self.raKey_A, Vector),
301 (self.decKey_B, Vector),
302 (self.raKey_B, Vector),
303 )
305 def __call__(self, data: KeyedData, **kwargs) -> Vector:
306 ra_A = np.array(data[self.raKey_A])
307 dec_A = np.array(data[self.decKey_A])
308 ra_B = np.array(data[self.raKey_B])
309 dec_B = np.array(data[self.decKey_B])
310 coord_A = SkyCoord(ra_A * u.degree, dec_A * u.degree)
311 coord_B = SkyCoord(ra_B * u.degree, dec_B * u.degree)
312 return coord_A.separation(coord_B).to(u.Unit(self.outputUnit)).value
315# Statistical vectorActions
318class PerGroupStatistic(VectorAction):
319 """Compute per-group statistic values and return result as a vector with
320 one element per group. The computed statistic can be any function accepted
321 by pandas DataFrameGroupBy.aggregate passed in as a string function name.
322 """
324 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index")
325 buildAction = ConfigurableActionField[VectorAction](doc="Action to build vector", default=LoadVector)
326 func = Field[str](doc="Name of function to be applied per group")
328 def getInputSchema(self) -> KeyedDataSchema:
329 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
331 def __call__(self, data: KeyedData, **kwargs) -> Vector:
332 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": self.buildAction(data, **kwargs)})
333 result = df.groupby("groupKey")["value"].aggregate(self.func)
334 return np.array(result)
337class ResidualWithPerGroupStatistic(VectorAction):
338 """Compute residual between individual elements of group and the per-group
339 statistic."""
341 groupKey = Field[str](doc="Column key to use for forming groups", default="obj_index")
342 buildAction = ConfigurableActionField(doc="Action to build vector", default=LoadVector)
343 func = Field[str](doc="Name of function to be applied per group", default="mean")
345 def getInputSchema(self) -> KeyedDataSchema:
346 return tuple(self.buildAction.getInputSchema()) + ((self.groupKey, Vector),)
348 def __call__(self, data: KeyedData, **kwargs) -> Vector:
349 values = self.buildAction(data, **kwargs)
350 df = pd.DataFrame({"groupKey": data[self.groupKey], "value": values})
351 result = df.groupby("groupKey")["value"].aggregate(self.func)
353 joinedDf = df.join(result, on="groupKey", validate="m:1", lsuffix="_individual", rsuffix="_group")
355 result = joinedDf["value_individual"] - joinedDf["value_group"]
356 return np.array(result)