"""Node2vec random walk"""
from .._ffi.function import _init_api
from .. import backend as F
from .. import ndarray as nd
from .. import utils
# pylint: disable=invalid-name
__all__ = ['node2vec_random_walk']
[docs]def node2vec_random_walk(g, nodes, p, q, walk_length, prob=None, return_eids=False):
"""
Generate random walk traces from an array of starting nodes based on the node2vec model.
Paper: `node2vec: Scalable Feature Learning for Networks
<https://arxiv.org/abs/1607.00653>`__.
The returned traces all have length ``walk_length + 1``, where the first node
is the starting node itself.
Note that if a random walk stops in advance, DGL pads the trace with -1 to have the same
length.
Parameters
----------
g : DGLGraph
The graph. Must be on CPU.
Note that node2vec only support homogeneous graph.
nodes : Tensor
Node ID tensor from which the random walk traces starts.
The tensor must be on CPU, and must have the same dtype as the ID type
of the graph.
p: float
Likelihood of immediately revisiting a node in the walk.
q: float
Control parameter to interpolate between breadth-first strategy and depth-first strategy.
walk_length: int
Length of random walks.
prob : str, optional
The name of the edge feature tensor on the graph storing the (unnormalized)
probabilities associated with each edge for choosing the next node.
The feature tensor must be non-negative and the sum of the probabilities
must be positive for the outbound edges of all nodes (although they don't have
to sum up to one). The result will be undefined otherwise.
If omitted, DGL assumes that the neighbors are picked uniformly.
return_eids : bool, optional
If True, additionally return the edge IDs traversed.
Default: False.
Returns
-------
traces : Tensor
A 2-dimensional node ID tensor with shape ``(num_seeds, walk_length + 1)``.
eids : Tensor, optional
A 2-dimensional edge ID tensor with shape ``(num_seeds, length)``.
Only returned if :attr:`return_eids` is True.
Examples
--------
>>> g1 = dgl.graph(([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]))
>>> dgl.sampling.node2vec_random_walk(g1, [0, 1, 2, 0], 1, 1, walk_length=4)
tensor([[0, 1, 3, 0, 1],
[1, 2, 0, 1, 3],
[2, 0, 1, 3, 0],
[0, 1, 2, 0, 1]])
>>> dgl.sampling.node2vec_random_walk(g1, [0, 1, 2, 0], 1, 1, walk_length=4, return_eids=True)
(tensor([[0, 1, 3, 0, 1],
[1, 2, 0, 1, 2],
[2, 0, 1, 2, 0],
[0, 1, 2, 0, 1]]),
tensor([[0, 2, 4, 0],
[1, 3, 0, 1],
[3, 0, 1, 3],
[0, 1, 3, 0]]))
"""
assert g.device == F.cpu(), "Graph must be on CPU."
gidx = g._graph
nodes = F.to_dgl_nd(utils.prepare_tensor(g, nodes, 'nodes'))
if prob is None:
prob_nd = nd.array([], ctx=nodes.ctx)
else:
prob_nd = F.to_dgl_nd(g.edata[prob])
traces, eids = _CAPI_DGLSamplingNode2vec(gidx, nodes, p, q, walk_length, prob_nd)
traces = F.from_dgl_nd(traces)
eids = F.from_dgl_nd(eids)
return (traces, eids) if return_eids else traces
_init_api('dgl.sampling.randomwalks', __name__)