fix egnn + equivariance test

Author: GiovanniCanali

pina/model/block/message_passing/en_equivariant_network_block.py

@@ -102,14 +102,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.
@@ -145,13 +155,21 @@ def message(self, x_i, x_j, pos_i, pos_j, edge_attr):
         :return: The message to be passed.
         :rtype: 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 save them for feature update
+        self._m_ij = self.message_net(input_)
+
+        # Rescale the message by the euclidean distance
+        return diff * self.update_pos_net(self._m_ij)
 
     def update(self, message, x, pos, edge_index):
         """
@@ -167,10 +185,14 @@ def update(self, message, x, pos, edge_index):
         :return: The updated node features and node positions.
         :rtype: tuple(torch.Tensor, torch.Tensor)
         """
+        # Sum the incoming messages for each node (m_i = sum_j m_ij)
+        m_sum = torch.zeros(x.size(0), self._m_ij.shape[-1], device=x.device)
+        m_sum.index_add_(0, edge_index[1], self._m_ij)
+
         # Update the node features
-        x = self.update_net(torch.cat((x, message), dim=-1))
+        x = self.update_feat_net(torch.cat((x, m_sum), dim=-1))
 
         # Update the node positions
-        c = degree(edge_index[0], pos.shape[0]).unsqueeze(-1)
+        c = degree(edge_index[1], pos.shape[0]).unsqueeze(-1)
         pos = pos + message / c
         return x, pos

tests/test_messagepassing/test_equivariant_network_block.py

@@ -128,3 +128,37 @@ 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()
+
+    _, pos1 = model(edge_index=edge_index, x=x, pos=pos)
+    _, 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)