102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
import torch
|
|
import pytest
|
|
from pina.data.dataset import PinaDatasetFactory, PinaGraphDataset
|
|
from pina.graph import KNNGraph
|
|
from torch_geometric.data import Data
|
|
|
|
x = torch.rand((100, 20, 10))
|
|
pos = torch.rand((100, 20, 2))
|
|
input_ = KNNGraph(x=x, pos=pos, k=3, build_edge_attr=True)
|
|
output_ = torch.rand((100, 20, 10))
|
|
|
|
x_2 = torch.rand((50, 20, 10))
|
|
pos_2 = torch.rand((50, 20, 2))
|
|
input_2_ = KNNGraph(x=x_2, pos=pos_2, k=3, build_edge_attr=True)
|
|
output_2_ = torch.rand((50, 20, 10))
|
|
|
|
|
|
# Problem with a single condition
|
|
conditions_dict_single = {
|
|
'data': {
|
|
'input_points': input_.data,
|
|
'output_points': output_,
|
|
}
|
|
}
|
|
max_conditions_lengths_single = {
|
|
'data': 100
|
|
}
|
|
|
|
# Problem with multiple conditions
|
|
conditions_dict_single_multi = {
|
|
'data_1': {
|
|
'input_points': input_.data,
|
|
'output_points': output_,
|
|
},
|
|
'data_2': {
|
|
'input_points': input_2_.data,
|
|
'output_points': output_2_,
|
|
}
|
|
}
|
|
|
|
max_conditions_lengths_multi = {
|
|
'data_1': 100,
|
|
'data_2': 50
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"conditions_dict, max_conditions_lengths",
|
|
[
|
|
(conditions_dict_single, max_conditions_lengths_single),
|
|
(conditions_dict_single_multi, max_conditions_lengths_multi)
|
|
]
|
|
)
|
|
def test_constructor(conditions_dict, max_conditions_lengths):
|
|
dataset = PinaDatasetFactory(conditions_dict,
|
|
max_conditions_lengths=max_conditions_lengths,
|
|
automatic_batching=True)
|
|
assert isinstance(dataset, PinaGraphDataset)
|
|
assert len(dataset) == 100
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"conditions_dict, max_conditions_lengths",
|
|
[
|
|
(conditions_dict_single, max_conditions_lengths_single),
|
|
(conditions_dict_single_multi, max_conditions_lengths_multi)
|
|
]
|
|
)
|
|
def test_getitem(conditions_dict, max_conditions_lengths):
|
|
dataset = PinaDatasetFactory(conditions_dict,
|
|
max_conditions_lengths=max_conditions_lengths,
|
|
automatic_batching=True)
|
|
data = dataset[50]
|
|
assert isinstance(data, dict)
|
|
assert all([isinstance(d['input_points'], Data)
|
|
for d in data.values()])
|
|
assert all([isinstance(d['output_points'], torch.Tensor)
|
|
for d in data.values()])
|
|
assert all([d['input_points'].x.shape == torch.Size((20, 10))
|
|
for d in data.values()])
|
|
assert all([d['output_points'].shape == torch.Size((20, 10))
|
|
for d in data.values()])
|
|
assert all([d['input_points'].edge_index.shape ==
|
|
torch.Size((2, 60)) for d in data.values()])
|
|
assert all([d['input_points'].edge_attr.shape[0]
|
|
== 60 for d in data.values()])
|
|
|
|
data = dataset.fetch_from_idx_list([i for i in range(20)])
|
|
assert isinstance(data, dict)
|
|
assert all([isinstance(d['input_points'], Data)
|
|
for d in data.values()])
|
|
assert all([isinstance(d['output_points'], torch.Tensor)
|
|
for d in data.values()])
|
|
assert all([d['input_points'].x.shape == torch.Size((400, 10))
|
|
for d in data.values()])
|
|
assert all([d['output_points'].shape == torch.Size((400, 10))
|
|
for d in data.values()])
|
|
assert all([d['input_points'].edge_index.shape ==
|
|
torch.Size((2, 1200)) for d in data.values()])
|
|
assert all([d['input_points'].edge_attr.shape[0]
|
|
== 1200 for d in data.values()])
|