Python networkx.bfs_edges() Examples

The following are 25 code examples of networkx.bfs_edges(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module networkx , or try the search function .
Example #1
Source File: greedy_coloring.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 6 votes vote down vote up
def strategy_connected_sequential(G, colors, traversal='bfs'):
    """
    Connected sequential ordering (CS). Yield nodes in such an order, that
    each node, except the first one, has at least one neighbour in the
    preceeding sequence. The sequence can be generated using both BFS and
    DFS search (using the strategy_connected_sequential_bfs and
    strategy_connected_sequential_dfs method). The default is bfs.
    """
    for component_graph in nx.connected_component_subgraphs(G):
        source = component_graph.nodes()[0]

        yield source  # Pick the first node as source

        if traversal == 'bfs':
            tree = nx.bfs_edges(component_graph, source)
        elif traversal == 'dfs':
            tree = nx.dfs_edges(component_graph, source)
        else:
            raise nx.NetworkXError(
                'Please specify bfs or dfs for connected sequential ordering')

        for (_, end) in tree:
            # Then yield nodes in the order traversed by either BFS or DFS
            yield end 
Example #2
Source File: decomposers.py    From dwave-hybrid with Apache License 2.0 5 votes vote down vote up
def _bfs_nodes(cls, graph, source, size, **kwargs):
        """Traverse `graph` with BFS starting from `source`, up to `size` nodes.
        Return an iterator of subgraph nodes (including source node).
        """
        if size < 1:
            return iter(())

        return itertools.chain(
            (source,),
            itertools.islice((v for u, v in nx.bfs_edges(graph, source)), size-1)
        ) 
Example #3
Source File: test_bfs.py    From aws-kube-codesuite with Apache License 2.0 5 votes vote down vote up
def test_bfs_edges_reverse(self):
        D = nx.DiGraph()
        D.add_edges_from([(0, 1), (1, 2), (1, 3), (2, 4), (3, 4)])
        edges = nx.bfs_edges(D, source=4, reverse=True)
        assert_equal(list(edges), [(4, 2), (4, 3), (2, 1), (1, 0)]) 
Example #4
Source File: test_bfs.py    From aws-kube-codesuite with Apache License 2.0 5 votes vote down vote up
def test_bfs_edges(self):
        edges = nx.bfs_edges(self.G, source=0)
        assert_equal(list(edges), [(0, 1), (1, 2), (1, 3), (2, 4)]) 
Example #5
Source File: breadth_first_search.py    From aws-kube-codesuite with Apache License 2.0 5 votes vote down vote up
def bfs_predecessors(G, source):
    """Returns an iterator of predecessors in breadth-first-search from source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    Returns
    -------
    pred: iterator
        (node, predecessors) iterator where predecessors is the list of
        predecessors of the node.

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(dict(nx.bfs_predecessors(G, 0)))
    {1: 0, 2: 1}
    >>> H = nx.Graph()
    >>> H.add_edges_from([(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)])
    >>> dict(nx.bfs_predecessors(H, 0))
    {1: 0, 2: 0, 3: 1, 4: 1, 5: 2, 6: 2}

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    for s, t in bfs_edges(G, source):
        yield (t, s) 
Example #6
Source File: breadth_first_search.py    From aws-kube-codesuite with Apache License 2.0 5 votes vote down vote up
def bfs_tree(G, source, reverse=False):
    """Return an oriented tree constructed from of a breadth-first-search
    starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    Returns
    -------
    T: NetworkX DiGraph
       An oriented tree

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(list(nx.bfs_tree(G,1).edges()))
    [(1, 0), (1, 2)]

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    T = nx.DiGraph()
    T.add_node(source)
    T.add_edges_from(bfs_edges(G, source, reverse=reverse))
    return T 
Example #7
Source File: greedy_coloring.py    From aws-kube-codesuite with Apache License 2.0 5 votes vote down vote up
def strategy_connected_sequential(G, colors, traversal='bfs'):
    """Returns an iterable over nodes in ``G`` in the order given by a
    breadth-first or depth-first traversal.

    ``traversal`` must be one of the strings ``'dfs'`` or ``'bfs'``,
    representing depth-first traversal or breadth-first traversal,
    respectively.

    The generated sequence has the property that for each node except
    the first, at least one neighbor appeared earlier in the sequence.

    ``G`` is a NetworkX graph. ``colors`` is ignored.

    """
    if traversal == 'bfs':
        traverse = nx.bfs_edges
    elif traversal == 'dfs':
        traverse = nx.dfs_edges
    else:
        raise nx.NetworkXError("Please specify one of the strings 'bfs' or"
                               " 'dfs' for connected sequential ordering")
    for component in nx.connected_component_subgraphs(G):
        source = arbitrary_element(component)
        # Yield the source node, then all the nodes in the specified
        # traversal order.
        yield source
        for (_, end) in traverse(component, source):
            yield end 
Example #8
Source File: test_bfs.py    From Carnets with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def bfs_test_edges(self):
        edges = nx.bfs_edges(self.G, source=9, depth_limit=4)
        assert_equal(list(edges), [(9, 8), (9, 10), (8, 7),
                                   (7, 2), (2, 1), (2, 3)]) 
Example #9
Source File: test_bfs.py    From Carnets with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def test_bfs_edges_reverse(self):
        D = nx.DiGraph()
        D.add_edges_from([(0, 1), (1, 2), (1, 3), (2, 4), (3, 4)])
        edges = nx.bfs_edges(D, source=4, reverse=True)
        assert_equal(list(edges), [(4, 2), (4, 3), (2, 1), (1, 0)]) 
Example #10
Source File: greedy_coloring.py    From Carnets with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def strategy_connected_sequential(G, colors, traversal='bfs'):
    """Returns an iterable over nodes in ``G`` in the order given by a
    breadth-first or depth-first traversal.

    ``traversal`` must be one of the strings ``'dfs'`` or ``'bfs'``,
    representing depth-first traversal or breadth-first traversal,
    respectively.

    The generated sequence has the property that for each node except
    the first, at least one neighbor appeared earlier in the sequence.

    ``G`` is a NetworkX graph. ``colors`` is ignored.

    """
    if traversal == 'bfs':
        traverse = nx.bfs_edges
    elif traversal == 'dfs':
        traverse = nx.dfs_edges
    else:
        raise nx.NetworkXError("Please specify one of the strings 'bfs' or"
                               " 'dfs' for connected sequential ordering")
    for component in nx.connected_component_subgraphs(G):
        source = arbitrary_element(component)
        # Yield the source node, then all the nodes in the specified
        # traversal order.
        yield source
        for (_, end) in traverse(component, source):
            yield end 
Example #11
Source File: test_bfs.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 5 votes vote down vote up
def test_bfs_edges_reverse(self):
        D = nx.DiGraph()
        D.add_edges_from([(0, 1), (1, 2), (1, 3), (2, 4), (3, 4)])
        edges = nx.bfs_edges(D, source=4, reverse=True)
        assert_equal(list(edges), [(4, 2), (4, 3), (2, 1), (1, 0)]) 
Example #12
Source File: breadth_first_search.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 5 votes vote down vote up
def bfs_successors(G, source):
    """Return dictionary of successors in breadth-first-search from source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    Returns
    -------
    succ: dict
       A dictionary with nodes as keys and list of succssors nodes as values.

    Examples
    --------
    >>> G = nx.Graph()
    >>> G.add_path([0,1,2])
    >>> print(nx.bfs_successors(G,0))
    {0: [1], 1: [2]}

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    d = defaultdict(list)
    for s,t in bfs_edges(G,source):
        d[s].append(t)
    return dict(d) 
Example #13
Source File: breadth_first_search.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 5 votes vote down vote up
def bfs_predecessors(G, source):
    """Return dictionary of predecessors in breadth-first-search from source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    Returns
    -------
    pred: dict
       A dictionary with nodes as keys and predecessor nodes as values.

    Examples
    --------
    >>> G = nx.Graph()
    >>> G.add_path([0,1,2])
    >>> print(nx.bfs_predecessors(G,0))
    {1: 0, 2: 1}

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    return dict((t,s) for s,t in bfs_edges(G,source)) 
Example #14
Source File: breadth_first_search.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 5 votes vote down vote up
def bfs_tree(G, source, reverse=False):
    """Return an oriented tree constructed from of a breadth-first-search
    starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    Returns
    -------
    T: NetworkX DiGraph
       An oriented tree

    Examples
    --------
    >>> G = nx.Graph()
    >>> G.add_path([0,1,2])
    >>> print(list(nx.bfs_edges(G,0)))
    [(0, 1), (1, 2)]

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    T = nx.DiGraph()
    T.add_node(source)
    T.add_edges_from(bfs_edges(G,source,reverse=reverse))
    return T 
Example #15
Source File: network.py    From ditto with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def bfs_order(self, source="sourcebus"):
        start_node = self.graph[source]
        return set(nx.bfs_edges(self.graph, source)) 
Example #16
Source File: test_highlevel.py    From tskit with MIT License 5 votes vote down vote up
def verify_nx_algorithm_equivalence(self, tree, g):
        for root in tree.roots:
            self.assertTrue(nx.is_directed_acyclic_graph(g))

            # test descendants
            self.assertSetEqual(
                {u for u in tree.nodes() if tree.is_descendant(u, root)},
                set(nx.descendants(g, root)) | {root},
            )

            # test MRCA
            if tree.num_nodes < 20:
                for u, v in itertools.combinations(tree.nodes(), 2):
                    mrca = nx.lowest_common_ancestor(g, u, v)
                    if mrca is None:
                        mrca = -1
                    self.assertEqual(tree.mrca(u, v), mrca)

            # test node traversal modes
            self.assertEqual(
                list(tree.nodes(root=root, order="breadthfirst")),
                [root] + [v for u, v in nx.bfs_edges(g, root)],
            )
            self.assertEqual(
                list(tree.nodes(root=root, order="preorder")),
                list(nx.dfs_preorder_nodes(g, root)),
            ) 
Example #17
Source File: decomposers.py    From dwave-hybrid with Apache License 2.0 5 votes vote down vote up
def next(self, state, **runopts):
        CG = self.constraint_graph
        size = self.size
        constraints = self.constraints
        bqm = state.problem

        # get a random constraint to start with
        n = random.choice(list(CG.nodes))

        if len(constraints[n]) > size:
            raise NotImplementedError

        # starting from our constraint, do a breadth-first search adding constraints until our max
        # size is reached
        variables = set(constraints[n])
        for _, ci in nx.bfs_edges(CG, n):
            proposed = [v for v in constraints[ci] if v not in variables]
            if len(proposed) + len(variables) <= size:
                variables.union(proposed)
            if len(variables) == size:
                # can exit early
                break

        sample = state.samples.change_vartype(bqm.vartype).first.sample
        subbqm = bqm_induced_by(bqm, variables, sample)
        return state.updated(subproblem=subbqm) 
Example #18
Source File: breadth_first_search.py    From Carnets with BSD 3-Clause "New" or "Revised" License 4 votes vote down vote up
def bfs_edges(G, source, reverse=False, depth_limit=None):
    """Iterate over edges in a breadth-first-search starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    depth_limit : int, optional(default=len(G))
        Specify the maximum search depth

    Returns
    -------
    edges: generator
       A generator of edges in the breadth-first-search.

    Examples
    --------
    To get the edges in a breadth-first search::

        >>> G = nx.path_graph(3)
        >>> list(nx.bfs_edges(G, 0))
        [(0, 1), (1, 2)]
        >>> list(nx.bfs_edges(G, source=0, depth_limit=1))
        [(0, 1)]

    To get the nodes in a breadth-first search order::

        >>> G = nx.path_graph(3)
        >>> root = 2
        >>> edges = nx.bfs_edges(G, root)
        >>> nodes = [root] + [v for u, v in edges]
        >>> nodes
        [2, 1, 0]

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py.
    by D. Eppstein, July 2004. The modifications
    to allow depth limits based on the Wikipedia article
    "`Depth-limited-search`_".

    .. _Depth-limited-search: https://en.wikipedia.org/wiki/Depth-limited_search
    """
    if reverse and G.is_directed():
        successors = G.predecessors
    else:
        successors = G.neighbors
    # TODO In Python 3.3+, this should be `yield from ...`
    for e in generic_bfs_edges(G, source, successors, depth_limit):
        yield e 
Example #19
Source File: breadth_first_search.py    From Carnets with BSD 3-Clause "New" or "Revised" License 4 votes vote down vote up
def bfs_tree(G, source, reverse=False, depth_limit=None):
    """Returns an oriented tree constructed from of a breadth-first-search
    starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    depth_limit : int, optional(default=len(G))
        Specify the maximum search depth

    Returns
    -------
    T: NetworkX DiGraph
       An oriented tree

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(list(nx.bfs_tree(G,1).edges()))
    [(1, 0), (1, 2)]
    >>> H = nx.Graph()
    >>> nx.add_path(H, [0, 1, 2, 3, 4, 5, 6])
    >>> nx.add_path(H, [2, 7, 8, 9, 10])
    >>> print(sorted(list(nx.bfs_tree(H, source=3, depth_limit=3).edges())))
    [(1, 0), (2, 1), (2, 7), (3, 2), (3, 4), (4, 5), (5, 6), (7, 8)]


    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004. The modifications
    to allow depth limits based on the Wikipedia article
    "`Depth-limited-search`_".

    .. _Depth-limited-search: https://en.wikipedia.org/wiki/Depth-limited_search
    """
    T = nx.DiGraph()
    T.add_node(source)
    edges_gen = bfs_edges(G, source, reverse=reverse, depth_limit=depth_limit)
    T.add_edges_from(edges_gen)
    return T 
Example #20
Source File: breadth_first_search.py    From Carnets with BSD 3-Clause "New" or "Revised" License 4 votes vote down vote up
def bfs_predecessors(G, source, depth_limit=None):
    """Returns an iterator of predecessors in breadth-first-search from source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    depth_limit : int, optional(default=len(G))
        Specify the maximum search depth

    Returns
    -------
    pred: iterator
        (node, predecessors) iterator where predecessors is the list of
        predecessors of the node.

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(dict(nx.bfs_predecessors(G, 0)))
    {1: 0, 2: 1}
    >>> H = nx.Graph()
    >>> H.add_edges_from([(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)])
    >>> print(dict(nx.bfs_predecessors(H, 0)))
    {1: 0, 2: 0, 3: 1, 4: 1, 5: 2, 6: 2}
    >>> M = nx.Graph()
    >>> nx.add_path(M, [0, 1, 2, 3, 4, 5, 6])
    >>> nx.add_path(M, [2, 7, 8, 9, 10])
    >>> print(sorted(nx.bfs_predecessors(M, source=1, depth_limit=3)))
    [(0, 1), (2, 1), (3, 2), (4, 3), (7, 2), (8, 7)]


    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004. The modifications
    to allow depth limits based on the Wikipedia article
    "`Depth-limited-search`_".

    .. _Depth-limited-search: https://en.wikipedia.org/wiki/Depth-limited_search
    """
    for s, t in bfs_edges(G, source, depth_limit=depth_limit):
        yield (t, s) 
Example #21
Source File: breadth_first_search.py    From Carnets with BSD 3-Clause "New" or "Revised" License 4 votes vote down vote up
def bfs_successors(G, source, depth_limit=None):
    """Returns an iterator of successors in breadth-first-search from source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    depth_limit : int, optional(default=len(G))
        Specify the maximum search depth

    Returns
    -------
    succ: iterator
       (node, successors) iterator where successors is the list of
       successors of the node.

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(dict(nx.bfs_successors(G,0)))
    {0: [1], 1: [2]}
    >>> H = nx.Graph()
    >>> H.add_edges_from([(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)])
    >>> print(dict(nx.bfs_successors(H, 0)))
    {0: [1, 2], 1: [3, 4], 2: [5, 6]}
    >>> G = nx.Graph()
    >>> nx.add_path(G, [0, 1, 2, 3, 4, 5, 6])
    >>> nx.add_path(G, [2, 7, 8, 9, 10])
    >>> print(dict(nx.bfs_successors(G, source=1, depth_limit=3)))
    {1: [0, 2], 2: [3, 7], 3: [4], 7: [8]}


    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.The modifications
    to allow depth limits based on the Wikipedia article
    "`Depth-limited-search`_".

    .. _Depth-limited-search: https://en.wikipedia.org/wiki/Depth-limited_search
    """
    parent = source
    children = []
    for p, c in bfs_edges(G, source, depth_limit=depth_limit):
        if p == parent:
            children.append(c)
            continue
        yield (parent, children)
        children = [c]
        parent = p
    yield (parent, children) 
Example #22
Source File: breadth_first_search.py    From qgisSpaceSyntaxToolkit with GNU General Public License v3.0 4 votes vote down vote up
def bfs_edges(G, source, reverse=False):
    """Produce edges in a breadth-first-search starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    Returns
    -------
    edges: generator
       A generator of edges in the breadth-first-search.

    Examples
    --------
    >>> G = nx.Graph()
    >>> G.add_path([0,1,2])
    >>> print(list(nx.bfs_edges(G,0)))
    [(0, 1), (1, 2)]

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    if reverse and isinstance(G, nx.DiGraph):
        neighbors = G.predecessors_iter
    else:
        neighbors = G.neighbors_iter
    visited = set([source])
    queue = deque([(source, neighbors(source))])
    while queue:
        parent, children = queue[0]
        try:
            child = next(children)
            if child not in visited:
                yield parent, child
                visited.add(child)
                queue.append((child, neighbors(child)))
        except StopIteration:
            queue.popleft() 
Example #23
Source File: breadth_first_search.py    From aws-kube-codesuite with Apache License 2.0 4 votes vote down vote up
def generic_bfs_edges(G, source, neighbors=None):
    """Iterate over edges in a breadth-first search.

    The breadth-first search begins at `source` and enqueues the
    neighbors of newly visited nodes specified by the `neighbors`
    function.

    Parameters
    ----------
    G : NetworkX graph

    source : node
        Starting node for the breadth-first search; this function
        iterates over only those edges in the component reachable from
        this node.

    neighbors : function
        A function that takes a newly visited node of the graph as input
        and returns an *iterator* (not just a list) of nodes that are
        neighbors of that node. If not specified, this is just the
        ``G.neighbors`` method, but in general it can be any function
        that returns an iterator over some or all of the neighbors of a
        given node, in any order.

    Yields
    ------
    edge
        Edges in the breadth-first search starting from `source`.

    Examples
    --------
    >>> G = nx.path_graph(3)
    >>> print(list(nx.bfs_edges(G,0)))
    [(0, 1), (1, 2)]

    Notes
    -----
    This implementation is from `PADS`_, which was in the public domain
    when it was first accessed in July, 2004.

    .. _PADS: http://www.ics.uci.edu/~eppstein/PADS/BFS.py

    """
    visited = {source}
    queue = deque([(source, neighbors(source))])
    while queue:
        parent, children = queue[0]
        try:
            child = next(children)
            if child not in visited:
                yield parent, child
                visited.add(child)
                queue.append((child, neighbors(child)))
        except StopIteration:
            queue.popleft() 
Example #24
Source File: breadth_first_search.py    From aws-kube-codesuite with Apache License 2.0 4 votes vote down vote up
def bfs_edges(G, source, reverse=False):
    """Iterate over edges in a breadth-first-search starting at source.

    Parameters
    ----------
    G : NetworkX graph

    source : node
       Specify starting node for breadth-first search and return edges in
       the component reachable from source.

    reverse : bool, optional
       If True traverse a directed graph in the reverse direction

    Returns
    -------
    edges: generator
       A generator of edges in the breadth-first-search.

    Examples
    --------
    To get the edges in a breadth-first search::

        >>> G = nx.path_graph(3)
        >>> list(nx.bfs_edges(G, 0))
        [(0, 1), (1, 2)]

    To get the nodes in a breadth-first search order::

        >>> G = nx.path_graph(3)
        >>> root = 2
        >>> edges = nx.bfs_edges(G, root)
        >>> nodes = [root] + [v for u, v in edges]
        >>> nodes
        [2, 1, 0]

    Notes
    -----
    Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
    by D. Eppstein, July 2004.
    """
    if reverse and G.is_directed():
        successors = G.predecessors
    else:
        successors = G.neighbors
    # TODO In Python 3.3+, this should be `yield from ...`
    for e in generic_bfs_edges(G, source, successors):
        yield e 
Example #25
Source File: test_traversal.py    From dgl with Apache License 2.0 4 votes vote down vote up
def test_bfs(index_dtype, n=100):
    def _bfs_nx(g_nx, src):
        edges = nx.bfs_edges(g_nx, src)
        layers_nx = [set([src])]
        edges_nx = []
        frontier = set()
        edge_frontier = set()
        for u, v in edges:
            if u in layers_nx[-1]:
                frontier.add(v)
                edge_frontier.add(g.edge_id(u, v))
            else:
                layers_nx.append(frontier)
                edges_nx.append(edge_frontier)
                frontier = set([v])
                edge_frontier = set([g.edge_id(u, v)])
        # avoids empty successors
        if len(frontier) > 0 and len(edge_frontier) > 0:
            layers_nx.append(frontier)
            edges_nx.append(edge_frontier)
        return layers_nx, edges_nx

    g = dgl.DGLGraph()
    a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
    g.from_scipy_sparse_matrix(a)
    if index_dtype == 'int32':
        g = dgl.graph(g.edges()).int()
    else:
        g = dgl.graph(g.edges()).long()

    g_nx = g.to_networkx()
    src = random.choice(range(n))
    layers_nx, _ = _bfs_nx(g_nx, src)
    layers_dgl = dgl.bfs_nodes_generator(g, src)
    assert len(layers_dgl) == len(layers_nx)
    assert all(toset(x) == y for x, y in zip(layers_dgl, layers_nx))

    g_nx = nx.random_tree(n, seed=42)
    g = dgl.DGLGraph()
    g.from_networkx(g_nx)
    if index_dtype == 'int32':
        g = dgl.graph(g.edges()).int()
    else:
        g = dgl.graph(g.edges()).long()

    src = 0
    _, edges_nx = _bfs_nx(g_nx, src)
    edges_dgl = dgl.bfs_edges_generator(g, src)
    assert len(edges_dgl) == len(edges_nx)
    assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))