Coverage for python/lsst/pipe/base/quantum_graph_skeleton.py: 98%
208 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-29 22:16 +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 extract_overall_inputs(self) -> dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef]:
566 """Find overall input datasets.
568 Returns
569 -------
570 datasets : `dict` [ `DatasetKey` or `PrerequisiteDatasetKey`, \
571 `~lsst.daf.butler.DatasetRef` ]
572 Overall-input datasets, including prerequisites and init-inputs.
573 """
574 result = {}
575 for generation in networkx.algorithms.topological_generations(self._xgraph): 575 ↛ 585line 575 didn't jump to line 585 because the loop on line 575 didn't complete
576 for dataset_key in generation:
577 if dataset_key.is_task:
578 continue
579 if (ref := self.get_dataset_ref(dataset_key)) is None:
580 raise AssertionError(
581 f"Logic bug in QG generation: dataset {dataset_key} was never resolved."
582 )
583 result[dataset_key] = ref
584 break
585 return result
587 def set_dataset_ref(
588 self, ref: DatasetRef, key: DatasetKey | PrerequisiteDatasetKey | None = None
589 ) -> None:
590 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef`
591 instance.
593 Parameters
594 ----------
595 ref : `~lsst.daf.butler.DatasetRef`
596 `~lsst.daf.butler.DatasetRef` to associate with the node.
597 key : `DatasetKey` or `PrerequisiteDatasetKey`, optional
598 Identifier for the graph node. If not provided, a `DatasetKey`
599 is constructed from the dataset type name and data ID of ``ref``.
600 """
601 if key is None:
602 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
603 self._xgraph.nodes[key]["ref"] = ref
605 def set_output_for_skip(self, ref: DatasetRef) -> None:
606 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
607 represents an existing output in a collection where such outputs can
608 cause a quantum to be skipped.
610 Parameters
611 ----------
612 ref : `~lsst.daf.butler.DatasetRef`
613 `~lsst.daf.butler.DatasetRef` to associate with the node.
614 """
615 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
616 self._xgraph.nodes[key]["output_for_skip"] = ref
618 def set_output_in_the_way(self, ref: DatasetRef) -> None:
619 """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
620 represents an existing output in the output RUN collection.
622 Parameters
623 ----------
624 ref : `~lsst.daf.butler.DatasetRef`
625 `~lsst.daf.butler.DatasetRef` to associate with the node.
626 """
627 key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
628 self._xgraph.nodes[key]["output_in_the_way"] = ref
630 def get_dataset_ref(self, key: DatasetKey | PrerequisiteDatasetKey) -> DatasetRef | None:
631 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
632 node.
634 This does not return "output for skip" and "output in the way"
635 datasets.
637 Parameters
638 ----------
639 key : `DatasetKey` or `PrerequisiteDatasetKey`
640 Identifier for the graph node.
642 Returns
643 -------
644 ref : `~lsst.daf.butler.DatasetRef` or `None`
645 Dataset reference associated with the node.
646 """
647 return self._xgraph.nodes[key].get("ref")
649 def get_output_for_skip(self, key: DatasetKey) -> DatasetRef | None:
650 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
651 node in a collection where it could lead to a quantum being skipped.
653 Parameters
654 ----------
655 key : `DatasetKey`
656 Identifier for the graph node.
658 Returns
659 -------
660 ref : `~lsst.daf.butler.DatasetRef` or `None`
661 Dataset reference associated with the node.
662 """
663 return self._xgraph.nodes[key].get("output_for_skip")
665 def get_output_in_the_way(self, key: DatasetKey) -> DatasetRef | None:
666 """Return the `~lsst.daf.butler.DatasetRef` associated with the given
667 node in the output RUN collection.
669 Parameters
670 ----------
671 key : `DatasetKey`
672 Identifier for the graph node.
674 Returns
675 -------
676 ref : `~lsst.daf.butler.DatasetRef` or `None`
677 Dataset reference associated with the node.
678 """
679 return self._xgraph.nodes[key].get("output_in_the_way")
681 def discard_output_in_the_way(self, key: DatasetKey) -> None:
682 """Drop any `~lsst.daf.butler.DatasetRef` associated with this node in
683 the output RUN collection.
685 Does nothing if there is no such `~lsst.daf.butler.DatasetRef`.
687 Parameters
688 ----------
689 key : `DatasetKey`
690 Identifier for the graph node.
691 """
692 self._xgraph.nodes[key].pop("output_in_the_way", None)
694 def set_data_id(self, key: Key, data_id: DataCoordinate) -> None:
695 """Set the data ID associated with a node.
697 This updates the data ID in any `~lsst.daf.butler.DatasetRef` objects
698 associated with the node via `set_ref`, `set_output_for_skip`, or
699 `set_output_in_the_way` as well, assuming it is an expanded version
700 of the original data ID.
702 Parameters
703 ----------
704 key : `Key`
705 Identifier for the graph node.
706 data_id : `~lsst.daf.butler.DataCoordinate`
707 Data ID for the node.
708 """
709 state: MutableMapping[str, Any] = self._xgraph.nodes[key]
710 state["data_id"] = data_id
711 ref: DatasetRef | None
712 if (ref := state.get("ref")) is not None:
713 state["ref"] = ref.expanded(data_id)
714 output_for_skip: DatasetRef | None
715 if (output_for_skip := state.get("output_for_skip")) is not None: 715 ↛ 716line 715 didn't jump to line 716 because the condition on line 715 was never true
716 state["output_for_skip"] = output_for_skip.expanded(data_id)
717 output_in_the_way: DatasetRef | None
718 if (output_in_the_way := state.get("output_in_the_way")) is not None:
719 state["output_in_the_way"] = output_in_the_way.expanded(data_id)
721 def get_data_id(self, key: Key) -> DataCoordinate:
722 """Return the full data ID for a quantum or dataset, if available.
724 Parameters
725 ----------
726 key : `Key`
727 Identifier for the graph node.
729 Returns
730 -------
731 data_id : `~lsst.daf.butler.DataCoordinate`
732 Expanded data ID for the node, if one is available.
734 Raises
735 ------
736 KeyError
737 Raised if this node does not have an expanded data ID.
738 """
739 return self._xgraph.nodes[key]["data_id"]
741 def attach_dimension_records(
742 self,
743 butler: Butler,
744 dimensions: DimensionGroup,
745 dimension_records: Iterable[DimensionRecordSet] = (),
746 ) -> None:
747 """Attach dimension records to the data IDs in the skeleton.
749 This both attaches records to data IDs in the skeleton and aggregates
750 any existing records on data IDS, so `get_dimension_data` returns all
751 dimension records used in the skeleton. It can be called multiple
752 times.
754 Parameters
755 ----------
756 butler : `lsst.daf.butler.Butler`
757 Butler to use to query for missing dimension records.
758 dimensions : `lsst.daf.butler.DimensionGroup`
759 Superset of all of the dimensions of all data IDs.
760 dimension_records : `~collections.abc.Iterable` [ \
761 `lsst.daf.butler.DimensionRecordSet` ], optional
762 Iterable of sets of dimension records to attach.
763 """
764 for record_set in dimension_records:
765 self._dimension_data.setdefault(
766 record_set.element.name, DimensionRecordSet(record_set.element)
767 ).update(record_set)
768 # Group all nodes by data ID (and dimensions of data ID).
769 data_ids_to_expand: defaultdict[DimensionGroup, defaultdict[DataCoordinate, list[Key]]] = defaultdict(
770 lambda: defaultdict(list)
771 )
772 extractor = DimensionDataExtractor.from_dimension_group(dimensions)
773 data_id: DataCoordinate | None
774 for node_key in self:
775 if data_id := self[node_key].get("data_id"):
776 if data_id.hasRecords():
777 extractor.update([data_id])
778 else:
779 data_ids_to_expand[data_id.dimensions][data_id].append(node_key)
780 # Add records we extracted from data IDs that were already expanded, in
781 # case other nodes want them.
782 for record_set in extractor.records.values():
783 self._dimension_data.setdefault(
784 record_set.element.name, DimensionRecordSet(record_set.element)
785 ).update(record_set)
786 attacher = DimensionDataAttacher(records=self._dimension_data.values(), dimensions=dimensions)
787 for dimensions, data_ids in data_ids_to_expand.items():
788 with butler.query() as query:
789 # Butler query will be used as-needed to get dimension records
790 # (from prerequisites) we didn't fetch in advance. These are
791 # cached in the attacher so we don't look them up multiple
792 # times.
793 expanded_data_ids = attacher.attach(dimensions, data_ids.keys(), query=query)
794 for expanded_data_id, node_keys in zip(expanded_data_ids, data_ids.values()):
795 for node_key in node_keys:
796 self.set_data_id(node_key, expanded_data_id)
797 # Hold on to any records that we had to query for or extracted.
798 self._dimension_data = attacher.records
800 def get_dimension_data(self) -> list[DimensionRecordSet]:
801 """Return the dimension records attached to data IDs."""
802 return list(self._dimension_data.values())