pyg_lib.sampler
- neighbor_sample(rowptr: Tensor, col: Tensor, seed: Tensor, num_neighbors: List[int], node_time: Optional[Tensor] = None, edge_time: Optional[Tensor] = None, seed_time: Optional[Tensor] = None, edge_weight: Optional[Tensor] = None, csc: bool = False, replace: bool = False, directed: bool = True, disjoint: bool = False, temporal_strategy: str = 'uniform', return_edge_id: bool = True) Tuple[Tensor, Tensor, Tensor, Optional[Tensor], List[int], List[int]] [source]
Recursively samples neighbors from all node indices in
seed
in the graph given by(rowptr, col)
.Note
For temporal sampling, the
col
vector needs to be sorted according totime
within individual neighborhoods since we use binary search to find neighbors that fulfill temporal constraints.- Parameters:
rowptr (torch.Tensor) – Compressed source node indices.
col (torch.Tensor) – Target node indices.
seed (torch.Tensor) – The seed node indices.
num_neighbors (List[int]) – The number of neighbors to sample for each node in each iteration. If an entry is set to
-1
, all neighbors will be included.node_time (torch.Tensor, optional) – Timestamps for the nodes in the graph. If set, temporal sampling will be used such that neighbors are guaranteed to fulfill temporal constraints, i.e. sampled nodes have an earlier or equal timestamp than the seed node. If used, the
col
vector needs to be sorted according to time within individual neighborhoods. Requiresdisjoint=True
. Only eithernode_time
oredge_time
can be specified. (default:None
)edge_time (torch.Tensor, optional) – Timestamps for the edges in the graph. If set, temporal sampling will be used such that neighbors are guaranteed to fulfill temporal constraints, i.e. sampled edges have an earlier or equal timestamp than the seed node. If used, the
col
vector needs to be sorted according to time within individual neighborhoods. Requiresdisjoint=True
. Only eithernode_time
oredge_time
can be specified. (default:None
)seed_time (torch.Tensor, optional) – Optional values to override the timestamp for seed nodes. If not set, will use timestamps in
node_time
as default for seed nodes. Needs to be specified in case edge-level sampling is used viaedge_time
. (default:None
)edge-weight (torch.Tensor, optional) – If given, will perform biased sampling based on the weight of each edge. (default:
None
)csc (bool, optional) – If set to
True
, assumes that the graph is given in CSC format(colptr, row)
. (default:False
)replace (bool, optional) – If set to
True
, will sample with replacement. (default:False
)directed (bool, optional) – If set to
False
, will include all edges between all sampled nodes. (default:True
)disjoint (bool, optional) – If set to
True
, will create disjoint subgraphs for every seed node. (default:False
)temporal_strategy (string, optional) – The sampling strategy when using temporal sampling (
"uniform"
,"last"
). (default:"uniform"
)return_edge_id (bool, optional) – If set to
False
, will not return the indices of edges of the original graph. (default: :obj: True)
- Returns:
(torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], List[int], List[int]): Row indices, col indices of the returned subtree/subgraph, as well as original node indices for all nodes sampled. In addition, may return the indices of edges of the original graph. Lastly, returns information about the sampled amount of nodes and edges per hop.
- hetero_neighbor_sample(rowptr_dict: Dict[Tuple[str, str, str], Tensor], col_dict: Dict[Tuple[str, str, str], Tensor], seed_dict: Dict[str, Tensor], num_neighbors_dict: Dict[Tuple[str, str, str], List[int]], node_time_dict: Optional[Dict[str, Tensor]] = None, edge_time_dict: Optional[Dict[Tuple[str, str, str], Tensor]] = None, seed_time_dict: Optional[Dict[str, Tensor]] = None, edge_weight_dict: Optional[Dict[Tuple[str, str, str], Tensor]] = None, csc: bool = False, replace: bool = False, directed: bool = True, disjoint: bool = False, temporal_strategy: str = 'uniform', return_edge_id: bool = True) Tuple[Dict[Tuple[str, str, str], Tensor], Dict[Tuple[str, str, str], Tensor], Dict[str, Tensor], Optional[Dict[Tuple[str, str, str], Tensor]], Dict[str, List[int]], Dict[Tuple[str, str, str], List[int]]] [source]
Recursively samples neighbors from all node indices in
seed_dict
in the heterogeneous graph given by(rowptr_dict, col_dict)
.Note
Similar to
neighbor_sample()
, but expects a dictionary of node types (str
) and edge types (Tuple[str, str, str]
) for each non-boolean argument.- Parameters:
kwargs – Arguments of
neighbor_sample()
.
- subgraph(rowptr: Tensor, col: Tensor, nodes: Tensor, return_edge_id: bool = True) Tuple[Tensor, Tensor, Optional[Tensor]] [source]
Returns the induced subgraph of the graph given by
(rowptr, col)
, containing only the nodes innodes
.- Parameters:
rowptr (torch.Tensor) – Compressed source node indices.
col (torch.Tensor) – Target node indices.
nodes (torch.Tensor) – Node indices of the induced subgraph.
return_edge_id (bool, optional) – If set to
False
, will not return the indices of edges of the original graph contained in the induced subgraph. (default:True
)
- Returns:
Compressed source node indices and target node indices of the induced subgraph. In addition, may return the indices of edges of the original graph.
- Return type:
(torch.Tensor, torch.Tensor, Optional[torch.Tensor])
- random_walk(rowptr: Tensor, col: Tensor, seed: Tensor, walk_length: int, p: float = 1.0, q: float = 1.0) Tensor [source]
Samples random walks of length
walk_length
from all node indices inseed
in the graph given by(rowptr, col)
, as described in the “node2vec: Scalable Feature Learning for Networks” paper.- Parameters:
rowptr (torch.Tensor) – Compressed source node indices.
col (torch.Tensor) – Target node indices.
seed (torch.Tensor) – Seed node indices from where random walks start.
walk_length (int) – The walk length of a random walk.
p (float, optional) – Likelihood of immediately revisiting a node in the walk. (default:
1.0
)q (float, optional) – Control parameter to interpolate between breadth-first strategy and depth-first strategy. (default:
1.0
)
- Returns:
A tensor of shape
[seed.size(0), walk_length + 1]
, holding the nodes indices of each walk for each seed node.- Return type: