nx-parallel/nx_parallel/algorithms/shortest_paths/unweighted.py at e9c6c066d07d1128c997e8771d1ba844eee308f1 · Schefflera-Arboricola/nx-parallel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# TODO : update docstrings

"""
Shortest path parallel algorithms for unweighted graphs.
"""

from joblib import Parallel, delayed
import nx_parallel as nxp
from networkx.algorithms.shortest_paths.unweighted import (
    single_source_shortest_path_length,
    single_source_shortest_path,
)

__all__ = [
    "all_pairs_shortest_path",
    "all_pairs_shortest_path_length",
]


def all_pairs_shortest_path_length(G, cutoff=None):
    """Computes the shortest path lengths between all nodes in `G`."""
    if hasattr(G, "graph_object"):
        G = G.graph_object

    length = single_source_shortest_path_length

    def _calculate_all_pairs_shortest_path_length_subset(n):
        return (n, (length(G, n, cutoff=cutoff)))

    total_cores = nxp.cpu_count()

    return Parallel(n_jobs=total_cores, return_as="generator")(
        delayed(_calculate_all_pairs_shortest_path_length_subset)(n) for n in G
    )


def all_pairs_shortest_path(G, cutoff=None):
    """Compute shortest paths between all nodes."""
    if hasattr(G, "graph_object"):
        G = G.graph_object

    def _calculate_all_pairs_shortest_path_subset(n):
        return (n, (single_source_shortest_path(G, n, cutoff=cutoff)))

    total_cores = nxp.cpu_count()

    return Parallel(n_jobs=total_cores, return_as="generator")(
        delayed(_calculate_all_pairs_shortest_path_subset)(n) for n in G
    )