graclus.py 1.17 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
2
import torch_cluster.graclus_cpu
rusty1s's avatar
rusty1s committed
3
4

if torch.cuda.is_available():
5
    import torch_cluster.graclus_cuda
rusty1s's avatar
rusty1s committed
6
7
8


def graclus_cluster(row, col, weight=None, num_nodes=None):
rusty1s's avatar
rusty1s committed
9
    """A greedy clustering algorithm of picking an unmarked vertex and matching
rusty1s's avatar
typo  
rusty1s committed
10
    it with one its unmarked neighbors (that maximizes its edge weight).
rusty1s's avatar
rusty1s committed
11
12
13
14
15
16
17

    Args:
        row (LongTensor): Source nodes.
        col (LongTensor): Target nodes.
        weight (Tensor, optional): Edge weights. (default: :obj:`None`)
        num_nodes (int, optional): The number of nodes. (default: :obj:`None`)

rusty1s's avatar
docs  
rusty1s committed
18
19
    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
20
21
    Examples::

rusty1s's avatar
rusty1s committed
22
23
        >>> row = torch.tensor([0, 1, 1, 2])
        >>> col = torch.tensor([1, 0, 2, 1])
rusty1s's avatar
rusty1s committed
24
25
26
        >>> weight = torch.Tensor([1, 1, 1, 1])
        >>> cluster = graclus_cluster(row, col, weight)
    """
rusty1s's avatar
rusty1s committed
27

rusty1s's avatar
rusty1s committed
28
29
    if num_nodes is None:
        num_nodes = max(row.max().item(), col.max().item()) + 1
rusty1s's avatar
rusty1s committed
30

31
32
33
34
    if row.is_cuda:
        op = torch_cluster.graclus_cuda
    else:
        op = torch_cluster.graclus_cpu
rusty1s's avatar
rusty1s committed
35

rusty1s's avatar
rusty1s committed
36
37
38
39
    if weight is None:
        cluster = op.graclus(row, col, num_nodes)
    else:
        cluster = op.weighted_graclus(row, col, weight, num_nodes)
rusty1s's avatar
rusty1s committed
40
41

    return cluster