dynamic gni

Author: aditya0by0

chebai_graph/models/__init__.py

@@ -4,6 +4,7 @@
     ResGatedAugNodePoolGraphPred,
     ResGatedGraphNodeFGNodePoolGraphPred,
 )
+from .dynamic_gni import ResGatedDynamicGNIGraphPred
 from .gat import GATGraphPred
 from .resgated import ResGatedGraphPred
 
@@ -14,4 +15,5 @@
     "GATGraphPred",
     "GATAugNodePoolGraphPred",
     "GATGraphNodeFGNodePoolGraphPred",
+    "ResGatedDynamicGNIGraphPred",
 ]

chebai_graph/models/dynamic_gni.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import ELU
+from torch_geometric.data import Data as GraphData
+from torch_geometric.nn.models.basic_gnn import BasicGNN
+
+from chebai_graph.preprocessing.reader import RandomFeatureInitializationReader
+
+from .base import GraphModelBase, GraphNetWrapper
+from .resgated import ResGatedModel
+
+
+class ResGatedDynamicGNI(GraphModelBase):
+    """
+    Base model class for applying ResGatedGraphConv layers to graph-structured data
+    with dynamic initialization of features for nodes and edges.
+
+    Args:
+        config (dict): Configuration dictionary containing model hyperparameters.
+        **kwargs: Additional keyword arguments for parent class.
+    """
+
+    def __init__(self, config: dict[str, Any], **kwargs: Any):
+        super().__init__(config=config, **kwargs)
+        self.activation = ELU()  # Instantiate ELU once for reuse.
+        distribution = config.get("distribution", "normal")
+        assert distribution in ["normal", "uniform", "xavier_normal", "xavier_uniform"]
+        self.distribution = distribution
+
+        self.resgated: BasicGNN = ResGatedModel(
+            in_channels=self.in_channels,
+            hidden_channels=self.hidden_channels,
+            out_channels=self.out_channels,
+            num_layers=self.num_layers,
+            edge_dim=self.edge_dim,
+            act=self.activation,
+        )
+
+    def forward(self, batch: dict[str, Any]) -> Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            batch (dict): A batch containing graph input features under the key "features".
+
+        Returns:
+            Tensor: The output node-level embeddings after the final activation.
+        """
+        graph_data = batch["features"][0]
+        assert isinstance(graph_data, GraphData), "Expected GraphData instance"
+
+        random_x = torch.empty(graph_data.x.shape[0], graph_data.x.shape[1])
+        RandomFeatureInitializationReader.random_gni(random_x, self.distribution)
+        random_edge_attr = torch.empty(
+            graph_data.edge_attr.shape[0], graph_data.edge_attr.shape[1]
+        )
+        RandomFeatureInitializationReader.random_gni(
+            random_edge_attr, self.distribution
+        )
+
+        out = self.resgated(
+            x=graph_data.x.float(),
+            edge_index=graph_data.edge_index.long(),
+            edge_attr=graph_data.edge_attr,
+        )
+
+        return self.activation(out)
+
+
+class ResGatedDynamicGNIGraphPred(GraphNetWrapper):
+    """
+    Wrapper for graph-level prediction using ResGatedDynamicGNI.
+
+    This class instantiates the core GNN model using the provided config.
+    """
+
+    def _get_gnn(self, config: dict[str, Any]) -> ResGatedDynamicGNI:
+        """
+        Returns the core ResGated GNN model.
+
+        Args:
+            config (dict): Configuration dictionary for the GNN model.
+
+        Returns:
+            ResGatedDynamicGNI: The core graph convolutional network.
+        """
+        return ResGatedDynamicGNI(config=config)

chebai_graph/preprocessing/datasets/chebi.py

@@ -36,7 +36,7 @@
     GN_WithAtoms_FG_WithAtoms_NoFGE,
     GraphPropertyReader,
     GraphReader,
-    RandomNodeInitializationReader,
+    RandomFeatureInitializationReader,
 )
 
 from .utils import resolve_property
@@ -518,7 +518,7 @@ def _merge_props_into_base(
 
 
 class ChEBI50_StaticGNI(DataPropertiesSetter, ChEBIOver50):
-    READER = RandomNodeInitializationReader
+    READER = RandomFeatureInitializationReader
 
     def _setup_properties(self): ...
 

chebai_graph/preprocessing/reader/__init__.py

@@ -10,7 +10,7 @@
     GN_WithAtoms_FG_WithAtoms_NoFGE,
 )
 from .reader import GraphPropertyReader, GraphReader
-from .static_gni import RandomNodeInitializationReader
+from .static_gni import RandomFeatureInitializationReader
 
 __all__ = [
     "GraphReader",
@@ -20,7 +20,7 @@
     "AtomFGReader_NoFGEdges_WithGraphNode",
     "AtomFGReader_WithFGEdges_NoGraphNode",
     "AtomFGReader_WithFGEdges_WithGraphNode",
-    "RandomNodeInitializationReader",
+    "RandomFeatureInitializationReader",
     "GN_WithAtoms_FG_WithAtoms_FGE",
     "GN_WithAtoms_FG_WithAtoms_NoFGE",
     "GN_WithAllNodes_FG_WithAtoms_FGE",

chebai_graph/preprocessing/reader/static_gni.py

@@ -12,7 +12,7 @@
 from .reader import GraphPropertyReader
 
 
-class RandomNodeInitializationReader(GraphPropertyReader):
+class RandomFeatureInitializationReader(GraphPropertyReader):
     def __init__(
         self,
         num_node_properties: int,
@@ -46,24 +46,9 @@ def _read_data(self, raw_data):
         )
         random_molecule_properties = torch.empty(1, self.num_molecule_properties)
 
-        if self.distribution == "normal":
-            torch.nn.init.normal_(random_x)
-            torch.nn.init.normal_(random_edge_attr)
-            torch.nn.init.normal_(random_molecule_properties)
-        elif self.distribution == "uniform":
-            torch.nn.init.uniform_(random_x, a=-1.0, b=1.0)
-            torch.nn.init.uniform_(random_edge_attr, a=-1.0, b=1.0)
-            torch.nn.init.uniform_(random_molecule_properties, a=-1.0, b=1.0)
-        elif self.distribution == "xavier_normal":
-            torch.nn.init.xavier_normal_(random_x)
-            torch.nn.init.xavier_normal_(random_edge_attr)
-            torch.nn.init.xavier_normal_(random_molecule_properties)
-        elif self.distribution == "xavier_uniform":
-            torch.nn.init.xavier_uniform_(random_x)
-            torch.nn.init.xavier_uniform_(random_edge_attr)
-            torch.nn.init.xavier_uniform_(random_molecule_properties)
-        else:
-            raise ValueError("Unknown distribution type")
+        self.random_gni(random_x, self.distribution)
+        self.random_gni(random_edge_attr, self.distribution)
+        self.random_gni(random_molecule_properties, self.distribution)
 
         data.x = random_x
         data.edge_attr = random_edge_attr
@@ -73,3 +58,16 @@ def _read_data(self, raw_data):
     def read_property(self, *args, **kwargs) -> Exception:
         """This reader does not support reading specific properties."""
         raise NotImplementedError("This reader only performs random initialization.")
+
+    @staticmethod
+    def random_gni(tensor: torch.Tensor, distribution: str) -> None:
+        if distribution == "normal":
+            torch.nn.init.normal_(tensor)
+        elif distribution == "uniform":
+            torch.nn.init.uniform_(tensor, a=-1.0, b=1.0)
+        elif distribution == "xavier_normal":
+            torch.nn.init.xavier_normal_(tensor)
+        elif distribution == "xavier_uniform":
+            torch.nn.init.xavier_uniform_(tensor)
+        else:
+            raise ValueError("Unknown distribution type")

configs/model/resgated_dynamic_gni.yml

@@ -0,0 +1,13 @@
+class_path: chebai_graph.models.ResGatedDynamicGNIGraphPred
+init_args:
+  optimizer_kwargs:
+    lr: 1e-3
+  config:
+    in_channels: 158 # number of node/atom properties
+    hidden_channels: 256
+    out_channels: 512
+    num_layers: 4
+    edge_dim: 7 # number of bond properties
+    dropout: 0
+  n_molecule_properties: 0
+  n_linear_layers: 1