add deep tensor network block

Author: GiovanniCanali

pina/model/block/message_passing/__init__.py

@@ -0,0 +1,9 @@
+"""Module for the message passing blocks of the graph neural models."""
+
+__all__ = [
+    "InteractionNetworkBlock",
+    "DeepTensorNetworkBlock",
+]
+
+from .interaction_network_block import InteractionNetworkBlock
+from .deep_tensor_network_block import DeepTensorNetworkBlock

pina/model/block/message_passing/deep_tensor_network_block.py

@@ -0,0 +1,128 @@
+"""Module for the Deep Tensor Network block."""
+
+import torch
+from torch_geometric.nn import MessagePassing
+
+
+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.
+
+    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).
+
+    The update step is performed by a simple addition of the incoming messages
+    to the node features.
+
+    .. seealso::
+
+        **Original reference**: Schutt, K., Arbabzadah, F., Chmiela, S. et al.
+        *Quantum-Chemical Insights from Deep Tensor Neural Networks*.
+        Nature Communications 8, 13890 (2017).
+        DOI: `<https://doi.org/10.1038/ncomms13890>_`
+    """
+
+    def __init__(
+        self,
+        node_feature_dim,
+        edge_feature_dim,
+        activation=torch.nn.Tanh,
+        aggr="add",
+        node_dim=-2,
+        flow="source_to_target",
+    ):
+        """
+        Initialization of the :class:`AVNOBDeepTensorNetworkBlocklock` class.
+
+        :param int node_feature_dim: The dimension of the node features.
+        :param int edge_feature_dim: The dimension of the edge features.
+        :param torch.nn.Module activation: The activation function.
+            Default is :class:`torch.nn.Tanh`.
+        :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 node_dim: The axis along which to propagate. Default is -2.
+        :param str flow: The direction of message passing.
+            See :class:`torch_geometric.nn.MessagePassing` for more details.
+            Default is "source_to_target".
+        """
+        super().__init__(aggr=aggr, node_dim=node_dim, flow=flow)
+
+        self.node_feature_dim = node_feature_dim
+        self.edge_feature_dim = edge_feature_dim
+        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,
+            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,
+            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,
+            bias=False,
+        )
+
+    def forward(self, x, edge_index, edge_attr):
+        """
+        Forward pass of the block. It performs a message-passing operation
+        between nodes and edges.
+
+        :param x: The node features.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor edge_index: The edge indeces.
+        :param edge_attr: The edge attributes.
+        :type edge_attr: torch.Tensor | LabelTensor
+        :return: The updated node features.
+        :rtype: torch.Tensor
+        """
+        return self.propagate(edge_index=edge_index, x=x, edge_attr=edge_attr)
+
+    def message(self, x_j, edge_attr):
+        """
+        Compute the message to be passed between nodes and edges.
+
+        :param x_j: The node features of the sender nodes.
+        :type x_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
+        """
+        # Process node and edge features
+        filter_node = self.node_layer(x_j)
+        filter_edge = self.edge_layer(edge_attr)
+
+        # Compute the message to be passed
+        message = self.message_layer(filter_node * filter_edge)
+
+        return self.activation(message)
+
+    def update(self, message, x):
+        """
+        Update the node features with the received messages.
+
+        :param torch.Tensor message: The message to be passed.
+        :param x: The node features.
+        :type x: torch.Tensor | LabelTensor
+        :return: The updated node features.
+        :rtype: torch.Tensor
+        """
+        return x + message

pina/model/block/message_passing/egnn_block.py

@@ -1,61 +1,97 @@
+"""Module for the E(n) Equivariant Graph Neural Network block."""
+
 import torch
-import torch.nn as nn
 from torch_geometric.nn import MessagePassing
 from torch_geometric.utils import degree
-from ....utils import check_consistency
 
 
 class EnEquivariantGraphBlock(MessagePassing):
-    def __init__(self,
-                 channels_h,
-                 channels_m, 
-                 channels_a,
-                 aggr: str = 'add', 
-                 hidden_channels: int = 64,
-                 **kwargs):
+    """
+    TODO
+    """
+
+    def __init__(
+        self,
+        channels_h,
+        channels_m,
+        channels_a,
+        aggr: str = "add",
+        hidden_channels: int = 64,
+        **kwargs,
+    ):
+        """
+        TODO
+        """
         super().__init__(aggr=aggr, **kwargs)
 
-        self.phi_e = nn.Sequential(
-                nn.Linear(2 * channels_h + 1 + channels_a, hidden_channels),
-                nn.LayerNorm(hidden_channels),
-                nn.SiLU(),
-                nn.Linear(hidden_channels, channels_m),
-                nn.LayerNorm(channels_m),
-                nn.SiLU()
+        self.phi_e = torch.nn.Sequential(
+            torch.nn.Linear(2 * channels_h + 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_x = nn.Sequential(
-                nn.Linear(channels_m, hidden_channels),
-                nn.LayerNorm(hidden_channels),
-                nn.SiLU(),
-                nn.Linear(hidden_channels, 1),
+        self.phi_x = torch.nn.Sequential(
+            torch.nn.Linear(channels_m, hidden_channels),
+            torch.nn.LayerNorm(hidden_channels),
+            torch.nn.SiLU(),
+            torch.nn.Linear(hidden_channels, 1),
+        )
+        self.phi_h = torch.nn.Sequential(
+            torch.nn.Linear(channels_h + channels_m, hidden_channels),
+            torch.nn.LayerNorm(hidden_channels),
+            torch.nn.SiLU(),
+            torch.nn.Linear(hidden_channels, channels_h),
         )
-        self.phi_h = nn.Sequential(
-                nn.Linear(channels_h + channels_m, hidden_channels),
-                nn.LayerNorm(hidden_channels),
-                nn.SiLU(),
-                nn.Linear(hidden_channels, channels_h),
-            )
 
     def forward(self, x, h, edge_attr, edge_index, c=None):
+        """
+        TODO
+        """
         if c is None:
             c = degree(edge_index[0], x.shape[0]).unsqueeze(-1)
-        return self.propagate(edge_index=edge_index, x=x, h=h, edge_attr=edge_attr, c=c)
+        return self.propagate(
+            edge_index=edge_index, x=x, h=h, edge_attr=edge_attr, c=c
+        )
 
     def message(self, x_i, x_j, h_i, h_j, edge_attr):
-        mh_ij = self.phi_e(torch.cat([h_i, h_j, torch.norm(x_i - x_j, dim=-1, keepdim=True)**2, edge_attr], dim=-1))
+        """
+        TODO
+        """
+        mh_ij = self.phi_e(
+            torch.cat(
+                [
+                    h_i,
+                    h_j,
+                    torch.norm(x_i - x_j, dim=-1, keepdim=True) ** 2,
+                    edge_attr,
+                ],
+                dim=-1,
+            )
+        )
         mx_ij = (x_i - x_j) * self.phi_x(mh_ij)
         return torch.cat((mx_ij, mh_ij), dim=-1)
 
     def update(self, aggr_out, x, h, edge_attr, c):
-        m_x, m_h = aggr_out[:, :self.m_len], aggr_out[:, self.m_len:]
+        """
+        TODO
+        """
+        m_x, m_h = aggr_out[:, : self.m_len], aggr_out[:, self.m_len :]
         h_l1 = self.phi_h(torch.cat([h, m_h], dim=-1))
         x_l1 = x + (m_x / c)
         return x_l1, h_l1
 
     @property
     def edge_function(self):
+        """
+        TODO
+        """
         return self._edge_function
 
     @property
     def attribute_function(self):
+        """
+        TODO
+        """
         return self._attribute_function

pina/model/block/message_passing/interaction_network_block.py

@@ -0,0 +1,10 @@
+"""Module for the Interaction Network block."""
+
+import torch
+from torch_geometric.nn import MessagePassing
+
+
+class InteractionNetworkBlock(MessagePassing):
+    """
+    TODO
+    """