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

Simplify Graph class (#459)

Browse Source 
* Simplifying Graph class and adjust tests

---------

Co-authored-by: Dario Coscia <dariocos99@gmail.com>
This commit is contained in:
Filippo Olivo
2025-03-03 09:30:44 +01:00
committed by Nicola Demo
parent 4c3e305b09
commit ab6ca78d85
7 changed files with 909 additions and 719 deletions
Download Patch File Download Diff File Expand all files Collapse all files

186
tests/test_data/test_data_module.py
View File

@@ -15,16 +15,15 @@ output_tensor = torch.rand((100, 2))
x = torch.rand((100, 50, 10))
pos = torch.rand((100, 50, 2))
input_graph = RadiusGraph(x, pos, r=.1, build_edge_attr=True)
input_graph = [
RadiusGraph(x=x_, pos=pos_, radius=0.2) for x_, pos_, in zip(x, pos)
]
output_graph = torch.rand((100, 50, 10))
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
def test_constructor(input_, output_):
problem = SupervisedProblem(input_=input_, output_=output_)
@@ -33,22 +32,16 @@ def test_constructor(input_, output_):
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
@pytest.mark.parametrize(
"train_size, val_size, test_size",
[
(.7, .2, .1),
(.7, .3, 0)
]
"train_size, val_size, test_size", [(0.7, 0.2, 0.1), (0.7, 0.3, 0)]
)
def test_setup_train(input_, output_, train_size, val_size, test_size):
problem = SupervisedProblem(input_=input_, output_=output_)
dm = PinaDataModule(problem, train_size=train_size,
val_size=val_size, test_size=test_size)
dm = PinaDataModule(
problem, train_size=train_size, val_size=val_size, test_size=test_size
)
dm.setup()
assert hasattr(dm, "train_dataset")
if isinstance(input_, torch.Tensor):
@@ -71,23 +64,17 @@ def test_setup_train(input_, output_, train_size, val_size, test_size):
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
@pytest.mark.parametrize(
"train_size, val_size, test_size",
[
(.7, .2, .1),
(0., 0., 1.)
]
"train_size, val_size, test_size", [(0.7, 0.2, 0.1), (0.0, 0.0, 1.0)]
)
def test_setup_test(input_, output_, train_size, val_size, test_size):
problem = SupervisedProblem(input_=input_, output_=output_)
dm = PinaDataModule(problem, train_size=train_size,
val_size=val_size, test_size=test_size)
dm.setup(stage='test')
dm = PinaDataModule(
problem, train_size=train_size, val_size=val_size, test_size=test_size
)
dm.setup(stage="test")
if train_size > 0:
assert hasattr(dm, "train_dataset")
assert dm.train_dataset is None
@@ -109,16 +96,14 @@ def test_setup_test(input_, output_, train_size, val_size, test_size):
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
def test_dummy_dataloader(input_, output_):
problem = SupervisedProblem(input_=input_, output_=output_)
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
trainer = Trainer(solver, batch_size=None, train_size=.7,
val_size=.3, test_size=0.)
trainer = Trainer(
solver, batch_size=None, train_size=0.7, val_size=0.3, test_size=0.0
)
dm = trainer.data_module
dm.setup()
dm.trainer = trainer
@@ -128,11 +113,11 @@ def test_dummy_dataloader(input_, output_):
data = next(dataloader)
assert isinstance(data, list)
assert isinstance(data[0], tuple)
if isinstance(input_, RadiusGraph):
assert isinstance(data[0][1]['input_points'], Batch)
if isinstance(input_, list):
assert isinstance(data[0][1]["input_points"], 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_points"], torch.Tensor)
assert isinstance(data[0][1]["output_points"], torch.Tensor)
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DummyDataloader)
@@ -140,31 +125,29 @@ def test_dummy_dataloader(input_, output_):
data = next(dataloader)
assert isinstance(data, list)
assert isinstance(data[0], tuple)
if isinstance(input_, RadiusGraph):
assert isinstance(data[0][1]['input_points'], Batch)
if isinstance(input_, list):
assert isinstance(data[0][1]["input_points"], 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_points"], torch.Tensor)
assert isinstance(data[0][1]["output_points"], torch.Tensor)
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
)
@pytest.mark.parametrize(
"automatic_batching",
[
True, False
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
@pytest.mark.parametrize("automatic_batching", [True, False])
def test_dataloader(input_, output_, automatic_batching):
problem = SupervisedProblem(input_=input_, output_=output_)
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
trainer = Trainer(solver, batch_size=10, train_size=.7, val_size=.3,
test_size=0., automatic_batching=automatic_batching)
trainer = Trainer(
solver,
batch_size=10,
train_size=0.7,
val_size=0.3,
test_size=0.0,
automatic_batching=automatic_batching,
)
dm = trainer.data_module
dm.setup()
dm.trainer = trainer
@@ -173,51 +156,53 @@ def test_dataloader(input_, output_, automatic_batching):
assert len(dataloader) == 7
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, RadiusGraph):
assert isinstance(data['data']['input_points'], Batch)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
else:
assert isinstance(data['data']['input_points'], torch.Tensor)
assert isinstance(data['data']['output_points'], torch.Tensor)
assert isinstance(data["data"]["input_points"], torch.Tensor)
assert isinstance(data["data"]["output_points"], torch.Tensor)
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DataLoader)
assert len(dataloader) == 3
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, RadiusGraph):
assert isinstance(data['data']['input_points'], Batch)
if isinstance(input_, list):
assert isinstance(data["data"]["input_points"], Batch)
else:
assert isinstance(data['data']['input_points'], torch.Tensor)
assert isinstance(data['data']['output_points'], torch.Tensor)
assert isinstance(data["data"]["input_points"], torch.Tensor)
assert isinstance(data["data"]["output_points"], torch.Tensor)
from pina import LabelTensor
input_tensor = LabelTensor(torch.rand((100, 3)), ['u', 'v', 'w'])
output_tensor = LabelTensor(torch.rand((100, 3)), ['u', 'v', 'w'])
input_tensor = LabelTensor(torch.rand((100, 3)), ["u", "v", "w"])
output_tensor = LabelTensor(torch.rand((100, 3)), ["u", "v", "w"])
x = LabelTensor(torch.rand((100, 50, 3)), ["u", "v", "w"])
pos = LabelTensor(torch.rand((100, 50, 2)), ["x", "y"])
input_graph = [
RadiusGraph(x=x[i], pos=pos[i], radius=0.1) for i in range(len(x))
]
output_graph = LabelTensor(torch.rand((100, 50, 3)), ["u", "v", "w"])
x = LabelTensor(torch.rand((100, 50, 3)), ['u', 'v', 'w'])
pos = LabelTensor(torch.rand((100, 50, 2)), ['x', 'y'])
input_graph = RadiusGraph(x, pos, r=.1, build_edge_attr=True)
output_graph = LabelTensor(torch.rand((100, 50, 3)), ['u', 'v', 'w'])
@pytest.mark.parametrize(
"input_, output_",
[
(input_tensor, output_tensor),
(input_graph, output_graph)
]
)
@pytest.mark.parametrize(
"automatic_batching",
[
True, False
]
[(input_tensor, output_tensor), (input_graph, output_graph)],
)
@pytest.mark.parametrize("automatic_batching", [True, False])
def test_dataloader_labels(input_, output_, automatic_batching):
problem = SupervisedProblem(input_=input_, output_=output_)
solver = SupervisedSolver(problem=problem, model=torch.nn.Linear(10, 10))
trainer = Trainer(solver, batch_size=10, train_size=.7, val_size=.3,
test_size=0., automatic_batching=automatic_batching)
trainer = Trainer(
solver,
batch_size=10,
train_size=0.7,
val_size=0.3,
test_size=0.0,
automatic_batching=automatic_batching,
)
dm = trainer.data_module
dm.setup()
dm.trainer = trainer
@@ -226,31 +211,30 @@ def test_dataloader_labels(input_, output_, automatic_batching):
assert len(dataloader) == 7
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, RadiusGraph):
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']
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']
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"]
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"]
dataloader = dm.val_dataloader()
assert isinstance(dataloader, DataLoader)
assert len(dataloader) == 3
data = next(iter(dataloader))
assert isinstance(data, dict)
if isinstance(input_, RadiusGraph):
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']
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"]
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']
test_dataloader_labels(input_graph, output_graph, True)
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"]