21 from __future__
import annotations
23 """Module defining quantum graph classes and related methods.
25 There could be different representations of the quantum graph depending
26 on the client needs. Presently this module contains graph implementation
27 which is based on requirements of command-line environment. In the future
28 we could add other implementations and methods to convert between those
33 __all__ = [
"QuantumGraph",
"QuantumGraphTaskNodes",
"QuantumIterData"]
38 from itertools
import chain
39 from dataclasses
import dataclass
40 from typing
import List, FrozenSet, Mapping
45 from .pipeline
import TaskDef
46 from .pipeTools
import orderPipeline
47 from lsst.daf.butler
import DatasetRef, DatasetType, NamedKeyDict, Quantum
60 """Helper class for iterating over quanta in a graph.
62 The `QuantumGraph.traverse` method needs to return topologically ordered
63 Quanta together with their dependencies. This class is used as a value
64 for the iterator, it contains enumerated Quantum and its dependencies.
67 __slots__ = [
"index",
"quantum",
"taskDef",
"dependencies"]
70 """Index of this Quantum, a unique but arbitrary integer."""
73 """Quantum corresponding to a graph node."""
76 """Task class to be run on this quantum, and corresponding label and
80 dependencies: FrozenSet(int)
81 """Possibly empty set of indices of dependencies for this Quantum.
82 Dependencies include other nodes in the graph; they do not reflect data
83 already in butler (there are no graph nodes for those).
89 """QuantumGraphTaskNodes represents a bunch of nodes in an quantum graph
90 corresponding to a single task.
92 The node in quantum graph is represented by the `PipelineTask` and a
93 single `~lsst.daf.butler.Quantum` instance. One possible representation
94 of the graph is just a list of nodes without edges (edges can be deduced
95 from nodes' quantum inputs and outputs if needed). That representation can
96 be reduced to the list of PipelineTasks (or their corresponding TaskDefs)
97 and the corresponding list of Quanta. This class is used in this reduced
98 representation for a single task, and full `QuantumGraph` is a sequence of
99 tinstances of this class for one or more tasks.
101 Different frameworks may use different graph representation, this
102 representation was based mostly on requirements of command-line
103 executor which does not need explicit edges information.
107 """Task defintion for this set of nodes."""
109 quanta: List[Quantum]
110 """List of quanta corresponding to the task."""
112 initInputs: Mapping[DatasetType, DatasetRef]
113 """Datasets that must be loaded or created to construct this task."""
115 initOutputs: Mapping[DatasetType, DatasetRef]
116 """Datasets that may be written after constructing this task."""
120 """QuantumGraph is a sequence of `QuantumGraphTaskNodes` objects.
122 Typically the order of the tasks in the list will be the same as the
123 order of tasks in a pipeline (obviously depends on the code which
128 iterable : iterable of `QuantumGraphTaskNodes`, optional
129 Initial sequence of per-task nodes.
132 list.__init__(self, iterable
or [])
137 initInputs: NamedKeyDict
138 """Datasets that must be provided to construct one or more Tasks in this
139 graph, and must be obtained from the data repository.
141 This is disjoint with both `initIntermediates` and `initOutputs`.
144 initIntermediates: NamedKeyDict
145 """Datasets that must be provided to construct one or more Tasks in this
146 graph, but are also produced after constructing a Task in this graph.
148 This is disjoint with both `initInputs` and `initOutputs`.
151 initOutputs: NamedKeyDict
152 """Datasets that are produced after constructing a Task in this graph,
153 and are not used to construct any other Task in this graph.
155 This is disjoint from both `initInputs` and `initIntermediates`.
159 """Iterator over quanta in a graph.
161 Quanta are returned in unspecified order.
166 Task definition for a Quantum.
167 quantum : `~lsst.daf.butler.Quantum`
170 for taskNodes
in self:
171 taskDef = taskNodes.taskDef
172 for quantum
in taskNodes.quanta:
173 yield taskDef, quantum
176 """Iterator over quanta in a graph.
178 QuantumGraph containing individual quanta are returned.
182 graph : `QuantumGraph`
184 for taskDef, quantum
in self.
quanta():
186 quantum.initInputs, quantum.outputs)
191 """Return total count of quanta in a graph.
196 Number of quanta in a graph.
198 return sum(len(taskNodes.quanta)
for taskNodes
in self)
201 """Return topologically ordered Quanta and their dependencies.
203 This method iterates over all Quanta in topological order, enumerating
204 them during iteration. Returned `QuantumIterData` object contains
205 Quantum instance, its ``index`` and the ``index`` of all its
206 prerequsites (Quanta that produce inputs for this Quantum):
208 - the ``index`` values are generated by an iteration of a
209 QuantumGraph, and are not intrinsic to the QuantumGraph
210 - during iteration, each ID will appear in index before it ever
211 appears in dependencies.
215 quantumData : `QuantumIterData`
218 def orderedTaskNodes(graph):
219 """Return topologically ordered task nodes.
223 nodes : `QuantumGraphTaskNodes`
228 nodesMap = {id(item.taskDef): item
for item
in graph}
230 for taskDef
in pipeline:
231 yield nodesMap[id(taskDef)]
235 for nodes
in orderedTaskNodes(self):
236 for quantum
in nodes.quanta:
240 for dataRef
in chain.from_iterable(quantum.predictedInputs.values()):
242 if dataRef.id
is None:
244 name, component = dataRef.datasetType.nameAndComponent()
245 key = (name, dataRef.dataId)
247 prereq.append(outputs[key])
255 if not (len(self) == 1
and len(self[0].quanta) == 1):
259 for dataRef
in chain.from_iterable(quantum.outputs.values()):
260 key = (dataRef.datasetType.name, dataRef.dataId)
263 yield QuantumIterData(index=index, quantum=quantum, taskDef=nodes.taskDef,
264 dependencies=frozenset(prereq))