Coverage for python / lsst / analysis / tools / actions / keyedData / keyedDataActions.py: 34%

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 

22 

23__all__ = ( 

24 "ChainedKeyedDataActions", 

25 "AddComputedVector", 

26 "KeyedDataSelectorAction", 

27 "KeyedScalars", 

28 "KeyedDataKeyAccessAction", 

29 "KeyedDataUnitAccessAction", 

30) 

31 

32from typing import Any, cast 

33 

34import numpy as np 

35 

36from lsst.pex.config import Field 

37from lsst.pex.config.configurableActions import ConfigurableActionField, ConfigurableActionStructField 

38from lsst.pex.config.listField import ListField 

39 

40from ...interfaces import ( 

41 KeyedData, 

42 KeyedDataAction, 

43 KeyedDataSchema, 

44 Scalar, 

45 ScalarAction, 

46 Vector, 

47 VectorAction, 

48) 

49 

50 

51class ChainedKeyedDataActions(KeyedDataAction): 

52 r"""Run a series of `KeyedDataAction`\ s and accumulated their output into 

53 one KeyedData result. 

54 """ 

55 

56 keyedDataActions = ConfigurableActionStructField[KeyedDataAction]( 

57 doc="Set of KeyedData actions to run, results will be concatenated into a final output KeyedData" 

58 "object" 

59 ) 

60 

61 def getInputSchema(self) -> KeyedDataSchema: 

62 for action in self.keyedDataActions: 

63 yield from action.getInputSchema() 

64 

65 def getOutputSchema(self) -> KeyedDataSchema: 

66 for action in self.keyedDataActions: 

67 output = action.getOutputSchema() 

68 if output is not None: 

69 yield from output 

70 

71 def __call__(self, data: KeyedData, **kwargs) -> KeyedData: 

72 result: KeyedData = {} # type:ignore 

73 for action in self.keyedDataActions: 

74 for column, values in action(data, **kwargs).items(): 

75 result[column] = values 

76 return result 

77 

78 

79class AddComputedVector(KeyedDataAction): 

80 """Compute a `Vector` from the specified `VectorAction` and add it to a 

81 copy of the KeyedData, returning the result. 

82 """ 

83 

84 action = ConfigurableActionField[VectorAction](doc="Action to use to compute Vector") 

85 keyName = Field[str](doc="Key name to add to KeyedData") 

86 

87 def getInputSchema(self) -> KeyedDataSchema: 

88 yield from self.action.getInputSchema() 

89 

90 def getOutputSchema(self) -> KeyedDataSchema: 

91 return ((self.keyName, Vector),) 

92 

93 def __call__(self, data: KeyedData, **kwargs) -> KeyedData: 

94 result = dict(data) 

95 result[self.keyName.format(**kwargs)] = self.action(data, **kwargs) 

96 return result 

97 

98 

99class KeyedDataSelectorAction(KeyedDataAction): 

100 """Extract Vector specified by ``vectorKeys`` from input KeyedData and 

101 optionally apply selectors to down select extracted vectors. 

102 

103 Note this action will not work with keyed scalars, see `getInputSchema` for 

104 expected schema. 

105 """ 

106 

107 vectorKeys = ListField[str](doc="Keys to extract from KeyedData and return", default=[]) 

108 

109 selectors = ConfigurableActionStructField[VectorAction]( 

110 doc="Selectors for selecting rows, will be AND together", 

111 ) 

112 

113 def getInputSchema(self) -> KeyedDataSchema: 

114 yield from ((column, Vector | Scalar) for column in self.vectorKeys) # type: ignore 

115 for action in self.selectors: 

116 yield from action.getInputSchema() 

117 

118 def getOutputSchema(self) -> KeyedDataSchema: 

119 return ((column, Vector | Scalar) for column in self.vectorKeys) # type: ignore 

120 

121 def __call__(self, data: KeyedData, **kwargs) -> KeyedData: 

122 mask: np.ndarray | None = None 

123 for selector in self.selectors: 

124 subMask = selector(data, **kwargs) 

125 if mask is None: 

126 mask = subMask 

127 else: 

128 mask *= subMask # type: ignore 

129 result = {} 

130 for key in self.vectorKeys: 

131 key_arg = key.format_map(kwargs) 

132 result[f"{self.identity}_{key_arg}"] = data[key_arg] 

133 if mask is not None: 

134 return {key: cast(Vector, col)[mask] for key, col in result.items()} 

135 else: 

136 return result 

137 

138 

139class KeyedScalars(KeyedDataAction): 

140 """Creates an output of type KeyedData, where the keys are given by the 

141 identifiers in `scalarActions` and the values are the results of the 

142 corresponding `ScalarAction`. 

143 """ 

144 

145 scalarActions = ConfigurableActionStructField[ScalarAction]( 

146 doc="Create a KeyedData of individual ScalarActions" 

147 ) 

148 

149 def getInputSchema(self) -> KeyedDataSchema: 

150 for action in self.scalarActions: 

151 yield from action.getInputSchema() 

152 

153 def getOutputSchema(self) -> KeyedDataSchema: 

154 for name in self.scalarActions.fieldNames: 

155 yield (name, Scalar) 

156 

157 def __call__(self, data: KeyedData, **kwargs) -> KeyedData: 

158 result: KeyedData = {} # type: ignore 

159 for name, action in self.scalarActions.items(): 

160 result[name] = action(data, **kwargs) 

161 return result 

162 

163 

164class KeyedDataKeyAccessAction(KeyedDataAction): 

165 """Fetches the value for a given top-level key in a KeyedData. 

166 

167 Note that the returned value may optionally be a KeyedData itself. 

168 """ 

169 

170 topLevelKey = Field[str](doc="The top-level key in the KeyedData structure to extract data from.") 

171 

172 def getInputSchema(self) -> KeyedDataSchema: 

173 yield from [(self.topLevelKey, Any)] 

174 

175 def __call__(self, data: KeyedData, **kwargs) -> Any: 

176 return data[self.topLevelKey] 

177 

178 

179class KeyedDataUnitAccessAction(KeyedDataAction): 

180 """Fetches the unit for a given key in a KeyedData as a string. 

181 

182 If the requested key doesn't have a unit, an empty string is returned. 

183 """ 

184 

185 key = Field[str](doc="The key in the KeyedData structure to extract the unit from.") 

186 

187 def getInputSchema(self) -> KeyedDataSchema: 

188 yield from [(self.key, Any)] 

189 

190 def __call__(self, data: KeyedData, **kwargs) -> str: 

191 if data[self.key.format_map(kwargs)].unit is not None: 

192 return data[self.key.format_map(kwargs)].unit.to_string() 

193 else: 

194 return ""