doc model resgated and gat

Author: aditya0by0

chebai_graph/models/gat.py

@@ -7,11 +7,28 @@
 
 
 class GATGraphConvNetBase(GraphModelBase):
-    def __init__(self, config, **kwargs):
+    """
+    Graph Attention Network (GAT) base module for graph convolution.
+
+    Uses PyTorch Geometric's `GAT` implementation to process atomic node features
+    and bond edge attributes through multiple attention heads and layers.
+    """
+
+    def __init__(self, config: dict, **kwargs):
+        """
+        Initialize the GATGraphConvNetBase.
+
+        Args:
+            config (dict): Model configuration containing:
+                - 'heads' (int): Number of attention heads.
+                - 'v2' (bool): Whether to use the GATv2 variant.
+                - Other required GraphModelBase parameters.
+            **kwargs: Additional arguments for the base class.
+        """
         super().__init__(config=config, **kwargs)
         self.heads = int(config["heads"])
         self.v2 = bool(config["v2"])
-        self.activation = ELU()  # instantiate once
+        self.activation = ELU()  # Instantiate ELU once for reuse.
         self.gat = GAT(
             in_channels=self.n_atom_properties,
             hidden_channels=self.hidden_length,
@@ -24,20 +41,49 @@ def __init__(self, config, **kwargs):
         )
 
     def forward(self, batch: dict) -> torch.Tensor:
+        """
+        Forward pass through the GAT network.
+
+        Processes atomic node features and edge attributes, and applies
+        an ELU activation to the output.
+
+        Args:
+            batch (dict): Input batch containing:
+                - 'features': A list with a `GraphData` object as its first element.
+
+        Returns:
+            torch.Tensor: Node embeddings after GAT and activation.
+        """
         graph_data = batch["features"][0]
         assert isinstance(graph_data, GraphData)
 
-        a = self.gat(
+        out = self.gat(
             x=graph_data.x.float(),
             edge_index=graph_data.edge_index,
             edge_attr=graph_data.edge_attr,
         )
 
-        return self.activation(a)
+        return self.activation(out)
 
 
 class GATGraphPred(GraphNetWrapper):
+    """
+    GAT-based graph prediction model.
+
+    Uses a `GATGraphConvNetBase` as the GNN backbone for generating node embeddings,
+    which are then pooled by the GraphNetWrapper for final prediction.
+    """
+
     NAME = "GATGraphPred"
 
-    def _get_gnn(self, config):
+    def _get_gnn(self, config: dict) -> GATGraphConvNetBase:
+        """
+        Instantiate the GAT graph convolutional network base.
+
+        Args:
+            config (dict): Model configuration.
+
+        Returns:
+            GATGraphConvNetBase: The initialized GNN.
+        """
         return GATGraphConvNetBase(config=config)

chebai_graph/models/resgated.py

@@ -8,20 +8,36 @@
 
 
 class ResGatedGraphConvNetBase(GraphModelBase):
-    """GNN that supports edge attributes"""
+    """
+    Residual Gated Graph Convolutional Network with edge attributes support.
+
+    This model uses a stack of `ResGatedGraphConv` layers from PyTorch Geometric,
+    allowing edge attributes as part of message passing. A final projection layer maps
+    to the hidden length specified for downstream graph prediction tasks.
+    """
 
     NAME = "ResGatedGraphConvNetBase"
 
-    def __init__(self, config, **kwargs):
+    def __init__(self, config: dict, **kwargs):
+        """
+        Initialize the ResGatedGraphConvNetBase.
+
+        Args:
+            config (dict): Configuration dictionary with keys:
+                - 'in_length' (int): Intermediate feature length used in GNN layers.
+                - Other parameters inherited from GraphModelBase.
+            **kwargs: Additional keyword arguments passed to GraphModelBase.
+        """
         super().__init__(config=config, **kwargs)
         self.in_length = int(config["in_length"])
 
         self.activation = F.elu
         self.dropout = nn.Dropout(self.dropout_rate)
 
-        self.convs = torch.nn.ModuleList([])
+        self.convs = torch.nn.ModuleList()
         for i in range(self.n_conv_layers):
             if i == 0:
+                # Initial layer uses atom features as input
                 self.convs.append(
                     tgnn.ResGatedGraphConv(
                         self.n_atom_properties,
@@ -30,36 +46,66 @@ def __init__(self, config, **kwargs):
                         edge_dim=self.n_bond_properties,
                     )
                 )
+            # Intermediate layers
             self.convs.append(
                 tgnn.ResGatedGraphConv(
                     self.in_length, self.in_length, edge_dim=self.n_bond_properties
                 )
             )
+
+        # Final projection layer to hidden dimension
         self.final_conv = tgnn.ResGatedGraphConv(
             self.in_length, self.hidden_length, edge_dim=self.n_bond_properties
         )
 
-    def forward(self, batch):
+    def forward(self, batch: dict) -> torch.Tensor:
+        """
+        Forward pass through residual gated GNN layers.
+
+        Args:
+            batch (dict): A batch containing:
+                - 'features': A list with a `GraphData` instance as the first element.
+
+        Returns:
+            torch.Tensor: Node-level embeddings of shape [num_nodes, hidden_length].
+        """
         graph_data = batch["features"][0]
         assert isinstance(graph_data, GraphData)
-        a = graph_data.x.float()
-        # a = self.embedding(a)
+
+        x = graph_data.x.float()  # Atom features
 
         for conv in self.convs:
             assert isinstance(conv, tgnn.ResGatedGraphConv)
-            a = self.activation(
-                conv(a, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr)
+            x = self.activation(
+                conv(x, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr)
             )
-        a = self.activation(
+
+        x = self.activation(
             self.final_conv(
-                a, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr
+                x, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr
             )
         )
-        return a
+
+        return x
 
 
 class ResGatedGraphPred(GraphNetWrapper):
+    """
+    Residual Gated GNN for Graph Prediction.
+
+    Uses `ResGatedGraphConvNetBase` as the GNN encoder to compute node embeddings.
+    """
+
     NAME = "ResGatedGraphPred"
 
-    def _get_gnn(self, config):
+    def _get_gnn(self, config: dict) -> ResGatedGraphConvNetBase:
+        """
+        Instantiate the residual gated GNN backbone.
+
+        Args:
+            config (dict): Model configuration.
+
+        Returns:
+            ResGatedGraphConvNetBase: The GNN encoder.
+        """
         return ResGatedGraphConvNetBase(config=config)