Utilities

Utility functions.

resolve_device(device=None)

Resolve the device to use.

Parameters

device (Union[None, str, device]) – the device hint

Return type

device

Returns

the resolved device

prepare_num_nodes(edge_index, num_nodes=None)

Prepare the number of nodes.

If an explicit number is given, this number will be used. Otherwise, infers the number of nodes as the maximum id in the edge index.

Parameters

• edge_index (Tensor) – shape: (2, m) the edge index

• num_nodes (Optional[int]) – the number of nodes. If None, it is inferred from edge_index.

Return type

int

Returns

the number of nodes

edge_index_to_sparse_matrix(edge_index, num_nodes=None)

Convert an edge index to a sparse matrix.

Uses the edge index for non-zero entries, and fills in 1 as entries.

Parameters

• edge_index (LongTensor) – shape: (2, m) the edge index

• num_nodes (Optional[int]) – the number of nodes used to determine the shape of the matrix. If None, it is inferred from edge_index.

Return type

Tensor

Returns

shape: (n, n) the adjacency matrix as a sparse tensor, cf. torch.sparse_coo_tensor().

prepare_page_rank_adjacency(adj=None, edge_index=None, num_nodes=None, add_identity=False)

Prepare the page-rank adjacency matrix.

If no explicit adjacency is given, the methods first creates an adjacency matrix from the edge index, cf. edge_index_to_sparse_matrix(). Next, the matrix is symmetrized as

\[A := A + A^T\]

Finally, the matrix is normalized such that the columns sum to one.

Parameters

• adj (Optional[Tensor]) – shape: (n, n) the adjacency matrix

• edge_index (Optional[LongTensor]) – shape: (2, m) the edge index

• num_nodes (Optional[int]) – the number of nodes used to determine the shape of the adjacency matrix. If None, and adj is not already provided, it is inferred from edge_index.

• add_identity (bool) – whether to add an identity matrix to A to ensure that each node has a degree of at least one.

Raises

ValueError – if neither is provided, or the adjacency matrix is invalid

Return type

Tensor

Returns

shape: (n, n) the symmetric, normalized, and sparse adjacency matrix

validate_x(x, n=None)

Validate a (batched) page-rank vector.

In particular, the method checks that

• the tensor dimension is (n,) or (n, batch_size)

• all entries are between 0 and 1

• the entries sum to 1 (along the first dimension)

Parameters

• x (Tensor) – the initial value.

• n (Optional[int]) – the number of nodes.

Raises

ValueError – if the input is invalid.

Return type

None

prepare_x0(x0=None, indices=None, n=None)

Prepare a start value.

The following precedence order is used:

1. an explicit start value, via x0. If present, this tensor is passed through without further modification.

2. a one-hot matrix created via indices. The matrix is of shape (n, len(indices)) and has a single 1 per column at the given indices.

3. a uniform 1/n vector of shape (n,)

Parameters

• x0 (Optional[Tensor]) – the start value.

• indices (Optional[Collection[int]]) – a non-zero indices

• n (Optional[int]) – the number of nodes

Raises

ValueError – if neither x0 nor n are provided

Return type

Tensor

Returns

shape: (n,) or (n, batch_size) the initial value x

power_iteration(adj, x0, alpha=0.05, max_iter=1000, use_tqdm=False, epsilon=0.0001, device=None)

Perform the power iteration.

\[\mathbf{x}^{(i+1)} = (1 - \alpha) \cdot \mathbf{A} \mathbf{x}^{(i)} + \alpha \mathbf{x}^{(0)}\]

Parameters

• adj (Tensor) – shape: (n, n) the (sparse) adjacency matrix

• x0 (Tensor) – shape: (n,), or (n, batch_size) the initial value for x.

• alpha (float) – 0 < alpha < 1 the smoothing value / teleport probability

• max_iter (int) – 0 < max_iter the maximum number of iterations

• epsilon (float) – epsilon > 0 a (small) constant to check for convergence

• use_tqdm (bool) – whether to use a tqdm progress bar

• device (Union[None, str, device]) – the device to use, or a hint thereof, cf. resolve_device()

Return type

Tensor

Returns

shape: (n,) or (n, batch_size) the x value after convergence (or maximum number of iterations).

batched_personalized_page_rank(adj, indices, batch_size, **kwargs)

Batch-wise PPR computation with automatic memory optimization.

Parameters

• adj (Tensor) – shape: (n, n) the adjacency matrix.

• indices (Tensor) – shape: k the indices for which to compute PPR

• batch_size (int) – batch_size > 0 the batch size. Will be reduced if necessary

• kwargs – additional keyword-based parameters passed to power_iteration()

Return type

Tensor

Returns

shape: (n, k) the PPR vectors for each node index

sparse_normalize(matrix, dim=0)

Normalize a sparse matrix to row/column sum of 1.

Parameters

• matrix (Tensor) – the sparse matrix

• dim (int) – the dimension along which to normalize, either 0 for rows or 1 for columns

Return type

Tensor

Returns

the normalized sparse matrix