Coverage for python/lsst/pipe/base/quantum_graph_skeleton.py: 98%
235 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-09 09:08 +0000
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""An under-construction version of QuantumGraph and various helper
29classes.
30"""
32from __future__ import annotations
34__all__ = (
35 "DatasetKey",
36 "PrerequisiteDatasetKey",
37 "QuantumGraphSkeleton",
38 "QuantumKey",
39 "TaskInitKey",
40)
42import dataclasses
43from collections import defaultdict
44from collections.abc import Iterable, Iterator, MutableMapping, Set
45from typing import TYPE_CHECKING, Any, ClassVar, Literal
47import networkx
49from lsst.daf.butler import (
50 Butler,
51 DataCoordinate,
52 DataIdValue,
53 DatasetRef,
54 DimensionDataAttacher,
55 DimensionDataExtractor,
56 DimensionGroup,
57 DimensionRecordSet,
58)
59from lsst.utils.logging import getLogger
61if TYPE_CHECKING:
62 pass
64_LOG = getLogger(__name__)
67@dataclasses.dataclass(slots=True, eq=True, frozen=True)
68class QuantumKey:
69 """Identifier type for quantum keys in a `QuantumGraphSkeleton`."""
71 task_label: str
72 """Label of the task in the pipeline."""
74 data_id_values: tuple[DataIdValue, ...]
75 """Data ID values of the quantum.
77 Note that keys are fixed given `task_label`, so using only the values here
78 speeds up comparisons.
79 """
81 is_task: ClassVar[Literal[True]] = True
82 """Whether this node represents a quantum or task initialization rather
83 than a dataset (always `True`).
84 """
87@dataclasses.dataclass(slots=True, eq=True, frozen=True)
88class TaskInitKey:
89 """Identifier type for task init keys in a `QuantumGraphSkeleton`."""
91 task_label: str
92 """Label of the task in the pipeline."""
94 is_task: ClassVar[Literal[True]] = True
95 """Whether this node represents a quantum or task initialization rather
96 than a dataset (always `True`).
97 """
100@dataclasses.dataclass(slots=True, eq=True, frozen=True)
101class DatasetKey:
102 """Identifier type for dataset keys in a `QuantumGraphSkeleton`."""
104 parent_dataset_type_name: str
105 """Name of the dataset type (never a component)."""
107 data_id_values: tuple[DataIdValue, ...]
108 """Data ID values of the dataset.
110 Note that keys are fixed given `parent_dataset_type_name`, so using only
111 the values here speeds up comparisons.
112 """
114 is_task: ClassVar[Literal[False]] = False
115 """Whether this node represents a quantum or task initialization rather
116 than a dataset (always `False`).
117 """
119 is_prerequisite: ClassVar[Literal[False]] = False
122@dataclasses.dataclass(slots=True, eq=True, frozen=True)
123class PrerequisiteDatasetKey:
124 """Identifier type for prerequisite dataset keys in a
125 `QuantumGraphSkeleton`.
127 Unlike regular datasets, prerequisites are not actually required to come
128 from a find-first search of `input_collections`, so we don't want to
129 assume that the same data ID implies the same dataset. Happily we also
130 don't need to search for them by data ID in the graph, so we can use the
131 dataset ID (UUID) instead.
132 """
134 parent_dataset_type_name: str
135 """Name of the dataset type (never a component)."""
137 dataset_id_bytes: bytes
138 """Dataset ID (UUID) as raw bytes."""
140 is_task: ClassVar[Literal[False]] = False
141 """Whether this node represents a quantum or task initialization rather
142 than a dataset (always `False`).
143 """
145 is_prerequisite: ClassVar[Literal[True]] = True
148type Key = QuantumKey | TaskInitKey | DatasetKey | PrerequisiteDatasetKey
151class QuantumGraphSkeleton:
152 """An under-construction quantum graph.
154 QuantumGraphSkeleton is intended for use inside `QuantumGraphBuilder` and
155 its subclasses.
157 Parameters
158 ----------
159 task_labels : `~collections.abc.Iterable` [ `str` ]
160 The labels of all tasks whose quanta may be included in the graph, in
161 topological order.
163 Notes
164 -----
165 QuantumGraphSkeleton models a bipartite version of the quantum graph, in
166 which both quanta and datasets are represented as nodes and each type of
167 node only has edges to the other type.
169 Square-bracket (`getitem`) indexing returns a mutable mapping of a node's
170 flexible attributes.
172 The details of the `QuantumGraphSkeleton` API (e.g. which operations
173 operate on multiple nodes vs. a single node) are set by what's actually
174 needed by current quantum graph generation algorithms. New variants can be
175 added as needed, but adding all operations that *might* be useful for some
176 future algorithm seems premature.
177 """
179 def __init__(self, task_labels: Iterable[str]):
180 self._tasks: dict[str, tuple[TaskInitKey, set[QuantumKey]]] = {}
181 self._xgraph: networkx.DiGraph = networkx.DiGraph()
182 self._global_init_outputs: set[DatasetKey] = set()
183 self._dimension_data: dict[str, DimensionRecordSet] = {}
184 for task_label in task_labels:
185 task_init_key = TaskInitKey(task_label)
186 self._tasks[task_label] = (task_init_key, set())
187 self._xgraph.add_node(task_init_key)
189 def __contains__(self, key: Key) -> bool:
190 return key in self._xgraph.nodes
192 def __getitem__(self, key: Key) -> MutableMapping[str, Any]:
193 return self._xgraph.nodes[key]
195 def __iter__(self) -> Iterator[Key]:
196 return iter(self._xgraph.nodes)
198 @property
199 def n_nodes(self) -> int:
200 """The total number of nodes of all types."""
201 return len(self._xgraph.nodes)
203 @property
204 def n_edges(self) -> int:
205 """The total number of edges."""
206 return len(self._xgraph.edges)
208 @property
209 def has_any_quanta(self) -> bool:
210 """Test whether this graph has any quanta."""
211 for _ in self.iter_all_quanta():
212 return True
213 return False
215 def has_task(self, task_label: str) -> bool:
216 """Test whether the given task is in this skeleton.
218 Tasks are only added to the skeleton at initialization, but may be
219 removed by `remove_task` if they end up having no quanta.
221 Parameters
222 ----------
223 task_label : `str`
224 Task to check for.
226 Returns
227 -------
228 has : `bool`
229 `True` if the task is in this skeleton.
230 """
231 return task_label in self._tasks
233 def get_task_init_node(self, task_label: str) -> TaskInitKey:
234 """Return the graph node that represents a task's initialization.
236 Parameters
237 ----------
238 task_label : `str`
239 The task label to use.
241 Returns
242 -------
243 node : `TaskInitKey`
244 The graph node representing this task's initialization.
245 """
246 return self._tasks[task_label][0]
248 def get_quanta(self, task_label: str) -> Set[QuantumKey]:
249 """Return the quanta for the given task label.
251 Parameters
252 ----------
253 task_label : `str`
254 Label for the task.
256 Returns
257 -------
258 quanta : `~collections.abc.Set` [ `QuantumKey` ]
259 A set-like object with the identifiers of all quanta for the given
260 task. *The skeleton object's set of quanta must not be modified
261 while iterating over this container; make a copy if mutation during
262 iteration is necessary*.
263 """
264 return self._tasks[task_label][1]
266 @property
267 def global_init_outputs(self) -> Set[DatasetKey]:
268 """The set of dataset nodes that are not associated with any task."""
269 return self._global_init_outputs
271 def iter_all_quanta(self) -> Iterator[QuantumKey]:
272 """Iterate over all quanta from any task, in topological (but otherwise
273 unspecified) order.
274 """
275 for _, quanta in self._tasks.values():
276 yield from quanta
278 def iter_outputs_of(self, quantum_key: QuantumKey | TaskInitKey) -> Iterator[DatasetKey]:
279 """Iterate over the datasets produced by the given quantum.
281 Parameters
282 ----------
283 quantum_key : `QuantumKey` or `TaskInitKey`
284 Quantum to iterate over.
286 Returns
287 -------
288 datasets : `~collections.abc.Iterator` of `DatasetKey`
289 Datasets produced by the given quanta.
290 """
291 return self._xgraph.successors(quantum_key)
293 def iter_inputs_of(
294 self, quantum_key: QuantumKey | TaskInitKey
295 ) -> Iterator[DatasetKey | PrerequisiteDatasetKey]:
296 """Iterate over the datasets consumed by the given quantum.
298 Parameters
299 ----------
300 quantum_key : `QuantumKey` or `TaskInitKey`
301 Quantum to iterate over.
303 Returns
304 -------
305 datasets : `~collections.abc.Iterator` of `DatasetKey` \
306 or `PrequisiteDatasetKey`
307 Datasets consumed by the given quanta.
308 """
309 return self._xgraph.predecessors(quantum_key)
311 def get_producer_of(self, dataset_key: DatasetKey) -> QuantumKey | TaskInitKey | None:
312 """Return the quantum that produces the given dataset, or ``None``.
314 Parameters
315 ----------
316 dataset_key : `DatasetKey`
317 Dataset to look up the producer for.
319 Returns
320 -------
321 producer : `QuantumKey`, `TaskInitKey`, or `None`
322 The quantum or init task that produces the dataset, or ``None``
323 if the dataset has no producer in this graph (an external input).
324 """
325 return next(iter(self._xgraph.predecessors(dataset_key)), None)
327 def update(self, other: QuantumGraphSkeleton) -> None:
328 """Copy all nodes from ``other`` to ``self``.
330 Parameters
331 ----------
332 other : `QuantumGraphSkeleton`
333 Source of nodes. The tasks in ``other`` must be a subset of the
334 tasks in ``self`` (this method is expected to be used to populate
335 a skeleton for a full from independent-subgraph skeletons).
336 """
337 for task_label, (_, quanta) in other._tasks.items():
338 self._tasks[task_label][1].update(quanta)
339 self._xgraph.update(other._xgraph)
340 for record_set in other._dimension_data.values():
341 self._dimension_data.setdefault(
342 record_set.element.name, DimensionRecordSet(record_set.element)
343 ).update(record_set)
345 def add_quantum_node(self, task_label: str, data_id: DataCoordinate, **attrs: Any) -> QuantumKey:
346 """Add a new node representing a quantum.
348 Parameters
349 ----------
350 task_label : `str`
351 Name of task.
352 data_id : `~lsst.daf.butler.DataCoordinate`
353 The data ID of the quantum.
354 **attrs : `~typing.Any`
355 Additional attributes.
356 """
357 key = QuantumKey(task_label, data_id.required_values)
358 self._xgraph.add_node(key, data_id=data_id, **attrs)
359 self._tasks[key.task_label][1].add(key)
360 return key
362 def add_dataset_node(
363 self,
364 parent_dataset_type_name: str,
365 data_id: DataCoordinate,
366 is_global_init_output: bool = False,
367 **attrs: Any,
368 ) -> DatasetKey:
369 """Add a new node representing a dataset.
371 Parameters
372 ----------
373 parent_dataset_type_name : `str`
374 Name of the parent dataset type.
375 data_id : `~lsst.daf.butler.DataCoordinate`
376 The dataset data ID.
377 is_global_init_output : `bool`, optional
378 Whether this dataset is a global init output.
379 **attrs : `~typing.Any`
380 Additional attributes for the node.
381 """
382 key = DatasetKey(parent_dataset_type_name, data_id.required_values)
383 self._xgraph.add_node(key, data_id=data_id, **attrs)
384 if is_global_init_output:
385 assert isinstance(key, DatasetKey), str(key)
386 self._global_init_outputs.add(key)
387 return key
389 def add_prerequisite_node(
390 self,
391 ref: DatasetRef,
392 **attrs: Any,
393 ) -> PrerequisiteDatasetKey:
394 """Add a new node representing a prerequisite input dataset.
396 Parameters
397 ----------
398 ref : `~lsst.daf.butler.DatasetRef`
399 The dataset ref of the prerequisite.
400 **attrs : `~typing.Any`
401 Additional attributes for the node.
402 """
403 key = PrerequisiteDatasetKey(ref.datasetType.name, ref.id.bytes)
404 self._xgraph.add_node(key, data_id=ref.dataId, ref=ref, **attrs)
405 return key
407 def remove_quantum_node(self, key: QuantumKey, remove_outputs: bool) -> None:
408 """Remove a node representing a quantum.
410 Parameters
411 ----------
412 key : `QuantumKey`
413 Identifier for the node.
414 remove_outputs : `bool`
415 If `True`, also remove all dataset nodes produced by this quantum.
416 If `False`, any such dataset nodes will become overall inputs.
417 """
418 _, quanta = self._tasks[key.task_label]
419 quanta.remove(key)
420 if remove_outputs:
421 to_remove = list(self._xgraph.successors(key))
422 to_remove.append(key)
423 self._xgraph.remove_nodes_from(to_remove)
424 else:
425 self._xgraph.remove_node(key)
427 def remove_dataset_nodes(self, keys: Iterable[DatasetKey | PrerequisiteDatasetKey]) -> None:
428 """Remove nodes representing datasets.
430 Parameters
431 ----------
432 keys : `~collections.abc.Iterable` of `DatasetKey`\
433 or `PrerequisiteDatasetKey`
434 Nodes to remove.
435 """
436 self._xgraph.remove_nodes_from(keys)
438 def remove_task(self, task_label: str) -> None:
439 """Fully remove a task from the skeleton.
441 All init-output datasets and quanta for the task must already have been
442 removed.
444 Parameters
445 ----------
446 task_label : `str`
447 Name of task to remove.
448 """
449 task_init_key, quanta = self._tasks.pop(task_label)
450 assert not quanta, "Cannot remove task unless all quanta have already been removed."
451 assert not list(self._xgraph.successors(task_init_key))
452 self._xgraph.remove_node(task_init_key)
454 def add_input_edges(
455 self,
456 task_key: QuantumKey | TaskInitKey,
457 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
458 ) -> None:
459 """Add edges connecting datasets to a quantum that consumes them.
461 Parameters
462 ----------
463 task_key : `QuantumKey` or `TaskInitKey`
464 Quantum to connect.
465 dataset_keys : `~collections.abc.Iterable` of `DatasetKey`\
466 or `PrequisiteDatasetKey`
467 Datasets to join to the quantum.
469 Notes
470 -----
471 This must only be called if the task node has already been added.
472 Use `add_input_edge` if this cannot be assumed.
474 Dataset nodes that are not already present will be created.
475 """
476 assert task_key in self._xgraph, str(task_key)
477 self._xgraph.add_edges_from((dataset_key, task_key) for dataset_key in dataset_keys)
479 def remove_input_edges(
480 self,
481 task_key: QuantumKey | TaskInitKey,
482 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
483 ) -> None:
484 """Remove edges connecting datasets to a quantum that consumes them.
486 Parameters
487 ----------
488 task_key : `QuantumKey` or `TaskInitKey`
489 Quantum to disconnect.
490 dataset_keys : `~collections.abc.Iterable` of `DatasetKey`\
491 or `PrequisiteDatasetKey`
492 Datasets to remove from the quantum.
493 """
494 self._xgraph.remove_edges_from((dataset_key, task_key) for dataset_key in dataset_keys)
496 def add_input_edge(
497 self,
498 task_key: QuantumKey | TaskInitKey,
499 dataset_key: DatasetKey | PrerequisiteDatasetKey,
500 ignore_unrecognized_quanta: bool = False,
501 ) -> bool:
502 """Add an edge connecting a dataset to a quantum that consumes it.
504 Parameters
505 ----------
506 task_key : `QuantumKey` or `TaskInitKey`
507 Identifier for the quantum node.
508 dataset_key : `DatasetKey` or `PrerequisiteKey`
509 Identifier for the dataset node.
510 ignore_unrecognized_quanta : `bool`, optional
511 If `False`, do nothing if the quantum node is not already present.
512 If `True`, the quantum node is assumed to be present.
514 Returns
515 -------
516 added : `bool`
517 `True` if an edge was actually added, `False` if the quantum was
518 not recognized and the edge was not added as a result.
520 Notes
521 -----
522 Dataset nodes that are not already present will be created.
523 """
524 if ignore_unrecognized_quanta and task_key not in self._xgraph: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 return False
526 self._xgraph.add_edge(dataset_key, task_key)
527 return True
529 def add_output_edge(self, task_key: QuantumKey | TaskInitKey, dataset_key: DatasetKey) -> None:
530 """Add an edge connecting a dataset to the quantum that produces it.
532 Parameters
533 ----------
534 task_key : `QuantumKey` or `TaskInitKey`
535 Identifier for the quantum node. Must identify a node already
536 present in the graph.
537 dataset_key : `DatasetKey`
538 Identifier for the dataset node. Must identify a node already
539 present in the graph.
540 """
541 assert task_key in self._xgraph, str(task_key)
542 assert dataset_key in self._xgraph, str(dataset_key)
543 self._xgraph.add_edge(task_key, dataset_key)
545 def remove_output_edge(self, dataset_key: DatasetKey) -> None:
546 """Remove the edge connecting a dataset to the quantum that produces
547 it.
549 Parameters
550 ----------
551 dataset_key : `DatasetKey`
552 Identifier for the dataset node. Must identify a node already
553 present in the graph.
554 """
555 (task_key,) = self._xgraph.predecessors(dataset_key)
556 assert dataset_key in self._xgraph, str(dataset_key)
557 self._xgraph.remove_edge(task_key, dataset_key)
559 def remove_orphan_datasets(self) -> None:
560 """Remove any dataset nodes that do not have any edges."""
561 for orphan in list(networkx.isolates(self._xgraph)):
562 if not orphan.is_task and orphan not in self._global_init_outputs:
563 self._xgraph.remove_node(orphan)
565 def remove_unanchored_quanta(self, source_label: str, anchor_label: str) -> dict[str, int]:
566 """Remove unanchored source quanta and their entire downstream chain.
568 A source quantum is considered unanchored if no quantum with
569 ``anchor_label`` is reachable along directed edges from it.
570 Unanchored source quanta and every descendant reachable from them are
571 removed unconditionally, regardless of what else may reach it.
573 Parameters
574 ----------
575 source_label : `str`
576 Task label of the source task whose unanchored quanta to remove.
577 anchor_label : `str`
578 Task label that must appear downstream of a source quantum for that
579 quantum to be considered anchored.
581 Returns
582 -------
583 removed : `dict` [`str`, `int`]
584 Mapping of task label to number of quanta removed for that task.
585 Empty if nothing was removed.
586 """
587 if not self.has_task(source_label):
588 return {}
589 source_quanta = set(self.get_quanta(source_label))
590 anchor_quanta = set(self.get_quanta(anchor_label)) if self.has_task(anchor_label) else set()
591 reachable = set()
592 for quantum in anchor_quanta:
593 reachable.update(networkx.ancestors(self._xgraph, quantum))
595 unanchored = source_quanta - reachable
596 if not unanchored:
597 return {}
599 # to_remove collects unanchored source quanta and all their
600 # descendants. It has both QuantumKey and DatasetKey nodes.
601 to_remove = set(unanchored)
602 for quantum in unanchored:
603 to_remove.update(networkx.descendants(self._xgraph, quantum))
604 removed: dict[str, int] = {}
605 for node in to_remove:
606 if isinstance(node, QuantumKey):
607 removed[node.task_label] = removed.get(node.task_label, 0) + 1
608 _, quanta = self._tasks[node.task_label]
609 quanta.remove(node)
610 self._xgraph.remove_nodes_from(to_remove)
611 # For any task with no quanta remaining, remove its TaskInitKey and
612 # any init-output dataset nodes attached to it, then drop the task.
613 for label in removed:
614 task_init_key, remaining = self._tasks[label]
615 if not remaining:
616 self._xgraph.remove_nodes_from(list(self._xgraph.successors(task_init_key)))
617 self.remove_task(label)
618 return removed
620 def extract_overall_inputs(self) -> dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef]:
621 """Find overall input datasets.
623 Returns
624 -------
625 datasets : `dict` [ `DatasetKey` or `PrerequisiteDatasetKey`, \
626 `~lsst.daf.butler.DatasetRef` ]
627 Overall-input datasets, including prerequisites and init-inputs.
628 """
629 result = {}
630 for generation in networkx.algorithms.topological_generations(self._xgraph): 630 ↛ 640line 630 didn't jump to line 640 because the loop on line 630 didn't complete
631 for dataset_key in generation:
632 if dataset_key.is_task:
633 continue
634 if (ref := self.get_dataset_ref(dataset_key)) is None:
635 raise AssertionError(
636 f"Logic bug in QG generation: dataset {dataset_key} was never resolved."
637 )
638 result[dataset_key] = ref
639 break
640 return result
642 def set_dataset_ref(
643 self, ref: DatasetRef, key: DatasetKey | PrerequisiteDatasetKey | None = None
644 ) -> None:
645 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef`
646 instance.
648 Parameters
649 ----------
650 ref : `~lsst.daf.butler.DatasetRef`
651 `~lsst.daf.butler.DatasetRef` to associate with the node.
652 key : `DatasetKey` or `PrerequisiteDatasetKey`, optional
653 Identifier for the graph node. If not provided, a `DatasetKey`
654 is constructed from the dataset type name and data ID of ``ref``.
655 """
656 if key is None:
657 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
658 self._xgraph.nodes[key]["ref"] = ref
660 def set_output_for_skip(self, ref: DatasetRef) -> None:
661 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
662 represents an existing output in a collection where such outputs can
663 cause a quantum to be skipped.
665 Parameters
666 ----------
667 ref : `~lsst.daf.butler.DatasetRef`
668 `~lsst.daf.butler.DatasetRef` to associate with the node.
669 """
670 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
671 self._xgraph.nodes[key]["output_for_skip"] = ref
673 def set_output_in_the_way(self, ref: DatasetRef) -> None:
674 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
675 represents an existing output in the output RUN collection.
677 Parameters
678 ----------
679 ref : `~lsst.daf.butler.DatasetRef`
680 `~lsst.daf.butler.DatasetRef` to associate with the node.
681 """
682 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
683 self._xgraph.nodes[key]["output_in_the_way"] = ref
685 def get_dataset_ref(self, key: DatasetKey | PrerequisiteDatasetKey) -> DatasetRef | None:
686 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
687 node.
689 This does not return "output for skip" and "output in the way"
690 datasets.
692 Parameters
693 ----------
694 key : `DatasetKey` or `PrerequisiteDatasetKey`
695 Identifier for the graph node.
697 Returns
698 -------
699 ref : `~lsst.daf.butler.DatasetRef` or `None`
700 Dataset reference associated with the node.
701 """
702 return self._xgraph.nodes[key].get("ref")
704 def get_output_for_skip(self, key: DatasetKey) -> DatasetRef | None:
705 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
706 node in a collection where it could lead to a quantum being skipped.
708 Parameters
709 ----------
710 key : `DatasetKey`
711 Identifier for the graph node.
713 Returns
714 -------
715 ref : `~lsst.daf.butler.DatasetRef` or `None`
716 Dataset reference associated with the node.
717 """
718 return self._xgraph.nodes[key].get("output_for_skip")
720 def get_output_in_the_way(self, key: DatasetKey) -> DatasetRef | None:
721 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
722 node in the output RUN collection.
724 Parameters
725 ----------
726 key : `DatasetKey`
727 Identifier for the graph node.
729 Returns
730 -------
731 ref : `~lsst.daf.butler.DatasetRef` or `None`
732 Dataset reference associated with the node.
733 """
734 return self._xgraph.nodes[key].get("output_in_the_way")
736 def discard_output_in_the_way(self, key: DatasetKey) -> None:
737 """Drop any `~lsst.daf.butler.DatasetRef` associated with this node in
738 the output RUN collection.
740 Does nothing if there is no such `~lsst.daf.butler.DatasetRef`.
742 Parameters
743 ----------
744 key : `DatasetKey`
745 Identifier for the graph node.
746 """
747 self._xgraph.nodes[key].pop("output_in_the_way", None)
749 def set_data_id(self, key: Key, data_id: DataCoordinate) -> None:
750 """Set the data ID associated with a node.
752 This updates the data ID in any `~lsst.daf.butler.DatasetRef` objects
753 associated with the node via `set_ref`, `set_output_for_skip`, or
754 `set_output_in_the_way` as well, assuming it is an expanded version
755 of the original data ID.
757 Parameters
758 ----------
759 key : `Key`
760 Identifier for the graph node.
761 data_id : `~lsst.daf.butler.DataCoordinate`
762 Data ID for the node.
763 """
764 state: MutableMapping[str, Any] = self._xgraph.nodes[key]
765 state["data_id"] = data_id
766 ref: DatasetRef | None
767 if (ref := state.get("ref")) is not None:
768 state["ref"] = ref.expanded(data_id)
769 output_for_skip: DatasetRef | None
770 if (output_for_skip := state.get("output_for_skip")) is not None: 770 ↛ 771line 770 didn't jump to line 771 because the condition on line 770 was never true
771 state["output_for_skip"] = output_for_skip.expanded(data_id)
772 output_in_the_way: DatasetRef | None
773 if (output_in_the_way := state.get("output_in_the_way")) is not None:
774 state["output_in_the_way"] = output_in_the_way.expanded(data_id)
776 def get_data_id(self, key: Key) -> DataCoordinate:
777 """Return the full data ID for a quantum or dataset, if available.
779 Parameters
780 ----------
781 key : `Key`
782 Identifier for the graph node.
784 Returns
785 -------
786 data_id : `~lsst.daf.butler.DataCoordinate`
787 Expanded data ID for the node, if one is available.
789 Raises
790 ------
791 KeyError
792 Raised if this node does not have an expanded data ID.
793 """
794 return self._xgraph.nodes[key]["data_id"]
796 def attach_dimension_records(
797 self,
798 butler: Butler,
799 dimensions: DimensionGroup,
800 dimension_records: Iterable[DimensionRecordSet] = (),
801 ) -> None:
802 """Attach dimension records to the data IDs in the skeleton.
804 This both attaches records to data IDs in the skeleton and aggregates
805 any existing records on data IDS, so `get_dimension_data` returns all
806 dimension records used in the skeleton. It can be called multiple
807 times.
809 Parameters
810 ----------
811 butler : `lsst.daf.butler.Butler`
812 Butler to use to query for missing dimension records.
813 dimensions : `lsst.daf.butler.DimensionGroup`
814 Superset of all of the dimensions of all data IDs.
815 dimension_records : `~collections.abc.Iterable` [ \
816 `lsst.daf.butler.DimensionRecordSet` ], optional
817 Iterable of sets of dimension records to attach.
818 """
819 for record_set in dimension_records:
820 self._dimension_data.setdefault(
821 record_set.element.name, DimensionRecordSet(record_set.element)
822 ).update(record_set)
823 # Group all nodes by data ID (and dimensions of data ID).
824 data_ids_to_expand: defaultdict[DimensionGroup, defaultdict[DataCoordinate, list[Key]]] = defaultdict(
825 lambda: defaultdict(list)
826 )
827 extractor = DimensionDataExtractor.from_dimension_group(dimensions)
828 data_id: DataCoordinate | None
829 for node_key in self:
830 if data_id := self[node_key].get("data_id"):
831 if data_id.hasRecords():
832 extractor.update([data_id])
833 else:
834 data_ids_to_expand[data_id.dimensions][data_id].append(node_key)
835 # Add records we extracted from data IDs that were already expanded, in
836 # case other nodes want them.
837 for record_set in extractor.records.values():
838 self._dimension_data.setdefault(
839 record_set.element.name, DimensionRecordSet(record_set.element)
840 ).update(record_set)
841 attacher = DimensionDataAttacher(records=self._dimension_data.values(), dimensions=dimensions)
842 for dimensions, data_ids in data_ids_to_expand.items():
843 with butler.query() as query:
844 # Butler query will be used as-needed to get dimension records
845 # (from prerequisites) we didn't fetch in advance. These are
846 # cached in the attacher so we don't look them up multiple
847 # times.
848 expanded_data_ids = attacher.attach(dimensions, data_ids.keys(), query=query)
849 for expanded_data_id, node_keys in zip(expanded_data_ids, data_ids.values()):
850 for node_key in node_keys:
851 self.set_data_id(node_key, expanded_data_id)
852 # Hold on to any records that we had to query for or extracted.
853 self._dimension_data = attacher.records
855 def get_dimension_data(self) -> list[DimensionRecordSet]:
856 """Return the dimension records attached to data IDs."""
857 return list(self._dimension_data.values())