global fixes

Author: GiovanniCanali

pina/model/block/message_passing/__init__.py

@@ -3,7 +3,13 @@
 __all__ = [
     "InteractionNetworkBlock",
     "DeepTensorNetworkBlock",
+    "EnEquivariantNetworkBlock",
+    "RadialFieldNetworkBlock",
+    "SchnetBlock",
 ]
 
 from .interaction_network_block import InteractionNetworkBlock
 from .deep_tensor_network_block import DeepTensorNetworkBlock
+from .egnn_block import EnEquivariantNetworkBlock
+from .radial_field_network_block import RadialFieldNetworkBlock
+from .schnet_block import SchnetBlock

pina/model/block/message_passing/deep_tensor_network_block.py

@@ -2,16 +2,17 @@
 
 import torch
 from torch_geometric.nn import MessagePassing
-from ....utils import check_consistency
+from ....utils import check_positive_integer
 
 
 class DeepTensorNetworkBlock(MessagePassing):
     """
     Implementation of the Deep Tensor Network block.
 
     This block is used to perform message-passing between nodes and edges in a
-    graph neural network, following the scheme proposed by Schutt et al. (2017).
-    It serves as an inner block in a larger graph neural network architecture.
+    graph neural network, following the scheme proposed by Schutt et al. in
+    2017. It serves as an inner block in a larger graph neural network
+    architecture.
 
     The message between two nodes connected by an edge is computed by applying a
     linear transformation to the sender node features and the edge features,
@@ -24,7 +25,7 @@ class DeepTensorNetworkBlock(MessagePassing):
     .. seealso::
 
         **Original reference**: Schutt, K., Arbabzadah, F., Chmiela, S. et al.
-        *Quantum-Chemical Insights from Deep Tensor Neural Networks*.
+        (2017). *Quantum-Chemical Insights from Deep Tensor Neural Networks*.
         Nature Communications 8, 13890 (2017).
         DOI: `<https://doi.org/10.1038/ncomms13890>_`.
     """
