folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve conditions and refactor dataset classes (#475)

Browse Source 
* Reimplement conditions

* Refactor datasets and implement LabelBatch

---------

Co-authored-by: Dario Coscia <dariocos99@gmail.com>
This commit is contained in:
Filippo Olivo
2025-03-07 11:24:09 +01:00
committed by Nicola Demo
parent bdad144461
commit a0cbf1c44a
40 changed files with 943 additions and 550 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
tests/test_data/test_data_module.py
View File

@@ -114,10 +114,10 @@ def test_dummy_dataloader(input_, output_):
assert isinstance(data, list)
assert isinstance(data[0], tuple)
if isinstance(input_, list):
assert isinstance(data[0][1]["input_points"], Batch)
assert isinstance(data[0][1]["input"], Batch)
else:
assert isinstance(data[0][1]["input_points"], torch.Tensor)
assert isinstance(data[0][1]["output_points"], torch.Tensor)
assert isinstance(data[0][1]["input"], torch.Tensor)
assert isinstance(data[0][1]["target"], torch.Tensor)
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DummyDataloader)
@@ -126,10 +126,10 @@ def test_dummy_dataloader(input_, output_):
assert isinstance(data, list)
assert isinstance(data[0], tuple)
if isinstance(input_, list):
assert isinstance(data[0][1]["input_points"], Batch)
assert isinstance(data[0][1]["input"], Batch)
else:
assert isinstance(data[0][1]["input_points"], torch.Tensor)
assert isinstance(data[0][1]["output_points"], torch.Tensor)
assert isinstance(data[0][1]["input"], torch.Tensor)
assert isinstance(data[0][1]["target"], torch.Tensor)
@pytest.mark.parametrize(
@@ -157,10 +157,10 @@ def test_dataloader(input_, output_, automatic_batching):
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
assert isinstance(data["data"]["input"], Batch)
else:
assert isinstance(data["data"]["input_points"], torch.Tensor)
assert isinstance(data["data"]["output_points"], torch.Tensor)
assert isinstance(data["data"]["input"], torch.Tensor)
assert isinstance(data["data"]["target"], torch.Tensor)
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DataLoader)
@@ -168,10 +168,10 @@ def test_dataloader(input_, output_, automatic_batching):
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
assert isinstance(data["data"]["input"], Batch)
else:
assert isinstance(data["data"]["input_points"], torch.Tensor)
assert isinstance(data["data"]["output_points"], torch.Tensor)
assert isinstance(data["data"]["input"], torch.Tensor)
assert isinstance(data["data"]["target"], torch.Tensor)
from pina import LabelTensor
@@ -212,15 +212,15 @@ def test_dataloader_labels(input_, output_, automatic_batching):
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
assert isinstance(data["data"]["input_points"].x, LabelTensor)
assert data["data"]["input_points"].x.labels == ["u", "v", "w"]
assert data["data"]["input_points"].pos.labels == ["x", "y"]
assert isinstance(data["data"]["input"], Batch)
assert isinstance(data["data"]["input"].x, LabelTensor)
assert data["data"]["input"].x.labels == ["u", "v", "w"]
assert data["data"]["input"].pos.labels == ["x", "y"]
else:
assert isinstance(data["data"]["input_points"], LabelTensor)
assert data["data"]["input_points"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["output_points"], LabelTensor)
assert data["data"]["output_points"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["input"], LabelTensor)
assert data["data"]["input"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["target"], LabelTensor)
assert data["data"]["target"].labels == ["u", "v", "w"]
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DataLoader)
@@ -228,13 +228,13 @@ def test_dataloader_labels(input_, output_, automatic_batching):
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
assert isinstance(data["data"]["input_points"].x, LabelTensor)
assert data["data"]["input_points"].x.labels == ["u", "v", "w"]
assert data["data"]["input_points"].pos.labels == ["x", "y"]
assert isinstance(data["data"]["input"], Batch)
assert isinstance(data["data"]["input"].x, LabelTensor)
assert data["data"]["input"].x.labels == ["u", "v", "w"]
assert data["data"]["input"].pos.labels == ["x", "y"]
else:
assert isinstance(data["data"]["input_points"], torch.Tensor)
assert isinstance(data["data"]["input_points"], LabelTensor)
assert data["data"]["input_points"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["output_points"], torch.Tensor)
assert data["data"]["output_points"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["input"], torch.Tensor)
assert isinstance(data["data"]["input"], LabelTensor)
assert data["data"]["input"].labels == ["u", "v", "w"]
assert isinstance(data["data"]["target"], torch.Tensor)
assert data["data"]["target"].labels == ["u", "v", "w"]

36
tests/test_data/test_graph_dataset.py
View File

