[Transform] Add add_reverse_edges (#1936)

* Add add_reverse_edges and is_simple_graph
Update to_bidirected

* remove is_simple_graph

* Fix lint

* to bidirected only support cpu

* support copy ndata

* upd
Co-authored-by: Ubuntu <ubuntu@ip-172-31-51-214.ec2.internal>

python/dgl/partition.py

@@ -177,7 +177,7 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False):
     # The METIS runs on the symmetric graph to generate the node assignment to partitions.
     from .transform import to_bidirected # avoid cyclic import
     start = time.time()
-    sym_g = to_bidirected(g, copy_ndata=False)
+    sym_g = to_bidirected(g)
     print('Convert a graph into a bidirected graph: {:.3f} seconds'.format(
         time.time() - start))
     vwgt = []

python/dgl/transform.py

@@ -28,6 +28,7 @@ __all__ = [
     'reverse',
     'to_bidirected',
     'to_bidirected_stale',
+    'add_reverse_edges',
     'laplacian_lambda_max',
     'knn_graph',
     'segmented_knn_graph',
@@ -147,19 +148,99 @@ def segmented_knn_graph(x, k, segs):
     g = convert.graph(adj)
     return g
 
-def to_bidirected(g, readonly=None, copy_ndata=True,
+def to_bidirected(g, readonly=None, copy_ndata=False):
+    r""" Convert the graph to a bidirected one.
+
+    The function generates a new graph with no edge features.
+    If g has an edge for i->j but no edge for j->i, then the
+    returned graph will have both i->j and j->i.
+
+    For a heterograph with multiple edge types, we can treat edges corresponding
+    to each type as a separate graph and convert the graph to a bidirected one
+    for each of them.
+
+    Since **to_bidirected is not well defined for unidirectional bipartite graphs**,
+    an error will be raised if an edge type of the input heterograph is for a
+    unidirectional bipartite graph.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The input graph.
+    readonly : bool, default to be True
+        Deprecated. There will be no difference between readonly and non-readonly
+    copy_ndata: bool, optional
+        If True, the node features of the bidirected graph are copied from the
+        original graph. If False, the bidirected graph will not have any node features.
+        (Default: False)
+
+    Notes
+    -----
+    Please make sure g is a single graph.
+
+    Returns
+    -------
+    dgl.DGLGraph
+        The bidirected graph
+
+    Examples
+    --------
+    The following examples use PyTorch backend.
+
+    >>> import dgl
+    >>> import torch as th
+    >>> g = dgl.graph((th.tensor([0, 1, 2]), th.tensor([1, 2, 0])))
+    >>> bg1 = dgl.to_bidirected(g)
+    >>> bg1.edges()
+    (tensor([0, 1, 2, 1, 2, 0]), tensor([1, 2, 0, 0, 1, 2]))
+
+    The graph already have i->j and j->i
+
+    >>> g = dgl.graph((th.tensor([0, 1, 2, 0]), th.tensor([1, 2, 0, 2])))
+    >>> bg1 = dgl.to_bidirected(g)
+    >>> bg1.edges()
+    (tensor([0, 1, 2, 1, 2, 0]), tensor([1, 2, 0, 0, 1, 2]))
+
+    **Heterographs with Multiple Edge Types**
+
+    >>> g = dgl.heterograph({
+    >>>     ('user', 'wins', 'user'): (th.tensor([0, 2, 0, 2]), th.tensor([1, 1, 2, 0])),
+    >>>     ('user', 'follows', 'user'): (th.tensor([1, 2, 1]), th.tensor([2, 1, 1]))
+    >>> })
+    >>> bg1 = dgl.to_bidirected(g)
+    >>> bg1.edges(etype='wins')
+    (tensor([0, 0, 1, 1, 2, 2]), tensor([1, 2, 0, 2, 0, 1]))
+    >>> bg1.edges(etype='follows')
+    (tensor([1, 1, 2]), tensor([1, 2, 1]))
+    """
+    if readonly is not None:
+        dgl_warning("Parameter readonly is deprecated" \
+            "There will be no difference between readonly and non-readonly DGLGraph")
+
+    for c_etype in g.canonical_etypes:
+        if c_etype[0] != c_etype[2]:
+            assert False, "to_bidirected is not well defined for " \
+                "unidirectional bipartite graphs" \
+                ", but {} is unidirectional bipartite".format(c_etype)
+
+    assert g.is_multigraph is False, "to_bidirected only support simple graph"
+
+    g = add_reverse_edges(g, copy_ndata=copy_ndata, copy_edata=False)
+    g = to_simple(g, return_counts=None, copy_ndata=copy_ndata, copy_edata=False)
+    return g
+
+def add_reverse_edges(g, readonly=None, copy_ndata=True,
                       copy_edata=False, ignore_bipartite=False):
-    r"""Convert the graph to a bidirected one.
+    r"""Add reverse edges to a graph
 
     For a graph with edges :math:`(i_1, j_1), \cdots, (i_n, j_n)`, this
     function creates a new graph with edges
     :math:`(i_1, j_1), \cdots, (i_n, j_n), (j_1, i_1), \cdots, (j_n, i_n)`.
 
     For a heterograph with multiple edge types, we can treat edges corresponding
-    to each type as a separate graph and convert the graph to a bidirected one
-    for each of them.
+    to each type as a separate graph and add reverse edges for each of them.
 
-    Since **to_bidirected is not well defined for unidirectional bipartite graphs**,
+    Since **add_reverse_edges is not well defined for unidirectional bipartite graphs**,
     an error will be raised if an edge type of the input heterograph is for a
     unidirectional bipartite graph. We can simply skip the edge types corresponding
     to unidirectional bipartite graphs by specifying ``ignore_bipartite=True``.
@@ -171,14 +252,14 @@ def to_bidirected(g, readonly=None, copy_ndata=True,
     readonly : bool, default to be True
         Deprecated. There will be no difference between readonly and non-readonly
     copy_ndata: bool, optional
-        If True, the node features of the bidirected graph are copied from
-        the original graph. If False, the bidirected
-        graph will not have any node features.
+        If True, the node features of the new graph are copied from
+        the original graph. If False, the new graph will not have any
+        node features.
         (Default: True)
     copy_edata: bool, optional
         If True, the features of the reversed edges will be identical to
         the original ones."
-        If False, the bidirected graph will not have any edge
+        If False, the new graph will not have any edge
         features.
         (Default: False)
     ignore_bipartite: bool, optional
@@ -189,8 +270,8 @@ def to_bidirected(g, readonly=None, copy_ndata=True,
 
     Returns
     -------
-    DGLGraph
-        The bidirected graph
+    dgl.DGLGraph
+        The graph with reversed edges added.
 
     Notes
     -----
@@ -208,7 +289,7 @@ def to_bidirected(g, readonly=None, copy_ndata=True,
     **Homographs**
 
     >>> g = dgl.graph(th.tensor([0, 0]), th.tensor([0, 1]))
-    >>> bg1 = dgl.to_bidirected(g)
+    >>> bg1 = dgl.add_reverse_edges(g)
     >>> bg1.edges()
     (tensor([0, 0, 0, 1]), tensor([0, 1, 0, 0]))
 
@@ -224,14 +305,14 @@ def to_bidirected(g, readonly=None, copy_ndata=True,
     >>> g.nodes['game'].data['hv'] = th.ones(3, 1)
     >>> g.edges['wins'].data['h'] = th.tensor([0, 1, 2, 3, 4])
 
-    The to_bidirected operation is applied to the subgraph
+    The add_reverse_edges operation is applied to the subgraph
     corresponding to ('user', 'wins', 'user') and the
     subgraph corresponding to ('user', 'follows', 'user).
     The unidirectional bipartite subgraph ('user', 'plays', 'game')
     is ignored. Both the node features and edge features
     are shared.
 
-    >>> bg = dgl.to_bidirected(g, copy_ndata=True,
+    >>> bg = dgl.add_reverse_edges(g, copy_ndata=True,
                                copy_edata=True, ignore_bipartite=True)
     >>> bg.edges(('user', 'wins', 'user'))
     (tensor([0, 2, 0, 2, 2, 1, 1, 2, 1, 0]), tensor([1, 1, 2, 1, 0, 0, 2, 0, 2, 2]))
@@ -254,7 +335,7 @@ def to_bidirected(g, readonly=None, copy_ndata=True,
         subgs = {}
         for c_etype in canonical_etypes:
             if c_etype[0] != c_etype[2]:
-                assert False, "to_bidirected is not well defined for " \
+                assert False, "add_reverse_edges is not well defined for " \
                     "unidirectional bipartite graphs" \
                     ", but {} is unidirectional bipartite".format(c_etype)
 

tests/compute/test_transform.py

@@ -234,12 +234,62 @@ def test_reverse_shared_frames(idtype):
     assert F.allclose(g.edges[[0, 2], [1, 1]].data['h'],
                       rg.edges[[1, 1], [0, 2]].data['h'])
 
+@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
 def test_to_bidirected():
+    # homogeneous graph
+    elist = [(0, 0), (0, 1), (1, 0),
+                (1, 1), (2, 1), (2, 2)]
+    num_edges = 7
+    g = dgl.graph(elist)
+    elist.append((1, 2))
+    elist = set(elist)
+    big = dgl.to_bidirected(g)
+    assert big.number_of_edges() == num_edges
+    src, dst = big.edges()
+    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
+    assert eset == set(elist)
+
+    # heterogeneous graph
+    elist1 = [(0, 0), (0, 1), (1, 0),
+                (1, 1), (2, 1), (2, 2)]
+    elist2 = [(0, 0), (0, 1)]
+    g = dgl.heterograph({
+        ('user', 'wins', 'user'): elist1,
+        ('user', 'follows', 'user'): elist2
+    })
+    g.nodes['user'].data['h'] = F.ones((3, 1))
+    elist1.append((1, 2))
+    elist1 = set(elist1)
+    elist2.append((1, 0))
+    elist2 = set(elist2)
+    big = dgl.to_bidirected(g)
+    assert big.number_of_edges('wins') == 7
+    assert big.number_of_edges('follows') == 3
+    src, dst = big.edges(etype='wins')
+    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
+    assert eset == set(elist1)
+    src, dst = big.edges(etype='follows')
+    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
+    assert eset == set(elist2)
+
+    big = dgl.to_bidirected(g, copy_ndata=True)
+    assert F.array_equal(g.nodes['user'].data['h'], big.nodes['user'].data['h'])
+
+    # test multigraph
+    g = dgl.graph((F.tensor([0, 1, 3, 1]), F.tensor([1, 2, 0, 2])))
+    raise_error = False
+    try:
+        big = dgl.to_bidirected(g)
+    except:
+        raise_error = True
+    assert raise_error
+
+def test_add_reverse_edges():
     # homogeneous graph
     g = dgl.graph((F.tensor([0, 1, 3, 1]), F.tensor([1, 2, 0, 2])))
     g.ndata['h'] = F.tensor([[0.], [1.], [2.], [1.]])
     g.edata['h'] = F.tensor([[3.], [4.], [5.], [6.]])
-    bg = dgl.to_bidirected(g, copy_ndata=True, copy_edata=True)
+    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True)
     u, v = g.edges()
     ub, vb = bg.edges()
     assert F.array_equal(F.cat([u, v], dim=0), ub)
@@ -252,7 +302,7 @@ def test_to_bidirected():
     assert ('hh' in g.edata) is False
 
     # donot share ndata and edata
-    bg = dgl.to_bidirected(g, copy_ndata=False, copy_edata=False)
+    bg = dgl.add_reverse_edges(g, copy_ndata=False, copy_edata=False)
     ub, vb = bg.edges()
     assert F.array_equal(F.cat([u, v], dim=0), ub)
     assert F.array_equal(F.cat([v, u], dim=0), vb)
@@ -261,7 +311,7 @@ def test_to_bidirected():
 
     # zero edge graph
     g = dgl.graph([])
-    bg = dgl.to_bidirected(g, copy_ndata=True, copy_edata=True)
+    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True)
 
     # heterogeneous graph
     g = dgl.heterograph({
@@ -272,7 +322,7 @@ def test_to_bidirected():
     g.nodes['game'].data['hv'] = F.ones((3, 1))
     g.nodes['user'].data['hv'] = F.ones((3, 1))
     g.edges['wins'].data['h'] = F.tensor([0, 1, 2, 3, 4])
-    bg = dgl.to_bidirected(g, copy_ndata=True, copy_edata=True, ignore_bipartite=True)
+    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True, ignore_bipartite=True)
     assert F.array_equal(g.nodes['game'].data['hv'], bg.nodes['game'].data['hv'])
     assert F.array_equal(g.nodes['user'].data['hv'], bg.nodes['user'].data['hv'])
     u, v = g.all_edges(order='eid', etype=('user', 'wins', 'user'))
@@ -293,7 +343,7 @@ def test_to_bidirected():
     assert len(bg.edges['follows'].data) == 0
 
     # donot share ndata and edata
-    bg = dgl.to_bidirected(g, copy_ndata=False, copy_edata=False, ignore_bipartite=True)
+    bg = dgl.add_reverse_edges(g, copy_ndata=False, copy_edata=False, ignore_bipartite=True)
     assert len(bg.edges['wins'].data) == 0
     assert len(bg.edges['plays'].data) == 0
     assert len(bg.edges['follows'].data) == 0