Small benchmarking for triangle calculations

test_triangles.py

from timeit import timeit

from networkx import __version__, erdos_renyi_graph, triangles

print(__version__)
#3.0rc2.dev0

def _triangles(G, nodes=None):
    """Compute the number of triangles.

    Finds the number of triangles that include a node as one vertex.

    Parameters
    ----------
    G : graph
       A networkx graph

    nodes : container of nodes, optional (default = compute for all nodes in G)
       Compute triangles for nodes in this container.

    Returns
    -------
    out : dictionary or int
       Number of triangles keyed by node label (dict) if nodes is a container of nodes.
       Number of triangles for a specific node (int) if nodes is a specific node.

    Examples
    --------
    >>> G = nx.complete_graph(5)
    >>> print(nx.triangles(G, 0))
    6
    >>> print(nx.triangles(G))
    {0: 6, 1: 6, 2: 6, 3: 6, 4: 6}
    >>> print(list(nx.triangles(G, (0, 1)).values()))
    [6, 6]

    Notes
    -----
    Self loops are ignored.

    """

    # if nodes is None, then compute triangles for the complete graph
    if nodes is None:
        return _triangle_count(G)

    # if specific nodes are provided, then for efficiency create a subgraph
    # that preserves the edges among the alters for each node in the `nodes`
    from networkx import ancestors

    if _return_single_value := not hasattr(nodes, "__contains__"):
        _return_node = nodes
        nodes = [nodes]

    subset_nodes = set(nodes).union(*(ancestors(G, n) for n in nodes))
    G_subgraph = G.subgraph(subset_nodes)

    # compute the triangles for the subgraph
    # note that the triangle counts will be correct only for nodes in the `nodes`
    # since the other edges are not preserved when G_subgraph is created
    triangle_counts = {
        node: n_triangles
        for node, n_triangles in _triangle_count(G_subgraph).items()
        if node in nodes
    }

    # if only one node is passed, then the returned value should be the triangle count
    if _return_single_value:
        return triangle_counts[_return_node]
    else:
        return triangle_counts


def _triangle_count(G):
    """This is a utility function to calculate triangles via a single pass."""

    # dict used to avoid visiting the same nodes twice
    # this allows calculating/counting each triangle only once
    later_neighbors = {}

    # iterate over the nodes in a graph
    for node, neighbors in G.adjacency():
        later_neighbors[node] = {
            n for n in neighbors if n not in later_neighbors and n is not node
        }

    # instantiate with a zero count for each node
    # add 1 to the count if a nodes neighbor's neighbor is also a neighbor
    triangle_counts = {node: 0 for node in G}
    for node1, neighbors in later_neighbors.items():
        for node2 in neighbors:
            for node3 in neighbors & later_neighbors[node2]:
                triangle_counts[node1] += 1
                triangle_counts[node2] += 1
                triangle_counts[node3] += 1

    return triangle_counts

setup = """
from networkx import triangles, erdos_renyi_graph
G = erdos_renyi_graph(1500,0.1)
"""
time_current = timeit(stmt="triangles(G)", setup=setup, number=10)

setup = """
from networkx import erdos_renyi_graph
G = erdos_renyi_graph(1500,0.1)
"""
time_proposed = timeit(
    stmt="_triangles(G)",
    setup=setup,
    number=10,
    globals=globals(),
)

print(time_proposed, time_current)
print(time_proposed/time_current)

# a few checks on correctness
G = erdos_renyi_graph(500, 0.1)
assert _triangles(G) == triangles(G)
_nodes = [1, 2, 3]
assert _triangles(G, _nodes) == triangles(G, _nodes)
_nodes = 1
assert _triangles(G, _nodes) == triangles(G, _nodes)