@@ -24,8 +24,8 @@ output_2_ = torch.rand((50, 20, 10))
# Problem with a single condition
conditions_dict_single = {
"data": {
"input_points": input_,
"output_points": output_,
"input": input_,
"target": output_,
}
}
max_conditions_lengths_single = {"data": 100}
@@ -33,12 +33,12 @@ max_conditions_lengths_single = {"data": 100}
# Problem with multiple conditions
conditions_dict_single_multi = {
"data_1": {
"input_points": input_,
"output_points": output_,
"input": input_,
"target": output_,
},
"data_2": {
"input_points": input_2_,
"output_points": output_2_,
"input": input_2_,
"target": output_2_,
},
}
@@ -77,56 +77,56 @@ def test_getitem(conditions_dict, max_conditions_lengths):
)
data = dataset[50]
assert isinstance(data, dict)
assert all([isinstance(d["input_points"], Data) for d in data.values()])
assert all([isinstance(d["input"], Data) for d in data.values()])
assert all(
[isinstance(d["output_points"], torch.Tensor) for d in data.values()]
[isinstance(d["target"], torch.Tensor) for d in data.values()]
)
assert all(
[
d["input_points"].x.shape == torch.Size((20, 10))
d["input"].x.shape == torch.Size((20, 10))
for d in data.values()
]
)
assert all(
[
d["output_points"].shape == torch.Size((20, 10))
d["target"].shape == torch.Size((20, 10))
for d in data.values()
]
)
assert all(
[
d["input_points"].edge_index.shape == torch.Size((2, 60))
d["input"].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()]
[d["input"].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["input"], Data) for d in data.values()])
assert all(
[isinstance(d["output_points"], torch.Tensor) for d in data.values()]
[isinstance(d["target"], torch.Tensor) for d in data.values()]
)
assert all(
[
d["input_points"].x.shape == torch.Size((400, 10))
d["input"].x.shape == torch.Size((400, 10))
for d in data.values()
]
)
assert all(
[
d["output_points"].shape == torch.Size((400, 10))
d["target"].shape == torch.Size((400, 10))
for d in data.values()
]
)
assert all(
[
d["input_points"].edge_index.shape == torch.Size((2, 1200))
d["input"].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()]
[d["input"].edge_attr.shape[0] == 1200 for d in data.values()]
)

42
tests/test_data/test_tensor_dataset.py
View File

@@ -10,19 +10,19 @@ output_tensor_2 = torch.rand((50, 2))
conditions_dict_single = {
'data': {
'input_points': input_tensor,
'output_points': output_tensor,
'input': input_tensor,
'target': output_tensor,
}
}
conditions_dict_single_multi = {
'data_1': {
'input_points': input_tensor,
'output_points': output_tensor,
'input': input_tensor,
'target': output_tensor,
},
'data_2': {
'input_points': input_tensor_2,
'output_points': output_tensor_2,
'input': input_tensor_2,
'target': output_tensor_2,
}
}
@@ -59,11 +59,11 @@ def test_getitem_single():
assert isinstance(tensors, dict)
assert list(tensors.keys()) == ['data']
assert sorted(list(tensors['data'].keys())) == [
'input_points', 'output_points']
assert isinstance(tensors['data']['input_points'], torch.Tensor)
assert tensors['data']['input_points'].shape == torch.Size((70, 10))
assert isinstance(tensors['data']['output_points'], torch.Tensor)
assert tensors['data']['output_points'].shape == torch.Size((70, 2))
'input', 'target']
assert isinstance(tensors['data']['input'], torch.Tensor)
assert tensors['data']['input'].shape == torch.Size((70, 10))
assert isinstance(tensors['data']['target'], torch.Tensor)
assert tensors['data']['target'].shape == torch.Size((70, 2))
def test_getitem_multi():
@@ -74,15 +74,15 @@ def test_getitem_multi():
assert isinstance(tensors, dict)
assert list(tensors.keys()) == ['data_1', 'data_2']
assert sorted(list(tensors['data_1'].keys())) == [
'input_points', 'output_points']
assert isinstance(tensors['data_1']['input_points'], torch.Tensor)
assert tensors['data_1']['input_points'].shape == torch.Size((70, 10))
assert isinstance(tensors['data_1']['output_points'], torch.Tensor)
assert tensors['data_1']['output_points'].shape == torch.Size((70, 2))
'input', 'target']
assert isinstance(tensors['data_1']['input'], torch.Tensor)
assert tensors['data_1']['input'].shape == torch.Size((70, 10))
assert isinstance(tensors['data_1']['target'], torch.Tensor)
assert tensors['data_1']['target'].shape == torch.Size((70, 2))
assert sorted(list(tensors['data_2'].keys())) == [
'input_points', 'output_points']
assert isinstance(tensors['data_2']['input_points'], torch.Tensor)
assert tensors['data_2']['input_points'].shape == torch.Size((50, 10))
assert isinstance(tensors['data_2']['output_points'], torch.Tensor)
assert tensors['data_2']['output_points'].shape == torch.Size((50, 2))
'input', 'target']
assert isinstance(tensors['data_2']['input'], torch.Tensor)
assert tensors['data_2']['input'].shape == torch.Size((50, 10))
assert isinstance(tensors['data_2']['target'], torch.Tensor)
assert tensors['data_2']['target'].shape == torch.Size((50, 2))