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, Optional, cast 

33 

34import numpy as np 

35from lsst.pex.config import Field 

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

37from lsst.pex.config.listField import ListField 

38 

39from ...interfaces import ( 

40 KeyedData, 

41 KeyedDataAction, 

42 KeyedDataSchema, 

43 Scalar, 

44 ScalarAction, 

45 Vector, 

46 VectorAction, 

47) 

48 

49 

50class ChainedKeyedDataActions(KeyedDataAction): 

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

52 one KeyedData result. 

53 """ 

54 

55 keyedDataActions = ConfigurableActionStructField[KeyedDataAction]( 

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

57 "object" 

58 ) 

59 

60 def getInputSchema(self) -> KeyedDataSchema: 

61 for action in self.keyedDataActions: 

62 yield from action.getInputSchema() 

63 

64 def getOutputSchema(self) -> KeyedDataSchema: 

65 for action in self.keyedDataActions: 

66 output = action.getOutputSchema() 

67 if output is not None: 

68 yield from output 

69 

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

71 result: KeyedData = {} # type:ignore 

72 for action in self.keyedDataActions: 

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

74 result[column] = values 

75 return result 

76 

77 

78class AddComputedVector(KeyedDataAction): 

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

80 copy of the KeyedData, returning the result. 

81 """ 

82 

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

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

85 

86 def getInputSchema(self) -> KeyedDataSchema: 

87 yield from self.action.getInputSchema() 

88 

89 def getOutputSchema(self) -> KeyedDataSchema: 

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

91 

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

93 result = dict(data) 

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

95 return result 

96 

97 

98class KeyedDataSelectorAction(KeyedDataAction): 

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

100 optionally apply selectors to down select extracted vectors. 

101 

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

103 expected schema. 

104 """ 

105 

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

107 

108 selectors = ConfigurableActionStructField[VectorAction]( 

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

110 ) 

111 

112 def getInputSchema(self) -> KeyedDataSchema: 

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

114 for action in self.selectors: 

115 yield from action.getInputSchema() 

116 

117 def getOutputSchema(self) -> KeyedDataSchema: 

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

119 

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

121 mask: Optional[np.ndarray] = None 

122 for selector in self.selectors: 

123 subMask = selector(data, **kwargs) 

124 if mask is None: 

125 mask = subMask 

126 else: 

127 mask *= subMask # type: ignore 

128 result = {} 

129 for key in self.vectorKeys: 

130 key_arg = key.format_map(kwargs) 

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

132 if mask is not None: 

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

134 else: 

135 return result 

136 

137 

138class KeyedScalars(KeyedDataAction): 

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

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

141 corresponding `ScalarAction`. 

142 """ 

143 

144 scalarActions = ConfigurableActionStructField[ScalarAction]( 

145 doc="Create a KeyedData of individual ScalarActions" 

146 ) 

147 

148 def getInputSchema(self) -> KeyedDataSchema: 

149 for action in self.scalarActions: 

150 yield from action.getInputSchema() 

151 

152 def getOutputSchema(self) -> KeyedDataSchema: 

153 for name in self.scalarActions.fieldNames: 

154 yield (name, Scalar) 

155 

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

157 result: KeyedData = {} # type: ignore 

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

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

160 return result 

161 

162 

163class KeyedDataKeyAccessAction(KeyedDataAction): 

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

165 

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

167 """ 

168 

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

170 

171 def getInputSchema(self) -> KeyedDataSchema: 

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

173 

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

175 return data[self.topLevelKey] 

176 

177 

178class KeyedDataUnitAccessAction(KeyedDataAction): 

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

180 

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

182 """ 

183 

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

185 

186 def getInputSchema(self) -> KeyedDataSchema: 

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

188 

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

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

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

192 else: 

193 return ""