wehub-resource-sync
132 lines
4.2 KiB
Python
132 lines
4.2 KiB
Python
"""Module for various graph generator functions."""
|
|
|
|
from . import backend as F, convert, random
|
|
|
|
__all__ = ["rand_graph", "rand_bipartite"]
|
|
|
|
|
|
def rand_graph(num_nodes, num_edges, idtype=F.int64, device=F.cpu()):
|
|
"""Generate a random graph of the given number of nodes/edges and return.
|
|
|
|
It uniformly chooses ``num_edges`` from all possible node pairs and form a graph.
|
|
The random choice is without replacement, which means there will be no multi-edge
|
|
in the resulting graph.
|
|
|
|
To control the randomness, set the random seed via :func:`dgl.seed`.
|
|
|
|
Parameters
|
|
----------
|
|
num_nodes : int
|
|
The number of nodes
|
|
num_edges : int
|
|
The number of edges
|
|
idtype : int32, int64, optional
|
|
The data type for storing the structure-related graph information
|
|
such as node and edge IDs. It should be a framework-specific data type object
|
|
(e.g., torch.int32). By default, DGL uses int64.
|
|
device : Device context, optional
|
|
The device of the resulting graph. It should be a framework-specific device
|
|
object (e.g., torch.device). By default, DGL stores the graph on CPU.
|
|
|
|
Returns
|
|
-------
|
|
DGLGraph
|
|
The generated random graph.
|
|
|
|
See Also
|
|
--------
|
|
rand_bipartite
|
|
|
|
Examples
|
|
--------
|
|
>>> import dgl
|
|
>>> dgl.rand_graph(100, 10)
|
|
Graph(num_nodes=100, num_edges=10,
|
|
ndata_schemes={}
|
|
edata_schemes={})
|
|
"""
|
|
# TODO(minjie): support RNG as one of the arguments.
|
|
eids = random.choice(num_nodes * num_nodes, num_edges, replace=False)
|
|
eids = F.zerocopy_to_numpy(eids)
|
|
rows = F.zerocopy_from_numpy(eids // num_nodes)
|
|
cols = F.zerocopy_from_numpy(eids % num_nodes)
|
|
rows = F.copy_to(F.astype(rows, idtype), device)
|
|
cols = F.copy_to(F.astype(cols, idtype), device)
|
|
return convert.graph(
|
|
(rows, cols), num_nodes=num_nodes, idtype=idtype, device=device
|
|
)
|
|
|
|
|
|
def rand_bipartite(
|
|
utype,
|
|
etype,
|
|
vtype,
|
|
num_src_nodes,
|
|
num_dst_nodes,
|
|
num_edges,
|
|
idtype=F.int64,
|
|
device=F.cpu(),
|
|
):
|
|
"""Generate a random uni-directional bipartite graph and return.
|
|
|
|
It uniformly chooses ``num_edges`` from all possible node pairs and form a graph.
|
|
The random choice is without replacement, which means there will be no multi-edge
|
|
in the resulting graph.
|
|
|
|
To control the randomness, set the random seed via :func:`dgl.seed`.
|
|
|
|
Parameters
|
|
----------
|
|
utype : str, optional
|
|
The name of the source node type.
|
|
etype : str, optional
|
|
The name of the edge type.
|
|
vtype : str, optional
|
|
The name of the destination node type.
|
|
num_src_nodes : int
|
|
The number of source nodes.
|
|
num_dst_nodes : int
|
|
The number of destination nodes.
|
|
num_edges : int
|
|
The number of edges
|
|
idtype : int32, int64, optional
|
|
The data type for storing the structure-related graph information
|
|
such as node and edge IDs. It should be a framework-specific data type object
|
|
(e.g., torch.int32). By default, DGL uses int64.
|
|
device : Device context, optional
|
|
The device of the resulting graph. It should be a framework-specific device
|
|
object (e.g., torch.device). By default, DGL stores the graph on CPU.
|
|
|
|
Returns
|
|
-------
|
|
DGLGraph
|
|
The generated random bipartite graph.
|
|
|
|
See Also
|
|
--------
|
|
rand_graph
|
|
|
|
Examples
|
|
--------
|
|
>>> import dgl
|
|
>>> dgl.rand_bipartite('user', 'buys', 'game', 50, 100, 10)
|
|
Graph(num_nodes={'game': 100, 'user': 50},
|
|
num_edges={('user', 'buys', 'game'): 10},
|
|
metagraph=[('user', 'game', 'buys')])
|
|
"""
|
|
# TODO(minjie): support RNG as one of the arguments.
|
|
eids = random.choice(
|
|
num_src_nodes * num_dst_nodes, num_edges, replace=False
|
|
)
|
|
eids = F.zerocopy_to_numpy(eids)
|
|
rows = F.zerocopy_from_numpy(eids // num_dst_nodes)
|
|
cols = F.zerocopy_from_numpy(eids % num_dst_nodes)
|
|
rows = F.copy_to(F.astype(rows, idtype), device)
|
|
cols = F.copy_to(F.astype(cols, idtype), device)
|
|
return convert.heterograph(
|
|
{(utype, etype, vtype): (rows, cols)},
|
|
{utype: num_src_nodes, vtype: num_dst_nodes},
|
|
idtype=idtype,
|
|
device=device,
|
|
)
|