revert mol props

Author: aditya0by0

chebai_graph/preprocessing/reader/static_gni.py

@@ -17,15 +17,15 @@ def __init__(
         self,
         num_node_properties: int,
         num_bond_properties: int,
-        # num_molecule_properties: int,
+        num_molecule_properties: int,
         distribution: str = "normal",
         *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
+        self.num_molecule_properties = num_molecule_properties
         assert distribution in ["normal", "uniform", "xavier_normal", "xavier_uniform"]
         self.distribution = distribution
 
@@ -44,30 +44,30 @@ def _read_data(self, raw_data):
         random_edge_attr = torch.empty(
             data.edge_index.shape[1], self.num_bond_properties
         )
-        # random_molecule_properties = torch.empty(1, self.num_molecule_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)
+            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)
+            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)
+            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)
+            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
+        data.molecule_attr = random_molecule_properties
         return data
 
     def read_property(self, *args, **kwargs) -> Exception:

configs/data/chebi50_static_gni.yml

@@ -3,3 +3,4 @@ init_args:
   reader_kwargs:
     num_node_properties: 158
     num_bond_properties: 7
+    num_molecule_properties: 0