Coverage for python/lsst/analysis/tools/interfaces/_stages.py: 17%

119 statements  

« prev     ^ index     » next       coverage.py v7.2.4, created at 2023-04-30 03:04 -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 

22 

23__all__ = ("BasePrep", "BaseProcess", "BaseMetricAction", "BaseProduce") 

24 

25from collections import abc 

26from typing import Any, cast 

27 

28import astropy.units as apu 

29from lsst.pex.config import ListField 

30from lsst.pex.config.configurableActions import ConfigurableActionStructField 

31from lsst.pex.config.dictField import DictField 

32from lsst.verify import Measurement 

33 

34from ._actions import ( 

35 AnalysisAction, 

36 JointAction, 

37 KeyedDataAction, 

38 MetricAction, 

39 MetricResultType, 

40 NoPlot, 

41 VectorAction, 

42) 

43from ._interfaces import KeyedData, KeyedDataSchema, KeyedDataTypes, Scalar, Vector 

44 

45 

46class BasePrep(KeyedDataAction): 

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

48 

49 selectors = ConfigurableActionStructField[VectorAction]( 

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

51 ) 

52 

53 def getInputSchema(self) -> KeyedDataSchema: 

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

55 for action in self.selectors: 

56 yield from action.getInputSchema() 

57 

58 def getOutputSchema(self) -> KeyedDataSchema: 

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

60 

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

62 mask: Vector | None = None 

63 for selector in self.selectors: 

64 subMask = selector(data, **kwargs) 

65 if mask is None: 

66 mask = subMask 

67 else: 

68 mask *= subMask # type: ignore 

69 result: dict[str, Any] = {} 

70 for key in self.vectorKeys: 

71 formattedKey = key.format_map(kwargs) 

72 result[formattedKey] = cast(Vector, data[formattedKey]) 

73 if mask is not None: 

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

75 else: 

76 return result 

77 

78 def addInputSchema(self, inputSchema: KeyedDataSchema) -> None: 

79 self.vectorKeys = [name for name, _ in inputSchema] 

80 

81 

82class BaseProcess(KeyedDataAction): 

83 buildActions = ConfigurableActionStructField[VectorAction | KeyedDataAction]( 

84 doc="Actions which compute a Vector which will be added to results" 

85 ) 

86 filterActions = ConfigurableActionStructField[VectorAction | KeyedDataAction]( 

87 doc="Actions which filter one or more input or build Vectors into shorter vectors" 

88 ) 

89 calculateActions = ConfigurableActionStructField[AnalysisAction]( 

90 doc="Actions which compute quantities from the input or built data" 

91 ) 

92 

93 def getInputSchema(self) -> KeyedDataSchema: 

94 inputSchema: KeyedDataTypes = {} # type: ignore 

95 buildOutputSchema: KeyedDataTypes = {} # type: ignore 

96 filterOutputSchema: KeyedDataTypes = {} # type: ignore 

97 action: AnalysisAction 

98 

99 for fieldName, action in self.buildActions.items(): 

100 for name, typ in action.getInputSchema(): 

101 inputSchema[name] = typ 

102 if isinstance(action, KeyedDataAction): 

103 buildOutputSchema.update(action.getOutputSchema() or {}) 

104 else: 

105 buildOutputSchema[fieldName] = Vector 

106 

107 for fieldName, action in self.filterActions.items(): 

108 for name, typ in action.getInputSchema(): 

109 if name not in buildOutputSchema: 

110 inputSchema[name] = typ 

111 if isinstance(action, KeyedDataAction): 

112 filterOutputSchema.update(action.getOutputSchema() or {}) 

113 else: 

114 filterOutputSchema[fieldName] = Vector 

115 

116 for calcAction in self.calculateActions: 

117 for name, typ in calcAction.getInputSchema(): 

118 if name not in buildOutputSchema and name not in filterOutputSchema: 

119 inputSchema[name] = typ 

120 return ((name, typ) for name, typ in inputSchema.items()) 

121 

122 def getOutputSchema(self) -> KeyedDataSchema: 

123 for action in self.buildActions: 

124 if isinstance(action, KeyedDataAction): 

125 outSchema = action.getOutputSchema() 

126 if outSchema is not None: 

127 yield from outSchema 

128 

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

130 action: AnalysisAction 

131 results = {} 

132 data = dict(data) 

133 for name, action in self.buildActions.items(): 

134 match action(data, **kwargs): 

135 case abc.Mapping() as item: 

136 for key, result in item.items(): 

137 results[key] = result 

138 case item: 

139 results[name] = item 

140 view1 = data | results 

141 for name, action in self.filterActions.items(): 

142 match action(view1, **kwargs): 

143 case abc.Mapping() as item: 

144 for key, result in item.items(): 

145 results[key] = result 

146 case item: 

147 results[name] = item 

148 

149 view2 = data | results 

150 for name, calcAction in self.calculateActions.items(): 

151 match calcAction(view2, **kwargs): 

152 case abc.Mapping() as item: 

153 for key, result in item.items(): 

154 results[key] = result 

155 case item: 

156 results[name] = item 

157 return results 

158 

159 

160class BaseMetricAction(MetricAction): 

161 units = DictField[str, str](doc="Mapping of scalar key to astropy unit string", default={}) 

162 newNames = DictField[str, str]( 

163 doc="Mapping of key to new name if needed prior to creating metric", 

164 default={}, 

165 ) 

166 

167 def getInputSchema(self) -> KeyedDataSchema: 

168 # Something is wrong with the typing for DictField key iteration 

169 return [(key, Scalar) for key in self.units] # type: ignore 

170 

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

172 results = {} 

173 for key, unit in self.units.items(): 

174 formattedKey = key.format(**kwargs) 

175 if formattedKey not in data: 

176 raise ValueError(f"Key: {formattedKey} could not be found input data") 

177 value = data[formattedKey] 

178 if not isinstance(value, Scalar): 

179 raise ValueError(f"Data for key {key} is not a Scalar type") 

180 if newName := self.newNames.get(key): 

181 formattedKey = newName.format(**kwargs) 

182 notes = {"metric_tags": kwargs.get("metric_tags", [])} 

183 results[formattedKey] = Measurement(formattedKey, value * apu.Unit(unit), notes=notes) 

184 return results 

185 

186 

187class BaseProduce(JointAction): 

188 def setDefaults(self): 

189 super().setDefaults() 

190 self.metric = BaseMetricAction() 

191 self.plot = NoPlot