Group nodes together in a networkx.DiGraph

nx_coalesce_nodes.py

#!/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