include zero padding random init reader static method

Author: aditya0by0

chebai_graph/models/dynamic_gni.py

@@ -41,35 +41,35 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
 
         if not self.complete_randomness:
             assert (
-                "random_pad_node" in config or "random_pad_edge" in config
-            ), "Missing 'random_pad_node' or 'random_pad_edge' in config when complete_randomness is False"
-            self.random_pad_node = (
-                int(config["random_pad_node"])
-                if config.get("random_pad_node") is not None
+                "pad_node_features" in config or "pad_edge_features" in config
+            ), "Missing 'pad_node_features' or 'pad_edge_features' in config when complete_randomness is False"
+            self.pad_node_features = (
+                int(config["pad_node_features"])
+                if config.get("pad_node_features") is not None
                 else None
             )
-            if self.random_pad_node is not None:
+            if self.pad_node_features is not None:
                 print(
-                    f"[Info] Node features will be padded with {self.random_pad_node} "
+                    f"[Info] Node features will be padded with {self.pad_node_features} "
                     f"new set of random features from distribution {self.distribution} "
                     f"in each forward pass."
                 )
 
-            self.random_pad_edge = (
-                int(config["random_pad_edge"])
-                if config.get("random_pad_edge") is not None
+            self.pad_edge_features = (
+                int(config["pad_edge_features"])
+                if config.get("pad_edge_features") is not None
                 else None
             )
-            if self.random_pad_edge is not None:
+            if self.pad_edge_features is not None:
                 print(
-                    f"[Info] Edge features will be padded with {self.random_pad_edge} "
+                    f"[Info] Edge features will be padded with {self.pad_edge_features} "
                     f"new set of random features from distribution {self.distribution} "
                     f"in each forward pass."
                 )
 
             assert (
-                self.random_pad_node > 0 or self.random_pad_edge > 0
-            ), "'random_pad_node' or 'random_pad_edge' must be positive integers"
+                self.pad_node_features > 0 or self.pad_edge_features > 0
+            ), "'pad_node_features' or 'pad_edge_features' must be positive integers"
 
         self.resgated: BasicGNN = ResGatedModel(
             in_channels=self.in_channels,
@@ -111,21 +111,21 @@ def forward(self, batch: dict[str, Any]) -> Tensor:
                 new_edge_attr, self.distribution
             )
         else:
-            if self.random_pad_node is not None:
+            if self.pad_node_features is not None:
                 pad_node = torch.empty(
                     graph_data.x.shape[0],
-                    self.random_pad_node,
+                    self.pad_node_features,
                     device=self.device,
                 )
                 RandomFeatureInitializationReader.random_gni(
                     pad_node, self.distribution
                 )
                 new_x = torch.cat((graph_data.x, pad_node), dim=1)
 
-            if self.random_pad_edge is not None:
+            if self.pad_edge_features is not None:
                 pad_edge = torch.empty(
                     graph_data.edge_attr.shape[0],
-                    self.random_pad_edge,
+                    self.pad_edge_features,
                     device=self.device,
                 )
                 RandomFeatureInitializationReader.random_gni(

chebai_graph/preprocessing/datasets/chebi.py

@@ -189,45 +189,29 @@ def __init__(
         self,
         properties=None,
         transform=None,
-        zero_pad_node: int = None,
-        zero_pad_edge: int = None,
-        random_pad_node: int = None,
-        random_pad_edge: int = None,
+        pad_node_features: int = None,
+        pad_edge_features: int = None,
         distribution: str = "normal",
         **kwargs,
     ):
         super().__init__(properties, transform, **kwargs)
-        self.zero_pad_node = int(zero_pad_node) if zero_pad_node else None
-        if self.zero_pad_node:
-            print(
-                f"[Info] Node-level features will be zero-padded with "
-                f"{self.zero_pad_node} additional dimensions."
-            )
-
-        self.zero_pad_edge = int(zero_pad_edge) if zero_pad_edge else None
-        if self.zero_pad_edge:
-            print(
-                f"[Info] Edge-level features will be zero-padded with "
-                f"{self.zero_pad_edge} additional dimensions."
-            )
-
-        self.random_pad_edge = int(random_pad_edge) if random_pad_edge else None
-        self.random_pad_node = int(random_pad_node) if random_pad_node else None
-        if self.random_pad_node or self.random_pad_edge:
+        self.pad_edge_features = int(pad_edge_features) if pad_edge_features else None
+        self.pad_node_features = int(pad_node_features) if pad_node_features else None
+        if self.pad_node_features or self.pad_edge_features:
             assert (
                 distribution is not None
                 and distribution in RandomFeatureInitializationReader.DISTRIBUTIONS
-            ), "When using random padding, a valid distribution must be specified."
+            ), "When using padding for features, a valid distribution must be specified."
             self.distribution = distribution
