fix egnn + equivariance/invariance tests

Co-authored-by: Dario Coscia <[email protected]>

Author: GiovanniCanali

pina/model/block/message_passing/en_equivariant_network_block.py

@@ -10,7 +10,6 @@
 class EnEquivariantNetworkBlock(MessagePassing):
     """
     Implementation of the E(n) Equivariant Graph Neural Network block.
-
     This block is used to perform message-passing between nodes and edges in a
     graph neural network, following the scheme proposed by Satorras et al. in
     2021. It serves as an inner block in a larger graph neural network
@@ -102,14 +101,24 @@ def __init__(
         )
 
         # Layer for updating the node features
-        self.update_net = FeedForward(
+        self.update_feat_net = FeedForward(
             input_dimensions=node_feature_dim + pos_dim,
             output_dimensions=node_feature_dim,
             inner_size=hidden_dim,
             n_layers=n_update_layers,
             func=activation,
         )
 
+        # Layer for updating the node positions
+        # The output dimension is set to 1 for equivariant updates
+        self.update_pos_net = FeedForward(
+            input_dimensions=pos_dim,
+            output_dimensions=1,
+            inner_size=hidden_dim,
+            n_layers=n_update_layers,
+            func=activation,
+        )
+
     def forward(self, x, pos, edge_index, edge_attr=None):
         """
         Forward pass of the block, triggering the message-passing routine.
@@ -143,22 +152,62 @@ def message(self, x_i, x_j, pos_i, pos_j, edge_attr):
         :param edge_attr: The edge attributes.
         :type edge_attr: torch.Tensor | LabelTensor
         :return: The message to be passed.
-        :rtype: torch.Tensor
+        :rtype: tuple(torch.Tensor, torch.Tensor)
         """
-        dist = torch.norm(pos_i - pos_j, dim=-1, keepdim=True) ** 2
+        # Compute the euclidean distance between the sender and recipient nodes
+        diff = pos_i - pos_j
+        dist = torch.norm(diff, dim=-1, keepdim=True) ** 2
+
+        # Compute the message input
         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_)
+        # Compute the messages and their equivariant counterpart
+        m_ij = self.message_net(input_)
+        message = diff * self.update_pos_net(m_ij)
+
+        return message, m_ij
 
-    def update(self, message, x, pos, edge_index):
+    def aggregate(self, inputs, index, ptr=None, dim_size=None):
+        """
+        Aggregate the messages at the nodes during message passing.
+
+        This method receives a tuple of tensors corresponding to the messages
+        to be aggregated. Both messages are aggregated separately according to
+        the specified aggregation scheme.
+
+        :param tuple(torch.Tensor) inputs: Tuple containing two messages to
+            aggregate.
+        :param index: The indices of target nodes for each message. This tensor
+            specifies which node each message is aggregated into.
+        :type index: torch.Tensor | LabelTensor
+        :param ptr: Optional tensor to specify the slices of messages for each
+            node (used in some aggregation strategies). Default is None.
+        :type ptr: torch.Tensor | LabelTensor
+        :param int dim_size: Optional size of the output dimension, i.e.,
+            number of nodes. Default is None.
+        :return: Tuple of aggregated tensors corresponding to (aggregated
+            messages for position updates, aggregated messages for feature
+            updates).
+        :rtype: tuple(torch.Tensor, torch.Tensor)
+        """
+        # Unpack the messages from the inputs
+        message, m_ij = inputs
+
+        # Aggregate messages as usual using self.aggr method
+        agg_message = super().aggregate(message, index, ptr, dim_size)
+        agg_m_ij = super().aggregate(m_ij, index, ptr, dim_size)
+
+        return agg_message, agg_m_ij
+
+    def update(self, aggregated_inputs, x, pos, edge_index):
         """
         Update the node features and the node coordinates with the received
         messages.
 
-        :param torch.Tensor message: The message to be passed.
+        :param tuple(torch.Tensor) aggregated_inputs: The messages to be passed.
         :param x: The node features.
         :type x: torch.Tensor | LabelTensor
         :param pos: The euclidean coordinates of the nodes.