@@ -57,51 +58,36 @@ def __init__(
             flow means that messages are sent from the target node to the
             source node. See :class:`torch_geometric.nn.MessagePassing` for more
             details. Default is "source_to_target".
-        :raises ValueError: If `node_feature_dim` is not a positive integer.
-        :raises ValueError: If `edge_feature_dim` is not a positive integer.
+        :raises AssertionError: If `node_feature_dim` is not a positive integer.
+        :raises AssertionError: If `edge_feature_dim` is not a positive integer.
         """
         super().__init__(aggr=aggr, node_dim=node_dim, flow=flow)
 
-        # Check consistency
-        check_consistency(node_feature_dim, int)
-        check_consistency(edge_feature_dim, int)
-
         # Check values
-        if node_feature_dim <= 0:
-            raise ValueError(
-                "`node_feature_dim` must be a positive integer,"
-                f" got {node_feature_dim}."
-            )
-
-        if edge_feature_dim <= 0:
-            raise ValueError(
-                "`edge_feature_dim` must be a positive integer,"
-                f" got {edge_feature_dim}."
-            )
-
-        # Initialize parameters
-        self.node_feature_dim = node_feature_dim
-        self.edge_feature_dim = edge_feature_dim
+        check_positive_integer(node_feature_dim, strict=True)
+        check_positive_integer(edge_feature_dim, strict=True)
+
+        # Activation function
         self.activation = activation
 
         # Layer for processing node features
         self.node_layer = torch.nn.Linear(
-            in_features=self.node_feature_dim,
-            out_features=self.node_feature_dim,
+            in_features=node_feature_dim,
+            out_features=node_feature_dim,
             bias=True,
         )
 
         # Layer for processing edge features
         self.edge_layer = torch.nn.Linear(
-            in_features=self.edge_feature_dim,
-            out_features=self.node_feature_dim,
+            in_features=edge_feature_dim,
+            out_features=node_feature_dim,
             bias=True,
         )
 
         # Layer for computing the message
         self.message_layer = torch.nn.Linear(
-            in_features=self.node_feature_dim,
-            out_features=self.node_feature_dim,
+            in_features=node_feature_dim,
+            out_features=node_feature_dim,
             bias=False,
         )
 

pina/model/block/message_passing/egnn_block.py

@@ -3,135 +3,170 @@
 import torch
 from torch_geometric.nn import MessagePassing
 from torch_geometric.utils import degree
+from ....utils import check_positive_integer
+from ....model import FeedForward
 
 
-class EnEquivariantGraphBlock(MessagePassing):
+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. (2021).
-    It serves as an inner block in a larger graph neural network architecture.
+    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
+    architecture.
 
     The message between two nodes connected by an edge is computed by applying a
     linear transformation to the sender node features and the edge features,
-    followed by a non-linear activation function. Messages are then aggregated
-    using an aggregation scheme (e.g., sum, mean, min, max, or product).
+    together with the squared euclidean distance between the sender and
+    recipient node positions, followed by a non-linear activation function.
+    Messages are then aggregated using an aggregation scheme (e.g., sum, mean,
+    min, max, or product).
 
-    The update step is performed by a simple addition of the incoming messages
-    to the node features.
+    The update step is performed by applying another MLP to the concatenation of
+    the incoming messages and the node features. Here, also the node
+    positions are updated by adding the incoming messages divided by the
+    degree of the recipient node.
 
     .. seealso::
 
-        **Original reference** Satorras, V. G., Hoogeboom, E., & Welling, M. (2021, July). 
-        E (n) equivariant graph neural networks. 
-        In International conference on machine learning (pp. 9323-9332). PMLR.
+        **Original reference** Satorras, V. G., Hoogeboom, E., Welling, M.
+        (2021). *E(n) Equivariant Graph Neural Networks.*
+        In International Conference on Machine Learning.
+        DOI: `<https://doi.org/10.48550/arXiv.2102.09844>_`.
     """
 
     def __init__(
         self,
-        channels_x,
-        channels_m,
-        channels_a,
-        aggr: str = "add",
-        hidden_channels: int = 64,
-        **kwargs,
+        node_feature_dim,
+        edge_feature_dim,
+        pos_dim,
+        hidden_dim=64,
+        n_message_layers=2,
+        n_update_layers=2,
+        activation=torch.nn.SiLU,
+        aggr="add",
+        node_dim=-2,
+        flow="source_to_target",
     ):
         """
-        Initialization of the :class:`EnEquivariantGraphBlock` class.
-
-        :param int channels_x: The dimension of the node features.
-        :param int channels_m: The dimension of the Euclidean coordinates (should be =3).
-        :param int channels_a: The dimension of the edge features.
+        Initialization of the :class:`EnEquivariantNetworkBlock` class.
+
+        :param int node_feature_dim: The dimension of the node features.
+        :param int edge_feature_dim: The dimension of the edge features.
+        :param int pos_dim: The dimension of the position features.
+        :param int hidden_dim: The dimension of the hidden features.
+            Default is 64.
+        :param int n_message_layers: The number of layers in the message
+            network. Default is 2.
+        :param int n_update_layers: The number of layers in the update network.
+            Default is 2.
+        :param torch.nn.Module activation: The activation function.
+            Default is :class:`torch.nn.SiLU`.
         :param str aggr: The aggregation scheme to use for message passing.
             Available options are "add", "mean", "min", "max", "mul".
             See :class:`torch_geometric.nn.MessagePassing` for more details.
             Default is "add".
-        :param int hidden_channels_dim: The hidden dimension in each MLPs initialized in the block.
+        :param int node_dim: The axis along which to propagate. Default is -2.
+        :param str flow: The direction of message passing. Available options
+            are "source_to_target" and "target_to_source".
+            The "source_to_target" flow means that messages are sent from
+            the source node to the target node, while the "target_to_source"
+            flow means that messages are sent from the target node to the
+            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 `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.
+        :raises AssertionError: If `n_update_layers` is not a positive integer.
         """
-        super().__init__(aggr=aggr, **kwargs)
-
-        self.phi_e = torch.nn.Sequential(
-            torch.nn.Linear(2 * channels_x + 1 + channels_a, hidden_channels),
-            torch.nn.LayerNorm(hidden_channels),
-            torch.nn.SiLU(),
-            torch.nn.Linear(hidden_channels, channels_m),
-            torch.nn.LayerNorm(channels_m),
-            torch.nn.SiLU(),
-        )
-        self.phi_pos = torch.nn.Sequential(
-            torch.nn.Linear(channels_m, hidden_channels),
-            torch.nn.LayerNorm(hidden_channels),
-            torch.nn.SiLU(),
-            torch.nn.Linear(hidden_channels, 1),
+        super().__init__(aggr=aggr, node_dim=node_dim, flow=flow)
+
+        # Check values
+        check_positive_integer(node_feature_dim, strict=True)
+        check_positive_integer(edge_feature_dim, strict=True)
+        check_positive_integer(pos_dim, strict=True)
+        check_positive_integer(hidden_dim, strict=True)
+        check_positive_integer(n_message_layers, strict=True)
+        check_positive_integer(n_update_layers, strict=True)
+
+        # Layer for computing the message
+        self.message_net = FeedForward(
+            input_dimensions=2 * node_feature_dim + edge_feature_dim + 1,
+            output_dimensions=pos_dim,
+            inner_size=hidden_dim,
+            n_layers=n_message_layers,
+            func=activation,
         )
-        self.phi_x = torch.nn.Sequential(
-            torch.nn.Linear(channels_x + channels_m, hidden_channels),
-            torch.nn.LayerNorm(hidden_channels),
-            torch.nn.SiLU(),
-            torch.nn.Linear(hidden_channels, channels_x),
+
+        # Layer for updating the node features
+        self.update_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,
         )
 
-    def forward(self, x, pos, edge_attr, edge_index, c=None):
+    def forward(self, x, pos, edge_index, edge_attr):
         """
         Forward pass of the block, triggering the message-passing routine.
 
         :param x: The node features.
         :type x: torch.Tensor | LabelTensor
-        :param pos_i: 3D Euclidean coordinates.
-        :type pos_i: 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. 
-        :return: The updated node features.
-        :rtype: torch.Tensor
+        :param pos: The euclidean coordinates of the nodes.
+        :type pos: torch.Tensor | LabelTensor
+        :param torch.Tensor edge_index: The edge indices.
+        :param edge_attr: The edge attributes. Default is None.
+        :type edge_attr: torch.Tensor | LabelTensor
+        :return: The updated node features and node positions.
+        :rtype: tuple(torch.Tensor, torch.Tensor)
         """
-        if c is None:
-            c = degree(edge_index[0], pos.shape[0]).unsqueeze(-1)
         return self.propagate(
-            edge_index=edge_index, x=x, pos=pos, edge_attr=edge_attr, c=c
+            edge_index=edge_index, x=x, pos=pos, edge_attr=edge_attr
         )
 
     def message(self, x_i, x_j, pos_i, pos_j, edge_attr):
         """
         Compute the message to be passed between nodes and edges.
 
-        :param x_i: Node features of the sender nodes.
+        :param x_i: The node features of the recipient nodes.
         :type x_i: torch.Tensor | LabelTensor
-        :param pos_i: 3D Euclidean coordinates of the sender nodes.
+        :param x_j: The node features of the sender nodes.
+        :type x_j: torch.Tensor | LabelTensor
+        :param pos_i: The node coordinates of the recipient nodes.
         :type pos_i: torch.Tensor | LabelTensor
+        :param pos_j: The node coordinates of the sender nodes.
+        :type pos_j: torch.Tensor | LabelTensor
         :param edge_attr: The edge attributes.
         :type edge_attr: torch.Tensor | LabelTensor
         :return: The message to be passed.
         :rtype: torch.Tensor
         """
-        mpos_ij = self.phi_e(
-            torch.cat(
-                [
-                    x_i,
-                    x_j,
-                    torch.norm(pos_i - pos_j, dim=-1, keepdim=True) ** 2,
-                    edge_attr,
-                ],
-                dim=-1,
-            )
-        )
-        mpos_ij = (pos_i - pos_j) * self.phi_pos(mpos_ij)
-        return  mpos_ij
+        dist = torch.norm(pos_i - pos_j, dim=-1, keepdim=True) ** 2
+        input_ = torch.cat((x_i, x_j, dist, edge_attr), dim=-1)
+        return self.message_net(input_)
 
-    def update(self, message, x, pos, c):
+    def update(self, message, x, pos, edge_index):
         """
-        Update the node features with the received messages.
+        Update the node features and the node coordinates with the received
+        messages.
 
         :param torch.Tensor message: The message to be passed.
         :param x: The node features.
         :type x: torch.Tensor | LabelTensor
-        :param pos: The 3D Euclidean coordinates of the nodes.
+        :param pos: The euclidean coordinates of the nodes.
         :type pos: torch.Tensor | LabelTensor
-        :param c: the constant that divides the aggregated message (it should be (M-1), where M is the number of nodes)
-        :type pos: torch.Tensor 
-        :return: The concatenation of the update position features and the updated node features.
-        :rtype: torch.Tensor
+        :param torch.Tensor edge_index: The edge indices.
+        :return: The updated node features and node positions.
+        :rtype: tuple(torch.Tensor, torch.Tensor)
         """
-        x = self.phi_x(torch.cat([x, message], dim=-1))
-        pos = pos + (message / c)
+        # Update the node features
+        x = self.update_net(torch.cat((x, message), dim=-1))
+
+        # Update the node positions
+        c = degree(edge_index[0], pos.shape[0]).unsqueeze(-1)
+        pos = pos + message / c
         return pos, x