Source code for torch_geometric.utils.random

import warnings
from typing import List, Union

import numpy as np
import torch

from torch_geometric.utils import remove_self_loops, to_undirected


[docs]def erdos_renyi_graph( num_nodes: int, edge_prob: float, directed: bool = False, ) -> torch.Tensor: r"""Returns the :obj:`edge_index` of a random Erdos-Renyi graph. Args: num_nodes (int): The number of nodes. edge_prob (float): Probability of an edge. directed (bool, optional): If set to :obj:`True`, will return a directed graph. (default: :obj:`False`) Examples: >>> erdos_renyi_graph(5, 0.2, directed=False) tensor([[0, 1, 1, 4], [1, 0, 4, 1]]) >>> erdos_renyi_graph(5, 0.2, directed=True) tensor([[0, 1, 3, 3, 4, 4], [4, 3, 1, 2, 1, 3]]) """ if directed: idx = torch.arange((num_nodes - 1) * num_nodes) idx = idx.view(num_nodes - 1, num_nodes) idx = idx + torch.arange(1, num_nodes).view(-1, 1) idx = idx.view(-1) else: warnings.filterwarnings('ignore', '.*pass the indexing argument.*') idx = torch.combinations(torch.arange(num_nodes), r=2) # Filter edges. mask = torch.rand(idx.size(0)) < edge_prob idx = idx[mask] if directed: row = idx.div(num_nodes, rounding_mode='floor') col = idx % num_nodes edge_index = torch.stack([row, col], dim=0) else: edge_index = to_undirected(idx.t(), num_nodes=num_nodes) return edge_index
[docs]def stochastic_blockmodel_graph( block_sizes: Union[List[int], torch.Tensor], edge_probs: Union[List[List[float]], torch.Tensor], directed: bool = False, ) -> torch.Tensor: r"""Returns the :obj:`edge_index` of a stochastic blockmodel graph. Args: block_sizes ([int] or LongTensor): The sizes of blocks. edge_probs ([[float]] or FloatTensor): The density of edges going from each block to each other block. Must be symmetric if the graph is undirected. directed (bool, optional): If set to :obj:`True`, will return a directed graph. (default: :obj:`False`) Examples: >>> block_sizes = [2, 2, 4] >>> edge_probs = [[0.25, 0.05, 0.02], ... [0.05, 0.35, 0.07], ... [0.02, 0.07, 0.40]] >>> stochastic_blockmodel_graph(block_sizes, edge_probs, ... directed=False) tensor([[2, 4, 4, 5, 5, 6, 7, 7], [5, 6, 7, 2, 7, 4, 4, 5]]) >>> stochastic_blockmodel_graph(block_sizes, edge_probs, ... directed=True) tensor([[0, 2, 3, 4, 4, 5, 5], [3, 4, 1, 5, 6, 6, 7]]) """ size, prob = block_sizes, edge_probs if not isinstance(size, torch.Tensor): size = torch.tensor(size, dtype=torch.long) if not isinstance(prob, torch.Tensor): prob = torch.tensor(prob, dtype=torch.float) assert size.dim() == 1 assert prob.dim() == 2 and prob.size(0) == prob.size(1) assert size.size(0) == prob.size(0) if not directed: assert torch.allclose(prob, prob.t()) node_idx = torch.cat([size.new_full((b, ), i) for i, b in enumerate(size)]) num_nodes = node_idx.size(0) if directed: idx = torch.arange((num_nodes - 1) * num_nodes) idx = idx.view(num_nodes - 1, num_nodes) idx = idx + torch.arange(1, num_nodes).view(-1, 1) idx = idx.view(-1) row = idx.div(num_nodes, rounding_mode='floor') col = idx % num_nodes else: row, col = torch.combinations(torch.arange(num_nodes), r=2).t() mask = torch.bernoulli(prob[node_idx[row], node_idx[col]]).to(torch.bool) edge_index = torch.stack([row[mask], col[mask]], dim=0) if not directed: edge_index = to_undirected(edge_index, num_nodes=num_nodes) return edge_index
[docs]def barabasi_albert_graph(num_nodes: int, num_edges: int) -> torch.Tensor: r"""Returns the :obj:`edge_index` of a Barabasi-Albert preferential attachment model, where a graph of :obj:`num_nodes` nodes grows by attaching new nodes with :obj:`num_edges` edges that are preferentially attached to existing nodes with high degree. Args: num_nodes (int): The number of nodes. num_edges (int): The number of edges from a new node to existing nodes. Example: >>> barabasi_albert_graph(num_nodes=4, num_edges=3) tensor([[0, 0, 0, 1, 1, 2, 2, 3], [1, 2, 3, 0, 2, 0, 1, 0]]) """ assert num_edges > 0 and num_edges < num_nodes row, col = torch.arange(num_edges), torch.randperm(num_edges) for i in range(num_edges, num_nodes): row = torch.cat([row, torch.full((num_edges, ), i, dtype=torch.long)]) choice = np.random.choice(torch.cat([row, col]).numpy(), num_edges) col = torch.cat([col, torch.from_numpy(choice)]) edge_index = torch.stack([row, col], dim=0) edge_index, _ = remove_self_loops(edge_index) edge_index = to_undirected(edge_index, num_nodes=num_nodes) return edge_index