[Doc] refine doc and optimize metis part for dgl.reorder() (#3090)

Co-authored-by: Tong He <hetong007@gmail.com>

python/dgl/transform.py

@@ -2895,48 +2895,55 @@ def sort_in_edges(g, tag, tag_offset_name='_TAG_OFFSET'):
 
 def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
     r"""Return a new graph which re-order and re-label the nodes
-        according to the specified permute algo.
+    according to the specified permute algorithm.
 
-    Homogeneous graph is supported only.
+    Support homogeneous graph only for the moment.
 
-    This API is basically implemented by leveraging dgl.node_subgraph(),
+    This API is basically implemented by leveraging :func:`~dgl.node_subgraph`,
     so the function signature is similar and raw IDs could be stored
-    in dgl.NID and dgl.EID.
+    in ``dgl.NID`` and ``dgl.EID``.
+
+    Please note that edges are re-ordered/re-labeled according to re-ordered
+    ``'src'`` nodes. This behavior is realized in :func:`dgl.node_subgraph`.
+    What's more, if user wants to re-order/re-label according to ``'dst'`` nodes
+    or any other algorithms, please use :func:`dgl.edge_subgraph` with new edge
+    permutation.
 
     Parameters
     ----------
     g : DGLGraph
         The homogeneous graph.
     permute_algo: str, optional
-        can be ``'rcmk'`` or ``'metis'`` or ``'custom'``. ``'rcmk'`` is the default algo.
-        * ``'rcmk'``: The Reverse Cuthill–McKee algorithm is an algorithm to permute
-        a sparse matrix that has a symmetric sparsity pattern into a band matrix form
-        with a small bandwidth. The resulting index numbers is reversed.
-        * ``'metis'``: METIS is a set of serial algorithms for partitioning graphs,
-        partitioning finite element meshes, and producing fill reducing orderings
-        for sparse matrices. This algorithm has already available in DGL:
-        ``'dgl.partition.metis_partition_assignment'``.
+        can be ``'rcmk'`` or ``'metis'`` or ``'custom'``. ``'rcmk'`` is the default value.
+
+        * ``'rcmk'``: Call `Reverse Cuthill–McKee <https://docs.scipy.org/doc/scipy/reference/
+          generated/scipy.sparse.csgraph.reverse_cuthill_mckee.html#
+          scipy-sparse-csgraph-reverse-cuthill-mckee>`__ from ``'scipy'`` to generate nodes
+          permutation and pass it into :func:`~dgl.node_subgraph` to generate new graph.
+        * ``'metis'``: Call :func:`~dgl.partition.metis_partition_assignment` from ``'DGL'``
+          to generate nodes permutation and pass it into :func:`~dgl.node_subgraph` to generate
+          new graph.
         * ``'custom'``: This enables user to pass in self-designed reorder algorithm.
-        User should pass in ``'nodes_perm'`` via another argument ``'permute_config'`` with
-        ``'custom'`` is specified here. By this way, can the graph be reordered according to
-        passed in nodes permutation.
+          User should pass in ``'nodes_perm'`` via another argument :attr:`permute_config` with
+          ``'custom'`` is specified here. By this way, can the graph be reordered according to
+          passed in nodes permutation.
     store_ids: bool, optional
-        It's passed into dgl.node_subgraph(). If True, it will store
-        the raw IDs of the extracted nodes and edges in the ndata
-        and edata of the resulting graph under name dgl.NID and
-        dgl.EID, respectively.
+        It is passed into :func:`~dgl.node_subgraph()`. If True, it will store
+        the raw IDs of the extracted nodes and edges in the ndata and edata of
+        the resulting graph under name ``'dgl.NID'`` and ``'dgl.EID'``, respectively.
     permute_config: dict, optional
-        additional config data for specified permute_algo.
-        * for ``'rcmk'``, this argument is not required.
-        * for ``'metis'``, partition part number ``'k'`` is required and specified in this
-        argument like this: {'k':10}.
-        * for ``'custom'``, ``'nodes_perm'`` should be specified in this argument like this:
-        {'nodes_perm':[1,2,3,0]}.
+        Additional config data for specified :attr:`permute_algo`.
 
-    Return
-    ------
+        * For ``'rcmk'``, this argument is not required.
+        * For ``'metis'``, partition part number ``'k'`` is required and specified in this
+          argument like this: {'k':10}.
+        * For ``'custom'``, ``'nodes_perm'`` should be specified in the format of
+          ``'Int Tensor'`` or ``'iterable[int]'`` like :attr:`nodes` in :func:`~dgl.node_subgraph`.
+
+    Returns
+    -------
     DGLGraph
-        The re-ordered graph
+        The re-ordered graph.
 
     Examples
     --------
@@ -2958,8 +2965,7 @@ def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
             [3],
             [4]])}
 
-    Reorder according to 'rcmk' permute_algo which is implemented in
-    scipy.sparse.csgraph.reverse_cuthill_mckee().
+    Reorder according to ``'rcmk'`` permute algorithm.
 
     >>> rg = dgl.reorder(g)
     >>> rg.ndata
@@ -2975,8 +2981,7 @@ def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
             [2],
             [0]]), '_ID': tensor([4, 3, 1, 2, 0])}
 
-    Reorder with according to 'metis' permute_algo which is implemented in
-    dgl.partition.metis_partition_assignment().
+    Reorder with according to ``'metis'`` permute algorithm.
 
     >>> rg = dgl.reorder(g, 'metis', permute_config={'k':2})
     >>> rg.ndata
@@ -2992,7 +2997,7 @@ def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
             [4],
             [3]]), '_ID': tensor([2, 1, 0, 4, 3])}
 
-    Reorder according to 'custom' permute_algo with user-provided nodes_perm.
+    Reorder according to ``'custom'`` permute algorithm with user-provided nodes_perm.
 
     >>> nodes_perm = torch.randperm(g.num_nodes())
     >>> nodes_perm
@@ -3019,12 +3024,12 @@ def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
         raise DGLError("Unexpected permute_algo is specified: {}. Expected algos: {}".format(
             permute_algo, expected_algo))
     if permute_algo == 'rcmk':
-        nodes_perm = RCMKPerm(g)
+        nodes_perm = rcmk_perm(g)
     elif permute_algo == 'metis':
         if permute_config is None or 'k' not in permute_config:
             raise DGLError(
                 "Partition parts 'k' is required for metis. Please specify in permute_config.")
-        nodes_perm = METISPerm(g, permute_config['k'])
+        nodes_perm = metis_perm(g, permute_config['k'])
     else:
         if permute_config is None or 'nodes_perm' not in permute_config:
             raise DGLError(
@@ -3040,30 +3045,43 @@ def reorder(g, permute_algo='rcmk', store_ids=True, permute_config=None):
 DGLHeteroGraph.reorder = utils.alias_func(reorder)
 
 
-def METISPerm(g, k):
-    """
+def metis_perm(g, k):
+    r"""Return nodes permutation according to ``'metis'`` algorithm.
+
     For internal use.
-    g: graph
-    k: partition parts number
-    return: permutation of node ids via metis partition and assignment
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The homogeneous graph.
+    k: int
+        The partition parts number.
+
+    Returns
+    -------
+    iterable[int]
+        The nodes permutation.
     """
     pids = metis_partition_assignment(
         g if g.device == F.cpu() else g.to(F.cpu()), k)
     pids = F.asnumpy(pids)
-    perm = np.zeros(pids.shape, np.int64)
-    bincnt = np.bincount(pids)
-    idcnt = np.cumsum(bincnt)
-    for i, e in enumerate(pids):
-        idcnt[e] -= 1
-        perm[idcnt[e]] = i
-    return perm
+    return np.argsort(pids).copy()
 
 
-def RCMKPerm(g):
-    """
+def rcmk_perm(g):
+    r"""Return nodes permutation according to ``'rcmk'`` algorithm.
+
     For internal use.
-    g: graph
-    return: permutation of node ids via RCMK algorithm
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The homogeneous graph.
+
+    Returns
+    -------
+    iterable[int]
+        The nodes permutation.
     """
     fmat = 'csr'
     allowed_fmats = sum(g.formats().values(), [])