Coverage for python/lsst/pipe/tasks/dataFrameActions/_actions.py: 57%

81 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-01 21:11 +0000

1from __future__ import annotations 

2 

3__all__ = ("SingleColumnAction", "MultiColumnAction", "CoordColumn", "MagColumnDN", "SumColumns", "AddColumn", 

4 "DivideColumns", "SubtractColumns", "MultiplyColumns", "MagColumnNanoJansky",) 

5 

6from typing import Iterable 

7 

8import numpy as np 

9import pandas as pd 

10from astropy import units 

11 

12from ..configurableActions import ConfigurableActionStructField, ConfigurableActionField 

13from ._baseDataFrameActions import DataFrameAction 

14from ._evalColumnExpression import makeColumnExpressionAction 

15 

16from lsst.pex.config import Field 

17 

18 

19class SingleColumnAction(DataFrameAction): 

20 column = Field(doc="Column to load for this action", dtype=str, optional=False) 

21 

22 @property 

23 def columns(self) -> Iterable[str]: 

24 return (self.column, ) 

25 

26 def __call__(self, df, **kwargs): 

27 return df[self.column] 

28 

29 

30class MultiColumnAction(DataFrameAction): 

31 actions = ConfigurableActionStructField(doc="Configurable actions to use in a joint action") 

32 

33 @property 

34 def columns(self) -> Iterable[str]: 

35 yield from (column for action in self.actions for column in action.columns) 

36 

37 

38class CoordColumn(SingleColumnAction): 

39 inRadians = Field(doc="Return the column in radians if true", default=True, dtype=bool) 

40 

41 def __call__(self, df): 

42 col = super().__call__(df) 

43 return col * 180 / np.pi if self.inRadians else col 

44 

45 

46class MagColumnDN(SingleColumnAction): 

47 coadd_zeropoint = Field(doc="Magnitude zero point", dtype=float, default=27) 

48 

49 def __call__(self, df: pd.DataFrame, **kwargs): 

50 if not (fluxMag0 := kwargs.get('fluxMag0')): 

51 fluxMag0 = 1/np.power(10, -0.4*self.coadd_zeropoint) 

52 

53 with np.warnings.catch_warnings(): 

54 np.warnings.filterwarnings('ignore', r'invalid value encountered') 

55 np.warnings.filterwarnings('ignore', r'divide by zero') 

56 return -2.5 * np.log10(df[self.column] / fluxMag0) 

57 

58 

59class MagColumnNanoJansky(SingleColumnAction): 

60 

61 def __call__(self, df: pd.DataFrame, **kwargs): 

62 

63 with np.warnings.catch_warnings(): 

64 np.warnings.filterwarnings('ignore', r'invalid value encountered') 

65 np.warnings.filterwarnings('ignore', r'divide by zero') 

66 return -2.5 * np.log10((df[self.column] * 1e-9) / 3631.0) 

67 

68 

69class NanoJansky(SingleColumnAction): 

70 ab_flux_scale = Field(doc="Scaling of ab flux", dtype=float, default=(0*units.ABmag).to_value(units.nJy)) 

71 coadd_zeropoint = Field(doc="Magnitude zero point", dtype=float, default=27) 

72 

73 def __call__(self, df, **kwargs): 

74 dataNumber = super().__call__(df, **kwargs) 

75 if not (fluxMag0 := kwargs.get('fluxMag0')): 

76 fluxMag0 = 1/np.power(10, -0.4*self.coadd_zeropoint) 

77 return self.ab_flux_scale * dataNumber / fluxMag0 

78 

79 def setDefaults(self): 

80 super().setDefaults() 

81 self.cache = True # cache this action for future calls 

82 

83 

84class NanoJanskyErr(SingleColumnAction): 

85 flux_mag_err = Field(doc="Error in the magnitude zeropoint", dtype=float, default=0) 

86 flux_action = ConfigurableActionField(doc="Action to use if flux is not provided to the call method", 

87 default=NanoJansky, dtype=DataFrameAction) 

88 

89 @property 

90 def columns(self): 

91 yield from zip((self.column,), self.flux_action.columns) 

92 

93 def __call__(self, df, flux_column=None, flux_mag_err=None, **kwargs): 

94 if flux_column is None: 

95 flux_column = self.flux_action(df, **kwargs) 

96 if flux_mag_err is None: 

97 flux_mag_err = self.flux_mag_err 

98 

99 

100_docs = """This is a `DataFrameAction` that is designed to add two columns 

101together and return the result. 

102""" 

103SumColumns = makeColumnExpressionAction("SumColumns", "colA+colB", 

104 exprDefaults={"colA": SingleColumnAction, 

105 "colB": SingleColumnAction}, 

106 docstring=_docs) 

107 

108_docs = """This is a `MultiColumnAction` that is designed to subtract two columns 

109together and return the result. 

110""" 

111SubtractColumns = makeColumnExpressionAction("SubtractColumns", "colA-colB", 

112 exprDefaults={"colA": SingleColumnAction, 

113 "colB": SingleColumnAction}, 

114 docstring=_docs) 

115 

116_docs = """This is a `MultiColumnAction` that is designed to multiply two columns 

117together and return the result. 

118""" 

119MultiplyColumns = makeColumnExpressionAction("MultiplyColumns", "colA*colB", 

120 exprDefaults={"colA": SingleColumnAction, 

121 "colB": SingleColumnAction}, 

122 docstring=_docs) 

123 

124_docs = """This is a `MultiColumnAction` that is designed to multiply two columns 

125together and return the result. 

126""" 

127DivideColumns = makeColumnExpressionAction("DivideColumns", "colA/colB", 

128 exprDefaults={"colA": SingleColumnAction, 

129 "colB": SingleColumnAction}, 

130 docstring=_docs) 

131 

132 

133class AddColumn(DataFrameAction): 

134 aggregator = ConfigurableActionField(doc="This is an instance of a Dataframe action that will be used " 

135 "to create a new column", dtype=DataFrameAction) 

136 newColumn = Field(doc="Name of the new column to add", dtype=str) 

137 

138 @property 

139 def columns(self) -> Iterable[str]: 

140 yield from self.aggregator.columns 

141 

142 def __call__(self, df, **kwargs) -> pd.DataFrame: 

143 # do your calculation and and 

144 df[self.newColumn] = self.aggregator(df, kwargs) 

145 return df