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
|
import torch
from torch_geometric.data import Data
from torch_geometric.transforms import FixedPoints
def test_fixed_points():
assert str(FixedPoints(1024)) == 'FixedPoints(1024, replace=True)'
data = Data(
pos=torch.randn(100, 3),
x=torch.randn(100, 16),
y=torch.randn(1),
edge_attr=torch.randn(100, 3),
num_nodes=100,
)
out = FixedPoints(50, replace=True)(data)
assert len(out) == 5
assert out.pos.size() == (50, 3)
assert out.x.size() == (50, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 50
out = FixedPoints(200, replace=True)(data)
assert len(out) == 5
assert out.pos.size() == (200, 3)
assert out.x.size() == (200, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 200
out = FixedPoints(50, replace=False, allow_duplicates=False)(data)
assert len(out) == 5
assert out.pos.size() == (50, 3)
assert out.x.size() == (50, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 50
out = FixedPoints(200, replace=False, allow_duplicates=False)(data)
assert len(out) == 5
assert out.pos.size() == (100, 3)
assert out.x.size() == (100, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 100
out = FixedPoints(50, replace=False, allow_duplicates=True)(data)
assert len(out) == 5
assert out.pos.size() == (50, 3)
assert out.x.size() == (50, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 50
out = FixedPoints(200, replace=False, allow_duplicates=True)(data)
assert len(out) == 5
assert out.pos.size() == (200, 3)
assert out.x.size() == (200, 16)
assert out.y.size() == (1, )
assert out.edge_attr.size() == (100, 3)
assert out.num_nodes == 200
|