tests and minor fixes

Author: GiovanniCanali

pina/model/block/message_passing/deep_tensor_network_block.py

@@ -40,7 +40,7 @@ def __init__(
         flow="source_to_target",
     ):
         """
-        Initialization of the :class:`DeepTensorNetworkBlocklock` class.
+        Initialization of the :class:`DeepTensorNetworkBlock` class.
 
         :param int node_feature_dim: The dimension of the node features.
         :param int edge_feature_dim: The dimension of the edge features.
@@ -68,7 +68,7 @@ def __init__(
         check_positive_integer(edge_feature_dim, strict=True)
 
         # Activation function
-        self.activation = activation
+        self.activation = activation()
 
         # Layer for processing node features
         self.node_layer = torch.nn.Linear(

pina/model/block/message_passing/egnn_block.py

@@ -76,7 +76,7 @@ def __init__(
             source node. See :class:`torch_geometric.nn.MessagePassing` for more
             details. Default is "source_to_target".
         :raises AssertionError: If `node_feature_dim` is not a positive integer.
-        :raises AssertionError: If `edge_feature_dim` is not a positive integer.
+        :raises AssertionError: If `edge_feature_dim` is a negative integer.
         :raises AssertionError: If `pos_dim` is not a positive integer.
         :raises AssertionError: If `hidden_dim` is not a positive integer.
         :raises AssertionError: If `n_message_layers` is not a positive integer.
@@ -86,7 +86,7 @@ def __init__(
 
         # Check values
         check_positive_integer(node_feature_dim, strict=True)
-        check_positive_integer(edge_feature_dim, strict=True)
+        check_positive_integer(edge_feature_dim, strict=False)
         check_positive_integer(pos_dim, strict=True)
         check_positive_integer(hidden_dim, strict=True)
         check_positive_integer(n_message_layers, strict=True)
@@ -110,7 +110,7 @@ def __init__(
             func=activation,
         )
 
-    def forward(self, x, pos, edge_index, edge_attr):
+    def forward(self, x, pos, edge_index, edge_attr=None):
         """
         Forward pass of the block, triggering the message-passing routine.
 
@@ -146,7 +146,11 @@ def message(self, x_i, x_j, pos_i, pos_j, edge_attr):
         :rtype: torch.Tensor
         """
         dist = torch.norm(pos_i - pos_j, dim=-1, keepdim=True) ** 2
-        input_ = torch.cat((x_i, x_j, dist, edge_attr), dim=-1)
+        if edge_attr is None:
+            input_ = torch.cat((x_i, x_j, dist), dim=-1)
+        else:
+            input_ = torch.cat((x_i, x_j, dist, edge_attr), dim=-1)
+
         return self.message_net(input_)
 
     def update(self, message, x, pos, edge_index):
@@ -169,4 +173,4 @@ def update(self, message, x, pos, edge_index):
         # Update the node positions
         c = degree(edge_index[0], pos.shape[0]).unsqueeze(-1)
         pos = pos + message / c
-        return pos, x
+        return x, pos

pina/model/block/message_passing/interaction_network_block.py

@@ -99,8 +99,8 @@ def __init__(
         # Update network
         self.update_net = FeedForward(
             input_dimensions=node_feature_dim + hidden_dim,
-            output_dimensions=hidden_dim,
-            inner_size=node_feature_dim,
+            output_dimensions=node_feature_dim,
+            inner_size=hidden_dim,
             n_layers=n_update_layers,
             func=activation,
         )

pina/model/block/message_passing/radial_field_network_block.py

@@ -97,7 +97,7 @@ def forward(self, x, edge_index):
         """
         return self.propagate(edge_index=edge_index, x=x)
 
-    def message(self, x_j, x_i):
+    def message(self, x_i, x_j):
         """
         Compute the message to be passed between nodes and edges.
 

pina/model/block/message_passing/schnet_block.py

@@ -120,8 +120,7 @@ def forward(self, x, pos, edge_index):
 
         :param x: The node features.
         :type x: torch.Tensor | LabelTensor
-        :param torch.Tensor edge_index: The edge indices. In the original formulation,
-        the messages are aggregated from all nodes, not only from the neighbours.
+        :param torch.Tensor edge_index: The edge indices.
         :return: The updated node features.
         :rtype: torch.Tensor
         """