@@ -167,10 +216,14 @@ def update(self, message, x, pos, edge_index):
         :return: The updated node features and node positions.
         :rtype: tuple(torch.Tensor, torch.Tensor)
         """
-        # Update the node features
-        x = self.update_net(torch.cat((x, message), dim=-1))
+        # aggregated_inputs is tuple (agg_message, agg_m_ij)
+        agg_message, agg_m_ij = aggregated_inputs
+
+        # Update node features with aggregated m_ij
+        x = self.update_feat_net(torch.cat((x, agg_m_ij), dim=-1))
+
+        # Degree for normalization of position updates
+        c = degree(edge_index[1], pos.shape[0]).unsqueeze(-1).clamp(min=1)
+        pos = pos + agg_message / c
 
-        # Update the node positions
-        c = degree(edge_index[0], pos.shape[0]).unsqueeze(-1)
-        pos = pos + message / c
         return x, pos

tests/test_messagepassing/test_equivariant_network_block.py

@@ -128,3 +128,38 @@ def test_backward(edge_feature_dim):
     loss.backward()
     assert x.grad.shape == x.shape
     assert pos.grad.shape == pos.shape
+
+
+def test_equivariance():
+
+    # Graph to be fully connected and undirected
+    edge_index = torch.combinations(torch.arange(x.shape[0]), r=2).T
+    edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    # Random rotation (det(rotation) should be 1)
+    rotation = torch.linalg.qr(torch.rand(pos.shape[-1], pos.shape[-1])).Q
+    if torch.det(rotation) < 0:
+        rotation[:, 0] *= -1
+
+    # Random translation
+    translation = torch.rand(1, pos.shape[-1])
+
+    model = EnEquivariantNetworkBlock(
+        node_feature_dim=x.shape[1],
+        edge_feature_dim=0,
+        pos_dim=pos.shape[1],
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+    ).eval()
+
+    h1, pos1 = model(edge_index=edge_index, x=x, pos=pos)
+    h2, pos2 = model(
+        edge_index=edge_index, x=x, pos=pos @ rotation.T + translation
+    )
+
+    # Transform model output
+    pos1_transformed = (pos1 @ rotation.T) + translation
+
+    assert torch.allclose(pos2, pos1_transformed, atol=1e-5)
+    assert torch.allclose(h1, h2, atol=1e-5)

tests/test_messagepassing/test_radial_field_network_block.py

@@ -65,3 +65,28 @@ def test_backward():
     loss = torch.mean(output_)
     loss.backward()
     assert x.grad.shape == x.shape
+
+
+def test_equivariance():
+
+    # Graph to be fully connected and undirected
+    edge_index = torch.combinations(torch.arange(x.shape[0]), r=2).T
+    edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    # Random rotation (det(rotation) should be 1)
+    rotation = torch.linalg.qr(torch.rand(x.shape[-1], x.shape[-1])).Q
+    if torch.det(rotation) < 0:
+        rotation[:, 0] *= -1
+
+    # Random translation
+    translation = torch.rand(1, x.shape[-1])
+
+    model = RadialFieldNetworkBlock(node_feature_dim=x.shape[1]).eval()
+
+    pos1 = model(edge_index=edge_index, x=x)
+    pos2 = model(edge_index=edge_index, x=x @ rotation.T + translation)
+
+    # Transform model output
+    pos1_transformed = (pos1 @ rotation.T) + translation
+
+    assert torch.allclose(pos2, pos1_transformed, atol=1e-5)

tests/test_messagepassing/test_schnet_block.py

@@ -71,3 +71,25 @@ def test_backward():
     loss = torch.mean(output_)
     loss.backward()
     assert x.grad.shape == x.shape
+
+
+def test_invariance():
+
+    # Graph to be fully connected and undirected
+    edge_index = torch.combinations(torch.arange(x.shape[0]), r=2).T
+    edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
+
+    # Random rotation (det(rotation) should be 1)
+    rotation = torch.linalg.qr(torch.rand(pos.shape[-1], pos.shape[-1])).Q
+    if torch.det(rotation) < 0:
+        rotation[:, 0] *= -1
+
+    # Random translation
+    translation = torch.rand(1, pos.shape[-1])
+
+    model = SchnetBlock(node_feature_dim=x.shape[1]).eval()
+
+    out1 = model(edge_index=edge_index, x=x, pos=pos)
+    out2 = model(edge_index=edge_index, x=x, pos=pos @ rotation.T + translation)
+
+    assert torch.allclose(out1, out2, atol=1e-5)