Add Graph class and tests for Graph and Collector + Dataloader refactoring

tests/test_graph.py

@@ -0,0 +1,143 @@
+import pytest
+import torch
+from pina import Graph
+from pina.graph import RadiusGraph, KNNGraph
+
+
+@pytest.mark.parametrize(
+    "x, pos",
+    [
+        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
+        ([torch.rand(10, 2) for _ in range(3)], [torch.rand(10, 3) for _ in range(3)]),
+        (torch.rand(3,10,2), torch.rand(3,10,3)),
+        (torch.rand(3,10,2), torch.rand(3,10,3)),
+    ]
+)
+def test_build_multiple_graph_multiple_val(x, pos):
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k = 3)
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+
+def test_build_single_graph_multiple_val():
+    x = torch.rand(10, 2)
+    pos = torch.rand(10, 3)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 1
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert len(graph.data) == 1
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+    x = torch.rand(10, 2)
+    pos = torch.rand(10, 3)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    assert len(graph.data) == 1
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    data = graph.data
+    assert len(graph.data) == 1
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos).all() for d_ in data)
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+
+@pytest.mark.parametrize(
+    "pos",
+    [
+        ([torch.rand(10, 3) for _ in range(3)]),
+        ([torch.rand(10, 3) for _ in range(3)]),
+        (torch.rand(3, 10, 3)),
+        (torch.rand(3, 10, 3))
+    ]
+)
+def test_build_single_graph_single_val(pos):
+    x = torch.rand(10, 2)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=False, r=.3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3)
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+    x = torch.rand(10, 2)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=False, k=3)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    graph = KNNGraph(x=x, pos=pos, build_edge_attr=True, k=3)
+    data = graph.data
+    assert all(torch.isclose(d.x, x).all() for d in data)
+    assert all(torch.isclose(d_.pos, pos_).all() for d_, pos_ in zip(data, pos))
+    assert all(len(d.edge_index) == 2 for d in data)
+    assert all(d.edge_attr is not None for d in data)
+    assert all([d.edge_index.shape[1] == d.edge_attr.shape[0]] for d in data)
+
+def test_additional_parameters_1():
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+    additional_parameters = {'y': torch.ones(3)}
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
+                  additional_params=additional_parameters)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(hasattr(d, 'y') for d in data)
+    assert all(d_.y == 1 for d_ in data)
+
+@pytest.mark.parametrize(
+    "additional_parameters",
+    [
+        ({'y': torch.rand(3,10,1)}),
+        ({'y': [torch.rand(10,1) for _ in range(3)]}),
+    ]
+)
+def test_additional_parameters_2(additional_parameters):
+    x = torch.rand(3, 10, 2)
+    pos = torch.rand(3, 10, 2)
+    graph = RadiusGraph(x=x, pos=pos, build_edge_attr=True, r=.3,
+                  additional_params=additional_parameters)
+    assert len(graph.data) == 3
+    data = graph.data
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))
+    assert all(hasattr(d, 'y') for d in data)
+    assert all(torch.isclose(d_.x, x_).all() for (d_, x_) in zip(data, x))