Coverage for python/lsst/analysis/tools/actions/keyedData/keyedDataActions.py: 29%
71 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-07 11:44 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-07 11:44 +0000
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 "ChainedKeyedDataActions",
25 "AddComputedVector",
26 "KeyedDataSelectorAction",
27 "KeyedScalars",
28)
30from typing import Optional, cast
32import numpy as np
33from lsst.pex.config import Field
34from lsst.pex.config.configurableActions import ConfigurableActionField, ConfigurableActionStructField
35from lsst.pex.config.listField import ListField
37from ...interfaces import (
38 KeyedData,
39 KeyedDataAction,
40 KeyedDataSchema,
41 Scalar,
42 ScalarAction,
43 Vector,
44 VectorAction,
45)
48class ChainedKeyedDataActions(KeyedDataAction):
49 r"""Run a series of `KeyedDataAction`\ s and accumulated their output into
50 one KeyedData result.
51 """
53 keyedDataActions = ConfigurableActionStructField[KeyedDataAction](
54 doc="Set of KeyedData actions to run, results will be concatenated into a final output KeyedData"
55 "object"
56 )
58 def getInputSchema(self) -> KeyedDataSchema:
59 for action in self.keyedDataActions:
60 yield from action.getInputSchema()
62 def getOutputSchema(self) -> KeyedDataSchema:
63 for action in self.keyedDataActions:
64 output = action.getOutputSchema()
65 if output is not None:
66 yield from output
68 def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
69 result: KeyedData = {} # type:ignore
70 for action in self.keyedDataActions:
71 for column, values in action(data, **kwargs).items():
72 result[column] = values
73 return result
76class AddComputedVector(KeyedDataAction):
77 """Compute a `Vector` from the specified `VectorAction` and add it to a
78 copy of the KeyedData, returning the result.
79 """
81 action = ConfigurableActionField[VectorAction](doc="Action to use to compute Vector")
82 keyName = Field[str](doc="Key name to add to KeyedData")
84 def getInputSchema(self) -> KeyedDataSchema:
85 yield from self.action.getInputSchema()
87 def getOutputSchema(self) -> KeyedDataSchema:
88 return ((self.keyName, Vector),)
90 def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
91 result = dict(data)
92 result[self.keyName.format(**kwargs)] = self.action(data, **kwargs)
93 return result
96class KeyedDataSelectorAction(KeyedDataAction):
97 """Extract Vector specified by ``vectorKeys`` from input KeyedData and
98 optionally apply selectors to down select extracted vectors.
100 Note this action will not work with keyed scalars, see `getInputSchema` for
101 expected schema.
102 """
104 vectorKeys = ListField[str](doc="Keys to extract from KeyedData and return", default=[])
106 selectors = ConfigurableActionStructField[VectorAction](
107 doc="Selectors for selecting rows, will be AND together",
108 )
110 def getInputSchema(self) -> KeyedDataSchema:
111 yield from ((column, Vector | Scalar) for column in self.vectorKeys) # type: ignore
112 for action in self.selectors:
113 yield from action.getInputSchema()
115 def getOutputSchema(self) -> KeyedDataSchema:
116 return ((column, Vector | Scalar) for column in self.vectorKeys) # type: ignore
118 def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
119 mask: Optional[np.ndarray] = None
120 for selector in self.selectors:
121 subMask = selector(data, **kwargs)
122 if mask is None:
123 mask = subMask
124 else:
125 mask *= subMask # type: ignore
126 result = {}
127 for key in self.vectorKeys:
128 key_arg = key.format_map(kwargs)
129 result[key_arg] = data[key_arg]
130 if mask is not None:
131 return {key: cast(Vector, col)[mask] for key, col in result.items()}
132 else:
133 return result
136class KeyedScalars(KeyedDataAction):
137 """Creates an output of type KeyedData, where the keys are given by the
138 identifiers in `scalarActions` and the values are the results of the
139 corresponding `ScalarAction`.
140 """
142 scalarActions = ConfigurableActionStructField[ScalarAction](
143 doc="Create a KeyedData of individual ScalarActions"
144 )
146 def getInputSchema(self) -> KeyedDataSchema:
147 for action in self.scalarActions:
148 yield from action.getInputSchema()
150 def getOutputSchema(self) -> KeyedDataSchema:
151 for name in self.scalarActions.fieldNames:
152 yield (name, Scalar)
154 def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
155 result: KeyedData = {} # type: ignore
156 for name, action in self.scalarActions.items():
157 result[name] = action(data, **kwargs)
158 return result