-            if self.random_pad_node:
+            if self.pad_node_features:
                 print(
-                    f"[Info] Node-level features will be padded with "
-                    f"{self.random_pad_node} additional dimensions initialized from {self.distribution} distribution."
+                    f"[Info] Node-level features will be padded with random"
+                    f"{self.pad_node_features} values from {self.distribution} distribution."
                 )
-            if self.random_pad_edge:
+            if self.pad_edge_features:
                 print(
-                    f"[Info] Edge-level features will be padded with "
-                    f"{self.random_pad_edge} additional dimensions initialized from {self.distribution} distribution."
+                    f"[Info] Edge-level features will be padded with random"
+                    f"{self.pad_edge_features} values from {self.distribution} distribution."
                 )
 
         if self.properties:
@@ -276,24 +260,19 @@ def _merge_props_into_base(self, row: pd.Series) -> GeomData:
             else:
                 raise TypeError(f"Unsupported property type: {type(property).__name__}")
 
-        if self.zero_pad_node:
-            x = torch.cat([x, torch.zeros((x.shape[0], self.zero_pad_node))], dim=1)
-
-        if self.zero_pad_edge:
-            edge_attr = torch.cat(
-                [edge_attr, torch.zeros((edge_attr.shape[0], self.zero_pad_edge))],
-                dim=1,
+        if self.pad_node_features:
+            padding_values = torch.empty((x.shape[0], self.pad_node_features))
+            RandomFeatureInitializationReader.random_gni(
+                padding_values, self.distribution
             )
+            x = torch.cat([x, padding_values], dim=1)
 
-        if self.random_pad_node:
-            random_pad = torch.empty((x.shape[0], self.random_pad_node))
-            RandomFeatureInitializationReader.random_gni(random_pad, self.distribution)
-            x = torch.cat([x, random_pad], dim=1)
-
-        if self.random_pad_edge:
-            random_pad = torch.empty((edge_attr.shape[0], self.random_pad_edge))
-            RandomFeatureInitializationReader.random_gni(random_pad, self.distribution)
-            edge_attr = torch.cat([edge_attr, random_pad], dim=1)
+        if self.pad_edge_features:
+            padding_values = torch.empty((edge_attr.shape[0], self.pad_edge_features))
+            RandomFeatureInitializationReader.random_gni(
+                padding_values, self.distribution
+            )
+            edge_attr = torch.cat([edge_attr, padding_values], dim=1)
 
         return GeomData(
             x=x,
@@ -350,13 +329,9 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
         )
 
         in_channels_str = ""
-        if self.zero_pad_node:
-            n_node_properties += self.zero_pad_node
-            in_channels_str += f" (with {self.zero_pad_node} padded zeros)"
-
-        if self.random_pad_node:
-            n_node_properties += self.random_pad_node
-            in_channels_str += f" (with {self.random_pad_node} random padded values from {self.distribution} distribution)"
+        if self.pad_node_features:
+            n_node_properties += self.pad_node_features
+            in_channels_str += f" (with {self.pad_node_features} padded random values from {self.distribution} distribution)"
 
         in_channels_str = f"in_channels: {n_node_properties}" + in_channels_str
 
@@ -367,14 +342,9 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
             if isinstance(p, BondProperty)
         )
         edge_dim_str = ""
-
-        if self.zero_pad_edge:
-            n_edge_properties += self.zero_pad_edge
-            edge_dim_str += f" (with {self.zero_pad_edge} padded zeros)"
-
-        if self.random_pad_edge:
-            n_edge_properties += self.random_pad_edge
-            edge_dim_str += f" (with {self.random_pad_edge} random padded values from {self.distribution} distribution)"
+        if self.pad_edge_features:
+            n_edge_properties += self.pad_edge_features
+            edge_dim_str += f" (with {self.pad_edge_features} padded random values from {self.distribution} distribution)"
 
         edge_dim_str = f"edge_dim: {n_edge_properties}" + edge_dim_str
 

chebai_graph/preprocessing/reader/static_gni.py

@@ -13,7 +13,7 @@
 
 
 class RandomFeatureInitializationReader(GraphPropertyReader):
-    DISTRIBUTIONS = ["normal", "uniform", "xavier_normal", "xavier_uniform"]
+    DISTRIBUTIONS = ["normal", "uniform", "xavier_normal", "xavier_uniform", "zeros"]
 
     def __init__(
         self,
@@ -74,5 +74,7 @@ def random_gni(tensor: torch.Tensor, distribution: str) -> None:
             torch.nn.init.xavier_normal_(tensor)
         elif distribution == "xavier_uniform":
             torch.nn.init.xavier_uniform_(tensor)
+        elif distribution == "zeros":
+            torch.nn.init.zeros_(tensor)
         else:
             raise ValueError("Unknown distribution type")