Source code for torch_geometric.utils.sparse

from typing import Any, Optional, Tuple, Union

import torch
from torch import Tensor

from torch_geometric.typing import SparseTensor


[docs]def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]: r"""Converts a dense adjacency matrix to a sparse adjacency matrix defined by edge indices and edge attributes. Args: adj (Tensor): The dense adjacency matrix of shape :obj:`[num_nodes, num_nodes]` or :obj:`[batch_size, num_nodes, num_nodes]`. :rtype: (:class:`LongTensor`, :class:`Tensor`) Examples: >>> # Forr a single adjacency matrix >>> adj = torch.tensor([[3, 1], ... [2, 0]]) >>> dense_to_sparse(adj) (tensor([[0, 0, 1], [0, 1, 0]]), tensor([3, 1, 2])) >>> # For two adjacency matrixes >>> adj = torch.tensor([[[3, 1], ... [2, 0]], ... [[0, 1], ... [0, 2]]]) >>> dense_to_sparse(adj) (tensor([[0, 0, 1, 2, 3], [0, 1, 0, 3, 3]]), tensor([3, 1, 2, 1, 2])) """ if adj.dim() < 2 or adj.dim() > 3: raise ValueError(f"Dense adjacency matrix 'adj' must be 2- or " f"3-dimensional (got {adj.dim()} dimensions)") edge_index = adj.nonzero().t() if edge_index.size(0) == 2: edge_attr = adj[edge_index[0], edge_index[1]] return edge_index, edge_attr else: edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]] row = edge_index[1] + adj.size(-2) * edge_index[0] col = edge_index[2] + adj.size(-1) * edge_index[0] return torch.stack([row, col], dim=0), edge_attr
[docs]def is_torch_sparse_tensor(src: Any) -> bool: """Returns :obj:`True` if the input :obj:`src` is a :class:`torch.sparse.Tensor` (in any sparse layout). Args: src (Any): The input object to be checked. """ return isinstance(src, Tensor) and src.is_sparse
[docs]def is_sparse(src: Any) -> bool: """Returns :obj:`True` if the input :obj:`src` is of type :class:`torch.sparse.Tensor` (in any sparse layout) or of type :class:`torch_sparse.SparseTensor`. Args: src (Any): The input object to be checked. """ return is_torch_sparse_tensor(src) or isinstance(src, SparseTensor)
[docs]def to_torch_coo_tensor( edge_index: Tensor, edge_attr: Optional[Tensor] = None, size: Optional[Union[int, Tuple[int, int]]] = None, ) -> Tensor: """Converts a sparse adjacency matrix defined by edge indices and edge attributes to a :class:`torch.sparse.Tensor`. Args: edge_index (LongTensor): The edge indices. edge_attr (Tensor, optional): The edge attributes. (default: :obj:`None`) size (int or (int, int), optional): The size of the sparse matrix. If given as an integer, will create a quadratic sparse matrix. If set to :obj:`None`, will infer a quadratic sparse matrix based on :obj:`edge_index.max() + 1`. (default: :obj:`None`) :rtype: :class:`torch.sparse.FloatTensor` Example: >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], ... [1, 0, 2, 1, 3, 2]]) >>> to_torch_coo_tensor(edge_index) tensor(indices=tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]]), values=tensor([1., 1., 1., 1., 1., 1.]), size=(4, 4), nnz=6, layout=torch.sparse_coo) """ if size is None: size = int(edge_index.max()) + 1 if not isinstance(size, (tuple, list)): size = (size, size) if edge_attr is None: edge_attr = torch.ones(edge_index.size(1), device=edge_index.device) size = tuple(size) + edge_attr.size()[1:] out = torch.sparse_coo_tensor(edge_index, edge_attr, size, device=edge_index.device) out = out.coalesce() return out