fix merge conflict

Author: sfluegel05

chebai_graph/models/base.py

@@ -74,11 +74,14 @@ class GraphNetWrapper(GraphBaseNet, ABC):
     """
 
     def __init__(
-        self, config: dict, n_linear_layers: int, n_molecule_properties: Optional[int] = 0, use_batch_norm: bool = False, **kwargs
-    ) -> None:
+        self,
+        config: dict,
+        n_linear_layers: int,
+        n_molecule_properties: Optional[int] = 0,
+        use_batch_norm: bool = False,
+        **kwargs,
+    ):
         """
-        Initialize the GNN and linear layers.
-
         Args:
             config (dict): Model configuration.
             n_linear_layers (int): Number of linear layers.
@@ -91,7 +94,9 @@ def __init__(
         self.activation = torch.nn.ELU
         self.lin_input_dim = self._get_lin_seq_input_dim(
             gnn_out_dim=gnn_out_dim,
-            n_molecule_properties=n_molecule_properties if n_molecule_properties is not None else 0,
+            n_molecule_properties=(
+                n_molecule_properties if n_molecule_properties is not None else 0
+            ),
         )
         self.use_batch_norm = use_batch_norm
         if self.use_batch_norm:

chebai_graph/models/dynamic_gni.py

@@ -33,25 +33,43 @@ def __init__(self, config: dict[str, Any], **kwargs: Any):
         )
         self.distribution = distribution
 
-        self.complete_randomness = config.get("complete_randomness", True)
+        self.complete_randomness = (
+            str(config.get("complete_randomness", "True")).lower() == "true"
+        )
+
+        print("Using complete randomness: ", self.complete_randomness)
 
         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
             )
-            self.random_pad_edge = (
-                int(config["random_pad_edge"])
-                if config.get("random_pad_edge") is not None
+            if self.pad_node_features is not None:
+                print(
+                    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.pad_edge_features = (
+                int(config["pad_edge_features"])
+                if config.get("pad_edge_features") is not None
                 else None
             )
+            if self.pad_edge_features is not None:
+                print(
+                    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,
@@ -93,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
 
@@ -388,32 +358,6 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
 
         return base_df[base_data[0].keys()].to_dict("records")
 
-    @property
-    def processed_file_names_dict(self) -> dict:
-        """
-        Returns a dictionary for the processed and tokenized data files.
-
-        Returns:
-            dict: A dictionary mapping dataset keys to their respective file names.
-                  For example, {"data": "data.pt"}.
-        """
-        if self.n_token_limit is not None:
-            return {"data": f"data_maxlen{self.n_token_limit}.pt"}
-
-        data_pt_filename = "data"
-        if self.zero_pad_node:
-            data_pt_filename += f"_zpn{self.zero_pad_node}"
-        if self.zero_pad_edge:
-            data_pt_filename += f"_zpe{self.zero_pad_edge}"
-        if self.random_pad_node:
-            data_pt_filename += f"_rpn{self.random_pad_node}"
-        if self.random_pad_edge:
-            data_pt_filename += f"_rpe{self.random_pad_edge}"
-        if self.random_pad_node or self.random_pad_edge:
-            data_pt_filename += f"_D{self.distribution}"
-
-        return {"data": data_pt_filename + ".pt"}
-
 
 class GraphPropAsPerNodeType(DataPropertiesSetter, ABC):
     def __init__(self, properties=None, transform=None, **kwargs):

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")