def level_order(forest, roots):
try:
edges = bfs_edges_generator(forest, roots)
_, leaves = forest.find_edges(edges[-1])
edges_back = bfs_edges_generator(forest, roots, reverse=True)
yield from reversed(edges_back)
yield from edges
except:
return None
def level_order(forest, roots):
device = forest.device
edges = list(bfs_edges_generator(forest, roots))
edges = [e.to(device) for e in edges]
_, leaves = forest.find_edges(edges[-1])
edges_back = list(bfs_edges_generator(forest, roots, reverse=True))
edges_back = [e.to(device) for e in edges_back]
yield from reversed(edges_back)
yield from edges
def level_order(forest, roots):
forest = tocpu(forest)
edges = bfs_edges_generator(forest, roots)
if len(edges) == 0:
# no edges in the tree; do not perform loopy BP
return
_, leaves = forest.find_edges(edges[-1])
edges_back = bfs_edges_generator(forest, roots, reverse=True)
yield from reversed(edges_back)
yield from edges
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))
