static gni

Author: aditya0by0

chebai_graph/preprocessing/datasets/__init__.py

@@ -2,6 +2,7 @@
     ChEBI50_Atom_WGNOnly_GraphProp,
     ChEBI50_NFGE_NGN_GraphProp,
     ChEBI50_NFGE_WGN_GraphProp,
+    ChEBI50_StaticGNI,
     ChEBI50_WFGE_NGN_GraphProp,
     ChEBI50_WFGE_WGN_GraphProp,
     ChEBI50GraphData,
@@ -19,4 +20,5 @@
     "ChEBI50_NFGE_WGN_GraphProp",
     "ChEBI50_WFGE_NGN_GraphProp",
     "ChEBI50_WFGE_WGN_GraphProp",
+    "ChEBI50_StaticGNI",
 ]

chebai_graph/preprocessing/datasets/chebi.py

@@ -31,6 +31,7 @@
     AtomsFGReader_NoFGEdges_NoGraphNode,
     GraphPropertyReader,
     GraphReader,
+    RandomNodeInitializationReader,
 )
 
 from .utils import resolve_property
@@ -188,9 +189,11 @@ def __init__(
                 f"[Info] Atom-level features will be zero-padded with "
                 f"{self.zero_pad_atom} additional dimensions."
             )
-        print(
-            f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}"
-        )
+
+        if self.properties:
+            print(
+                f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}"
+            )
 
     def _merge_props_into_base(self, row: pd.Series) -> GeomData:
         """
@@ -504,6 +507,24 @@ def _merge_props_into_base(
         )
 
 
+class ChEBI50_StaticGNI(DataPropertiesSetter, ChEBIOver50):
+    READER = RandomNodeInitializationReader
+
+    def _setup_properties(self): ...
+
+    def load_processed_data_from_file(self, filename):
+        base_data = super().load_processed_data_from_file(filename)
+        base_df = pd.DataFrame(base_data)
+
+        rank_zero_info(
+            f"Use following values for given parameters for model configuration: \n\t"
+            f"in_channels: {self.reader.num_node_properties} , "
+            f"edge_dim: {self.reader.num_bond_properties}, "
+            f"n_molecule_properties: {self.reader.num_molecule_properties}"
+        )
+        return base_df[base_data[0].keys()].to_dict("records")
+
+
 class ChEBI50GraphProperties(GraphPropertiesMixIn, ChEBIOver50):
     """ChEBIOver50 dataset with molecular property encodings."""
 

chebai_graph/preprocessing/reader/__init__.py

@@ -6,6 +6,7 @@
     AtomsFGReader_NoFGEdges_NoGraphNode,
 )
 from .reader import GraphPropertyReader, GraphReader
+from .static_gni import RandomNodeInitializationReader
 
 __all__ = [
     "GraphReader",
@@ -15,4 +16,5 @@
     "AtomFGReader_NoFGEdges_WithGraphNode",
     "AtomFGReader_WithFGEdges_NoGraphNode",
     "AtomFGReader_WithFGEdges_WithGraphNode",
+    "RandomNodeInitializationReader",
 ]

chebai_graph/preprocessing/reader/static_gni.py

@@ -0,0 +1,75 @@
+"""
+Abboud, Ralph, et al.
+"The surprising power of graph neural networks with random node initialization."
+arXiv preprint arXiv:2010.01179 (2020).
+
+Code Reference: https://github.com/ralphabb/GNN-RNI/blob/main/GNNHyb.py
+"""
+
+import torch
+from torch_geometric.data import Data as GeomData
+
+from .reader import GraphPropertyReader
+
+
+class RandomNodeInitializationReader(GraphPropertyReader):
+    def __init__(
+        self,
+        num_node_properties: int,
+        num_bond_properties: int,
+        num_molecule_properties: int,
+        distribution: str,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.num_node_properties = num_node_properties
+        self.num_bond_properties = num_bond_properties
+        self.num_molecule_properties = num_molecule_properties
+        assert distribution in ["normal", "uniform", "xavier_normal", "xavier_uniform"]
+        self.distribution = distribution
+
+    def name(self) -> str:
+        """
+        Get the name identifier of the reader.
+
+        Returns:
+            str: The name of the reader.
+        """
+        return f"gni-{self.distribution}-node{self.num_node_properties}-bond{self.num_bond_properties}-mol{self.num_molecule_properties}"
+
+    def _read_data(self, raw_data):
+        data: GeomData = super()._read_data(raw_data)
+        random_x = torch.empty(data.x.shape[0], self.num_node_properties)
+        random_edge_attr = torch.empty(
+            data.edge_index.shape[1], self.num_bond_properties
+        )
+        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")
+
+        data.x = random_x
+        data.edge_attr = random_edge_attr
+        data.molecule_attr = random_molecule_properties
+        return data
+
+    def read_property(self, *args, **kwargs) -> Exception:
+        """This reader does not support reading specific properties."""
+        raise NotImplementedError("This reader only performs random initialization.")

configs/data/chebi50_static_gni.yml

@@ -0,0 +1,7 @@
+class_path: chebai_graph.preprocessing.datasets.ChEBI50_StaticGNI
+init_args:
+  reader_kwargs:
+    num_node_properties: 158
+    num_bond_properties: 7
+    num_molecule_properties: 200
+    distribution: normal