Coverage for python/lsst/pipe/base/pipeline_graph/visualization/_merge.py: 87%
117 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-13 09:01 +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/>.
27from __future__ import annotations
29__all__ = (
30 "MergedNodeKey",
31 "merge_graph_input_trees",
32 "merge_graph_intermediates",
33 "merge_graph_output_trees",
34)
36import dataclasses
37import hashlib
38from collections import defaultdict
39from collections.abc import Iterable
40from functools import cached_property
41from typing import Any
43import networkx
44import networkx.algorithms.dag
45import networkx.algorithms.tree
47from lsst.daf.butler import DimensionGroup
49from .._nodes import NodeKey, NodeType
50from ._options import NodeAttributeOptions
53class MergedNodeKey(frozenset[NodeKey]):
54 """A key for NetworkX graph nodes that represent multiple similar tasks
55 or dataset types that have been merged to simplify graph visualization.
56 """
58 def __str__(self) -> str:
59 members = [str(k) for k in self]
60 members.sort(reverse=True)
61 return ", ".join(members)
63 @property
64 def node_type(self) -> NodeType:
65 """Enum value for whether this is a task, task initialization, or
66 dataset type node.
67 """
68 return next(iter(self)).node_type
70 @cached_property
71 def node_id(self) -> str:
72 """A unique string representation of the merged node key."""
73 hasher = hashlib.blake2b(digest_size=4)
74 for member in sorted(self):
75 hasher.update(str(member).encode())
76 return f"{hasher.digest().hex()}:{self.node_type.value}"
79def merge_graph_input_trees(
80 xgraph: networkx.DiGraph | networkx.MultiDiGraph, options: NodeAttributeOptions, depth: int
81) -> None:
82 """Merge trees of overall-input dataset type nodes and/or
83 beginning-of-pipeline task nodes that have similar properties and the same
84 structure.
86 Parameters
87 ----------
88 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
89 Graph to be processed; modified in-place.
90 options : `NodeAttributeOptions`
91 Properties of nodes that should be considered when determining whether
92 they are similar enough to be merged. Only the truthiness of
93 attributes is considered (e.g. ``options.dimensions == 'full'`` and
94 ``options.dimensions == 'concise'`` are both interpreted to mean "only
95 merge trees with the same dimensions"). This is typically the same set
96 of options that controls whether to display these attributes in the
97 graph visualization.
98 depth : `int`
99 How many nodes to traverse from the beginning of the graph before
100 terminating the merging algorithm.
101 """
102 groups = _make_tree_merge_groups(xgraph, options, depth)
103 _apply_tree_merges(xgraph, groups)
106def merge_graph_output_trees(
107 xgraph: networkx.DiGraph | networkx.MultiDiGraph, options: NodeAttributeOptions, depth: int
108) -> None:
109 """Merge trees of overall-output dataset type nodes and/or
110 end-of-pipeline task nodes that have similar properties and the same
111 structure.
113 Parameters
114 ----------
115 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
116 Graph to be processed; modified in-place.
117 options : `NodeAttributeOptions`
118 Properties of nodes that should be considered when determining whether
119 they are similar enough to be merged. Only the truthiness of
120 attributes is considered (e.g. ``options.dimensions == 'full'`` and
121 ``options.dimensions == 'concise'`` are both interpreted to mean "only
122 merge trees with the same dimensions"). This is typically the same set
123 of options that controls whether to display these attributes in the
124 graph visualization.
125 depth : `int`
126 How many nodes to traverse from the beginning of the graph before
127 terminating the merging algorithm.
128 """
129 groups = _make_tree_merge_groups(xgraph.reverse(copy=False), options, depth)
130 _apply_tree_merges(xgraph, groups)
133def merge_graph_intermediates(
134 xgraph: networkx.DiGraph | networkx.MultiDiGraph, options: NodeAttributeOptions
135) -> None:
136 """Merge parallel interior nodes of a graph with similar properties.
138 Parameters
139 ----------
140 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
141 Graph to be processed; modified in-place.
142 options : `NodeAttributeOptions`
143 Properties of nodes that should be considered when determining whether
144 they are similar enough to be merged. Only the truthiness of
145 attributes is considered (e.g. ``options.dimensions == 'full'`` and
146 ``options.dimensions == 'concise'`` are both interpreted to mean "only
147 merge trees with the same dimensions"). This is typically the same set
148 of options that controls whether to display these attributes in the
149 graph visualization.
151 Notes
152 -----
153 "Parallel" nodes here are nodes that have the exact same predecessor and
154 successors.
155 """
156 groups: dict[_MergeKey, set[NodeKey]] = defaultdict(set)
157 for node, state in xgraph.nodes.items():
158 merge_key = _MergeKey.from_node_state(
159 state,
160 xgraph.predecessors(node),
161 xgraph.successors(node),
162 options,
163 )
164 if merge_key.parents and merge_key.children:
165 groups[merge_key].add(node)
166 replacements: dict[NodeKey, MergedNodeKey] = {}
167 for merge_key, members in groups.items():
168 if len(members) < 2: 168 ↛ 170line 168 didn't jump to line 170 because the condition on line 168 was always true
169 continue
170 new_node_key = MergedNodeKey(frozenset(members))
171 xgraph.add_node(
172 new_node_key,
173 storage_class_name=merge_key.storage_class_name,
174 task_class_name=merge_key.task_class_name,
175 dimensions=merge_key.dimensions,
176 is_prerequisite=merge_key.is_prerequisite,
177 )
178 for parent in merge_key.parents:
179 _add_merged_edge(xgraph, replacements.get(parent, parent), new_node_key)
180 for child in merge_key.children:
181 _add_merged_edge(xgraph, new_node_key, replacements.get(child, child))
182 for member in members:
183 replacements[member] = new_node_key
184 xgraph.remove_nodes_from(members)
187@dataclasses.dataclass(frozen=True)
188class _MergeKey:
189 """A helper class for merge algorithms that is used as a dictionary key
190 when grouping nodes that may be merged by their attributes.
191 """
193 parents: frozenset[Any]
194 """Nodes of the original graph that are successors or predecessors of
195 the nodes being considered for merging.
196 """
198 dimensions: DimensionGroup | None
199 """Dimensions of the nodes being considered for merging, or `None` if
200 dimensions are not included in the similarity criteria.
201 """
203 storage_class_name: str | None
204 """Storage class of the nodes being considered for merging, or `None` if
205 storage classes are not included in the similarity criteria or this is a
206 task or task initialization node group.
207 """
209 task_class_name: str | None
210 """Name of the task class for the nodes being considered for merging, or
211 `None` if task classes are not included in the similarity criteria or
212 this is a dataset type node group.
213 """
215 is_prerequisite: bool
216 """Whether the nodes being considered for merging are prerequisite input
217 dataset types.
218 """
220 children: frozenset[Any]
221 """Nodes that are predecessors or successors (the opposite of ``parents``
222 of the nodes being considered for merging.
224 In the `merge_graph_intermediates` algorithm, these are regular unmerged
225 nodes. In the `merge_graph_input_trees` or `merge_graph_output_trees`
226 algorithms, these are more `_MergeKey` instances, representing
227 already-processed trees.
228 """
230 @classmethod
231 def from_node_state[P, C](
232 cls,
233 state: dict[str, Any],
234 parents: Iterable[P],
235 children: Iterable[C],
236 options: NodeAttributeOptions,
237 ) -> _MergeKey:
238 """Construct from a NetworkX node attribute state dictionary.
240 Parameters
241 ----------
242 state : `dict`
243 Dictionary used to hold NetworkX node attributes.
244 parents : `~collections.abc.Iterable` [ `NodeKey` ]
245 Predecessor or successor nodes (depending on the orientation of
246 the algorithm).
247 children : ~collections.abc.Iterable`
248 Successor or predecessor nodes (depending on the orientation of
249 the algorithm).
250 options : `NodeAttributeOptions`
251 Options for which node attributes to include in the new key.
252 """
253 return cls(
254 parents=frozenset(parents),
255 dimensions=state.get("dimensions"),
256 storage_class_name=(state.get("storage_class_name") if options.storage_classes else None),
257 task_class_name=(state.get("task_class_name") if options.task_classes else None),
258 is_prerequisite=state.get("is_prerequisite", False),
259 children=frozenset(children),
260 )
263def _add_merged_edge(
264 xgraph: networkx.DiGraph | networkx.MultiDiGraph,
265 a: NodeKey | MergedNodeKey,
266 b: NodeKey | MergedNodeKey,
267) -> None:
268 """Add an edge that involves at least one merged node.
270 Edges from prerequisite dataset type nodes are marked as prerequisite
271 edges, since the original edge attributes are not carried over when
272 nodes are merged.
273 """
274 xgraph.add_edge(a, b, is_prerequisite=xgraph.nodes[a].get("is_prerequisite", False))
277def _make_tree_merge_groups(
278 xgraph: networkx.DiGraph | networkx.MultiDiGraph,
279 options: NodeAttributeOptions,
280 depth: int,
281) -> list[dict[_MergeKey, set[NodeKey]]]:
282 """First-stage implementation of `merge_graph_input_trees` and
283 (when run on the reversed graph) `merge_graph_output_trees`.
284 """
285 # Our goal is to obtain mappings that groups trees of nodes by the
286 # attributes in a _TreeMergeKey. The nested dictionaries are the root of a
287 # tree and the nodes under that root, recursively (but not including the
288 # root). We nest these mappings inside a list, which each mapping
289 # corresponding to a different depth for the trees it represents. We start
290 # with a special empty dict for "0-depth trees", since that makes
291 # result[depth] valid and hence off-by-one errors less likely.
292 result: list[dict[_MergeKey, set[NodeKey]]] = [{}]
293 if depth == 0: 293 ↛ 294line 293 didn't jump to line 294 because the condition on line 293 was never true
294 return result
295 # We start with the nodes that have no predecessors in the graph.
296 # Ignore for now the fact that the 'current_candidates' data structure
297 # we process is actually a dict that associates each of those nodes
298 # with an empty dict. All of these initial nodes are valid trees,
299 # since they're just single nodes.
300 first_generation = next(networkx.algorithms.dag.topological_generations(xgraph))
301 current_candidates: dict[NodeKey, dict[NodeKey, _MergeKey]] = dict.fromkeys(first_generation, {})
302 # Set up an outer loop over tree depth; we'll construct a new set of
303 # candidates at each iteration.
304 while current_candidates:
305 # As we go, we'll remember nodes that have just one predecessor, as
306 # those predecessors might be the roots of slightly taller trees.
307 # We store the successors and their merge keys under them.
308 next_candidates: dict[NodeKey, dict[NodeKey, _MergeKey]] = defaultdict(dict)
309 # We also want to track the nodes the level up that are not trees
310 # because some node has both them and some other node as a
311 # predecessor.
312 nontrees: set[NodeKey] = set()
313 # Make a dictionary for the results at this depth, then start the
314 # inner iteration over candidates and (after the first iteration)
315 # their children.
316 result_for_depth: dict[_MergeKey, set[NodeKey]] = defaultdict(set)
317 for node, children in current_candidates.items():
318 # Make a _TreeMergeKey for this node and add it to the results for
319 # this depth. Two nodes with the same _TreeMergeKey are roots of
320 # isomorphic trees that have the same predecessor(s), and can be
321 # merged (with isomorphism defined as both both structure and
322 # whatever comparisons are in 'options').
323 merge_key = _MergeKey.from_node_state(
324 xgraph.nodes[node], xgraph.successors(node), children.values(), options
325 )
326 result_for_depth[merge_key].add(node)
327 if len(result) <= depth:
328 # See if this node's successor might be the root of a
329 # larger tree.
330 if len(merge_key.parents) == 1:
331 (parent,) = merge_key.parents
332 next_candidates[parent][node] = dataclasses.replace(merge_key, parents=frozenset())
333 else:
334 nontrees.update(merge_key.parents)
335 # Append the results for this depth.
336 result.append(result_for_depth)
337 # Trim out candidates that aren't trees after all.
338 for nontree_node in nontrees & next_candidates.keys():
339 del next_candidates[nontree_node]
340 current_candidates = next_candidates
341 return result
344def _apply_tree_merges(
345 xgraph: networkx.DiGraph | networkx.MultiDiGraph,
346 groups: list[dict[_MergeKey, set[NodeKey]]],
347) -> None:
348 """Second-stage implementation of `merge_graph_input_trees` and
349 `merge_graph_output_trees`.
350 """
351 replacements: dict[NodeKey, MergedNodeKey] = {}
352 for group in reversed(groups):
353 new_group: dict[_MergeKey, set[NodeKey]] = defaultdict(set)
354 for merge_key, members in group.items():
355 if merge_key.parents & replacements.keys(): 355 ↛ 356line 355 didn't jump to line 356 because the condition on line 355 was never true
356 replaced_parents = frozenset(replacements.get(p, p) for p in merge_key.parents)
357 new_group[dataclasses.replace(merge_key, parents=replaced_parents)].update(members)
358 else:
359 new_group[merge_key].update(members)
360 for merge_key, members in new_group.items():
361 if len(members) < 2:
362 continue
363 new_node_key = MergedNodeKey(frozenset(members))
364 new_edges: set[tuple[NodeKey | MergedNodeKey, NodeKey | MergedNodeKey]] = set()
365 for member_key in members:
366 replacements[member_key] = new_node_key
367 new_edges.update(
368 (replacements.get(a, a), replacements.get(b, b)) for a, b in xgraph.in_edges(member_key)
369 )
370 new_edges.update(
371 (replacements.get(a, a), replacements.get(b, b)) for a, b in xgraph.out_edges(member_key)
372 )
373 xgraph.add_node(
374 new_node_key,
375 storage_class_name=merge_key.storage_class_name,
376 task_class_name=merge_key.task_class_name,
377 dimensions=merge_key.dimensions,
378 is_prerequisite=merge_key.is_prerequisite,
379 )
380 for a, b in new_edges:
381 _add_merged_edge(xgraph, a, b)
382 xgraph.remove_nodes_from(replacements.keys())