test_basis.py 1.46 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
from itertools import product

import pytest
import torch
rusty1s's avatar
rusty1s committed
5
from torch_spline_conv import spline_basis
rusty1s's avatar
rusty1s committed
6

rusty1s's avatar
rusty1s committed
7
from .utils import dtypes, devices, tensor
rusty1s's avatar
rusty1s committed
8

rusty1s's avatar
rusty1s committed
9
tests = [{
rusty1s's avatar
rusty1s committed
10
    'pseudo': [[0], [0.0625], [0.25], [0.75], [0.9375], [1]],
rusty1s's avatar
rusty1s committed
11
12
13
14
15
    'kernel_size': [5],
    'is_open_spline': [1],
    'basis': [[1, 0], [0.75, 0.25], [1, 0], [1, 0], [0.25, 0.75], [1, 0]],
    'weight_index': [[0, 1], [0, 1], [1, 2], [3, 4], [3, 4], [4, 0]],
}, {
rusty1s's avatar
rusty1s committed
16
    'pseudo': [[0], [0.0625], [0.25], [0.75], [0.9375], [1]],
rusty1s's avatar
rusty1s committed
17
18
19
20
21
22
23
24
25
26
27
    'kernel_size': [4],
    'is_open_spline': [0],
    'basis': [[1, 0], [0.75, 0.25], [1, 0], [1, 0], [0.25, 0.75], [1, 0]],
    'weight_index': [[0, 1], [0, 1], [1, 2], [3, 0], [3, 0], [0, 1]],
}, {
    'pseudo': [[0.125, 0.5], [0.5, 0.5], [0.75, 0.125]],
    'kernel_size': [5, 5],
    'is_open_spline': [1, 1],
    'basis': [[0.5, 0.5, 0, 0], [1, 0, 0, 0], [0.5, 0, 0.5, 0]],
    'weight_index': [[10, 11, 15, 16], [12, 13, 17, 18], [3, 4, 8, 9]]
}]
rusty1s's avatar
rusty1s committed
28
29


rusty1s's avatar
rusty1s committed
30
31
32
33
34
@pytest.mark.parametrize('test,dtype,device', product(tests, dtypes, devices))
def test_spline_basis_forward(test, dtype, device):
    pseudo = tensor(test['pseudo'], dtype, device)
    kernel_size = tensor(test['kernel_size'], torch.long, device)
    is_open_spline = tensor(test['is_open_spline'], torch.uint8, device)
rusty1s's avatar
rusty1s committed
35
    degree = 1
rusty1s's avatar
rusty1s committed
36

rusty1s's avatar
rusty1s committed
37
38
    basis, weight_index = spline_basis(pseudo, kernel_size, is_open_spline,
                                       degree)
rusty1s's avatar
rusty1s committed
39
    assert basis.tolist() == test['basis']
rusty1s's avatar
rusty1s committed
40
    assert weight_index.tolist() == test['weight_index']