Last active
June 1, 2023 11:08
-
-
Save clbarnes/0fab4e03cefa1a1634b8bfcdb954f08d to your computer and use it in GitHub Desktop.
Group nodes together in a networkx.DiGraph
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
""" | |
Group nodes together in a networkx DiGraph. | |
Requires networkx. | |
""" | |
from typing import Hashable, Any, Callable, Optional | |
import networkx as nx | |
def default_node_data_fn( | |
data: list[tuple[Hashable, dict[Hashable, Any]]] | |
) -> Optional[dict[Hashable, Any]]: | |
return {"_old_data": data} | |
def default_edge_data_fn( | |
data: list[tuple[Hashable, Hashable, dict[Hashable, Any]]] | |
) -> Optional[dict[Hashable, Any]]: | |
return {"_old_data": data} | |
def coalesce_nodes_multi( | |
g: nx.DiGraph, | |
group_ids: dict[Hashable, Hashable], | |
node_data_fn: Callable[ | |
[list[tuple[Hashable, dict[Hashable, Any]]]], Optional[dict[Hashable, Any]] | |
] = default_node_data_fn, | |
edge_data_fn: Callable[ | |
[list[tuple[Hashable, Hashable, dict[Hashable, Any]]]], | |
Optional[dict[Hashable, Any]], | |
] = default_edge_data_fn, | |
keep_others: bool = True, | |
): | |
"""Coalesce groups of nodes into single nodes. | |
Coalesced node and edge data handling can be customised. | |
Node groupings are stored on the graph under "_node_provenance"; | |
a list where the last element is the {old_id: new_id} mapping | |
for the latest coalescing. | |
This allows for repeated coalescings without data loss. | |
Parameters | |
---------- | |
g : nx.DiGraph | |
Original graph. | |
group_ids : dict[Hashable, Hashable] | |
Mapping from original node ID to new coalesced node ID. | |
node_data_fn : Callable[ [list[tuple[Hashable, dict[Hashable, Any]]]], Optional[dict[Hashable, Any]] ], optional | |
How to coalesce node data. | |
A callable which takes a list of id-data tuples and returns None or a new data dict. | |
If none, node is not added to the output graph. | |
By default stores the input list under "_old_data". | |
edge_data_fn : Callable[ [list[tuple[Hashable, Hashable, dict[Hashable, Any]]]], Optional[dict[Hashable, Any]] ], optional | |
How to coalesce edge data. | |
A callable which takes a list of src-tgt-data tuples and returns None or a new data dict. | |
If None, edge is not added to the output graph. | |
By default stores the input list under "_old_data". | |
keep_others : bool, optional | |
Whether to keep nodes (and associated edges) not given in group_ids, by default True. | |
Returns | |
------- | |
nx.DiGraph | |
""" | |
out = nx.DiGraph() | |
out.graph.update(g.graph) | |
if keep_others: | |
g_ids = {n: group_ids.get(n, n) for n in g.nodes} | |
else: | |
g_ids = {n: new for n, new in group_ids.items() if n in g.nodes} | |
node_prov = out.graph.setdefault("_node_provenance", []) | |
node_prov.append(g_ids) | |
# new node ID to list of tuples of old node ID to old node data | |
nodes_to_add: dict[Hashable, list[tuple[Hashable, dict[Hashable, Any]]]] = dict() | |
for n, new_n in g_ids.items(): | |
nodes = nodes_to_add.setdefault(new_n, []) | |
nodes.append((n, g.nodes[n])) | |
for n, datas in nodes_to_add.items(): | |
new_data = node_data_fn(datas) | |
if new_data is not None: | |
out.add_node(n) | |
out.nodes[n].update(node_data_fn(datas)) | |
edges_to_add: dict[ | |
tuple[Hashable, Hashable], list[tuple[Hashable, Hashable, dict[Hashable, Any]]] | |
] = dict() | |
for src, tgt, data in g.edges(g_ids, True): | |
try: | |
new_src = g_ids[src] | |
new_tgt = g_ids[tgt] | |
except KeyError: | |
continue | |
edges = edges_to_add.setdefault((new_src, new_tgt), []) | |
edges.append((src, tgt, data)) | |
for (src, tgt), datas in edges_to_add.items(): | |
new_data = edge_data_fn(datas) | |
if new_data is not None: | |
out.add_edge(src, tgt) | |
out.edges[(src, tgt)].update(new_data) | |
return out |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment