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