pyg-lib/test/ops/test_cluster.py at 769d7433c99dd8f0ab75cb2c7be83008d2e4ed7c · pyg-team/pyg-lib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import pytest
import torch

import pyg_lib
from pyg_lib.testing import withCUDA


@withCUDA
@pytest.mark.parametrize('dtype', [torch.float, torch.double, torch.bfloat16])
def test_grid_cluster_2d(dtype: torch.dtype, device: torch.device) -> None:
    pos = torch.tensor(
        [[0.0, 0.0], [0.1, 0.1], [0.5, 0.5], [1.0, 1.0], [1.1, 1.1]],
        dtype=dtype, device=device)
    size = torch.tensor([0.5, 0.5], dtype=dtype, device=device)

    out = pyg_lib.ops.grid_cluster(pos, size)

    # Points (0,0) and (0.1,0.1) should be in the same cluster
    assert out[0] == out[1]
    # Point (0.5,0.5) should be in a different cluster from (0,0)
    assert out[0] != out[2]
    # Points (1.0,1.0) and (1.1,1.1) should be in the same cluster
    assert out[3] == out[4]


@withCUDA
@pytest.mark.parametrize('dtype', [torch.float, torch.double])
def test_grid_cluster_3d(dtype: torch.dtype, device: torch.device) -> None:
    pos = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.1, 0.1], [1.0, 1.0, 1.0]],
                       dtype=dtype, device=device)
    size = torch.tensor([0.5, 0.5, 0.5], dtype=dtype, device=device)

    out = pyg_lib.ops.grid_cluster(pos, size)

    assert out[0] == out[1]
    assert out[0] != out[2]


@withCUDA
@pytest.mark.parametrize('dtype', [torch.float, torch.double])
def test_grid_cluster_with_start_end(dtype: torch.dtype,
                                     device: torch.device) -> None:
    pos = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]], dtype=dtype,
                       device=device)
    size = torch.tensor([0.5, 0.5], dtype=dtype, device=device)
    start = torch.tensor([0.0, 0.0], dtype=dtype, device=device)
    end = torch.tensor([1.0, 1.0], dtype=dtype, device=device)

    out = pyg_lib.ops.grid_cluster(pos, size, start, end)

    assert out.shape == (3, )
    assert out.dtype == torch.long


@withCUDA
def test_grid_cluster_defaults_match_explicit(device: torch.device) -> None:
    pos = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]], device=device)
    size = torch.tensor([0.5, 0.5], device=device)

    out_default = pyg_lib.ops.grid_cluster(pos, size)
    out_explicit = pyg_lib.ops.grid_cluster(
        pos,
        size,
        start=pos.min(0).values,
        end=pos.max(0).values,
    )

    assert torch.equal(out_default, out_explicit)


@withCUDA
def test_grid_cluster_cpu_cuda_parity(device: torch.device) -> None:
    pos = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]])
    size = torch.tensor([0.5, 0.5])

    out_cpu = pyg_lib.ops.grid_cluster(pos, size)
    out_dev = pyg_lib.ops.grid_cluster(pos.to(device), size.to(device))

    assert torch.equal(out_cpu, out_dev.cpu())