pina/utils.py

@@ -202,7 +202,7 @@ def check_positive_integer(value, strict=True):
     :param int value: The value to check.
     :param bool strict: If True, the value must be strictly positive.
         Default is True.
-    :raises ValueError: If the value is not a positive integer.
+    :raises AssertionError: If the value is not a positive integer.
     """
     if strict:
         assert (

tests/test_messagepassing/test_deep_tensor_network_block.py

@@ -0,0 +1,59 @@
+import pytest
+import torch
+from pina.model.block.message_passing import DeepTensorNetworkBlock
+
+# Data for testing
+x = torch.rand(10, 3)
+edge_index = torch.randint(0, 10, (2, 20))
+edge_attr = torch.randn(20, 2)
+
+
+@pytest.mark.parametrize("node_feature_dim", [1, 3])
+@pytest.mark.parametrize("edge_feature_dim", [3, 5])
+def test_constructor(node_feature_dim, edge_feature_dim):
+
+    DeepTensorNetworkBlock(
+        node_feature_dim=node_feature_dim,
+        edge_feature_dim=edge_feature_dim,
+    )
+
+    # Should fail if node_feature_dim is negative
+    with pytest.raises(AssertionError):
+        DeepTensorNetworkBlock(
+            node_feature_dim=-1, edge_feature_dim=edge_feature_dim
+        )
+
+    # Should fail if edge_feature_dim is negative
+    with pytest.raises(AssertionError):
+        DeepTensorNetworkBlock(
+            node_feature_dim=node_feature_dim, edge_feature_dim=-1
+        )
+
+
+def test_forward():
+
+    model = DeepTensorNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_attr.shape[1],
+    )
+
+    output_ = model(edge_index=edge_index, x=x, edge_attr=edge_attr)
+    assert output_.shape == x.shape
+
+
+def test_backward():
+
+    model = DeepTensorNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_attr.shape[1],
+    )
+
+    output_ = model(
+        edge_index=edge_index,
+        x=x.requires_grad_(),
+        edge_attr=edge_attr.requires_grad_(),
+    )
+
+    loss = torch.mean(output_)
+    loss.backward()
+    assert x.grad.shape == x.shape

tests/test_messagepassing/test_equivariant_network_block.py

@@ -0,0 +1,130 @@
+import pytest
+import torch
+from pina.model.block.message_passing import EnEquivariantNetworkBlock
+
+# Data for testing
+x = torch.rand(10, 4)
+pos = torch.rand(10, 3)
+edge_index = torch.randint(0, 10, (2, 20))
+edge_attr = torch.randn(20, 2)
+
+
+@pytest.mark.parametrize("node_feature_dim", [1, 3])
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+@pytest.mark.parametrize("pos_dim", [2, 3])
+def test_constructor(node_feature_dim, edge_feature_dim, pos_dim):
+
+    EnEquivariantNetworkBlock(
+        node_feature_dim=node_feature_dim,
+        edge_feature_dim=edge_feature_dim,
+        pos_dim=pos_dim,
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    # Should fail if node_feature_dim is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=-1,
+            edge_feature_dim=edge_feature_dim,
+            pos_dim=pos_dim,
+        )
+
+    # Should fail if edge_feature_dim is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=node_feature_dim,
+            edge_feature_dim=-1,
+            pos_dim=pos_dim,
+        )
+
+    # Should fail if pos_dim is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=node_feature_dim,
+            edge_feature_dim=edge_feature_dim,
+            pos_dim=-1,
+        )
+
+    # Should fail if hidden_dim is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=node_feature_dim,
+            edge_feature_dim=edge_feature_dim,
+            pos_dim=pos_dim,
+            hidden_dim=-1,
+        )
+
+    # Should fail if n_message_layers is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=node_feature_dim,
+            edge_feature_dim=edge_feature_dim,
+            pos_dim=pos_dim,
+            n_message_layers=-1,
+        )
+
+    # Should fail if n_update_layers is negative
+    with pytest.raises(AssertionError):
+        EnEquivariantNetworkBlock(
+            node_feature_dim=node_feature_dim,
+            edge_feature_dim=edge_feature_dim,
+            pos_dim=pos_dim,
+            n_update_layers=-1,
+        )
+
+
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+def test_forward(edge_feature_dim):
+
+    model = EnEquivariantNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_feature_dim,
+        pos_dim=pos.shape[1],
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    if edge_feature_dim == 0:
+        output_ = model(edge_index=edge_index, x=x, pos=pos)
+    else:
+        output_ = model(
+            edge_index=edge_index, x=x, pos=pos, edge_attr=edge_attr
+        )
+
+    assert output_[0].shape == x.shape
+    assert output_[1].shape == pos.shape
+
+
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+def test_backward(edge_feature_dim):
+
+    model = EnEquivariantNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_feature_dim,
+        pos_dim=pos.shape[1],
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    if edge_feature_dim == 0:
+        output_ = model(
+            edge_index=edge_index,
+            x=x.requires_grad_(),
+            pos=pos.requires_grad_(),
+        )
+    else:
+        output_ = model(
+            edge_index=edge_index,
+            x=x.requires_grad_(),
+            pos=pos.requires_grad_(),
+            edge_attr=edge_attr.requires_grad_(),
+        )
+
+    loss = torch.mean(output_[0])
+    loss.backward()
+    assert x.grad.shape == x.shape
+    assert pos.grad.shape == pos.shape

tests/test_messagepassing/test_interaction_network_block.py

@@ -0,0 +1,84 @@
+import pytest
+import torch
+from pina.model.block.message_passing import InteractionNetworkBlock
+
+# Data for testing
+x = torch.rand(10, 3)
+edge_index = torch.randint(0, 10, (2, 20))
+edge_attr = torch.randn(20, 2)
+
+
+@pytest.mark.parametrize("node_feature_dim", [1, 3])
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+def test_constructor(node_feature_dim, edge_feature_dim):
+
+    InteractionNetworkBlock(
+        node_feature_dim=node_feature_dim,
+        edge_feature_dim=edge_feature_dim,
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    # Should fail if node_feature_dim is negative
+    with pytest.raises(AssertionError):
+        InteractionNetworkBlock(node_feature_dim=-1)
+
+    # Should fail if edge_feature_dim is negative
+    with pytest.raises(AssertionError):
+        InteractionNetworkBlock(node_feature_dim=3, edge_feature_dim=-1)
+
+    # Should fail if hidden_dim is negative
+    with pytest.raises(AssertionError):
+        InteractionNetworkBlock(node_feature_dim=3, hidden_dim=-1)
+
+    # Should fail if n_message_layers is negative
+    with pytest.raises(AssertionError):
+        InteractionNetworkBlock(node_feature_dim=3, n_message_layers=-1)
+
+    # Should fail if n_update_layers is negative
+    with pytest.raises(AssertionError):
+        InteractionNetworkBlock(node_feature_dim=3, n_update_layers=-1)
+
+
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+def test_forward(edge_feature_dim):
+
+    model = InteractionNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_feature_dim,
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    if edge_feature_dim == 0:
+        output_ = model(edge_index=edge_index, x=x)
+    else:
+        output_ = model(edge_index=edge_index, x=x, edge_attr=edge_attr)
+    assert output_.shape == x.shape
+
+
+@pytest.mark.parametrize("edge_feature_dim", [0, 2])
+def test_backward(edge_feature_dim):
+
+    model = InteractionNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=edge_feature_dim,
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    )
+
+    if edge_feature_dim == 0:
+        output_ = model(edge_index=edge_index, x=x.requires_grad_())
+    else:
+        output_ = model(
+            edge_index=edge_index,
+            x=x.requires_grad_(),
+            edge_attr=edge_attr.requires_grad_(),
+        )
+
+    loss = torch.mean(output_)
+    loss.backward()
+    assert x.grad.shape == x.shape