Function random_walk_pe

python/dgl/transforms/functional.py:3553–3611

r"""Random Walk Positional Encoding, as introduced in `Graph Neural Networks with Learnable Structural and Positional Representations `__ This function computes the random walk positional encodings as landing probabilities from 1-step to k-step, sta

(g, k, eweight_name=None)

def random_walk_pe(g, k, eweight_name=None):
    r"""Random Walk Positional Encoding, as introduced in
    `Graph Neural Networks with Learnable Structural and Positional Representations
    <https://arxiv.org/abs/2110.07875>`__

    This function computes the random walk positional encodings as landing probabilities
    from 1-step to k-step, starting from each node to itself.

    Parameters
    ----------
    g : DGLGraph
        The input graph. Must be homogeneous.
    k : int
        The number of random walk steps. The paper found the best value to be 16 and 20
        for two experiments.
    eweight_name : str, optional
        The name to retrieve the edge weights. Default: None, not using the edge weights.

    Returns
    -------
    Tensor
        The random walk positional encodings of shape :math:`(N, k)`, where :math:`N` is the
        number of nodes in the input graph.

    Example
    -------
    >>> import dgl
    >>> g = dgl.graph(([0,1,1], [1,1,0]))
    >>> dgl.random_walk_pe(g, 2)
    tensor([[0.0000, 0.5000],
            [0.5000, 0.7500]])
    """
    N = g.num_nodes()  # number of nodes
    M = g.num_edges()  # number of edges
    A = g.adj_external(scipy_fmt="csr")  # adjacency matrix
    if eweight_name is not None:
        # add edge weights if required
        W = sparse.csr_matrix(
            (g.edata[eweight_name].squeeze(), g.find_edges(list(range(M)))),
            shape=(N, N),
        )
        A = A.multiply(W)
    # 1-step transition probability
    if version.parse(scipy.__version__) < version.parse("1.11.0"):
        RW = np.array(A / (A.sum(1) + 1e-30))
    else:
        # Sparse matrix divided by a dense array returns a sparse matrix in
        # scipy since 1.11.0.
        RW = (A / (A.sum(1) + 1e-30)).toarray()

    # Iterate for k steps
    PE = [F.astype(F.tensor(np.array(RW.diagonal())), F.float32)]
    RW_power = RW
    for _ in range(k - 1):
        RW_power = RW_power @ RW
        PE.append(F.astype(F.tensor(np.array(RW_power.diagonal())), F.float32))
    PE = F.stack(PE, dim=-1)