From f48ce4f04566d8b783ee441ed955a4c24a0194c4 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 12:54:46 +0200 Subject: [PATCH 01/34] separate config folder for baselines --- .../model/{gat.yml => baselines/gat_baseline.yml} | 0 .../{resgated.yml => baselines/gine_baseline.yml} | 0 configs/model/baselines/rggcn_baseline.yml | 13 +++++++++++++ 3 files changed, 13 insertions(+) rename configs/model/{gat.yml => baselines/gat_baseline.yml} (100%) rename configs/model/{resgated.yml => baselines/gine_baseline.yml} (100%) create mode 100644 configs/model/baselines/rggcn_baseline.yml diff --git a/configs/model/gat.yml b/configs/model/baselines/gat_baseline.yml similarity index 100% rename from configs/model/gat.yml rename to configs/model/baselines/gat_baseline.yml diff --git a/configs/model/resgated.yml b/configs/model/baselines/gine_baseline.yml similarity index 100% rename from configs/model/resgated.yml rename to configs/model/baselines/gine_baseline.yml diff --git a/configs/model/baselines/rggcn_baseline.yml b/configs/model/baselines/rggcn_baseline.yml new file mode 100644 index 0000000..ccc6615 --- /dev/null +++ b/configs/model/baselines/rggcn_baseline.yml @@ -0,0 +1,13 @@ +class_path: chebai_graph.models.ResGatedGraphPred +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 From edd7c3064bb4eaf2cce69332108922c28dc65480 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 16:43:18 +0200 Subject: [PATCH 02/34] refine gine implementation --- README.md | 2 +- chebai_graph/models/gin_net.py | 72 +++++++++++++++---------------- configs/data/chebi50_baseline.yml | 12 ++++++ 3 files changed, 47 insertions(+), 39 deletions(-) create mode 100644 configs/data/chebi50_baseline.yml diff --git a/README.md b/README.md index 312f3df..b2a8ec0 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ The list can be found in the `configs/data/chebi50_graph_properties.yml` file. python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml ``` -## Augmented Graphs +## Augmented Graphs _See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_. Graph Neural Networks (GNNs) often fail to explicitly leverage the chemically meaningful substructures present within molecules (i.e. **functional groups (FGs)**). To make this implicit information explicitly accessible to GNNs, we augment molecular graphs with **artificial nodes** that represent these substructures. The resulting graph are referred to as **augmented graphs**. diff --git a/chebai_graph/models/gin_net.py b/chebai_graph/models/gin_net.py index 6fed4c6..72acce7 100644 --- a/chebai_graph/models/gin_net.py +++ b/chebai_graph/models/gin_net.py @@ -2,10 +2,11 @@ import torch import torch.nn.functional as F -import torch_geometric from torch_scatter import scatter_add +from torch_geometric.data import Data as GraphData -from chebai_graph.models.graph import GraphBaseNet +from chebai_graph.models.base import GraphModelBase, GraphNetWrapper +from torch_geometric import nn as tgnn class AggregateMLP(torch.nn.Module): @@ -24,7 +25,7 @@ def forward(self, x): return x -class GINEConvNet(GraphBaseNet): +class GINEConvNet(GraphModelBase): """Based on https://arxiv.org/pdf/1810.00826.pdf and https://arxiv.org/abs/1905.12265""" NAME = "GINEConvNet" @@ -32,42 +33,25 @@ class GINEConvNet(GraphBaseNet): def __init__(self, config: typing.Dict, **kwargs): super().__init__(**kwargs) - self.n_atom_properties = int(config["n_atom_properties"]) - self.n_bond_properties = int(config["n_bond_properties"]) - self.hidden_size = config["hidden_size"] - self.dropout_rate = config["dropout_rate"] - self.n_conv_layers = config["n_conv_layers"] if "n_conv_layers" in config else 5 - self.n_linear_layers = ( - config["n_linear_layers"] if "n_linear_layers" in config else 3 - ) - - self.dropout = torch.nn.Dropout(self.dropout_rate) - self.activation = F.relu + self.dropout_layer = torch.nn.Dropout(self.dropout) + self.activation = F.elu self.convs = torch.nn.ModuleList([]) # self.batch_norms = torch.nn.ModuleList([]) - for i in range(self.n_conv_layers): - in_length = self.n_atom_properties if i == 0 else self.hidden_size - out_length = self.hidden_size + for i in range(self.num_layers): self.convs.append( - torch_geometric.nn.GINEConv( - AggregateMLP(in_length, out_length, self.hidden_size), - edge_dim=self.n_bond_properties, + tgnn.GINEConv( + AggregateMLP( + self.in_channels, self.out_channels, self.hidden_channels + ), + edge_dim=self.edge_dim, ) ) # self.batch_norms.append(torch.nn.BatchNorm1d(out_length)) - self.linear_layers = torch.nn.ModuleList([]) - for i in range(self.n_linear_layers): - in_length = self.hidden_size - out_length = ( - self.out_dim if i == self.n_linear_layers - 1 else self.hidden_size - ) - self.linear_layers.append(torch.nn.Linear(in_length, out_length)) - def forward(self, batch): graph_data = batch["features"][0] - assert isinstance(graph_data, torch_geometric.data.Data) + assert isinstance(graph_data, GraphData) a = graph_data.x dropout_used = False # only apply dropout after first layer @@ -77,7 +61,7 @@ def forward(self, batch): conv(a, graph_data.edge_index.long(), graph_data.edge_attr) ) if not dropout_used: - a = self.dropout(a) + a = self.dropout_layer(a) dropout_used = True # a = norm(a) a = scatter_add(a, graph_data.batch, dim=0) @@ -85,12 +69,24 @@ def forward(self, batch): a = torch.cat(conv_out, dim=1) - for i in range(self.n_linear_layers): - if i != self.n_linear_layers - 1: - a = self.activation(self.linear_layers[i](a)) - else: - a = self.linear_layers[i](a) - if i == 0: - a = self.dropout(a) - return a + + +class GINEGraphPred(GraphNetWrapper): + """ + Wrapper for graph-level prediction using GINEConvNet. + + This class instantiates the core GNN model using the provided config. + """ + + def _get_gnn(self, config: dict[str, typing.Any]) -> GINEConvNet: + """ + Returns the core ResGated GNN model. + + Args: + config (dict): Configuration dictionary for the GNN model. + + Returns: + ResGatedGraphConvNetBase: The core graph convolutional network. + """ + return GINEConvNet(config=config) diff --git a/configs/data/chebi50_baseline.yml b/configs/data/chebi50_baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/chebi50_baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +init_args: + properties: + - chebai_graph.preprocessing.properties.AtomType + - chebai_graph.preprocessing.properties.NumAtomBonds + - chebai_graph.preprocessing.properties.AtomCharge + - chebai_graph.preprocessing.properties.AtomAromaticity + - chebai_graph.preprocessing.properties.AtomHybridization + - chebai_graph.preprocessing.properties.AtomNumHs + - chebai_graph.preprocessing.properties.BondType + - chebai_graph.preprocessing.properties.BondInRing + - chebai_graph.preprocessing.properties.BondAromaticity From 766e768868a8145107d5860afaffd3c2ff92c1a4 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 16:50:06 +0200 Subject: [PATCH 03/34] gine base config --- chebai_graph/models/__init__.py | 2 ++ configs/model/baselines/gine_baseline.yml | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 9e20b2d..f59e1c6 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -7,6 +7,7 @@ from .dynamic_gni import ResGatedDynamicGNIGraphPred from .gat import GATGraphPred from .resgated import ResGatedGraphPred +from .gin_net import GINEGraphPred __all__ = [ "ResGatedGraphPred", @@ -16,4 +17,5 @@ "GATAugNodePoolGraphPred", "GATGraphNodeFGNodePoolGraphPred", "ResGatedDynamicGNIGraphPred", + "GINEGraphPred", ] diff --git a/configs/model/baselines/gine_baseline.yml b/configs/model/baselines/gine_baseline.yml index ccc6615..8477ca6 100644 --- a/configs/model/baselines/gine_baseline.yml +++ b/configs/model/baselines/gine_baseline.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.ResGatedGraphPred +class_path: chebai_graph.models.GINEGraphPred init_args: optimizer_kwargs: lr: 1e-3 From 9a4afc65ffc1c39aaf1fcdb9cd827eeefe660bd9 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 20:56:12 +0200 Subject: [PATCH 04/34] move common data classes to base --- chebai_graph/preprocessing/datasets/base.py | 658 +++++++++++++++++++ chebai_graph/preprocessing/datasets/chebi.py | 651 +----------------- 2 files changed, 660 insertions(+), 649 deletions(-) create mode 100644 chebai_graph/preprocessing/datasets/base.py diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py new file mode 100644 index 0000000..abe13e4 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/base.py @@ -0,0 +1,658 @@ +import os +from abc import ABC +from collections.abc import Callable +from pprint import pformat +from typing import Optional + +import pandas as pd +import torch +import tqdm +from chebai.preprocessing.datasets.base import XYBaseDataModule +from lightning_utilities.core.rank_zero import rank_zero_info +from rdkit import Chem +from torch_geometric.data.data import Data as GeomData + +from chebai_graph.preprocessing.datasets.utils import resolve_property +from chebai_graph.preprocessing.properties import ( + AllNodeTypeProperty, + AtomNodeTypeProperty, + AtomProperty, + BondProperty, + FGNodeTypeProperty, + MolecularProperty, + MoleculeProperty, +) +from chebai_graph.preprocessing.reader import ( + GraphPropertyReader, + RandomFeatureInitializationReader, +) +from chebai_graph.preprocessing.reader.augmented_reader import _AugmentorReader + + +class DataPropertiesSetter(XYBaseDataModule, ABC): + """Mixin for adding molecular property encodings to graph-based given datasets.""" + + READER = GraphPropertyReader + + def __init__( + self, + properties: list | None = None, + transform: Callable | None = None, + **kwargs, + ): + """ + Initialize GraphPropertiesMixIn. + + Args: + properties: Optional list of MolecularProperty class paths or instances. + transform: Optional transformation applied to each data sample. + """ + super().__init__(**kwargs) + # atom_properties and bond_properties are given as lists containing class_paths + if properties is not None: + properties = [resolve_property(prop) for prop in properties] + properties = self._sort_properties(properties) + else: + properties = [] + self.properties: list[MolecularProperty] = properties + assert isinstance(self.properties, list) and all( + isinstance(p, MolecularProperty) for p in self.properties + ) + self.transform = transform + + def _sort_properties( + self, properties: list[MolecularProperty] + ) -> list[MolecularProperty]: + return sorted(properties, key=lambda prop: self.get_property_path(prop)) + + def _setup_properties(self) -> None: + """ + Process and cache molecular properties to disk. + + Returns: + None + """ + raw_data = [] + os.makedirs(self.processed_properties_dir, exist_ok=True) + + try: + file_names = self.processed_main_file_names + except NotImplementedError: + file_names = self.raw_file_names + + for file in file_names: + # processed_dir_main only exists for ChEBI datasets + path = os.path.join( + ( + self.processed_dir_main + if hasattr(self, "processed_dir_main") + else self.raw_dir + ), + file, + ) + raw_data += list(self._load_dict(path)) + + idents = [row["ident"] for row in raw_data] + features = [row["features"] for row in raw_data] + + # use vectorized version of encode function, apply only if value is present + def enc_if_not_none(encode, value): + return ( + [encode(v) for v in value] + if value is not None and len(value) > 0 + else None + ) + + if any( + not os.path.isfile(self.get_property_path(property)) + for property in self.properties + ): + # augment molecule graph if possible (this would also happen for the properties if needed, but this avoids redundancy) + if isinstance(self.reader, _AugmentorReader): + returned_results = [] + for mol in features: + try: + r = self.reader._create_augmented_graph(mol) + except Exception: + r = None + returned_results.append(r) + mols = [ + augmented_mol[1] if augmented_mol is not None else None + for augmented_mol in returned_results + ] + else: + mols = features + + for property in self.properties: + if not os.path.isfile(self.get_property_path(property)): + rank_zero_info(f"Processing property {property.name}") + # read all property values first, then encode + rank_zero_info(f"\tReading property values of {property.name}...") + property_values = [ + self.reader.read_property(mol, property) + if mol is not None + else None + for mol in tqdm.tqdm(mols) + ] + rank_zero_info(f"\tEncoding property values of {property.name}...") + property.encoder.on_start(property_values=property_values) + encoded_values = [ + enc_if_not_none(property.encoder.encode, value) + for value in tqdm.tqdm(property_values) + ] + assert len(encoded_values) == len(idents) == len(features) + torch.save( + [ + {property.name: torch.cat(feat), "ident": id} + for feat, id in zip(encoded_values, idents) + if feat is not None + ], + self.get_property_path(property), + ) + property.on_finish() + + @property + def processed_properties_dir(self) -> str: + return os.path.join(self.processed_dir, "properties") + + def get_property_path(self, property: MolecularProperty) -> str: + """ + Construct the cache path for a given molecular property. + + Args: + property: Instance of a MolecularProperty. + + Returns: + Path to the cached property file. + """ + return os.path.join( + self.processed_properties_dir, + f"{property.name}_{property.encoder.name}.pt", + ) + + def _after_setup(self, **kwargs) -> None: + """ + Finalize setup after ensuring properties are processed. + + Args: + **kwargs: Additional keyword arguments passed to superclass. + + Returns: + None + """ + self._setup_properties() + super()._after_setup(**kwargs) + + def _preprocess_smiles_for_pred( + self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None + ) -> Optional[dict]: + """Preprocess prediction data.""" + # Add dummy labels because the collate function requires them. + # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, + # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. + result = self.reader.to_data( + {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} + ) + # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object + if isinstance(result["features"], tuple): + result["features"], raw_data = result["features"] + if result is None or result["features"] is None: + return None + for property in self.properties: + property.encoder.eval = True + property_value = self.reader.read_property(raw_data, property) + if property_value is None or len(property_value) == 0: + encoded_value = None + else: + encoded_value = torch.stack( + [property.encoder.encode(v) for v in property_value] + ) + if len(encoded_value.shape) == 3: + encoded_value = encoded_value.squeeze(0) + result[property.name] = encoded_value + + result["features"] = self._prediction_merge_props_into_base_wrapper( + result, model_hparams + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + result["features"] = self.transform(result["features"]) + + return result + + def _prediction_merge_props_into_base_wrapper( + self, row: pd.Series | dict, model_hparams: Optional[dict] = None + ) -> GeomData: + """ + Wrapper to merge properties into base features for prediction. + + Args: + row: A dictionary or pd.Series containing 'features' and encoded properties. + Returns: + A GeomData object with merged features. + """ + return self._merge_props_into_base(row) + + +class GraphPropertiesMixIn(DataPropertiesSetter, ABC): + def __init__( + self, + properties=None, + transform=None, + pad_node_features: int | None = None, + pad_edge_features: int | None = None, + distribution: str = "normal", + **kwargs, + ): + super().__init__(properties, transform, **kwargs) + 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 padding for features, a valid distribution must be specified." + ) + self.distribution = distribution + if self.pad_node_features: + print( + f"[Info] Node-level features will be padded with random" + f"{self.pad_node_features} values from {self.distribution} distribution." + ) + if self.pad_edge_features: + print( + f"[Info] Edge-level features will be padded with random" + f"{self.pad_edge_features} values from {self.distribution} distribution." + ) + + 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 | dict) -> GeomData: + """ + Merge encoded molecular properties into the GeomData object. + + Args: + row: A dictionary containing 'features' and encoded properties. + + Returns: + A GeomData object with merged features. + """ + if isinstance(row["features"], tuple): + geom_data, _ = row[ + "features" + ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) + else: + geom_data = row["features"] + assert isinstance(geom_data, GeomData) + edge_attr = geom_data.edge_attr + x = geom_data.x + molecule_attr = torch.empty((1, 0)) + + for property in self.properties: + property_values = row[f"{property.name}"] + if isinstance(property_values, torch.Tensor): + if len(property_values.size()) == 0: + property_values = property_values.unsqueeze(0) + if len(property_values.size()) == 1: + property_values = property_values.unsqueeze(1) + else: + property_values = torch.zeros( + (0, property.encoder.get_encoding_length()) + ) + + if isinstance(property, AtomProperty): + x = torch.cat([x, property_values], dim=1) + elif isinstance(property, BondProperty): + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr = torch.cat( + [edge_attr, torch.cat([property_values, property_values], dim=0)], + dim=1, + ) + elif isinstance(property, MoleculeProperty): + molecule_attr = torch.cat([molecule_attr, property_values], dim=1) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + 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.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, + edge_index=geom_data.edge_index, + edge_attr=edge_attr, + molecule_attr=molecule_attr, + ) + + def load_processed_data( + self, kind: Optional[str] = None, filename: Optional[str] = None + ) -> list[dict]: + """ + Load dataset and merge cached properties into base features. + + Args: + filename: The path to the file to load. + + Returns: + List of data entries, each a dictionary. + """ + base_data = super().load_processed_data(kind, filename) + base_df = pd.DataFrame(base_data) + + for property in self.properties: + property_data = torch.load( + self.get_property_path(property), weights_only=False + ) + if len(property_data[0][property.name].shape) > 1: + property.encoder.set_encoding_length( + property_data[0][property.name].shape[1] + ) + + property_df = pd.DataFrame(property_data) + property_df.rename( + columns={property.name: f"{property.name}"}, inplace=True + ) + base_df = base_df.merge(property_df, on="ident", how="left") + + base_df["features"] = base_df.apply( + lambda row: self._merge_props_into_base(row), axis=1 + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + base_df["features"] = base_df["features"].apply(self.transform) + + prop_lengths = [ + (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties + ] + + # -------------------------- Count total node properties + n_node_properties = sum( + p.encoder.get_encoding_length() + for p in self.properties + if isinstance(p, AtomProperty) + ) + + in_channels_str = "" + 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 + + # -------------------------- Count total edge properties + n_edge_properties = sum( + p.encoder.get_encoding_length() + for p in self.properties + if isinstance(p, BondProperty) + ) + edge_dim_str = "" + 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 + + rank_zero_info( + f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" + f"Use following values for given parameters for model configuration: \n\t" + f"{in_channels_str} \n\t" + f"{edge_dim_str} \n\t" + f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}" + ) + + return base_df[base_data[0].keys()].to_dict("records") + + +class GraphPropAsPerNodeType(DataPropertiesSetter, ABC): + def __init__(self, properties=None, transform=None, **kwargs): + super().__init__(properties, transform, **kwargs) + # Sort properties so that AllNodeTypeProperty instances come first, rest of the properties order remain same + first = self._sort_properties( + [prop for prop in self.properties if isinstance(prop, AllNodeTypeProperty)] + ) + rest = self._sort_properties( + [ + prop + for prop in self.properties + if not isinstance(prop, AllNodeTypeProperty) + ] + ) + self.properties = first + rest + print( + "Properties are sorted so that `AllNodeTypeProperty` properties are first in sequence and rest of the order remains same\n", + f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}", + ) + + def load_processed_data( + self, kind: Optional[str] = None, filename: Optional[str] = None + ) -> list[dict]: + """ + Load dataset and merge cached properties into base features. + + Args: + filename: The path to the file to load. + + Returns: + List of data entries, each a dictionary. + """ + base_data = super().load_processed_data(kind, filename) + base_df = pd.DataFrame(base_data) + props_categories = { + "AllNodeTypeProperties": [], + "FGNodeTypeProperties": [], + "AtomNodeTypeProperties": [], + "GraphNodeTypeProperties": [], + "BondProperties": [], + } + n_atom_node_properties, n_fg_node_properties = 0, 0 + n_bond_properties, n_graph_node_properties = 0, 0 + prop_lengths = [] + for prop in self.properties: + prop_length = prop.encoder.get_encoding_length() + prop_name = prop.name + prop_lengths.append((prop_name, prop_length)) + if isinstance(prop, AllNodeTypeProperty): + n_atom_node_properties += prop_length + n_fg_node_properties += prop_length + n_graph_node_properties += prop_length + props_categories["AllNodeTypeProperties"].append(prop_name) + elif isinstance(prop, FGNodeTypeProperty): + n_fg_node_properties += prop_length + props_categories["FGNodeTypeProperties"].append(prop_name) + elif isinstance(prop, AtomNodeTypeProperty): + n_atom_node_properties += prop_length + props_categories["AtomNodeTypeProperties"].append(prop_name) + elif isinstance(prop, BondProperty): + n_bond_properties += prop_length + props_categories["BondProperties"].append(prop_name) + elif isinstance(prop, MoleculeProperty): + # molecule props will be used as graph node props + n_graph_node_properties += prop_length + props_categories["GraphNodeTypeProperties"].append(prop_name) + else: + raise TypeError(f"Unsupported property type: {type(prop).__name__}") + + n_node_properties = max( + n_atom_node_properties, n_fg_node_properties, n_graph_node_properties + ) + rank_zero_info( + f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" + f"Properties Categories:\n{pformat(props_categories)}\n\n" + f"n_atom_node_properties: {n_atom_node_properties}, " + f"n_fg_node_properties: {n_fg_node_properties}, " + f"n_bond_properties: {n_bond_properties}, " + f"n_graph_node_properties: {n_graph_node_properties}\n\n" + f"Use following values for given parameters for model configuration: \n\t" + f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0\n" + ) + + for property in self.properties: + rank_zero_info(f"Loading property {property.name}...") + property_data = torch.load( + self.get_property_path(property), weights_only=False + ) + if len(property_data[0][property.name].shape) > 1: + property.encoder.set_encoding_length( + property_data[0][property.name].shape[1] + ) + + property_df = pd.DataFrame(property_data) + property_df.rename( + columns={property.name: f"{property.name}"}, inplace=True + ) + base_df = base_df.merge(property_df, on="ident", how="left") + + base_df["features"] = base_df.apply( + lambda row: self._merge_props_into_base( + row, + max_len_node_properties=n_node_properties, + ), + axis=1, + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + base_df["features"] = base_df["features"].apply(self.transform) + + return base_df[base_data[0].keys()].to_dict("records") + + def _merge_props_into_base( + self, row: pd.Series, max_len_node_properties: int + ) -> GeomData: + """ + Merge encoded molecular properties into the GeomData object. + + Args: + row: A dictionary containing 'features' and encoded properties. + + Returns: + A GeomData object with merged features. + """ + geom_data = row["features"] + if geom_data is None: + return None + if isinstance(geom_data, tuple): + geom_data = geom_data[ + 0 + ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) + assert isinstance(geom_data, GeomData) + + is_atom_node = geom_data.is_atom_node + assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" + is_graph_node = geom_data.is_graph_node + assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" + + is_fg_node = ~is_atom_node & ~is_graph_node + num_nodes = geom_data.x.size(0) + edge_attr = geom_data.edge_attr + + # Initialize node feature matrix + assert max_len_node_properties is not None, ( + "Maximum len of node properties should not be None" + ) + x = torch.zeros((num_nodes, max_len_node_properties)) + + # Track column offsets for each node type + atom_offset, fg_offset, graph_offset = 0, 0, 0 + + for property in self.properties: + property_values = row[f"{property.name}"].to(dtype=torch.float32) + if isinstance(property_values, torch.Tensor): + if len(property_values.size()) == 0: + property_values = property_values.unsqueeze(0) + if len(property_values.size()) == 1: + property_values = property_values.unsqueeze(1) + else: + property_values = torch.zeros( + (0, property.encoder.get_encoding_length()) + ) + + enc_len = property_values.shape[1] + # -------------- Node properties --------------- + if isinstance(property, AllNodeTypeProperty): + x[:, atom_offset : atom_offset + enc_len] = property_values + atom_offset += enc_len + fg_offset += enc_len + graph_offset += enc_len + + elif isinstance(property, AtomNodeTypeProperty): + x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ + is_atom_node + ] + atom_offset += enc_len + + elif isinstance(property, FGNodeTypeProperty): + x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ + is_fg_node + ] + fg_offset += enc_len + + elif isinstance(property, MoleculeProperty): + x[is_graph_node, graph_offset : graph_offset + enc_len] = ( + property_values + ) + graph_offset += enc_len + + # ------------- Bond Properties -------------- + elif isinstance(property, BondProperty): + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr = torch.cat( + [edge_attr, torch.cat([property_values, property_values], dim=0)], + dim=1, + ) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + total_used_columns = max(atom_offset, fg_offset, graph_offset) + assert total_used_columns <= max_len_node_properties, ( + f"Used {total_used_columns} columns, but max allowed is {max_len_node_properties}" + ) + + return GeomData( + x=x, + edge_index=geom_data.edge_index, + edge_attr=edge_attr, + molecule_attr=torch.empty((1, 0)), # empty as not used for this class + is_atom_node=is_atom_node, + is_fg_node=is_fg_node, + is_graph_node=is_graph_node, + ) + + def _prediction_merge_props_into_base_wrapper( + self, row: pd.Series | dict, model_hparams: Optional[dict] = None + ) -> GeomData: + """ + Wrapper to merge properties into base features for prediction. + + Args: + row: A dictionary or pd.Series containing 'features' and encoded properties. + Returns: + A GeomData object with merged features. + """ + if ( + model_hparams is None + or "in_channels" not in model_hparams["config"] + or model_hparams["config"]["in_channels"] is None + ): + raise ValueError( + f"model_hparams must be provided for data class: {self.__class__.__name__}" + f" which should contain 'in_channels' key with valid value in 'config' dictionary." + ) + max_len_node_properties = int(model_hparams["config"]["in_channels"]) + return self._merge_props_into_base(row, max_len_node_properties) diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 2d03876..f6fa4d8 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -1,13 +1,6 @@ -import os from abc import ABC -from collections.abc import Callable -from pprint import pformat -from typing import Optional import pandas as pd -from chebai_graph.preprocessing.reader.augmented_reader import _AugmentorReader -import torch -import tqdm from chebai.preprocessing.datasets.chebi import ( ChEBIOver50, ChEBIOver100, @@ -16,17 +9,7 @@ ) from lightning_utilities.core.rank_zero import rank_zero_info from torch_geometric.data.data import Data as GeomData -from rdkit import Chem - -from chebai_graph.preprocessing.properties import ( - AllNodeTypeProperty, - AtomNodeTypeProperty, - AtomProperty, - BondProperty, - FGNodeTypeProperty, - MolecularProperty, - MoleculeProperty, -) + from chebai_graph.preprocessing.reader import ( AtomFGReader_NoFGEdges_WithGraphNode, AtomFGReader_WithFGEdges_NoGraphNode, @@ -37,12 +20,11 @@ GN_WithAllNodes_FG_WithAtoms_NoFGE, GN_WithAtoms_FG_WithAtoms_FGE, GN_WithAtoms_FG_WithAtoms_NoFGE, - GraphPropertyReader, GraphReader, RandomFeatureInitializationReader, ) -from chebai_graph.preprocessing.datasets.utils import resolve_property +from .base import DataPropertiesSetter, GraphPropAsPerNodeType, GraphPropertiesMixIn class ChEBI50GraphData(ChEBIOver50): @@ -54,635 +36,6 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -class DataPropertiesSetter(ChEBIOverX, ABC): - """Mixin for adding molecular property encodings to graph-based ChEBI datasets.""" - - READER = GraphPropertyReader - - def __init__( - self, - properties: list | None = None, - transform: Callable | None = None, - **kwargs, - ): - """ - Initialize GraphPropertiesMixIn. - - Args: - properties: Optional list of MolecularProperty class paths or instances. - transform: Optional transformation applied to each data sample. - """ - super().__init__(**kwargs) - # atom_properties and bond_properties are given as lists containing class_paths - if properties is not None: - properties = [resolve_property(prop) for prop in properties] - properties = self._sort_properties(properties) - else: - properties = [] - self.properties: list[MolecularProperty] = properties - assert isinstance(self.properties, list) and all( - isinstance(p, MolecularProperty) for p in self.properties - ) - self.transform = transform - - def _sort_properties( - self, properties: list[MolecularProperty] - ) -> list[MolecularProperty]: - return sorted(properties, key=lambda prop: self.get_property_path(prop)) - - def _setup_properties(self) -> None: - """ - Process and cache molecular properties to disk. - - Returns: - None - """ - raw_data = [] - os.makedirs(self.processed_properties_dir, exist_ok=True) - - try: - file_names = self.processed_main_file_names - except NotImplementedError: - file_names = self.raw_file_names - - for file in file_names: - # processed_dir_main only exists for ChEBI datasets - path = os.path.join( - ( - self.processed_dir_main - if hasattr(self, "processed_dir_main") - else self.raw_dir - ), - file, - ) - raw_data += list(self._load_dict(path)) - - idents = [row["ident"] for row in raw_data] - features = [row["features"] for row in raw_data] - - # use vectorized version of encode function, apply only if value is present - def enc_if_not_none(encode, value): - return ( - [encode(v) for v in value] - if value is not None and len(value) > 0 - else None - ) - - if any( - not os.path.isfile(self.get_property_path(property)) - for property in self.properties - ): - # augment molecule graph if possible (this would also happen for the properties if needed, but this avoids redundancy) - if isinstance(self.reader, _AugmentorReader): - returned_results = [] - for mol in features: - try: - r = self.reader._create_augmented_graph(mol) - except Exception: - r = None - returned_results.append(r) - mols = [ - augmented_mol[1] if augmented_mol is not None else None - for augmented_mol in returned_results - ] - else: - mols = features - - for property in self.properties: - if not os.path.isfile(self.get_property_path(property)): - rank_zero_info(f"Processing property {property.name}") - # read all property values first, then encode - rank_zero_info(f"\tReading property values of {property.name}...") - property_values = [ - self.reader.read_property(mol, property) - if mol is not None - else None - for mol in tqdm.tqdm(mols) - ] - rank_zero_info(f"\tEncoding property values of {property.name}...") - property.encoder.on_start(property_values=property_values) - encoded_values = [ - enc_if_not_none(property.encoder.encode, value) - for value in tqdm.tqdm(property_values) - ] - assert len(encoded_values) == len(idents) == len(features) - torch.save( - [ - {property.name: torch.cat(feat), "ident": id} - for feat, id in zip(encoded_values, idents) - if feat is not None - ], - self.get_property_path(property), - ) - property.on_finish() - - @property - def processed_properties_dir(self) -> str: - return os.path.join(self.processed_dir, "properties") - - def get_property_path(self, property: MolecularProperty) -> str: - """ - Construct the cache path for a given molecular property. - - Args: - property: Instance of a MolecularProperty. - - Returns: - Path to the cached property file. - """ - return os.path.join( - self.processed_properties_dir, - f"{property.name}_{property.encoder.name}.pt", - ) - - def _after_setup(self, **kwargs) -> None: - """ - Finalize setup after ensuring properties are processed. - - Args: - **kwargs: Additional keyword arguments passed to superclass. - - Returns: - None - """ - self._setup_properties() - super()._after_setup(**kwargs) - - def _preprocess_smiles_for_pred( - self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None - ) -> Optional[dict]: - """Preprocess prediction data.""" - # Add dummy labels because the collate function requires them. - # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, - # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. - result = self.reader.to_data( - {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} - ) - # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object - if isinstance(result["features"], tuple): - result["features"], raw_data = result["features"] - if result is None or result["features"] is None: - return None - for property in self.properties: - property.encoder.eval = True - property_value = self.reader.read_property(raw_data, property) - if property_value is None or len(property_value) == 0: - encoded_value = None - else: - encoded_value = torch.stack( - [property.encoder.encode(v) for v in property_value] - ) - if len(encoded_value.shape) == 3: - encoded_value = encoded_value.squeeze(0) - result[property.name] = encoded_value - - result["features"] = self._prediction_merge_props_into_base_wrapper( - result, model_hparams - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - result["features"] = self.transform(result["features"]) - - return result - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - return self._merge_props_into_base(row) - - -class GraphPropertiesMixIn(DataPropertiesSetter, ABC): - def __init__( - self, - properties=None, - transform=None, - pad_node_features: int = None, - pad_edge_features: int = None, - distribution: str = "normal", - **kwargs, - ): - super().__init__(properties, transform, **kwargs) - 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 padding for features, a valid distribution must be specified." - ) - self.distribution = distribution - if self.pad_node_features: - print( - f"[Info] Node-level features will be padded with random" - f"{self.pad_node_features} values from {self.distribution} distribution." - ) - if self.pad_edge_features: - print( - f"[Info] Edge-level features will be padded with random" - f"{self.pad_edge_features} values from {self.distribution} distribution." - ) - - 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 | dict) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - if isinstance(row["features"], tuple): - geom_data, _ = row[ - "features" - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - else: - geom_data = row["features"] - assert isinstance(geom_data, GeomData) - edge_attr = geom_data.edge_attr - x = geom_data.x - molecule_attr = torch.empty((1, 0)) - - for property in self.properties: - property_values = row[f"{property.name}"] - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - if isinstance(property, AtomProperty): - x = torch.cat([x, property_values], dim=1) - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - elif isinstance(property, MoleculeProperty): - molecule_attr = torch.cat([molecule_attr, property_values], dim=1) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - 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.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, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=molecule_attr, - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - - for property in self.properties: - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base(row), axis=1 - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - prop_lengths = [ - (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties - ] - - # -------------------------- Count total node properties - n_node_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, AtomProperty) - ) - - in_channels_str = "" - 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 - - # -------------------------- Count total edge properties - n_edge_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, BondProperty) - ) - edge_dim_str = "" - 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 - - rank_zero_info( - f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" - f"Use following values for given parameters for model configuration: \n\t" - f"{in_channels_str} \n\t" - f"{edge_dim_str} \n\t" - f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}" - ) - - return base_df[base_data[0].keys()].to_dict("records") - - -class GraphPropAsPerNodeType(DataPropertiesSetter, ABC): - def __init__(self, properties=None, transform=None, **kwargs): - super().__init__(properties, transform, **kwargs) - # Sort properties so that AllNodeTypeProperty instances come first, rest of the properties order remain same - first = self._sort_properties( - [prop for prop in self.properties if isinstance(prop, AllNodeTypeProperty)] - ) - rest = self._sort_properties( - [ - prop - for prop in self.properties - if not isinstance(prop, AllNodeTypeProperty) - ] - ) - self.properties = first + rest - print( - "Properties are sorted so that `AllNodeTypeProperty` properties are first in sequence and rest of the order remains same\n", - f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}", - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - props_categories = { - "AllNodeTypeProperties": [], - "FGNodeTypeProperties": [], - "AtomNodeTypeProperties": [], - "GraphNodeTypeProperties": [], - "BondProperties": [], - } - n_atom_node_properties, n_fg_node_properties = 0, 0 - n_bond_properties, n_graph_node_properties = 0, 0 - prop_lengths = [] - for prop in self.properties: - prop_length = prop.encoder.get_encoding_length() - prop_name = prop.name - prop_lengths.append((prop_name, prop_length)) - if isinstance(prop, AllNodeTypeProperty): - n_atom_node_properties += prop_length - n_fg_node_properties += prop_length - n_graph_node_properties += prop_length - props_categories["AllNodeTypeProperties"].append(prop_name) - elif isinstance(prop, FGNodeTypeProperty): - n_fg_node_properties += prop_length - props_categories["FGNodeTypeProperties"].append(prop_name) - elif isinstance(prop, AtomNodeTypeProperty): - n_atom_node_properties += prop_length - props_categories["AtomNodeTypeProperties"].append(prop_name) - elif isinstance(prop, BondProperty): - n_bond_properties += prop_length - props_categories["BondProperties"].append(prop_name) - elif isinstance(prop, MoleculeProperty): - # molecule props will be used as graph node props - n_graph_node_properties += prop_length - props_categories["GraphNodeTypeProperties"].append(prop_name) - else: - raise TypeError(f"Unsupported property type: {type(prop).__name__}") - - n_node_properties = max( - n_atom_node_properties, n_fg_node_properties, n_graph_node_properties - ) - rank_zero_info( - f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" - f"Properties Categories:\n{pformat(props_categories)}\n\n" - f"n_atom_node_properties: {n_atom_node_properties}, " - f"n_fg_node_properties: {n_fg_node_properties}, " - f"n_bond_properties: {n_bond_properties}, " - f"n_graph_node_properties: {n_graph_node_properties}\n\n" - f"Use following values for given parameters for model configuration: \n\t" - f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0\n" - ) - - for property in self.properties: - rank_zero_info(f"Loading property {property.name}...") - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base( - row, - max_len_node_properties=n_node_properties, - ), - axis=1, - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - return base_df[base_data[0].keys()].to_dict("records") - - def _merge_props_into_base( - self, row: pd.Series, max_len_node_properties: int - ) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - geom_data = row["features"] - if geom_data is None: - return None - if isinstance(geom_data, tuple): - geom_data = geom_data[ - 0 - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - assert isinstance(geom_data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" - - is_fg_node = ~is_atom_node & ~is_graph_node - num_nodes = geom_data.x.size(0) - edge_attr = geom_data.edge_attr - - # Initialize node feature matrix - assert max_len_node_properties is not None, ( - "Maximum len of node properties should not be None" - ) - x = torch.zeros((num_nodes, max_len_node_properties)) - - # Track column offsets for each node type - atom_offset, fg_offset, graph_offset = 0, 0, 0 - - for property in self.properties: - property_values = row[f"{property.name}"].to(dtype=torch.float32) - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - enc_len = property_values.shape[1] - # -------------- Node properties --------------- - if isinstance(property, AllNodeTypeProperty): - x[:, atom_offset : atom_offset + enc_len] = property_values - atom_offset += enc_len - fg_offset += enc_len - graph_offset += enc_len - - elif isinstance(property, AtomNodeTypeProperty): - x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ - is_atom_node - ] - atom_offset += enc_len - - elif isinstance(property, FGNodeTypeProperty): - x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ - is_fg_node - ] - fg_offset += enc_len - - elif isinstance(property, MoleculeProperty): - x[is_graph_node, graph_offset : graph_offset + enc_len] = ( - property_values - ) - graph_offset += enc_len - - # ------------- Bond Properties -------------- - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - total_used_columns = max(atom_offset, fg_offset, graph_offset) - assert total_used_columns <= max_len_node_properties, ( - f"Used {total_used_columns} columns, but max allowed is {max_len_node_properties}" - ) - - return GeomData( - x=x, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=torch.empty((1, 0)), # empty as not used for this class - is_atom_node=is_atom_node, - is_fg_node=is_fg_node, - is_graph_node=is_graph_node, - ) - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - if ( - model_hparams is None - or "in_channels" not in model_hparams["config"] - or model_hparams["config"]["in_channels"] is None - ): - raise ValueError( - f"model_hparams must be provided for data class: {self.__class__.__name__}" - f" which should contain 'in_channels' key with valid value in 'config' dictionary." - ) - max_len_node_properties = int(model_hparams["config"]["in_channels"]) - return self._merge_props_into_base(row, max_len_node_properties) - - class ChEBI50_StaticGNI(DataPropertiesSetter, ChEBIOver50): READER = RandomFeatureInitializationReader From 3c89f87430cc2efb22ef5bef0d7a46e4d3017e6d Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 21:31:19 +0200 Subject: [PATCH 05/34] create a separate aug base --- .../datasets/augmentation_base.py | 50 ++++++++++++++++++ chebai_graph/preprocessing/datasets/chebi.py | 51 ++----------------- 2 files changed, 54 insertions(+), 47 deletions(-) create mode 100644 chebai_graph/preprocessing/datasets/augmentation_base.py diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py new file mode 100644 index 0000000..9e2bf61 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -0,0 +1,50 @@ +from abc import ABC + +import pandas as pd +from torch_geometric.data.data import Data as GeomData + +from .base import GraphPropertiesMixIn + + +class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): + """Mixin for augmented graph data without additional graph nodes.""" + + READER = None + + def _merge_props_into_base(self, row: pd.Series) -> GeomData: + data = super()._merge_props_into_base(row) + geom_data = row["features"] + assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) + + is_atom_node = geom_data.is_atom_node + assert is_atom_node is not None, "is_atom_node must be set in the geom_data" + data.is_atom_node = is_atom_node + return data + + +class AugGraphPropMixIn_WithGraphNode(AugGraphPropMixIn_NoGraphNode, ABC): + """Mixin for augmented graph data with graph-level nodes.""" + + READER = None + + def _merge_props_into_base(self, row: pd.Series) -> GeomData: + data = super()._merge_props_into_base(row) + return self._add_graph_node_mask(data, row) + + def _add_graph_node_mask(self, data: GeomData, row: pd.Series) -> GeomData: + """ + Add a graph node mask to the GeomData object. + + Args: + data: A GeomData object with features. + row: A dictionary containing 'features' and other metadata. + + Returns: + Modified GeomData with graph node mask added. + """ + geom_data = row["features"] + assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) + is_graph_node = geom_data.is_graph_node + assert is_graph_node is not None, "is_graph_node must be set in the geom_data" + data.is_graph_node = is_graph_node + return data diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index f6fa4d8..59a3578 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -1,5 +1,3 @@ -from abc import ABC - import pandas as pd from chebai.preprocessing.datasets.chebi import ( ChEBIOver50, @@ -8,7 +6,6 @@ ChEBIOverXPartial, ) from lightning_utilities.core.rank_zero import rank_zero_info -from torch_geometric.data.data import Data as GeomData from chebai_graph.preprocessing.reader import ( AtomFGReader_NoFGEdges_WithGraphNode, @@ -24,6 +21,10 @@ RandomFeatureInitializationReader, ) +from .augmentation_base import ( + AugGraphPropMixIn_NoGraphNode, + AugGraphPropMixIn_WithGraphNode, +) from .base import DataPropertiesSetter, GraphPropAsPerNodeType, GraphPropertiesMixIn @@ -72,50 +73,6 @@ class ChEBI50GraphPropertiesPartial(ChEBI50GraphProperties, ChEBIOverXPartial): pass -class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): - """Mixin for augmented graph data without additional graph nodes.""" - - READER = None - - def _merge_props_into_base(self, row: pd.Series) -> GeomData: - data = super()._merge_props_into_base(row) - geom_data = row["features"] - assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "is_atom_node must be set in the geom_data" - data.is_atom_node = is_atom_node - return data - - -class AugGraphPropMixIn_WithGraphNode(AugGraphPropMixIn_NoGraphNode, ABC): - """Mixin for augmented graph data with graph-level nodes.""" - - READER = None - - def _merge_props_into_base(self, row: pd.Series) -> GeomData: - data = super()._merge_props_into_base(row) - return self._add_graph_node_mask(data, row) - - def _add_graph_node_mask(self, data: GeomData, row: pd.Series) -> GeomData: - """ - Add a graph node mask to the GeomData object. - - Args: - data: A GeomData object with features. - row: A dictionary containing 'features' and other metadata. - - Returns: - Modified GeomData with graph node mask added. - """ - geom_data = row["features"] - assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "is_graph_node must be set in the geom_data" - data.is_graph_node = is_graph_node - return data - - class ChEBI50_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): """ChEBIOver50 with with FG nodes and FG edges and graph node.""" From 5fad440c214288e7d6ce4e5bcac4a17c58f62213 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 16 Jul 2026 21:31:45 +0200 Subject: [PATCH 06/34] data classe for molecule net --- .../datasets/moleculeNet_classification.py | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 chebai_graph/preprocessing/datasets/moleculeNet_classification.py diff --git a/chebai_graph/preprocessing/datasets/moleculeNet_classification.py b/chebai_graph/preprocessing/datasets/moleculeNet_classification.py new file mode 100644 index 0000000..30c8998 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/moleculeNet_classification.py @@ -0,0 +1,65 @@ +from chebai.preprocessing.datasets.molecule_classfication import ( + BaceChem, + BBBPChem, + ClinToxChem, + HIVChem, + SiderChem, + Tox21Chem, +) + +from chebai_graph.preprocessing.datasets.base import ( + GraphPropAsPerNodeType, + GraphPropertiesMixIn, +) +from chebai_graph.preprocessing.reader.augmented_reader import ( + AtomFGReader_WithFGEdges_WithGraphNode, +) +from chebai_graph.preprocessing.reader.reader import GraphPropertyReader + + +class BaceChemDataset(GraphPropertiesMixIn, BaceChem): + READER = GraphPropertyReader + + +class BaceChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BaceChem): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class BBBPChemDataset(GraphPropertiesMixIn, BBBPChem): + READER = GraphPropertyReader + + +class BBBPChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BBBPChem): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ClinToxChemDataset(GraphPropertiesMixIn, ClinToxChem): + READER = GraphPropertyReader + + +class ClinToxChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ClinToxChem): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class HIVChemDataset(GraphPropertiesMixIn, HIVChem): + READER = GraphPropertyReader + + +class HIVChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIVChem): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class SiderChemDataset(GraphPropertiesMixIn, SiderChem): + READER = GraphPropertyReader + + +class SiderChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SiderChem): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Tox21ChemDataset(GraphPropertiesMixIn, Tox21Chem): + READER = GraphPropertyReader + + +class Tox21Chem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Tox21Chem): + READER = AtomFGReader_WithFGEdges_WithGraphNode From 49a9e8d7dec3a96a1c795e0c53f517be3b22b952 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 11:51:14 +0200 Subject: [PATCH 07/34] add requirements txt file --- pyproject.toml | 7 +-- requirements.txt | 120 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 3 deletions(-) create mode 100644 requirements.txt diff --git a/pyproject.toml b/pyproject.toml index a517cac..a1c3110 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,18 +13,19 @@ dependencies = [ # torch-geometric # torch_scatter ] -requires-python = ">=3.8" +requires-python = ">=3.10" [project.optional-dependencies] dev = [ "tox", + "omegaconf", + "chebi_utils", ] linters = [ - "isort", + "ruff", "pre-commit", - "black", ] [build-system] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0b7cd32 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,120 @@ +aiohappyeyeballs==2.7.1 +aiohttp==3.14.1 +aiosignal==1.4.0 +annotated-doc==0.0.4 +annotated-types==0.7.0 +antlr4-python3-runtime==4.9.3 +anyio==4.14.2 +async-timeout==5.0.1 +attrs==26.1.0 +cachetools==7.1.4 +certifi==2026.6.17 +cfgv==3.5.0 +chardet==7.4.3 +charset-normalizer==3.4.9 +chebai==1.3.0 +chebi-utils==0.2.1 +chembl-structure-pipeline==1.2.4 +click==8.4.2 +colorama==0.4.6 +descriptastorus==2.8.0 +distlib==0.4.3 +docstring-parser==0.18.0 +exceptiongroup==1.3.1 +fastobo==0.14.1 +filelock==3.30.2 +frozenlist==1.8.0 +fsspec==2025.12.0 +h11==0.16.0 +hf-xet==1.5.2 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.23.0 +identify==2.6.19 +idna==3.18 +importlib-resources==7.1.0 +iterative-stratification==0.1.9 +jinja2==3.1.6 +joblib==1.5.3 +jsonargparse==4.49.0 +lightning==2.5.1 +lightning-utilities==0.15.3 +markdown-it-py==4.2.0 +markupsafe==3.0.3 +mdurl==0.1.2 +mpmath==1.3.0 +multidict==6.7.1 +networkx==3.4.2 +nodeenv==1.10.0 +numpy==2.2.6 +nvidia-cublas-cu12==12.4.5.8 +nvidia-cuda-cupti-cu12==12.4.127 +nvidia-cuda-nvrtc-cu12==12.4.127 +nvidia-cuda-runtime-cu12==12.4.127 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cufft-cu12==11.2.1.3 +nvidia-curand-cu12==10.3.5.147 +nvidia-cusolver-cu12==11.6.1.9 +nvidia-cusparse-cu12==12.3.1.170 +nvidia-cusparselt-cu12==0.6.2 +nvidia-nccl-cu12==2.21.5 +nvidia-nvjitlink-cu12==12.4.127 +nvidia-nvtx-cu12==12.4.127 +omegaconf==2.3.1 +packaging==24.2 +pandas==2.3.3 +pandas-flavor==0.8.1 +pbr==7.0.3 +pillow==12.3.0 +platformdirs==4.10.0 +pluggy==1.6.0 +pre-commit==4.6.0 +propcache==0.5.2 +protobuf==7.35.1 +psutil==7.2.2 +pydantic==2.13.4 +pydantic-core==2.46.4 +pygments==2.20.0 +pyparsing==3.3.2 +pyproject-api==1.9.0 +pysmiles==1.1.2 +python-dateutil==2.9.0.post0 +python-discovery==1.4.4 +pytorch-lightning==2.6.5 +pytz==2026.2 +pyyaml==6.0.3 +rdkit==2024.3.6 +regex==2026.7.10 +requests==2.34.2 +rich==15.0.0 +safetensors==0.8.0 +scikit-learn==1.7.2 +scipy==1.15.3 +sentry-sdk==2.66.0 +setuptools==83.0.0 +shellingham==1.5.4 +six==1.17.0 +sympy==1.13.1 +threadpoolctl==3.6.0 +tokenizers==0.22.2 +tomli==2.4.1 +torch==2.6.0 +torch-geometric==2.8.0 +torch-scatter==2.1.2+pt26cu124 +torch-sparse==0.6.18+pt26cu124 +torchmetrics==1.9.0 +tox==4.27.0 +tqdm==4.68.4 +transformers==5.14.1 +triton==3.2.0 +typer==0.27.0 +typeshed-client==2.12.0 +typing-extensions==4.16.0 +typing-inspection==0.4.2 +tzdata==2026.3 +urllib3==2.7.0 +virtualenv==21.6.1 +wandb==0.28.1 +xarray==2025.6.1 +xxhash==3.8.1 +yarl==1.24.2 From e82ec951d28f5f5963ee72616e43ecbffbc5e2a7 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 11:52:15 +0200 Subject: [PATCH 08/34] fix data mol class --- .../datasets/moleculeNet_classification.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/moleculeNet_classification.py b/chebai_graph/preprocessing/datasets/moleculeNet_classification.py index 30c8998..b77bdd7 100644 --- a/chebai_graph/preprocessing/datasets/moleculeNet_classification.py +++ b/chebai_graph/preprocessing/datasets/moleculeNet_classification.py @@ -1,10 +1,10 @@ -from chebai.preprocessing.datasets.molecule_classfication import ( +from chebai.preprocessing.datasets.molecule_classification import ( BaceChem, BBBPChem, ClinToxChem, HIVChem, + MUVChem, SiderChem, - Tox21Chem, ) from chebai_graph.preprocessing.datasets.base import ( @@ -57,9 +57,15 @@ class SiderChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SiderChem): READER = AtomFGReader_WithFGEdges_WithGraphNode -class Tox21ChemDataset(GraphPropertiesMixIn, Tox21Chem): +class MUVChemDataset(GraphPropertiesMixIn, MUVChem): READER = GraphPropertyReader -class Tox21Chem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Tox21Chem): +class MUVChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, MUVChem): READER = AtomFGReader_WithFGEdges_WithGraphNode + + +if __name__ == "__main__": + dataset = BaceChemDataset() + dataset.prepare_data() + dataset.setup() From 971da2f381b6b1eb81e694b7c58cb19bb29c964b Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 11:56:30 +0200 Subject: [PATCH 09/34] wandb dep --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index a1c3110..c524311 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ linters = [ "pre-commit", ] +wandb = ["wandb"] + [build-system] build-backend = "flit_core.buildapi" requires = ["flit_core >=3.2,<4"] From e937089693beaf41d1b3a82b652bbd43d2bfad67 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 21:21:56 +0200 Subject: [PATCH 10/34] missing dropout to config --- configs/model/baselines/gat_baseline.yml | 1 + configs/model/gat_aug_aapool.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/configs/model/baselines/gat_baseline.yml b/configs/model/baselines/gat_baseline.yml index dda84dc..b7a921a 100644 --- a/configs/model/baselines/gat_baseline.yml +++ b/configs/model/baselines/gat_baseline.yml @@ -10,5 +10,6 @@ init_args: edge_dim: 7 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0 n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/gat_aug_aapool.yml b/configs/model/gat_aug_aapool.yml index fae47c3..1cdad92 100644 --- a/configs/model/gat_aug_aapool.yml +++ b/configs/model/gat_aug_aapool.yml @@ -10,5 +10,6 @@ init_args: edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0 n_molecule_properties: 0 n_linear_layers: 1 From 6a5aec178853e301f8482281a001139bd6473163 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 17 Jul 2026 21:41:48 +0200 Subject: [PATCH 11/34] update to correct props numbers --- configs/model/baselines/gat_baseline.yml | 4 ++-- configs/model/baselines/gine_baseline.yml | 4 ++-- configs/model/baselines/rggcn_baseline.yml | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/configs/model/baselines/gat_baseline.yml b/configs/model/baselines/gat_baseline.yml index b7a921a..1e96d57 100644 --- a/configs/model/baselines/gat_baseline.yml +++ b/configs/model/baselines/gat_baseline.yml @@ -3,11 +3,11 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 158 # number of node/atom properties + in_channels: 161 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 7 # number of bond properties + edge_dim: 8 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0 diff --git a/configs/model/baselines/gine_baseline.yml b/configs/model/baselines/gine_baseline.yml index 8477ca6..bd3d0c5 100644 --- a/configs/model/baselines/gine_baseline.yml +++ b/configs/model/baselines/gine_baseline.yml @@ -3,11 +3,11 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 158 # number of node/atom properties + in_channels: 161 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 7 # number of bond properties + edge_dim: 8 # number of bond properties dropout: 0 n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/baselines/rggcn_baseline.yml b/configs/model/baselines/rggcn_baseline.yml index ccc6615..863890d 100644 --- a/configs/model/baselines/rggcn_baseline.yml +++ b/configs/model/baselines/rggcn_baseline.yml @@ -3,11 +3,11 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 158 # number of node/atom properties + in_channels: 161 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 7 # number of bond properties + edge_dim: 8 # number of bond properties dropout: 0 n_molecule_properties: 0 n_linear_layers: 1 From 6e50bb59396cabbc5271c5e54ffbe93b266d1b75 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 18 Jul 2026 19:33:01 +0200 Subject: [PATCH 12/34] remove redundant code for mol propertiess addition to final classification linear layers --- chebai_graph/models/base.py | 122 ++++--------------- chebai_graph/models/graph.py | 11 +- chebai_graph/preprocessing/datasets/base.py | 3 +- chebai_graph/preprocessing/datasets/chebi.py | 1 - configs/model/baselines/gat_baseline.yml | 1 - configs/model/baselines/gine_baseline.yml | 1 - configs/model/baselines/rggcn_baseline.yml | 1 - configs/model/gat_aug_aapool.yml | 1 - configs/model/gat_aug_amgpool.yml | 1 - configs/model/gnn_res_gated.yml | 1 - configs/model/gnn_resgated_pretrain.yml | 1 - configs/model/res_aug_aapool.yml | 1 - configs/model/res_aug_amgpool.yml | 1 - configs/model/resgated_dynamic_gni.yml | 1 - 14 files changed, 26 insertions(+), 121 deletions(-) diff --git a/chebai_graph/models/base.py b/chebai_graph/models/base.py index 3226d21..56f16b1 100644 --- a/chebai_graph/models/base.py +++ b/chebai_graph/models/base.py @@ -77,7 +77,6 @@ def __init__( self, config: dict, n_linear_layers: int, - n_molecule_properties: Optional[int] = 0, use_batch_norm: bool = False, **kwargs, ): @@ -85,7 +84,6 @@ def __init__( Args: config (dict): Model configuration. n_linear_layers (int): Number of linear layers. - n_molecule_properties (int): Number of molecular-level features. **kwargs: Additional arguments. """ super().__init__(**kwargs) @@ -94,9 +92,6 @@ 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 - ), ) self.use_batch_norm = use_batch_norm if self.use_batch_norm: @@ -123,20 +118,17 @@ def _get_gnn(self, config: dict) -> torch.nn.Module: """ pass - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Compute input dimension for the linear layers. Args: gnn_out_dim (int): Output dimension of GNN. - n_molecule_properties (int): Number of molecule-level features. Returns: int: Total input dimension. """ - return gnn_out_dim + n_molecule_properties + return gnn_out_dim def _get_linear_module_list( self, @@ -188,7 +180,6 @@ def forward(self, batch: dict) -> torch.Tensor: assert isinstance(graph_data, GraphData) a = self.gnn(batch) a = scatter_add(a, graph_data.batch, dim=0) - a = torch.cat([a, graph_data.molecule_attr], dim=1) if self.use_batch_norm: a = self.batch_norm(a) return self.lin_sequential(a) @@ -204,25 +195,20 @@ class AugmentedNodePoolingNet(GraphNetWrapper, ABC): The concatenated vector is then passed through a linear sequential block. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Compute the input dimension for the final linear sequential block. Includes: - Atom embeddings - - Molecular attributes (if any) - Augmented node embeddings Args: gnn_out_dim (int): Dimension of the GNN output per node. - n_molecule_properties (int): Number of molecule-level attributes. - Returns: int: Total input dimension for the linear sequential block. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + return gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -234,7 +220,6 @@ def forward(self, batch: dict) -> torch.Tensor: 3. Aggregate embeddings for atoms and augmented nodes separately using scatter add. 4. Concatenate: - Atom nodes vector - - Molecular attributes - Augmented nodes vector 5. Pass the concatenated vector through the linear sequential block. @@ -265,9 +250,7 @@ def forward(self, batch: dict) -> torch.Tensor: ) # Concatenate all - graph_vector = torch.cat( - [atoms_vec, graph_data.molecule_attr, aug_nodes_vec], dim=1 - ) + graph_vector = torch.cat([atoms_vec, aug_nodes_vec], dim=1) return self.lin_sequential(graph_vector) @@ -282,25 +265,21 @@ class FGNodePoolingNet(GraphNetWrapper, ABC): The concatenated vector is then passed through a linear sequential block. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Computes the input dimension for the final linear sequential block. Combines: - All nodes embeddings except functional group nodes - - Molecular attributes - Functional group node embeddings Args: gnn_out_dim (int): Dimension of the GNN output per node. - n_molecule_properties (int): Number of molecule-level attributes. Returns: int: Total input dimension for the linear sequential block. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + return gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -344,9 +323,7 @@ def forward(self, batch: dict) -> torch.Tensor: fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) # Concatenate all - graph_vector = torch.cat( - [remaining_nodes_vec, graph_data.molecule_attr, fg_nodes_vec], dim=1 - ) + graph_vector = torch.cat([remaining_nodes_vec, fg_nodes_vec], dim=1) return self.lin_sequential(graph_vector) @@ -362,26 +339,22 @@ class GraphNodeFGNodePoolingNet(GraphNetWrapper, ABC): The concatenated vector is then passed through a linear sequential block. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Computes the input dimension for the final linear sequential block. Combines: - Atom embeddings - - Molecular attributes - Functional group node embeddings - Graph node embeddings Args: gnn_out_dim (int): Dimension of the GNN output per node. - n_molecule_properties (int): Number of molecule-level attributes. Returns: int: Total input dimension for the linear sequential block. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + gnn_out_dim + return gnn_out_dim + gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -427,9 +400,7 @@ def forward(self, batch: dict) -> torch.Tensor: fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) # Concatenate all - graph_vector = torch.cat( - [atoms_vec, graph_data.molecule_attr, fg_nodes_vec, graph_node_vec], dim=1 - ) + graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) return self.lin_sequential(graph_vector) @@ -439,9 +410,7 @@ class GraphNodePoolingNet(GraphNetWrapper, ABC): Pooling using non-graph nodes and graph node embeddings. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Return input dimension including graph node embeddings. - all_nodes_embeddings_except_graph_node + molecule attributes + graph_node_embedding @@ -449,7 +418,7 @@ def _get_lin_seq_input_dim( Returns: int: Total input dimension. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + return gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -478,9 +447,7 @@ def forward(self, batch: dict) -> torch.Tensor: remaining_nodes_embedding, remaining_nodes_batch, dim=0 ) - graph_vector = torch.cat( - [remaining_nodes_vec, graph_data.molecule_attr, graph_node_vec], dim=1 - ) + graph_vector = torch.cat([remaining_nodes_vec, graph_node_vec], dim=1) return self.lin_sequential(graph_vector) @@ -489,19 +456,16 @@ class FGNodePoolingNoGraphNodeNet(GraphNetWrapper, ABC): Graph Node not considered here in any computation. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Compute input dimension including: - atom_embeddings - - molecule attributes - functional_group_node_embeddings Returns: int: Total input dimension. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + return gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -531,9 +495,7 @@ def forward(self, batch: dict) -> torch.Tensor: atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - graph_vector = torch.cat( - [atoms_vec, graph_data.molecule_attr, fg_nodes_vec], dim=1 - ) + graph_vector = torch.cat([atoms_vec, fg_nodes_vec], dim=1) return self.lin_sequential(graph_vector) @@ -543,19 +505,16 @@ class GraphNodeNoFGNodePoolingNet(GraphNetWrapper, ABC): Functional Group Nodes not considered here in any computation. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ Compute input dimension including: - atom_embeddings - - molecule attributes - graph_node_embeddings Returns: int: Total input dimension. """ - return gnn_out_dim + n_molecule_properties + gnn_out_dim + return gnn_out_dim + gnn_out_dim def forward(self, batch: dict) -> torch.Tensor: """ @@ -584,9 +543,7 @@ def forward(self, batch: dict) -> torch.Tensor: graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) - graph_vector = torch.cat( - [atoms_vec, graph_data.molecule_attr, graph_node_vec], dim=1 - ) + graph_vector = torch.cat([atoms_vec, graph_node_vec], dim=1) return self.lin_sequential(graph_vector) @@ -596,17 +553,6 @@ class AugmentedOnlyPoolingNet(GraphNetWrapper, ABC): Only augmented node embeddings are pooled. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: - """ - Return input dimension using only augmented node embeddings. - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + n_molecule_properties - def forward(self, batch: dict) -> torch.Tensor: """ Forward pass pooling only augmented nodes. @@ -625,7 +571,7 @@ def forward(self, batch: dict) -> torch.Tensor: aug_nodes_vec = scatter_add( augmented_nodes_embeddings, augmented_nodes_batch, dim=0 ) - graph_vector = torch.cat([aug_nodes_vec, graph_data.molecule_attr], dim=1) + graph_vector = torch.cat([aug_nodes_vec], dim=1) return self.lin_sequential(graph_vector) @@ -635,17 +581,6 @@ class FGOnlyPoolingNet(GraphNetWrapper, ABC): Only functional group node embeddings are pooled. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: - """ - Return input dimension using only FG node embeddings. - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + n_molecule_properties - def forward(self, batch: dict) -> torch.Tensor: """ Forward pass pooling only functional group nodes. @@ -664,7 +599,7 @@ def forward(self, batch: dict) -> torch.Tensor: fg_nodes_batch = graph_data.batch[~is_fg_node] fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - graph_vector = torch.cat([fg_nodes_vec, graph_data.molecule_attr], dim=1) + graph_vector = torch.cat([fg_nodes_vec], dim=1) return self.lin_sequential(graph_vector) @@ -674,17 +609,6 @@ class GraphNodeOnlyPoolingNet(GraphNetWrapper, ABC): Only graph node embeddings are pooled. """ - def _get_lin_seq_input_dim( - self, gnn_out_dim: int, n_molecule_properties: int - ) -> int: - """ - Return input dimension using only graph node embeddings. - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + n_molecule_properties - def forward(self, batch: dict) -> torch.Tensor: """ Forward pass pooling only graph nodes. @@ -702,6 +626,6 @@ def forward(self, batch: dict) -> torch.Tensor: graph_node_batch = graph_data.batch[~is_graph_node] graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) - graph_vector = torch.cat([graph_node_vec, graph_data.molecule_attr], dim=1) + graph_vector = torch.cat([graph_node_vec], dim=1) return self.lin_sequential(graph_vector) diff --git a/chebai_graph/models/graph.py b/chebai_graph/models/graph.py index 3be0fdd..6037e93 100644 --- a/chebai_graph/models/graph.py +++ b/chebai_graph/models/graph.py @@ -88,11 +88,6 @@ def __init__(self, config: typing.Dict, **kwargs): self.n_bond_properties = ( int(config["n_bond_properties"]) if "n_bond_properties" in config else 7 ) - self.n_molecule_properties = ( - int(config["n_molecule_properties"]) - if "n_molecule_properties" in config - else 0 - ) self.activation = F.elu self.dropout = nn.Dropout(self.dropout_rate) @@ -158,7 +153,7 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length + (i == 0) * self.gnn.n_molecule_properties, + self.gnn.hidden_length, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) @@ -196,9 +191,7 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length - + (i == 0) * self.gnn.n_molecule_properties - + (i == 0) * self.gnn.hidden_length, + self.gnn.hidden_length + (i == 0) * self.gnn.hidden_length, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index abe13e4..13bacbb 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -413,7 +413,6 @@ def load_processed_data( f"Use following values for given parameters for model configuration: \n\t" f"{in_channels_str} \n\t" f"{edge_dim_str} \n\t" - f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}" ) return base_df[base_data[0].keys()].to_dict("records") @@ -499,7 +498,7 @@ def load_processed_data( f"n_bond_properties: {n_bond_properties}, " f"n_graph_node_properties: {n_graph_node_properties}\n\n" f"Use following values for given parameters for model configuration: \n\t" - f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0\n" + f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}\n" ) for property in self.properties: diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 59a3578..263747e 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -50,7 +50,6 @@ def load_processed_data_from_file(self, filename): 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") diff --git a/configs/model/baselines/gat_baseline.yml b/configs/model/baselines/gat_baseline.yml index b7a921a..5f53493 100644 --- a/configs/model/baselines/gat_baseline.yml +++ b/configs/model/baselines/gat_baseline.yml @@ -11,5 +11,4 @@ init_args: heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/baselines/gine_baseline.yml b/configs/model/baselines/gine_baseline.yml index 8477ca6..0fd7ac7 100644 --- a/configs/model/baselines/gine_baseline.yml +++ b/configs/model/baselines/gine_baseline.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 7 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/baselines/rggcn_baseline.yml b/configs/model/baselines/rggcn_baseline.yml index ccc6615..2400f45 100644 --- a/configs/model/baselines/rggcn_baseline.yml +++ b/configs/model/baselines/rggcn_baseline.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 7 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/gat_aug_aapool.yml b/configs/model/gat_aug_aapool.yml index 1cdad92..38a326c 100644 --- a/configs/model/gat_aug_aapool.yml +++ b/configs/model/gat_aug_aapool.yml @@ -11,5 +11,4 @@ init_args: heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/gat_aug_amgpool.yml b/configs/model/gat_aug_amgpool.yml index e596487..959642c 100644 --- a/configs/model/gat_aug_amgpool.yml +++ b/configs/model/gat_aug_amgpool.yml @@ -11,5 +11,4 @@ init_args: heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/gnn_res_gated.yml b/configs/model/gnn_res_gated.yml index 27d1e78..b8cfe54 100644 --- a/configs/model/gnn_res_gated.yml +++ b/configs/model/gnn_res_gated.yml @@ -10,4 +10,3 @@ init_args: n_linear_layers: 3 n_atom_properties: 158 n_bond_properties: 7 - n_molecule_properties: 200 diff --git a/configs/model/gnn_resgated_pretrain.yml b/configs/model/gnn_resgated_pretrain.yml index fad8c27..370fc9e 100644 --- a/configs/model/gnn_resgated_pretrain.yml +++ b/configs/model/gnn_resgated_pretrain.yml @@ -13,4 +13,3 @@ init_args: n_linear_layers: 3 n_atom_properties: 151 n_bond_properties: 7 - n_molecule_properties: 200 diff --git a/configs/model/res_aug_aapool.yml b/configs/model/res_aug_aapool.yml index de28d1c..f5c2e83 100644 --- a/configs/model/res_aug_aapool.yml +++ b/configs/model/res_aug_aapool.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 12 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/res_aug_amgpool.yml b/configs/model/res_aug_amgpool.yml index 9194cd7..9a59240 100644 --- a/configs/model/res_aug_amgpool.yml +++ b/configs/model/res_aug_amgpool.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 12 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/resgated_dynamic_gni.yml b/configs/model/resgated_dynamic_gni.yml index 4749795..ad55d88 100644 --- a/configs/model/resgated_dynamic_gni.yml +++ b/configs/model/resgated_dynamic_gni.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 7 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 From 38c25eea587a63e1a08fba1b768ecb74da694148 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 23 Jul 2026 12:25:17 +0200 Subject: [PATCH 13/34] refine gine implementation --- chebai_graph/models/__init__.py | 2 +- chebai_graph/models/base.py | 1 - chebai_graph/models/gin_net.py | 92 -------------- chebai_graph/models/gine.py | 145 ++++++++++++++++++++++ configs/model/baselines/gine_baseline.yml | 1 + 5 files changed, 147 insertions(+), 94 deletions(-) delete mode 100644 chebai_graph/models/gin_net.py create mode 100644 chebai_graph/models/gine.py diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index f59e1c6..cd67fb0 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -6,8 +6,8 @@ ) from .dynamic_gni import ResGatedDynamicGNIGraphPred from .gat import GATGraphPred +from .gine import GINEGraphPred from .resgated import ResGatedGraphPred -from .gin_net import GINEGraphPred __all__ = [ "ResGatedGraphPred", diff --git a/chebai_graph/models/base.py b/chebai_graph/models/base.py index 56f16b1..aaefc33 100644 --- a/chebai_graph/models/base.py +++ b/chebai_graph/models/base.py @@ -1,5 +1,4 @@ from abc import ABC, abstractmethod -from typing import Optional import torch from chebai.models.base import ChebaiBaseNet diff --git a/chebai_graph/models/gin_net.py b/chebai_graph/models/gin_net.py deleted file mode 100644 index 72acce7..0000000 --- a/chebai_graph/models/gin_net.py +++ /dev/null @@ -1,92 +0,0 @@ -import typing - -import torch -import torch.nn.functional as F -from torch_scatter import scatter_add -from torch_geometric.data import Data as GraphData - -from chebai_graph.models.base import GraphModelBase, GraphNetWrapper -from torch_geometric import nn as tgnn - - -class AggregateMLP(torch.nn.Module): - def __init__(self, in_channels, out_channels, hidden_channels): - super(AggregateMLP, self).__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.hidden_channels = hidden_channels - self.activation = F.relu - self.in_layer = torch.nn.Linear(in_channels, hidden_channels) - self.out_layer = torch.nn.Linear(hidden_channels, out_channels) - - def forward(self, x): - x = self.activation(self.in_layer(x)) - x = self.activation(self.out_layer(x)) - return x - - -class GINEConvNet(GraphModelBase): - """Based on https://arxiv.org/pdf/1810.00826.pdf and https://arxiv.org/abs/1905.12265""" - - NAME = "GINEConvNet" - - def __init__(self, config: typing.Dict, **kwargs): - super().__init__(**kwargs) - - self.dropout_layer = torch.nn.Dropout(self.dropout) - self.activation = F.elu - - self.convs = torch.nn.ModuleList([]) - # self.batch_norms = torch.nn.ModuleList([]) - for i in range(self.num_layers): - self.convs.append( - tgnn.GINEConv( - AggregateMLP( - self.in_channels, self.out_channels, self.hidden_channels - ), - edge_dim=self.edge_dim, - ) - ) - # self.batch_norms.append(torch.nn.BatchNorm1d(out_length)) - - def forward(self, batch): - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - a = graph_data.x - - dropout_used = False # only apply dropout after first layer - conv_out = [] - for conv in self.convs: # , norm in zip(self.convs, self.batch_norms): - a = self.activation( - conv(a, graph_data.edge_index.long(), graph_data.edge_attr) - ) - if not dropout_used: - a = self.dropout_layer(a) - dropout_used = True - # a = norm(a) - a = scatter_add(a, graph_data.batch, dim=0) - conv_out.append(a) - - a = torch.cat(conv_out, dim=1) - - return a - - -class GINEGraphPred(GraphNetWrapper): - """ - Wrapper for graph-level prediction using GINEConvNet. - - This class instantiates the core GNN model using the provided config. - """ - - def _get_gnn(self, config: dict[str, typing.Any]) -> GINEConvNet: - """ - Returns the core ResGated GNN model. - - Args: - config (dict): Configuration dictionary for the GNN model. - - Returns: - ResGatedGraphConvNetBase: The core graph convolutional network. - """ - return GINEConvNet(config=config) diff --git a/chebai_graph/models/gine.py b/chebai_graph/models/gine.py new file mode 100644 index 0000000..1885cce --- /dev/null +++ b/chebai_graph/models/gine.py @@ -0,0 +1,145 @@ +from typing import Any, Final + +from torch import Tensor +from torch.nn import ELU +from torch_geometric import nn as tgnn +from torch_geometric.data import Data as GraphData +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.models import MLP +from torch_geometric.nn.models.basic_gnn import BasicGNN + +from .base import GraphModelBase, GraphNetWrapper + + +class GINEModel(BasicGNN): + """ + A GIN-based GNN model based on PyG's BasicGNN, using GINEConv layers so that + edge (bond) features are incorporated into the message-passing step. + + See: + - https://pytorch-geometric.readthedocs.io/en/2.7.0/generated/torch_geometric.nn.conv.GINEConv.html + - https://arxiv.org/abs/1810.00826 (GIN) + - https://arxiv.org/abs/1905.12265 (GINE / edge-feature extension) + - https://github.com/pyg-team/pytorch_geometric/blob/master/examples/mutag_gin.py + + Attributes: + supports_edge_weight (bool): Indicates edge weights are not supported. + supports_edge_attr (bool): Indicates edge attributes are supported. + supports_norm_batch (bool): Indicates if batch normalization is supported. + """ + + supports_edge_weight: Final[bool] = False + supports_edge_attr: Final[bool] = True + supports_norm_batch: Final[bool] + + def init_conv( + self, in_channels: int | tuple[int, int], out_channels: int, **kwargs: Any + ) -> MessagePassing: + """ + Initializes a GINEConv layer. + + The inner network is a 2-layer MLP (Linear -> act -> Linear, no + activation on the last layer), matching both + `torch_geometric.nn.models.GIN.init_conv` and the GIN paper's + message-transform MLP. `edge_dim` (passed via **kwargs) lets + GINEConv linearly project bond features onto the node feature + space before adding them into the neighbor messages. + + Args: + in_channels (int or Tuple[int, int]): Number of input channels. + out_channels (int): Number of output channels. + **kwargs: Additional keyword arguments for the convolution layer + (e.g. `edge_dim`, `train_eps`). + + Returns: + MessagePassing: A GINEConv layer instance. + """ + mlp = MLP( + [in_channels, out_channels, out_channels], + act=self.act, + act_first=self.act_first, + norm=self.norm, + norm_kwargs=self.norm_kwargs, + ) + return tgnn.GINEConv(mlp, **kwargs) + + +class GINEConvNetBase(GraphModelBase): + """ + Base model class for applying GINEConv layers to graph-structured data. + + Based on: + - Xu et al., "How Powerful are Graph Neural Networks?" + (https://arxiv.org/abs/1810.00826) + - Hu et al., "Strategies for Pre-training Graph Neural Networks" + (https://arxiv.org/abs/1905.12265), reference implementation at + https://github.com/snap-stanford/pretrain-gnns/blob/master/chem/model.py + + Args: + config (dict): Configuration dictionary containing model hyperparameters. + Also supports an optional `train_eps` (bool, default True) key, + which makes GINEConv's epsilon a learnable parameter, as + recommended in the original GIN paper. + **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. + self.train_eps = bool(config.get("train_eps", True)) + + self.gine: BasicGNN = GINEModel( + in_channels=self.in_channels, + hidden_channels=self.hidden_channels, + out_channels=self.out_channels, + num_layers=self.num_layers, + dropout=self.dropout, + edge_dim=self.edge_dim, + train_eps=self.train_eps, + 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" + + out = self.gine( + x=graph_data.x.float(), + edge_index=graph_data.edge_index.long(), + edge_attr=graph_data.edge_attr, + ) + + return self.activation(out) + + +class GINEGraphPred(GraphNetWrapper): + """ + Wrapper for graph-level prediction using GINEConvNetBase. + + This class instantiates the core GNN model using the provided config. + Graph-level pooling (scatter-add over nodes) and the final linear + prediction head are handled by `GraphNetWrapper`, not here. + """ + + NAME = "GINEGraphPred" + + def _get_gnn(self, config: dict[str, Any]) -> GINEConvNetBase: + """ + Returns the core GINE GNN model. + + Args: + config (dict): Configuration dictionary for the GNN model. + + Returns: + GINEConvNetBase: The core graph convolutional network. + """ + return GINEConvNetBase(config=config) diff --git a/configs/model/baselines/gine_baseline.yml b/configs/model/baselines/gine_baseline.yml index 448e731..9d56063 100644 --- a/configs/model/baselines/gine_baseline.yml +++ b/configs/model/baselines/gine_baseline.yml @@ -9,4 +9,5 @@ init_args: num_layers: 4 edge_dim: 8 # number of bond properties dropout: 0 + train_eps: true n_linear_layers: 1 From 05e267958c31cda6800f28d1abdafa305cc04562 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 23 Jul 2026 15:38:51 +0200 Subject: [PATCH 14/34] remove redudant pooling types --- chebai_graph/models/base.py | 302 ------------------------------------ 1 file changed, 302 deletions(-) diff --git a/chebai_graph/models/base.py b/chebai_graph/models/base.py index aaefc33..8eb373f 100644 --- a/chebai_graph/models/base.py +++ b/chebai_graph/models/base.py @@ -188,7 +188,6 @@ class AugmentedNodePoolingNet(GraphNetWrapper, ABC): """ A pooling network that aggregates: - Atom node embeddings - - Molecular attributes (if provided else skipped) - Augmented node embeddings (FG nodes and graph node) The concatenated vector is then passed through a linear sequential block. @@ -254,84 +253,10 @@ def forward(self, batch: dict) -> torch.Tensor: return self.lin_sequential(graph_vector) -class FGNodePoolingNet(GraphNetWrapper, ABC): - """ - A pooling network that pools node embeddings by aggregating: - - All non-functional-group nodes' embeddings (atom and graph node) - - Molecular attributes - - Functional group node embeddings - - The concatenated vector is then passed through a linear sequential block. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Computes the input dimension for the final linear sequential block. - - Combines: - - All nodes embeddings except functional group nodes - - Functional group node embeddings - - Args: - gnn_out_dim (int): Dimension of the GNN output per node. - - Returns: - int: Total input dimension for the linear sequential block. - """ - return gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass for pooling node embeddings. - - Steps: - 1. Identify graph, atom, and functional group nodes. - 2. Aggregate embeddings for remaining nodes and functional group nodes separately. - 3. Concatenate: - - Remaining nodes vector - - Molecular attributes - - Functional group nodes vector - 4. Pass the concatenated vector through the linear sequential block. - - Args: - batch (dict): Batch containing graph data and features. - - Returns: - torch.Tensor: Output tensor after pooling and linear transformation. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - - is_graph_node = graph_data.is_graph_node.bool() - is_atom_node = graph_data.is_atom_node.bool() - is_fg_node = (~is_atom_node) & (~is_graph_node) - is_remaining_node = ~is_fg_node - - node_embeddings = self.gnn(batch) - - remaining_nodes_embedding = node_embeddings[is_remaining_node] - remaining_nodes_batch = graph_data.batch[is_remaining_node] - - fg_nodes_embeddings = node_embeddings[is_fg_node] - fg_nodes_batch = graph_data.batch[is_fg_node] - - # Scatter add separately - remaining_nodes_vec = scatter_add( - remaining_nodes_embedding, remaining_nodes_batch, dim=0 - ) - fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - - # Concatenate all - graph_vector = torch.cat([remaining_nodes_vec, fg_nodes_vec], dim=1) - - return self.lin_sequential(graph_vector) - - class GraphNodeFGNodePoolingNet(GraphNetWrapper, ABC): """ A pooling network that pools node embeddings by aggregating: - Atom nodes - - Molecular attributes - Functional group node embeddings - Graph node embeddings @@ -364,7 +289,6 @@ def forward(self, batch: dict) -> torch.Tensor: 2. Aggregate embeddings for each node type separately. 3. Concatenate: - Atom nodes vector - - Molecular attributes - Functional group nodes vector - Graph node vector 4. Pass the concatenated vector through the linear sequential block. @@ -402,229 +326,3 @@ def forward(self, batch: dict) -> torch.Tensor: graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) return self.lin_sequential(graph_vector) - - -class GraphNodePoolingNet(GraphNetWrapper, ABC): - """ - Pooling using non-graph nodes and graph node embeddings. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Return input dimension including graph node embeddings. - - all_nodes_embeddings_except_graph_node + molecule attributes + graph_node_embedding - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass with separate pooling for graph and other nodes. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - is_graph_node = graph_data.is_graph_node.bool() - is_not_graph_node = ~is_graph_node - - node_embeddings = self.gnn(batch) - graph_node_embedding = node_embeddings[is_graph_node] - graph_node_batch = graph_data.batch[is_graph_node] - - remaining_nodes_embedding = node_embeddings[is_not_graph_node] - remaining_nodes_batch = graph_data.batch[is_not_graph_node] - - graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) - remaining_nodes_vec = scatter_add( - remaining_nodes_embedding, remaining_nodes_batch, dim=0 - ) - - graph_vector = torch.cat([remaining_nodes_vec, graph_node_vec], dim=1) - return self.lin_sequential(graph_vector) - - -class FGNodePoolingNoGraphNodeNet(GraphNetWrapper, ABC): - """ - Graph Node not considered here in any computation. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Compute input dimension including: - - atom_embeddings - - functional_group_node_embeddings - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass pooling atoms and functional group nodes. - Graph nodes are ignored. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - is_graph_node = graph_data.is_graph_node.bool() - is_atom_node = graph_data.is_atom_node.bool() - is_fg_node = (~is_atom_node) & (~is_graph_node) - - node_embeddings = self.gnn(batch) - - atoms_embeddings = node_embeddings[is_atom_node] - atoms_batch = graph_data.batch[is_atom_node] - - fg_nodes_embeddings = node_embeddings[is_fg_node] - fg_nodes_batch = graph_data.batch[is_fg_node] - - atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) - fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - - graph_vector = torch.cat([atoms_vec, fg_nodes_vec], dim=1) - - return self.lin_sequential(graph_vector) - - -class GraphNodeNoFGNodePoolingNet(GraphNetWrapper, ABC): - """ - Functional Group Nodes not considered here in any computation. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Compute input dimension including: - - atom_embeddings - - graph_node_embeddings - - Returns: - int: Total input dimension. - """ - return gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass pooling atoms and graph nodes. - Functional group nodes are ignored. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - is_graph_node = graph_data.is_graph_node.bool() - is_atom_node = graph_data.is_atom_node.bool() - - node_embeddings = self.gnn(batch) - - graph_node_embedding = node_embeddings[is_graph_node] - graph_node_batch = graph_data.batch[is_graph_node] - - atoms_embeddings = node_embeddings[is_atom_node] - atoms_batch = graph_data.batch[is_atom_node] - - graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) - atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) - - graph_vector = torch.cat([atoms_vec, graph_node_vec], dim=1) - - return self.lin_sequential(graph_vector) - - -class AugmentedOnlyPoolingNet(GraphNetWrapper, ABC): - """ - Only augmented node embeddings are pooled. - """ - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass pooling only augmented nodes. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - is_atom_node = graph_data.is_atom_node.bool() - augmented_nodes_embeddings = self.gnn(batch)[~is_atom_node] - augmented_nodes_batch = graph_data.batch[~is_atom_node] - - aug_nodes_vec = scatter_add( - augmented_nodes_embeddings, augmented_nodes_batch, dim=0 - ) - graph_vector = torch.cat([aug_nodes_vec], dim=1) - - return self.lin_sequential(graph_vector) - - -class FGOnlyPoolingNet(GraphNetWrapper, ABC): - """ - Only functional group node embeddings are pooled. - """ - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass pooling only functional group nodes. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - is_graph_node = graph_data.is_graph_node.bool() - is_atom_node = graph_data.is_atom_node.bool() - is_fg_node = (~is_atom_node) & (~is_graph_node) - fg_nodes_embeddings = self.gnn(batch)[~is_fg_node] - fg_nodes_batch = graph_data.batch[~is_fg_node] - - fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - graph_vector = torch.cat([fg_nodes_vec], dim=1) - - return self.lin_sequential(graph_vector) - - -class GraphNodeOnlyPoolingNet(GraphNetWrapper, ABC): - """ - Only graph node embeddings are pooled. - """ - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass pooling only graph nodes. - - Args: - batch (dict): Input batch. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - is_graph_node = graph_data.is_graph_node.bool() - - graph_node_embedding = self.gnn(batch)[~is_graph_node] - graph_node_batch = graph_data.batch[~is_graph_node] - - graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) - graph_vector = torch.cat([graph_node_vec], dim=1) - - return self.lin_sequential(graph_vector) From 6bc6cf7ba2527339a84c94bc61a16b10507d36be Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 23 Jul 2026 15:44:32 +0200 Subject: [PATCH 15/34] move pooling types to separate file --- chebai_graph/models/augmented.py | 2 +- chebai_graph/models/base.py | 147 +----------------------------- chebai_graph/models/pooling.py | 151 +++++++++++++++++++++++++++++++ 3 files changed, 154 insertions(+), 146 deletions(-) create mode 100644 chebai_graph/models/pooling.py diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index fdb5388..3f43c2c 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,5 +1,5 @@ -from .base import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet from .gat import GATGraphPred +from .pooling import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet from .resgated import ResGatedGraphPred diff --git a/chebai_graph/models/base.py b/chebai_graph/models/base.py index 8eb373f..514e287 100644 --- a/chebai_graph/models/base.py +++ b/chebai_graph/models/base.py @@ -69,7 +69,8 @@ def __init__(self, config: dict, **kwargs) -> None: class GraphNetWrapper(GraphBaseNet, ABC): """ - Base wrapper class for GNNs with linear layers for property prediction. + Base wrapper class for GNNs with linear layers for graph classification + with standard pooling . """ def __init__( @@ -182,147 +183,3 @@ def forward(self, batch: dict) -> torch.Tensor: if self.use_batch_norm: a = self.batch_norm(a) return self.lin_sequential(a) - - -class AugmentedNodePoolingNet(GraphNetWrapper, ABC): - """ - A pooling network that aggregates: - - Atom node embeddings - - Augmented node embeddings (FG nodes and graph node) - - The concatenated vector is then passed through a linear sequential block. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Compute the input dimension for the final linear sequential block. - - Includes: - - Atom embeddings - - Augmented node embeddings - - Args: - gnn_out_dim (int): Dimension of the GNN output per node. - Returns: - int: Total input dimension for the linear sequential block. - """ - return gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass for pooling node embeddings. - - Steps: - 1. Identify atom nodes and augmented nodes. - 2. Compute node embeddings with the GNN. - 3. Aggregate embeddings for atoms and augmented nodes separately using scatter add. - 4. Concatenate: - - Atom nodes vector - - Augmented nodes vector - 5. Pass the concatenated vector through the linear sequential block. - - Args: - batch (dict): Input batch containing graph data and features. - - Returns: - torch.Tensor: Output tensor after pooling and linear transformation. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - - is_atom_node = graph_data.is_atom_node.bool() - is_augmented_node = ~is_atom_node - - node_embeddings = self.gnn(batch) - - atoms_embeddings = node_embeddings[is_atom_node] - atoms_batch = graph_data.batch[is_atom_node] - - augmented_nodes_embeddings = node_embeddings[is_augmented_node] - augmented_nodes_batch = graph_data.batch[is_augmented_node] - - # Scatter add separately - atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) - aug_nodes_vec = scatter_add( - augmented_nodes_embeddings, augmented_nodes_batch, dim=0 - ) - - # Concatenate all - graph_vector = torch.cat([atoms_vec, aug_nodes_vec], dim=1) - - return self.lin_sequential(graph_vector) - - -class GraphNodeFGNodePoolingNet(GraphNetWrapper, ABC): - """ - A pooling network that pools node embeddings by aggregating: - - Atom nodes - - Functional group node embeddings - - Graph node embeddings - - The concatenated vector is then passed through a linear sequential block. - """ - - def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: - """ - Computes the input dimension for the final linear sequential block. - - Combines: - - Atom embeddings - - Functional group node embeddings - - Graph node embeddings - - Args: - gnn_out_dim (int): Dimension of the GNN output per node. - - Returns: - int: Total input dimension for the linear sequential block. - """ - return gnn_out_dim + gnn_out_dim + gnn_out_dim - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass for pooling node embeddings. - - Steps: - 1. Identify graph, atom, and functional group nodes. - 2. Aggregate embeddings for each node type separately. - 3. Concatenate: - - Atom nodes vector - - Functional group nodes vector - - Graph node vector - 4. Pass the concatenated vector through the linear sequential block. - - Args: - batch (dict): Batch containing graph data and features. - - Returns: - torch.Tensor: Output tensor after pooling and linear transformation. - """ - graph_data = batch["features"][0] - assert isinstance(graph_data, GraphData) - - is_graph_node = graph_data.is_graph_node.bool() - is_atom_node = graph_data.is_atom_node.bool() - is_fg_node = (~is_atom_node) & (~is_graph_node) - - node_embeddings = self.gnn(batch) - - graph_node_embedding = node_embeddings[is_graph_node] - graph_node_batch = graph_data.batch[is_graph_node] - - atoms_embeddings = node_embeddings[is_atom_node] - atoms_batch = graph_data.batch[is_atom_node] - - fg_nodes_embeddings = node_embeddings[is_fg_node] - fg_nodes_batch = graph_data.batch[is_fg_node] - - # Scatter add separately - graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) - atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) - fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - - # Concatenate all - graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) - - return self.lin_sequential(graph_vector) diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py new file mode 100644 index 0000000..a69baf7 --- /dev/null +++ b/chebai_graph/models/pooling.py @@ -0,0 +1,151 @@ +from abc import ABC + +import torch +from torch_geometric.data import Data as GraphData +from torch_scatter import scatter_add + +from .base import GraphNetWrapper + + +class AugmentedNodePoolingNet(GraphNetWrapper, ABC): + """ + A pooling network that aggregates: + - Atom node embeddings + - Augmented node embeddings (FG nodes and graph node) + + The concatenated vector is then passed through a linear sequential block. + """ + + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: + """ + Compute the input dimension for the final linear sequential block. + + Includes: + - Atom embeddings + - Augmented node embeddings + + Args: + gnn_out_dim (int): Dimension of the GNN output per node. + Returns: + int: Total input dimension for the linear sequential block. + """ + return gnn_out_dim + gnn_out_dim + + def forward(self, batch: dict) -> torch.Tensor: + """ + Forward pass for pooling node embeddings. + + Steps: + 1. Identify atom nodes and augmented nodes. + 2. Compute node embeddings with the GNN. + 3. Aggregate embeddings for atoms and augmented nodes separately using scatter add. + 4. Concatenate: + - Atom nodes vector + - Augmented nodes vector + 5. Pass the concatenated vector through the linear sequential block. + + Args: + batch (dict): Input batch containing graph data and features. + + Returns: + torch.Tensor: Output tensor after pooling and linear transformation. + """ + graph_data = batch["features"][0] + assert isinstance(graph_data, GraphData) + + is_atom_node = graph_data.is_atom_node.bool() + is_augmented_node = ~is_atom_node + + node_embeddings = self.gnn(batch) + + atoms_embeddings = node_embeddings[is_atom_node] + atoms_batch = graph_data.batch[is_atom_node] + + augmented_nodes_embeddings = node_embeddings[is_augmented_node] + augmented_nodes_batch = graph_data.batch[is_augmented_node] + + # Scatter add separately + atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) + aug_nodes_vec = scatter_add( + augmented_nodes_embeddings, augmented_nodes_batch, dim=0 + ) + + # Concatenate all + graph_vector = torch.cat([atoms_vec, aug_nodes_vec], dim=1) + + return self.lin_sequential(graph_vector) + + +class GraphNodeFGNodePoolingNet(GraphNetWrapper, ABC): + """ + A pooling network that pools node embeddings by aggregating: + - Atom nodes + - Functional group node embeddings + - Graph node embeddings + + The concatenated vector is then passed through a linear sequential block. + """ + + def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: + """ + Computes the input dimension for the final linear sequential block. + + Combines: + - Atom embeddings + - Functional group node embeddings + - Graph node embeddings + + Args: + gnn_out_dim (int): Dimension of the GNN output per node. + + Returns: + int: Total input dimension for the linear sequential block. + """ + return gnn_out_dim + gnn_out_dim + gnn_out_dim + + def forward(self, batch: dict) -> torch.Tensor: + """ + Forward pass for pooling node embeddings. + + Steps: + 1. Identify graph, atom, and functional group nodes. + 2. Aggregate embeddings for each node type separately. + 3. Concatenate: + - Atom nodes vector + - Functional group nodes vector + - Graph node vector + 4. Pass the concatenated vector through the linear sequential block. + + Args: + batch (dict): Batch containing graph data and features. + + Returns: + torch.Tensor: Output tensor after pooling and linear transformation. + """ + graph_data = batch["features"][0] + assert isinstance(graph_data, GraphData) + + is_graph_node = graph_data.is_graph_node.bool() + is_atom_node = graph_data.is_atom_node.bool() + is_fg_node = (~is_atom_node) & (~is_graph_node) + + node_embeddings = self.gnn(batch) + + graph_node_embedding = node_embeddings[is_graph_node] + graph_node_batch = graph_data.batch[is_graph_node] + + atoms_embeddings = node_embeddings[is_atom_node] + atoms_batch = graph_data.batch[is_atom_node] + + fg_nodes_embeddings = node_embeddings[is_fg_node] + fg_nodes_batch = graph_data.batch[is_fg_node] + + # Scatter add separately + graph_node_vec = scatter_add(graph_node_embedding, graph_node_batch, dim=0) + atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) + fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) + + # Concatenate all + graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) + + return self.lin_sequential(graph_vector) From e04bb43cf849610e4e676c97dec88d81098bb0d3 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 23 Jul 2026 15:55:09 +0200 Subject: [PATCH 16/34] dedicate dir to core architectures --- chebai_graph/models/__init__.py | 6 +++--- chebai_graph/models/architectures/__init__.py | 0 chebai_graph/models/{ => architectures}/base.py | 0 chebai_graph/models/{ => architectures}/gat.py | 0 chebai_graph/models/{ => architectures}/gine.py | 0 chebai_graph/models/{ => architectures}/resgated.py | 0 chebai_graph/models/augmented.py | 4 ++-- chebai_graph/models/dynamic_gni.py | 4 ++-- chebai_graph/models/graph.py | 2 +- chebai_graph/models/pooling.py | 2 +- 10 files changed, 9 insertions(+), 9 deletions(-) create mode 100644 chebai_graph/models/architectures/__init__.py rename chebai_graph/models/{ => architectures}/base.py (100%) rename chebai_graph/models/{ => architectures}/gat.py (100%) rename chebai_graph/models/{ => architectures}/gine.py (100%) rename chebai_graph/models/{ => architectures}/resgated.py (100%) diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index cd67fb0..a0b518b 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -1,3 +1,6 @@ +from .architectures.gat import GATGraphPred +from .architectures.gine import GINEGraphPred +from .architectures.resgated import ResGatedGraphPred from .augmented import ( GATAugNodePoolGraphPred, GATGraphNodeFGNodePoolGraphPred, @@ -5,9 +8,6 @@ ResGatedGraphNodeFGNodePoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred -from .gat import GATGraphPred -from .gine import GINEGraphPred -from .resgated import ResGatedGraphPred __all__ = [ "ResGatedGraphPred", diff --git a/chebai_graph/models/architectures/__init__.py b/chebai_graph/models/architectures/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chebai_graph/models/base.py b/chebai_graph/models/architectures/base.py similarity index 100% rename from chebai_graph/models/base.py rename to chebai_graph/models/architectures/base.py diff --git a/chebai_graph/models/gat.py b/chebai_graph/models/architectures/gat.py similarity index 100% rename from chebai_graph/models/gat.py rename to chebai_graph/models/architectures/gat.py diff --git a/chebai_graph/models/gine.py b/chebai_graph/models/architectures/gine.py similarity index 100% rename from chebai_graph/models/gine.py rename to chebai_graph/models/architectures/gine.py diff --git a/chebai_graph/models/resgated.py b/chebai_graph/models/architectures/resgated.py similarity index 100% rename from chebai_graph/models/resgated.py rename to chebai_graph/models/architectures/resgated.py diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index 3f43c2c..c2b0a0c 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,6 +1,6 @@ -from .gat import GATGraphPred +from .architectures.gat import GATGraphPred +from .architectures.resgated import ResGatedGraphPred from .pooling import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet -from .resgated import ResGatedGraphPred class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): diff --git a/chebai_graph/models/dynamic_gni.py b/chebai_graph/models/dynamic_gni.py index 8cb6c7b..4fa1b1f 100644 --- a/chebai_graph/models/dynamic_gni.py +++ b/chebai_graph/models/dynamic_gni.py @@ -28,8 +28,8 @@ from chebai_graph.preprocessing.reader import RandomFeatureInitializationReader -from .base import GraphModelBase, GraphNetWrapper -from .resgated import ResGatedModel +from .architectures.base import GraphModelBase, GraphNetWrapper +from .architectures.resgated import ResGatedModel class ResGatedDynamicGNI(GraphModelBase): diff --git a/chebai_graph/models/graph.py b/chebai_graph/models/graph.py index 6037e93..69d27c0 100644 --- a/chebai_graph/models/graph.py +++ b/chebai_graph/models/graph.py @@ -10,7 +10,7 @@ from chebai_graph.loss.pretraining import MaskPretrainingLoss -from .base import GraphBaseNet +from .architectures.base import GraphBaseNet logging.getLogger("pysmiles").setLevel(logging.CRITICAL) diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py index a69baf7..c583d28 100644 --- a/chebai_graph/models/pooling.py +++ b/chebai_graph/models/pooling.py @@ -4,7 +4,7 @@ from torch_geometric.data import Data as GraphData from torch_scatter import scatter_add -from .base import GraphNetWrapper +from .architectures.base import GraphNetWrapper class AugmentedNodePoolingNet(GraphNetWrapper, ABC): From 542e55b07612fdb74aa63cc2c5099768dc878c9b Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 23 Jul 2026 16:30:35 +0200 Subject: [PATCH 17/34] pooling naming consistency --- chebai_graph/models/__init__.py | 8 ++++---- chebai_graph/models/augmented.py | 12 +++++------- chebai_graph/models/pooling.py | 4 ++-- .../amg_pool/gat.yml} | 0 configs/model/augmented/amg_pool/gine_baseline.yml | 13 +++++++++++++ configs/model/augmented/amg_pool/rggcn_baseline.yml | 12 ++++++++++++ 6 files changed, 36 insertions(+), 13 deletions(-) rename configs/model/{gat_aug_amgpool.yml => augmented/amg_pool/gat.yml} (100%) create mode 100644 configs/model/augmented/amg_pool/gine_baseline.yml create mode 100644 configs/model/augmented/amg_pool/rggcn_baseline.yml diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index a0b518b..7dcc154 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -2,19 +2,19 @@ from .architectures.gine import GINEGraphPred from .architectures.resgated import ResGatedGraphPred from .augmented import ( - GATAugNodePoolGraphPred, + GATAAPoolGraphPred, GATGraphNodeFGNodePoolGraphPred, - ResGatedAugNodePoolGraphPred, + ResGatedAAPoolGraphPred, ResGatedGraphNodeFGNodePoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred __all__ = [ "ResGatedGraphPred", - "ResGatedAugNodePoolGraphPred", + "ResGatedAAPoolGraphPred", "ResGatedGraphNodeFGNodePoolGraphPred", "GATGraphPred", - "GATAugNodePoolGraphPred", + "GATAAPoolGraphPred", "GATGraphNodeFGNodePoolGraphPred", "ResGatedDynamicGNIGraphPred", "GINEGraphPred", diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index c2b0a0c..3c3187f 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,9 +1,9 @@ from .architectures.gat import GATGraphPred from .architectures.resgated import ResGatedGraphPred -from .pooling import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet +from .pooling import AAPool, AMGPool -class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): +class ResGatedAAPoolGraphPred(AAPool, ResGatedGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -13,7 +13,7 @@ class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): ... -class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred): +class GATAAPoolGraphPred(AAPool, GATGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -23,9 +23,7 @@ class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred): ... -class ResGatedGraphNodeFGNodePoolGraphPred( - GraphNodeFGNodePoolingNet, ResGatedGraphPred -): +class ResGatedGraphNodeFGNodePoolGraphPred(AMGPool, ResGatedGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -35,7 +33,7 @@ class ResGatedGraphNodeFGNodePoolGraphPred( ... -class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, GATGraphPred): +class GATGraphNodeFGNodePoolGraphPred(AMGPool, GATGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py index c583d28..a1ee442 100644 --- a/chebai_graph/models/pooling.py +++ b/chebai_graph/models/pooling.py @@ -7,7 +7,7 @@ from .architectures.base import GraphNetWrapper -class AugmentedNodePoolingNet(GraphNetWrapper, ABC): +class AAPool(GraphNetWrapper, ABC): """ A pooling network that aggregates: - Atom node embeddings @@ -76,7 +76,7 @@ def forward(self, batch: dict) -> torch.Tensor: return self.lin_sequential(graph_vector) -class GraphNodeFGNodePoolingNet(GraphNetWrapper, ABC): +class AMGPool(GraphNetWrapper, ABC): """ A pooling network that pools node embeddings by aggregating: - Atom nodes diff --git a/configs/model/gat_aug_amgpool.yml b/configs/model/augmented/amg_pool/gat.yml similarity index 100% rename from configs/model/gat_aug_amgpool.yml rename to configs/model/augmented/amg_pool/gat.yml diff --git a/configs/model/augmented/amg_pool/gine_baseline.yml b/configs/model/augmented/amg_pool/gine_baseline.yml new file mode 100644 index 0000000..9d56063 --- /dev/null +++ b/configs/model/augmented/amg_pool/gine_baseline.yml @@ -0,0 +1,13 @@ +class_path: chebai_graph.models.GINEGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 161 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 8 # number of bond properties + dropout: 0 + train_eps: true + n_linear_layers: 1 diff --git a/configs/model/augmented/amg_pool/rggcn_baseline.yml b/configs/model/augmented/amg_pool/rggcn_baseline.yml new file mode 100644 index 0000000..d1847cc --- /dev/null +++ b/configs/model/augmented/amg_pool/rggcn_baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.models.ResGatedGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 161 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 8 # number of bond properties + dropout: 0 + n_linear_layers: 1 From 0622a685ca022007a8577f94faacbb882aaa1bf3 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 23 Jul 2026 17:14:45 +0200 Subject: [PATCH 18/34] amg pool rename --- chebai_graph/models/__init__.py | 8 ++++---- chebai_graph/models/augmented.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 7dcc154..581baa7 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -3,19 +3,19 @@ from .architectures.resgated import ResGatedGraphPred from .augmented import ( GATAAPoolGraphPred, - GATGraphNodeFGNodePoolGraphPred, + GATGraphAMGPoolGraphPred, ResGatedAAPoolGraphPred, - ResGatedGraphNodeFGNodePoolGraphPred, + ResGatedAMGPoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred __all__ = [ "ResGatedGraphPred", "ResGatedAAPoolGraphPred", - "ResGatedGraphNodeFGNodePoolGraphPred", + "ResGatedAMGPoolGraphPred", "GATGraphPred", "GATAAPoolGraphPred", - "GATGraphNodeFGNodePoolGraphPred", + "GATGraphAMGPoolGraphPred", "ResGatedDynamicGNIGraphPred", "GINEGraphPred", ] diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index 3c3187f..f7e1cb4 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -23,7 +23,7 @@ class GATAAPoolGraphPred(AAPool, GATGraphPred): ... -class ResGatedGraphNodeFGNodePoolGraphPred(AMGPool, ResGatedGraphPred): +class ResGatedAMGPoolGraphPred(AMGPool, ResGatedGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -33,7 +33,7 @@ class ResGatedGraphNodeFGNodePoolGraphPred(AMGPool, ResGatedGraphPred): ... -class GATGraphNodeFGNodePoolGraphPred(AMGPool, GATGraphPred): +class GATGraphAMGPoolGraphPred(AMGPool, GATGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). From 07a534ecf20f30438ca8ba99e0044f22799664bb Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 23 Jul 2026 17:22:18 +0200 Subject: [PATCH 19/34] update configs --- chebai_graph/models/__init__.py | 10 +++++--- chebai_graph/models/augmented.py | 23 ++++++++++++++++++- .../amg_pool => pooling/aa_pool}/gat.yml | 2 +- .../aa_pool/gine.yml} | 2 +- .../aa_pool/rggcn.yml} | 2 +- configs/model/pooling/amg_pool/gat.yml | 14 +++++++++++ configs/model/pooling/amg_pool/gine.yml | 13 +++++++++++ configs/model/pooling/amg_pool/rggcn.yml | 12 ++++++++++ .../standard/gat.yml} | 0 .../standard/gine.yml} | 0 .../standard/rggcn.yml} | 0 11 files changed, 71 insertions(+), 7 deletions(-) rename configs/model/{augmented/amg_pool => pooling/aa_pool}/gat.yml (88%) rename configs/model/{baselines/gine_baseline.yml => pooling/aa_pool/gine.yml} (84%) rename configs/model/{baselines/rggcn_baseline.yml => pooling/aa_pool/rggcn.yml} (81%) create mode 100644 configs/model/pooling/amg_pool/gat.yml create mode 100644 configs/model/pooling/amg_pool/gine.yml create mode 100644 configs/model/pooling/amg_pool/rggcn.yml rename configs/model/{baselines/gat_baseline.yml => pooling/standard/gat.yml} (100%) rename configs/model/{augmented/amg_pool/gine_baseline.yml => pooling/standard/gine.yml} (100%) rename configs/model/{augmented/amg_pool/rggcn_baseline.yml => pooling/standard/rggcn.yml} (100%) diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 581baa7..515a513 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -3,7 +3,9 @@ from .architectures.resgated import ResGatedGraphPred from .augmented import ( GATAAPoolGraphPred, - GATGraphAMGPoolGraphPred, + GATAMGPoolGraphPred, + GINEAAPoolGraphPred, + GINEAMGPoolGraphPred, ResGatedAAPoolGraphPred, ResGatedAMGPoolGraphPred, ) @@ -13,9 +15,11 @@ "ResGatedGraphPred", "ResGatedAAPoolGraphPred", "ResGatedAMGPoolGraphPred", + "ResGatedDynamicGNIGraphPred", "GATGraphPred", "GATAAPoolGraphPred", - "GATGraphAMGPoolGraphPred", - "ResGatedDynamicGNIGraphPred", + "GATAMGPoolGraphPred", "GINEGraphPred", + "GINEAAPoolGraphPred", + "GINEAMGPoolGraphPred", ] diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index f7e1cb4..87f1c30 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,4 +1,5 @@ from .architectures.gat import GATGraphPred +from .architectures.gine import GINEGraphPred from .architectures.resgated import ResGatedGraphPred from .pooling import AAPool, AMGPool @@ -23,6 +24,16 @@ class GATAAPoolGraphPred(AAPool, GATGraphPred): ... +class GINEAAPoolGraphPred(AAPool, GINEGraphPred): + """ + Combines: + - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). + - GINEGraphPred: Graph isomorphism network for final graph prediction. + """ + + ... + + class ResGatedAMGPoolGraphPred(AMGPool, ResGatedGraphPred): """ Combines: @@ -33,7 +44,7 @@ class ResGatedAMGPoolGraphPred(AMGPool, ResGatedGraphPred): ... -class GATGraphAMGPoolGraphPred(AMGPool, GATGraphPred): +class GATAMGPoolGraphPred(AMGPool, GATGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -41,3 +52,13 @@ class GATGraphAMGPoolGraphPred(AMGPool, GATGraphPred): """ ... + + +class GINEAMGPoolGraphPred(AMGPool, GINEGraphPred): + """ + Combines: + - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). + - GINEGraphPred: Graph isomorphism network for final graph prediction. + """ + + ... diff --git a/configs/model/augmented/amg_pool/gat.yml b/configs/model/pooling/aa_pool/gat.yml similarity index 88% rename from configs/model/augmented/amg_pool/gat.yml rename to configs/model/pooling/aa_pool/gat.yml index 959642c..36750e1 100644 --- a/configs/model/augmented/amg_pool/gat.yml +++ b/configs/model/pooling/aa_pool/gat.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.GATGraphNodeFGNodePoolGraphPred +class_path: chebai_graph.models.GATAAPoolGraphPred init_args: optimizer_kwargs: lr: 1e-3 diff --git a/configs/model/baselines/gine_baseline.yml b/configs/model/pooling/aa_pool/gine.yml similarity index 84% rename from configs/model/baselines/gine_baseline.yml rename to configs/model/pooling/aa_pool/gine.yml index 9d56063..7c60a73 100644 --- a/configs/model/baselines/gine_baseline.yml +++ b/configs/model/pooling/aa_pool/gine.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.GINEGraphPred +class_path: chebai_graph.models.GINEAAPoolGraphPred init_args: optimizer_kwargs: lr: 1e-3 diff --git a/configs/model/baselines/rggcn_baseline.yml b/configs/model/pooling/aa_pool/rggcn.yml similarity index 81% rename from configs/model/baselines/rggcn_baseline.yml rename to configs/model/pooling/aa_pool/rggcn.yml index d1847cc..4f6e94f 100644 --- a/configs/model/baselines/rggcn_baseline.yml +++ b/configs/model/pooling/aa_pool/rggcn.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.ResGatedGraphPred +class_path: chebai_graph.models.ResGatedAAPoolGraphPred init_args: optimizer_kwargs: lr: 1e-3 diff --git a/configs/model/pooling/amg_pool/gat.yml b/configs/model/pooling/amg_pool/gat.yml new file mode 100644 index 0000000..f88159d --- /dev/null +++ b/configs/model/pooling/amg_pool/gat.yml @@ -0,0 +1,14 @@ +class_path: chebai_graph.models.GATAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 12 # number of bond properties + heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) + v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0 + n_linear_layers: 1 diff --git a/configs/model/pooling/amg_pool/gine.yml b/configs/model/pooling/amg_pool/gine.yml new file mode 100644 index 0000000..87efd5a --- /dev/null +++ b/configs/model/pooling/amg_pool/gine.yml @@ -0,0 +1,13 @@ +class_path: chebai_graph.models.GINEAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 161 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 8 # number of bond properties + dropout: 0 + train_eps: true + n_linear_layers: 1 diff --git a/configs/model/pooling/amg_pool/rggcn.yml b/configs/model/pooling/amg_pool/rggcn.yml new file mode 100644 index 0000000..cb6b383 --- /dev/null +++ b/configs/model/pooling/amg_pool/rggcn.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.models.ResGatedAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 161 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 8 # number of bond properties + dropout: 0 + n_linear_layers: 1 diff --git a/configs/model/baselines/gat_baseline.yml b/configs/model/pooling/standard/gat.yml similarity index 100% rename from configs/model/baselines/gat_baseline.yml rename to configs/model/pooling/standard/gat.yml diff --git a/configs/model/augmented/amg_pool/gine_baseline.yml b/configs/model/pooling/standard/gine.yml similarity index 100% rename from configs/model/augmented/amg_pool/gine_baseline.yml rename to configs/model/pooling/standard/gine.yml diff --git a/configs/model/augmented/amg_pool/rggcn_baseline.yml b/configs/model/pooling/standard/rggcn.yml similarity index 100% rename from configs/model/augmented/amg_pool/rggcn_baseline.yml rename to configs/model/pooling/standard/rggcn.yml From f5af776754d23e01f57d26b1f9abd5eb804a65a5 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 23 Jul 2026 17:24:48 +0200 Subject: [PATCH 20/34] augmented config --- configs/model/{ => augmented}/pooling/aa_pool/gat.yml | 0 configs/model/{ => augmented}/pooling/aa_pool/gine.yml | 0 configs/model/{ => augmented}/pooling/aa_pool/rggcn.yml | 0 configs/model/{ => augmented}/pooling/amg_pool/gat.yml | 0 configs/model/{ => augmented}/pooling/amg_pool/gine.yml | 0 configs/model/{ => augmented}/pooling/amg_pool/rggcn.yml | 0 configs/model/{pooling/standard => baselines}/gat.yml | 0 configs/model/{pooling/standard => baselines}/gine.yml | 0 configs/model/{pooling/standard => baselines}/rggcn.yml | 0 9 files changed, 0 insertions(+), 0 deletions(-) rename configs/model/{ => augmented}/pooling/aa_pool/gat.yml (100%) rename configs/model/{ => augmented}/pooling/aa_pool/gine.yml (100%) rename configs/model/{ => augmented}/pooling/aa_pool/rggcn.yml (100%) rename configs/model/{ => augmented}/pooling/amg_pool/gat.yml (100%) rename configs/model/{ => augmented}/pooling/amg_pool/gine.yml (100%) rename configs/model/{ => augmented}/pooling/amg_pool/rggcn.yml (100%) rename configs/model/{pooling/standard => baselines}/gat.yml (100%) rename configs/model/{pooling/standard => baselines}/gine.yml (100%) rename configs/model/{pooling/standard => baselines}/rggcn.yml (100%) diff --git a/configs/model/pooling/aa_pool/gat.yml b/configs/model/augmented/pooling/aa_pool/gat.yml similarity index 100% rename from configs/model/pooling/aa_pool/gat.yml rename to configs/model/augmented/pooling/aa_pool/gat.yml diff --git a/configs/model/pooling/aa_pool/gine.yml b/configs/model/augmented/pooling/aa_pool/gine.yml similarity index 100% rename from configs/model/pooling/aa_pool/gine.yml rename to configs/model/augmented/pooling/aa_pool/gine.yml diff --git a/configs/model/pooling/aa_pool/rggcn.yml b/configs/model/augmented/pooling/aa_pool/rggcn.yml similarity index 100% rename from configs/model/pooling/aa_pool/rggcn.yml rename to configs/model/augmented/pooling/aa_pool/rggcn.yml diff --git a/configs/model/pooling/amg_pool/gat.yml b/configs/model/augmented/pooling/amg_pool/gat.yml similarity index 100% rename from configs/model/pooling/amg_pool/gat.yml rename to configs/model/augmented/pooling/amg_pool/gat.yml diff --git a/configs/model/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml similarity index 100% rename from configs/model/pooling/amg_pool/gine.yml rename to configs/model/augmented/pooling/amg_pool/gine.yml diff --git a/configs/model/pooling/amg_pool/rggcn.yml b/configs/model/augmented/pooling/amg_pool/rggcn.yml similarity index 100% rename from configs/model/pooling/amg_pool/rggcn.yml rename to configs/model/augmented/pooling/amg_pool/rggcn.yml diff --git a/configs/model/pooling/standard/gat.yml b/configs/model/baselines/gat.yml similarity index 100% rename from configs/model/pooling/standard/gat.yml rename to configs/model/baselines/gat.yml diff --git a/configs/model/pooling/standard/gine.yml b/configs/model/baselines/gine.yml similarity index 100% rename from configs/model/pooling/standard/gine.yml rename to configs/model/baselines/gine.yml diff --git a/configs/model/pooling/standard/rggcn.yml b/configs/model/baselines/rggcn.yml similarity index 100% rename from configs/model/pooling/standard/rggcn.yml rename to configs/model/baselines/rggcn.yml From 53514a46c45e433f26e4eb7520c0f8f0908a946b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 23 Jul 2026 19:41:58 +0200 Subject: [PATCH 21/34] config for final augmentation --- chebai_graph/models/architectures/gine.py | 1 + .../chebi50_final_augmented.yml} | 0 2 files changed, 1 insertion(+) rename configs/data/{chebi50_aug_prop_as_per_node.yml => augmented/chebi50_final_augmented.yml} (100%) diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py index 1885cce..f9e89a3 100644 --- a/chebai_graph/models/architectures/gine.py +++ b/chebai_graph/models/architectures/gine.py @@ -21,6 +21,7 @@ class GINEModel(BasicGNN): - https://arxiv.org/abs/1810.00826 (GIN) - https://arxiv.org/abs/1905.12265 (GINE / edge-feature extension) - https://github.com/pyg-team/pytorch_geometric/blob/master/examples/mutag_gin.py + - https://github.com/pyg-team/pytorch_geometric/issues/1311 Attributes: supports_edge_weight (bool): Indicates edge weights are not supported. diff --git a/configs/data/chebi50_aug_prop_as_per_node.yml b/configs/data/augmented/chebi50_final_augmented.yml similarity index 100% rename from configs/data/chebi50_aug_prop_as_per_node.yml rename to configs/data/augmented/chebi50_final_augmented.yml From f10b7ab7d70945351149849e3b84d94e701a4b28 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 12:20:38 +0200 Subject: [PATCH 22/34] update config --- configs/model/augmented/pooling/amg_pool/gine.yml | 4 ++-- configs/model/augmented/pooling/amg_pool/rggcn.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/configs/model/augmented/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml index 87efd5a..773088a 100644 --- a/configs/model/augmented/pooling/amg_pool/gine.yml +++ b/configs/model/augmented/pooling/amg_pool/gine.yml @@ -3,11 +3,11 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 161 # number of node/atom properties + in_channels: 203 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 8 # number of bond properties + edge_dim: 12 # number of bond properties dropout: 0 train_eps: true n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/amg_pool/rggcn.yml b/configs/model/augmented/pooling/amg_pool/rggcn.yml index cb6b383..b2657f2 100644 --- a/configs/model/augmented/pooling/amg_pool/rggcn.yml +++ b/configs/model/augmented/pooling/amg_pool/rggcn.yml @@ -3,10 +3,10 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 161 # number of node/atom properties + in_channels: 203 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 8 # number of bond properties + edge_dim: 12 # number of bond properties dropout: 0 n_linear_layers: 1 From cce6c6129e2f0ae99e9e009c24a1b123d60049df Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 14:45:28 +0200 Subject: [PATCH 23/34] update with new fg --- .../preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt index eae0a1d..273cc9b 100644 --- a/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt +++ b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt @@ -157,3 +157,4 @@ RING_71 RING_46 orthoester RING_55 +triiodomethyl From 32b863a7fe3ab6911668618f14cd33a8056cda39 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 14:47:53 +0200 Subject: [PATCH 24/34] update mol net classes --- .../preprocessing/datasets/__init__.py | 20 ++++++ .../datasets/moleculeNet_classification.py | 71 ------------------- .../datasets/molecule_net_classification.py | 60 ++++++++++++++++ .../data/augmented/BACE_final_augmented.yml | 24 +++++++ .../augmented/BBBP_final_augmented copy 2.yml | 24 +++++++ .../augmented/BBBP_final_augmented copy 3.yml | 24 +++++++ .../augmented/BBBP_final_augmented copy 4.yml | 24 +++++++ .../augmented/HIV_final_augmented copy 5.yml | 24 +++++++ .../augmented/Tox21_final_augmented copy.yml | 24 +++++++ .../ToxCast_final_augmented copy 6.yml | 24 +++++++ 10 files changed, 248 insertions(+), 71 deletions(-) delete mode 100644 chebai_graph/preprocessing/datasets/moleculeNet_classification.py create mode 100644 chebai_graph/preprocessing/datasets/molecule_net_classification.py create mode 100644 configs/data/augmented/BACE_final_augmented.yml create mode 100644 configs/data/augmented/BBBP_final_augmented copy 2.yml create mode 100644 configs/data/augmented/BBBP_final_augmented copy 3.yml create mode 100644 configs/data/augmented/BBBP_final_augmented copy 4.yml create mode 100644 configs/data/augmented/HIV_final_augmented copy 5.yml create mode 100644 configs/data/augmented/Tox21_final_augmented copy.yml create mode 100644 configs/data/augmented/ToxCast_final_augmented copy 6.yml diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 8708c28..38addf5 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -14,6 +14,17 @@ ChEBI50GraphProperties, ChEBI100GraphProperties, ) +from .molecule_net_classification import ( + Bace_WFGE_WGN_AsPerNodeType, + BBBP_WFGE_WGN_AsPerNodeType, + ClinTox_WFGE_WGN_AsPerNodeType, + HIV_WFGE_WGN_AsPerNodeType, + MUV_WFGE_WGN_AsPerNodeType, + PCBA_WFGE_WGN_AsPerNodeType, + Sider_WFGE_WGN_AsPerNodeType, + Tox21_WFGE_WGN_AsPerNodeType, + ToxCast_WFGE_WGN_AsPerNodeType, +) from .pubchem import PubChemGraphProperties __all__ = [ @@ -33,4 +44,13 @@ "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE", + "Bace_WFGE_WGN_AsPerNodeType", + "BBBP_WFGE_WGN_AsPerNodeType", + "ClinTox_WFGE_WGN_AsPerNodeType", + "HIV_WFGE_WGN_AsPerNodeType", + "MUV_WFGE_WGN_AsPerNodeType", + "Sider_WFGE_WGN_AsPerNodeType", + "Tox21_WFGE_WGN_AsPerNodeType", + "ToxCast_WFGE_WGN_AsPerNodeType", + "PCBA_WFGE_WGN_AsPerNodeType", ] diff --git a/chebai_graph/preprocessing/datasets/moleculeNet_classification.py b/chebai_graph/preprocessing/datasets/moleculeNet_classification.py deleted file mode 100644 index b77bdd7..0000000 --- a/chebai_graph/preprocessing/datasets/moleculeNet_classification.py +++ /dev/null @@ -1,71 +0,0 @@ -from chebai.preprocessing.datasets.molecule_classification import ( - BaceChem, - BBBPChem, - ClinToxChem, - HIVChem, - MUVChem, - SiderChem, -) - -from chebai_graph.preprocessing.datasets.base import ( - GraphPropAsPerNodeType, - GraphPropertiesMixIn, -) -from chebai_graph.preprocessing.reader.augmented_reader import ( - AtomFGReader_WithFGEdges_WithGraphNode, -) -from chebai_graph.preprocessing.reader.reader import GraphPropertyReader - - -class BaceChemDataset(GraphPropertiesMixIn, BaceChem): - READER = GraphPropertyReader - - -class BaceChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BaceChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class BBBPChemDataset(GraphPropertiesMixIn, BBBPChem): - READER = GraphPropertyReader - - -class BBBPChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BBBPChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class ClinToxChemDataset(GraphPropertiesMixIn, ClinToxChem): - READER = GraphPropertyReader - - -class ClinToxChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ClinToxChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class HIVChemDataset(GraphPropertiesMixIn, HIVChem): - READER = GraphPropertyReader - - -class HIVChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIVChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class SiderChemDataset(GraphPropertiesMixIn, SiderChem): - READER = GraphPropertyReader - - -class SiderChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SiderChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class MUVChemDataset(GraphPropertiesMixIn, MUVChem): - READER = GraphPropertyReader - - -class MUVChem_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, MUVChem): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -if __name__ == "__main__": - dataset = BaceChemDataset() - dataset.prepare_data() - dataset.setup() diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py new file mode 100644 index 0000000..4c7d399 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -0,0 +1,60 @@ +from chebai.preprocessing.datasets.molecule_net_classification import ( + BBBP, + HIV, + MUV, + PCBA, + Bace, + ClinTox, + Sider, + Tox21, + ToxCast, +) + +from chebai_graph.preprocessing.datasets.base import ( + GraphPropAsPerNodeType, +) +from chebai_graph.preprocessing.reader.augmented_reader import ( + AtomFGReader_WithFGEdges_WithGraphNode, +) + + +class PCBA_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, PCBA): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Bace_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Bace): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class BBBP_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BBBP): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ClinTox_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ClinTox): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class HIV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIV): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Sider_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Sider): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class MUV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, MUV): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Tox21_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Tox21): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ToxCast_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +if __name__ == "__main__": + dataset = Bace_WFGE_WGN_AsPerNodeType() + dataset.prepare_data() + dataset.setup() diff --git a/configs/data/augmented/BACE_final_augmented.yml b/configs/data/augmented/BACE_final_augmented.yml new file mode 100644 index 0000000..7ea67f0 --- /dev/null +++ b/configs/data/augmented/BACE_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BACE_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/BBBP_final_augmented copy 2.yml b/configs/data/augmented/BBBP_final_augmented copy 2.yml new file mode 100644 index 0000000..9a7ee90 --- /dev/null +++ b/configs/data/augmented/BBBP_final_augmented copy 2.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/BBBP_final_augmented copy 3.yml b/configs/data/augmented/BBBP_final_augmented copy 3.yml new file mode 100644 index 0000000..9a7ee90 --- /dev/null +++ b/configs/data/augmented/BBBP_final_augmented copy 3.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/BBBP_final_augmented copy 4.yml b/configs/data/augmented/BBBP_final_augmented copy 4.yml new file mode 100644 index 0000000..9a7ee90 --- /dev/null +++ b/configs/data/augmented/BBBP_final_augmented copy 4.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/HIV_final_augmented copy 5.yml b/configs/data/augmented/HIV_final_augmented copy 5.yml new file mode 100644 index 0000000..596fc19 --- /dev/null +++ b/configs/data/augmented/HIV_final_augmented copy 5.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.HIV_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/Tox21_final_augmented copy.yml b/configs/data/augmented/Tox21_final_augmented copy.yml new file mode 100644 index 0000000..1c88112 --- /dev/null +++ b/configs/data/augmented/Tox21_final_augmented copy.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/ToxCast_final_augmented copy 6.yml b/configs/data/augmented/ToxCast_final_augmented copy 6.yml new file mode 100644 index 0000000..4831ae5 --- /dev/null +++ b/configs/data/augmented/ToxCast_final_augmented copy 6.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType From 9c4c5ef6aa8e64aa29cf6679faa10b0dda44605a Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 25 Jul 2026 15:09:23 +0200 Subject: [PATCH 25/34] update mol net classes --- .../preprocessing/datasets/__init__.py | 8 +++---- .../datasets/molecule_net_classification.py | 10 ++++---- ...ed copy 2.yml => BBBP_final_augmented.yml} | 0 .../augmented/ClinTox_final_augmented.yml | 24 +++++++++++++++++++ ...ted copy 5.yml => HIV_final_augmented.yml} | 0 ...ted copy 4.yml => MUV_final_augmented.yml} | 2 +- ...ed copy 3.yml => PCBA_final_augmented.yml} | 2 +- .../data/augmented/SIDER_final_augmented.yml | 24 +++++++++++++++++++ ...ted copy.yml => Tox21_final_augmented.yml} | 0 ...copy 6.yml => ToxCast_final_augmented.yml} | 0 10 files changed, 59 insertions(+), 11 deletions(-) rename configs/data/augmented/{BBBP_final_augmented copy 2.yml => BBBP_final_augmented.yml} (100%) create mode 100644 configs/data/augmented/ClinTox_final_augmented.yml rename configs/data/augmented/{HIV_final_augmented copy 5.yml => HIV_final_augmented.yml} (100%) rename configs/data/augmented/{BBBP_final_augmented copy 4.yml => MUV_final_augmented.yml} (93%) rename configs/data/augmented/{BBBP_final_augmented copy 3.yml => PCBA_final_augmented.yml} (94%) create mode 100644 configs/data/augmented/SIDER_final_augmented.yml rename configs/data/augmented/{Tox21_final_augmented copy.yml => Tox21_final_augmented.yml} (100%) rename configs/data/augmented/{ToxCast_final_augmented copy 6.yml => ToxCast_final_augmented.yml} (100%) diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 38addf5..1dcb2f3 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -15,13 +15,13 @@ ChEBI100GraphProperties, ) from .molecule_net_classification import ( - Bace_WFGE_WGN_AsPerNodeType, + BACE_WFGE_WGN_AsPerNodeType, BBBP_WFGE_WGN_AsPerNodeType, ClinTox_WFGE_WGN_AsPerNodeType, HIV_WFGE_WGN_AsPerNodeType, MUV_WFGE_WGN_AsPerNodeType, PCBA_WFGE_WGN_AsPerNodeType, - Sider_WFGE_WGN_AsPerNodeType, + SIDER_WFGE_WGN_AsPerNodeType, Tox21_WFGE_WGN_AsPerNodeType, ToxCast_WFGE_WGN_AsPerNodeType, ) @@ -44,12 +44,12 @@ "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE", - "Bace_WFGE_WGN_AsPerNodeType", + "BACE_WFGE_WGN_AsPerNodeType", "BBBP_WFGE_WGN_AsPerNodeType", "ClinTox_WFGE_WGN_AsPerNodeType", "HIV_WFGE_WGN_AsPerNodeType", "MUV_WFGE_WGN_AsPerNodeType", - "Sider_WFGE_WGN_AsPerNodeType", + "SIDER_WFGE_WGN_AsPerNodeType", "Tox21_WFGE_WGN_AsPerNodeType", "ToxCast_WFGE_WGN_AsPerNodeType", "PCBA_WFGE_WGN_AsPerNodeType", diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py index 4c7d399..9afaf71 100644 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -1,11 +1,11 @@ from chebai.preprocessing.datasets.molecule_net_classification import ( + BACE, BBBP, HIV, MUV, PCBA, - Bace, + SIDER, ClinTox, - Sider, Tox21, ToxCast, ) @@ -22,7 +22,7 @@ class PCBA_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, PCBA): READER = AtomFGReader_WithFGEdges_WithGraphNode -class Bace_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Bace): +class BACE_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BACE): READER = AtomFGReader_WithFGEdges_WithGraphNode @@ -38,7 +38,7 @@ class HIV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIV): READER = AtomFGReader_WithFGEdges_WithGraphNode -class Sider_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Sider): +class SIDER_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SIDER): READER = AtomFGReader_WithFGEdges_WithGraphNode @@ -55,6 +55,6 @@ class ToxCast_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): if __name__ == "__main__": - dataset = Bace_WFGE_WGN_AsPerNodeType() + dataset = BACE_WFGE_WGN_AsPerNodeType() dataset.prepare_data() dataset.setup() diff --git a/configs/data/augmented/BBBP_final_augmented copy 2.yml b/configs/data/augmented/BBBP_final_augmented.yml similarity index 100% rename from configs/data/augmented/BBBP_final_augmented copy 2.yml rename to configs/data/augmented/BBBP_final_augmented.yml diff --git a/configs/data/augmented/ClinTox_final_augmented.yml b/configs/data/augmented/ClinTox_final_augmented.yml new file mode 100644 index 0000000..a4023b3 --- /dev/null +++ b/configs/data/augmented/ClinTox_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ClinTox_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/HIV_final_augmented copy 5.yml b/configs/data/augmented/HIV_final_augmented.yml similarity index 100% rename from configs/data/augmented/HIV_final_augmented copy 5.yml rename to configs/data/augmented/HIV_final_augmented.yml diff --git a/configs/data/augmented/BBBP_final_augmented copy 4.yml b/configs/data/augmented/MUV_final_augmented.yml similarity index 93% rename from configs/data/augmented/BBBP_final_augmented copy 4.yml rename to configs/data/augmented/MUV_final_augmented.yml index 9a7ee90..56a3e1e 100644 --- a/configs/data/augmented/BBBP_final_augmented copy 4.yml +++ b/configs/data/augmented/MUV_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.MUV_WFGE_WGN_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/BBBP_final_augmented copy 3.yml b/configs/data/augmented/PCBA_final_augmented.yml similarity index 94% rename from configs/data/augmented/BBBP_final_augmented copy 3.yml rename to configs/data/augmented/PCBA_final_augmented.yml index 9a7ee90..40688f4 100644 --- a/configs/data/augmented/BBBP_final_augmented copy 3.yml +++ b/configs/data/augmented/PCBA_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.PCBA_WFGE_WGN_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/SIDER_final_augmented.yml b/configs/data/augmented/SIDER_final_augmented.yml new file mode 100644 index 0000000..5eef128 --- /dev/null +++ b/configs/data/augmented/SIDER_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.SIDER_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/Tox21_final_augmented copy.yml b/configs/data/augmented/Tox21_final_augmented.yml similarity index 100% rename from configs/data/augmented/Tox21_final_augmented copy.yml rename to configs/data/augmented/Tox21_final_augmented.yml diff --git a/configs/data/augmented/ToxCast_final_augmented copy 6.yml b/configs/data/augmented/ToxCast_final_augmented.yml similarity index 100% rename from configs/data/augmented/ToxCast_final_augmented copy 6.yml rename to configs/data/augmented/ToxCast_final_augmented.yml From 8f207cc97333d4c28c97f3c606cabf8e16032609 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 1 Aug 2026 18:45:16 +0200 Subject: [PATCH 26/34] remove molecule attr from geom data --- chebai_graph/preprocessing/datasets/base.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 13bacbb..37e308e 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -336,7 +336,6 @@ def _merge_props_into_base(self, row: pd.Series | dict) -> GeomData: x=x, edge_index=geom_data.edge_index, edge_attr=edge_attr, - molecule_attr=molecule_attr, ) def load_processed_data( @@ -627,7 +626,6 @@ def _merge_props_into_base( x=x, edge_index=geom_data.edge_index, edge_attr=edge_attr, - molecule_attr=torch.empty((1, 0)), # empty as not used for this class is_atom_node=is_atom_node, is_fg_node=is_fg_node, is_graph_node=is_graph_node, From e4a6f771f9565d3c5aa04eba79ee333e64ee6cbe Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 2 Aug 2026 00:58:23 +0200 Subject: [PATCH 27/34] modularize _merge_props_into_base of graph per node type class --- chebai_graph/preprocessing/datasets/base.py | 117 ++++++++++++++------ 1 file changed, 83 insertions(+), 34 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 37e308e..792b9eb 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -559,6 +559,7 @@ def _merge_props_into_base( is_fg_node = ~is_atom_node & ~is_graph_node num_nodes = geom_data.x.size(0) edge_attr = geom_data.edge_attr + assert edge_attr is not None, "edge_attr must be set in the geom_data" # Initialize node feature matrix assert max_len_node_properties is not None, ( @@ -581,41 +582,24 @@ def _merge_props_into_base( (0, property.encoder.get_encoding_length()) ) - enc_len = property_values.shape[1] - # -------------- Node properties --------------- - if isinstance(property, AllNodeTypeProperty): - x[:, atom_offset : atom_offset + enc_len] = property_values - atom_offset += enc_len - fg_offset += enc_len - graph_offset += enc_len - - elif isinstance(property, AtomNodeTypeProperty): - x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ - is_atom_node - ] - atom_offset += enc_len - - elif isinstance(property, FGNodeTypeProperty): - x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ - is_fg_node - ] - fg_offset += enc_len - - elif isinstance(property, MoleculeProperty): - x[is_graph_node, graph_offset : graph_offset + enc_len] = ( - property_values - ) - graph_offset += enc_len + build_node_property_tensor_result = self._build_node_property_tensor( + node_tensor=x, + atom_offset=atom_offset, + fg_offset=fg_offset, + graph_offset=graph_offset, + property_values=property_values, + is_atom_node=is_atom_node, + is_fg_node=is_fg_node, + is_graph_node=is_graph_node, + ) + x = build_node_property_tensor_result["node_tensor"] + atom_offset = build_node_property_tensor_result["atom_offset"] + fg_offset = build_node_property_tensor_result["fg_offset"] + graph_offset = build_node_property_tensor_result["graph_offset"] - # ------------- Bond Properties -------------- - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") + edge_attr = self._build_edge_property_tensor( + edge_attr_tensor=edge_attr, property_values=property_values + ) total_used_columns = max(atom_offset, fg_offset, graph_offset) assert total_used_columns <= max_len_node_properties, ( @@ -631,6 +615,71 @@ def _merge_props_into_base( is_graph_node=is_graph_node, ) + def _build_node_property_tensor( + self, + node_tensor: torch.Tensor, + atom_offset: int, + fg_offset: int, + graph_offset: int, + property_values: torch.Tensor, + is_atom_node: torch.Tensor, + is_fg_node: torch.Tensor, + is_graph_node: torch.Tensor, + ) -> dict: + enc_len = property_values.shape[1] + # -------------- Node properties --------------- + if isinstance(property, AllNodeTypeProperty): + node_tensor[:, atom_offset : atom_offset + enc_len] = property_values + atom_offset += enc_len + fg_offset += enc_len + graph_offset += enc_len + + elif isinstance(property, AtomNodeTypeProperty): + node_tensor[is_atom_node, atom_offset : atom_offset + enc_len] = ( + property_values[is_atom_node] + ) + atom_offset += enc_len + + elif isinstance(property, FGNodeTypeProperty): + node_tensor[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ + is_fg_node + ] + fg_offset += enc_len + + elif isinstance(property, MoleculeProperty): + node_tensor[is_graph_node, graph_offset : graph_offset + enc_len] = ( + property_values + ) + graph_offset += enc_len + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + return { + "node_tensor": node_tensor, + "atom_offset": atom_offset, + "fg_offset": fg_offset, + "graph_offset": graph_offset, + } + + def _build_edge_property_tensor( + self, + edge_attr_tensor: torch.Tensor, + property_values: torch.Tensor, + ) -> torch.Tensor: + if isinstance(property, BondProperty): + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr_tensor = torch.cat( + [ + edge_attr_tensor, + torch.cat([property_values, property_values], dim=0), + ], + dim=1, + ) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + return edge_attr_tensor + def _prediction_merge_props_into_base_wrapper( self, row: pd.Series | dict, model_hparams: Optional[dict] = None ) -> GeomData: From b5e84cd22f53775d15a280748ac4c326b0658db0 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 2 Aug 2026 10:40:39 +0200 Subject: [PATCH 28/34] ablation for properties --- .../datasets/augmentation_base.py | 141 +++++++++++++++++- chebai_graph/preprocessing/datasets/base.py | 71 ++++++++- 2 files changed, 203 insertions(+), 9 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py index 9e2bf61..c12eb7a 100644 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -1,9 +1,10 @@ from abc import ABC import pandas as pd +import torch from torch_geometric.data.data import Data as GeomData -from .base import GraphPropertiesMixIn +from .base import GraphPropAsPerNodeType, GraphPropertiesMixIn class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): @@ -48,3 +49,141 @@ def _add_graph_node_mask(self, data: GeomData, row: pd.Series) -> GeomData: assert is_graph_node is not None, "is_graph_node must be set in the geom_data" data.is_graph_node = is_graph_node return data + + +class GraphPropForAtomLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_fg_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_fg_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_fg_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + def _fill_node_tensor_with_molecule_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_graph_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_graph_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + +class GraphPropForFGLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_atom_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_atom_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_atom_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + def _fill_node_tensor_with_molecule_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_graph_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_graph_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + +class GraphPropForGraphLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_atom_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_atom_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_atom_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + def _fill_node_tensor_with_fg_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_fg_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_fg_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + +class GraphPropForAtomAndFGLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_molecule_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_graph_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_graph_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + +class GraphPropForAtomAndGraphLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_fg_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_fg_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_fg_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor + + +class GraphPropForFGAndGraphLevelOnly(GraphPropAsPerNodeType): + def _fill_node_tensor_with_atom_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_atom_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( + torch.tensor.zeros( + property_values[is_atom_node].shape, dtype=property_values.dtype + ) + ) + return node_tensor diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 792b9eb..0103a12 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -629,26 +629,39 @@ def _build_node_property_tensor( enc_len = property_values.shape[1] # -------------- Node properties --------------- if isinstance(property, AllNodeTypeProperty): - node_tensor[:, atom_offset : atom_offset + enc_len] = property_values + node_tensor = self._fill_node_tensor_with_all_node_type_property( + node_tensor=node_tensor, + property_values=property_values, + offset=atom_offset, + ) atom_offset += enc_len fg_offset += enc_len graph_offset += enc_len elif isinstance(property, AtomNodeTypeProperty): - node_tensor[is_atom_node, atom_offset : atom_offset + enc_len] = ( - property_values[is_atom_node] + node_tensor = self._fill_node_tensor_with_atom_type_property( + node_tensor=node_tensor, + property_values=property_values, + offset=atom_offset, + is_atom_node=is_atom_node, ) atom_offset += enc_len elif isinstance(property, FGNodeTypeProperty): - node_tensor[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ - is_fg_node - ] + node_tensor = self._fill_node_tensor_with_fg_type_property( + node_tensor=node_tensor, + property_values=property_values, + offset=fg_offset, + is_fg_node=is_fg_node, + ) fg_offset += enc_len elif isinstance(property, MoleculeProperty): - node_tensor[is_graph_node, graph_offset : graph_offset + enc_len] = ( - property_values + node_tensor = self._fill_node_tensor_with_molecule_type_property( + node_tensor=node_tensor, + property_values=property_values, + offset=graph_offset, + is_graph_node=is_graph_node, ) graph_offset += enc_len else: @@ -661,6 +674,48 @@ def _build_node_property_tensor( "graph_offset": graph_offset, } + def _fill_node_tensor_with_all_node_type_property( + self, node_tensor: torch.Tensor, property_values: torch.Tensor, offset: int + ) -> torch.Tensor: + node_tensor[:, offset : offset + property_values.shape[1]] = property_values + return node_tensor + + def _fill_node_tensor_with_atom_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_atom_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( + property_values[is_atom_node] + ) + return node_tensor + + def _fill_node_tensor_with_fg_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_fg_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( + property_values[is_fg_node] + ) + return node_tensor + + def _fill_node_tensor_with_molecule_type_property( + self, + node_tensor: torch.Tensor, + property_values: torch.Tensor, + offset: int, + is_graph_node: torch.Tensor, + ) -> torch.Tensor: + node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( + property_values[is_graph_node] + ) + return node_tensor + def _build_edge_property_tensor( self, edge_attr_tensor: torch.Tensor, From 842f7b7a9c17c8e3e174bad51d3400b77821067d Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 2 Aug 2026 11:08:14 +0200 Subject: [PATCH 29/34] properties ablation config --- .../preprocessing/datasets/__init__.py | 12 +++++++ .../datasets/augmentation_base.py | 6 ++-- chebai_graph/preprocessing/datasets/chebi.py | 36 +++++++++++++++++++ .../{ => final}/BACE_final_augmented.yml | 0 .../{ => final}/BBBP_final_augmented.yml | 0 .../{ => final}/ClinTox_final_augmented.yml | 0 .../{ => final}/HIV_final_augmented.yml | 0 .../{ => final}/MUV_final_augmented.yml | 0 .../{ => final}/PCBA_final_augmented.yml | 0 .../{ => final}/SIDER_final_augmented.yml | 0 .../{ => final}/Tox21_final_augmented.yml | 0 .../{ => final}/ToxCast_final_augmented.yml | 0 .../{ => final}/chebi50_final_augmented.yml | 0 .../properties/chebi50_atom_fg_prop_only.yml | 24 +++++++++++++ .../chebi50_atom_graph_node_prop_only.yml | 24 +++++++++++++ .../properties/chebi50_atom_prop_only.yml | 24 +++++++++++++ .../chebi50_fg_graph_node_prop_only.yml | 24 +++++++++++++ .../properties/chebi50_fg_prop_only.yml | 24 +++++++++++++ .../chebi50_graph_node_prop_only.yml | 24 +++++++++++++ .../augmented/pooling/no_pooling/gat.yml | 14 ++++++++ 20 files changed, 209 insertions(+), 3 deletions(-) rename configs/data/augmented/{ => final}/BACE_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/BBBP_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/ClinTox_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/HIV_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/MUV_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/PCBA_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/SIDER_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/Tox21_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/ToxCast_final_augmented.yml (100%) rename configs/data/augmented/{ => final}/chebi50_final_augmented.yml (100%) create mode 100644 configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml create mode 100644 configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml create mode 100644 configs/data/augmented/properties/chebi50_atom_prop_only.yml create mode 100644 configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml create mode 100644 configs/data/augmented/properties/chebi50_fg_prop_only.yml create mode 100644 configs/data/augmented/properties/chebi50_graph_node_prop_only.yml create mode 100644 configs/model/augmented/pooling/no_pooling/gat.yml diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 1dcb2f3..44ea569 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -9,6 +9,12 @@ ChEBI50_StaticGNI, ChEBI50_WFGE_NGN_GraphProp, ChEBI50_WFGE_WGN_AsPerNodeType, + ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly, + ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly, + ChEBI50_WFGE_WGN_ForAtomLevelOnly, + ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly, + ChEBI50_WFGE_WGN_ForFGLevelOnly, + ChEBI50_WFGE_WGN_ForGraphNodeOnly, ChEBI50_WFGE_WGN_GraphProp, ChEBI50GraphData, ChEBI50GraphProperties, @@ -44,6 +50,12 @@ "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE", + "ChEBI50_WFGE_WGN_ForAtomLevelOnly", + "ChEBI50_WFGE_WGN_ForFGLevelOnly", + "ChEBI50_WFGE_WGN_ForGraphNodeOnly", + "ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly", + "ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly", + "ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly", "BACE_WFGE_WGN_AsPerNodeType", "BBBP_WFGE_WGN_AsPerNodeType", "ClinTox_WFGE_WGN_AsPerNodeType", diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py index c12eb7a..061813d 100644 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -111,7 +111,7 @@ def _fill_node_tensor_with_molecule_type_property( return node_tensor -class GraphPropForGraphLevelOnly(GraphPropAsPerNodeType): +class GraphPropForGraphNodeOnly(GraphPropAsPerNodeType): def _fill_node_tensor_with_atom_type_property( self, node_tensor: torch.Tensor, @@ -157,7 +157,7 @@ def _fill_node_tensor_with_molecule_type_property( return node_tensor -class GraphPropForAtomAndGraphLevelOnly(GraphPropAsPerNodeType): +class GraphPropForAtomLevelAndGraphNodeOnly(GraphPropAsPerNodeType): def _fill_node_tensor_with_fg_type_property( self, node_tensor: torch.Tensor, @@ -173,7 +173,7 @@ def _fill_node_tensor_with_fg_type_property( return node_tensor -class GraphPropForFGAndGraphLevelOnly(GraphPropAsPerNodeType): +class GraphPropForFGLevelAndGraphNodeOnly(GraphPropAsPerNodeType): def _fill_node_tensor_with_atom_type_property( self, node_tensor: torch.Tensor, diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 263747e..e8585e8 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -24,6 +24,12 @@ from .augmentation_base import ( AugGraphPropMixIn_NoGraphNode, AugGraphPropMixIn_WithGraphNode, + GraphPropForAtomAndFGLevelOnly, + GraphPropForAtomLevelAndGraphNodeOnly, + GraphPropForAtomLevelOnly, + GraphPropForFGLevelAndGraphNodeOnly, + GraphPropForFGLevelOnly, + GraphPropForGraphNodeOnly, ) from .base import DataPropertiesSetter, GraphPropAsPerNodeType, GraphPropertiesMixIn @@ -150,6 +156,36 @@ class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_WFGE_WGN_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_WFGE_WGN_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_WFGE_WGN_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly( + GraphPropForAtomAndFGLevelOnly, ChEBIOver50 +): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly( + GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly( + GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + class ChEBI100_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver100): READER = AtomFGReader_WithFGEdges_WithGraphNode diff --git a/configs/data/augmented/BACE_final_augmented.yml b/configs/data/augmented/final/BACE_final_augmented.yml similarity index 100% rename from configs/data/augmented/BACE_final_augmented.yml rename to configs/data/augmented/final/BACE_final_augmented.yml diff --git a/configs/data/augmented/BBBP_final_augmented.yml b/configs/data/augmented/final/BBBP_final_augmented.yml similarity index 100% rename from configs/data/augmented/BBBP_final_augmented.yml rename to configs/data/augmented/final/BBBP_final_augmented.yml diff --git a/configs/data/augmented/ClinTox_final_augmented.yml b/configs/data/augmented/final/ClinTox_final_augmented.yml similarity index 100% rename from configs/data/augmented/ClinTox_final_augmented.yml rename to configs/data/augmented/final/ClinTox_final_augmented.yml diff --git a/configs/data/augmented/HIV_final_augmented.yml b/configs/data/augmented/final/HIV_final_augmented.yml similarity index 100% rename from configs/data/augmented/HIV_final_augmented.yml rename to configs/data/augmented/final/HIV_final_augmented.yml diff --git a/configs/data/augmented/MUV_final_augmented.yml b/configs/data/augmented/final/MUV_final_augmented.yml similarity index 100% rename from configs/data/augmented/MUV_final_augmented.yml rename to configs/data/augmented/final/MUV_final_augmented.yml diff --git a/configs/data/augmented/PCBA_final_augmented.yml b/configs/data/augmented/final/PCBA_final_augmented.yml similarity index 100% rename from configs/data/augmented/PCBA_final_augmented.yml rename to configs/data/augmented/final/PCBA_final_augmented.yml diff --git a/configs/data/augmented/SIDER_final_augmented.yml b/configs/data/augmented/final/SIDER_final_augmented.yml similarity index 100% rename from configs/data/augmented/SIDER_final_augmented.yml rename to configs/data/augmented/final/SIDER_final_augmented.yml diff --git a/configs/data/augmented/Tox21_final_augmented.yml b/configs/data/augmented/final/Tox21_final_augmented.yml similarity index 100% rename from configs/data/augmented/Tox21_final_augmented.yml rename to configs/data/augmented/final/Tox21_final_augmented.yml diff --git a/configs/data/augmented/ToxCast_final_augmented.yml b/configs/data/augmented/final/ToxCast_final_augmented.yml similarity index 100% rename from configs/data/augmented/ToxCast_final_augmented.yml rename to configs/data/augmented/final/ToxCast_final_augmented.yml diff --git a/configs/data/augmented/chebi50_final_augmented.yml b/configs/data/augmented/final/chebi50_final_augmented.yml similarity index 100% rename from configs/data/augmented/chebi50_final_augmented.yml rename to configs/data/augmented/final/chebi50_final_augmented.yml diff --git a/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml new file mode 100644 index 0000000..a06a948 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml b/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml new file mode 100644 index 0000000..26ec990 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_atom_prop_only.yml b/configs/data/augmented/properties/chebi50_atom_prop_only.yml new file mode 100644 index 0000000..31821f1 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_atom_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForAtomLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml b/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml new file mode 100644 index 0000000..785f798 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_fg_prop_only.yml b/configs/data/augmented/properties/chebi50_fg_prop_only.yml new file mode 100644 index 0000000..10ed37d --- /dev/null +++ b/configs/data/augmented/properties/chebi50_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForFGLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml b/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml new file mode 100644 index 0000000..9827723 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/model/augmented/pooling/no_pooling/gat.yml b/configs/model/augmented/pooling/no_pooling/gat.yml new file mode 100644 index 0000000..f9d3575 --- /dev/null +++ b/configs/model/augmented/pooling/no_pooling/gat.yml @@ -0,0 +1,14 @@ +class_path: chebai_graph.models.GATGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 12 # number of bond properties + heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) + v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0 + n_linear_layers: 1 From 1da7014748f0675701a0dd613efcc25e33631d5e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 2 Aug 2026 11:59:07 +0200 Subject: [PATCH 30/34] fix prop error --- chebai_graph/preprocessing/datasets/base.py | 60 ++++++++++----------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 0103a12..6662abb 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -582,24 +582,28 @@ def _merge_props_into_base( (0, property.encoder.get_encoding_length()) ) - build_node_property_tensor_result = self._build_node_property_tensor( - node_tensor=x, - atom_offset=atom_offset, - fg_offset=fg_offset, - graph_offset=graph_offset, - property_values=property_values, - is_atom_node=is_atom_node, - is_fg_node=is_fg_node, - is_graph_node=is_graph_node, - ) - x = build_node_property_tensor_result["node_tensor"] - atom_offset = build_node_property_tensor_result["atom_offset"] - fg_offset = build_node_property_tensor_result["fg_offset"] - graph_offset = build_node_property_tensor_result["graph_offset"] + if isinstance(property, AtomProperty): + build_node_property_tensor_result = self._build_node_property_tensor( + node_tensor=x, + atom_offset=atom_offset, + fg_offset=fg_offset, + graph_offset=graph_offset, + property_values=property_values, + is_atom_node=is_atom_node, + is_fg_node=is_fg_node, + is_graph_node=is_graph_node, + ) + x = build_node_property_tensor_result["node_tensor"] + atom_offset = build_node_property_tensor_result["atom_offset"] + fg_offset = build_node_property_tensor_result["fg_offset"] + graph_offset = build_node_property_tensor_result["graph_offset"] - edge_attr = self._build_edge_property_tensor( - edge_attr_tensor=edge_attr, property_values=property_values - ) + elif isinstance(property, BondProperty): + edge_attr = self._build_edge_property_tensor( + edge_attr_tensor=edge_attr, property_values=property_values + ) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") total_used_columns = max(atom_offset, fg_offset, graph_offset) assert total_used_columns <= max_len_node_properties, ( @@ -664,8 +668,6 @@ def _build_node_property_tensor( is_graph_node=is_graph_node, ) graph_offset += enc_len - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") return { "node_tensor": node_tensor, @@ -721,18 +723,14 @@ def _build_edge_property_tensor( edge_attr_tensor: torch.Tensor, property_values: torch.Tensor, ) -> torch.Tensor: - if isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr_tensor = torch.cat( - [ - edge_attr_tensor, - torch.cat([property_values, property_values], dim=0), - ], - dim=1, - ) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr_tensor = torch.cat( + [ + edge_attr_tensor, + torch.cat([property_values, property_values], dim=0), + ], + dim=1, + ) return edge_attr_tensor def _prediction_merge_props_into_base_wrapper( From e0b0ad782e10f5cb838e19a0c96c6117e5c290a7 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 2 Aug 2026 12:23:49 +0200 Subject: [PATCH 31/34] property zeros fix --- .../datasets/augmentation_base.py | 45 ------------------- chebai_graph/preprocessing/datasets/base.py | 4 +- 2 files changed, 3 insertions(+), 46 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py index 061813d..d95e8ae 100644 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -59,11 +59,6 @@ def _fill_node_tensor_with_fg_type_property( offset: int, is_fg_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_fg_node].shape, dtype=property_values.dtype - ) - ) return node_tensor def _fill_node_tensor_with_molecule_type_property( @@ -73,11 +68,6 @@ def _fill_node_tensor_with_molecule_type_property( offset: int, is_graph_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_graph_node].shape, dtype=property_values.dtype - ) - ) return node_tensor @@ -89,11 +79,6 @@ def _fill_node_tensor_with_atom_type_property( offset: int, is_atom_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_atom_node].shape, dtype=property_values.dtype - ) - ) return node_tensor def _fill_node_tensor_with_molecule_type_property( @@ -103,11 +88,6 @@ def _fill_node_tensor_with_molecule_type_property( offset: int, is_graph_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_graph_node].shape, dtype=property_values.dtype - ) - ) return node_tensor @@ -119,11 +99,6 @@ def _fill_node_tensor_with_atom_type_property( offset: int, is_atom_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_atom_node].shape, dtype=property_values.dtype - ) - ) return node_tensor def _fill_node_tensor_with_fg_type_property( @@ -133,11 +108,6 @@ def _fill_node_tensor_with_fg_type_property( offset: int, is_fg_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_fg_node].shape, dtype=property_values.dtype - ) - ) return node_tensor @@ -149,11 +119,6 @@ def _fill_node_tensor_with_molecule_type_property( offset: int, is_graph_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_graph_node].shape, dtype=property_values.dtype - ) - ) return node_tensor @@ -165,11 +130,6 @@ def _fill_node_tensor_with_fg_type_property( offset: int, is_fg_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_fg_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_fg_node].shape, dtype=property_values.dtype - ) - ) return node_tensor @@ -181,9 +141,4 @@ def _fill_node_tensor_with_atom_type_property( offset: int, is_atom_node: torch.Tensor, ) -> torch.Tensor: - node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( - torch.tensor.zeros( - property_values[is_atom_node].shape, dtype=property_values.dtype - ) - ) return node_tensor diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 6662abb..3d8d517 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -582,8 +582,9 @@ def _merge_props_into_base( (0, property.encoder.get_encoding_length()) ) - if isinstance(property, AtomProperty): + if isinstance(property, (AtomProperty, MoleculeProperty)): build_node_property_tensor_result = self._build_node_property_tensor( + property=property, node_tensor=x, atom_offset=atom_offset, fg_offset=fg_offset, @@ -621,6 +622,7 @@ def _merge_props_into_base( def _build_node_property_tensor( self, + property: MolecularProperty, node_tensor: torch.Tensor, atom_offset: int, fg_offset: int, From 4623755770239a90a168a30fbbf7094c63f052e3 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 2 Aug 2026 13:17:06 +0200 Subject: [PATCH 32/34] graph node properties need no masking --- chebai_graph/preprocessing/datasets/base.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 3d8d517..35b92fc 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -691,6 +691,10 @@ def _fill_node_tensor_with_atom_type_property( offset: int, is_atom_node: torch.Tensor, ) -> torch.Tensor: + # We need a to mask property values + # node_tensor.shape : torch.Size([85, 203]) + # is_atom_node.shape : torch.Size([85]) + # property_values.shape : torch.Size([85, 1]) node_tensor[is_atom_node, offset : offset + property_values.shape[1]] = ( property_values[is_atom_node] ) @@ -715,8 +719,12 @@ def _fill_node_tensor_with_molecule_type_property( offset: int, is_graph_node: torch.Tensor, ) -> torch.Tensor: + # No mask for graph properties is required + # is_graph_node.shape : torch.Size([85]) + # node_tensor.shape : torch.Size([85, 203]) + # property_values.shape : torch.Size([1, 200]) node_tensor[is_graph_node, offset : offset + property_values.shape[1]] = ( - property_values[is_graph_node] + property_values ) return node_tensor From 0d76f65a176790e6f3bc1b35df832721f59fbc97 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 2 Aug 2026 15:24:26 +0200 Subject: [PATCH 33/34] configs for main augmentations --- chebai_graph/preprocessing/datasets/chebi.py | 53 ++++++++++--------- .../data/augmented/aug-ablation/FGN+E+WGN.yml | 12 +++++ configs/data/augmented/aug-ablation/FGN+E.yml | 12 +++++ .../data/augmented/aug-ablation/FGN+WGN.yml | 12 +++++ configs/data/augmented/aug-ablation/FGN.yml | 12 +++++ configs/data/augmented/aug-ablation/WGN.yml | 12 +++++ 6 files changed, 89 insertions(+), 24 deletions(-) create mode 100644 configs/data/augmented/aug-ablation/FGN+E+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/FGN+E.yml create mode 100644 configs/data/augmented/aug-ablation/FGN+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/FGN.yml create mode 100644 configs/data/augmented/aug-ablation/WGN.yml diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index e8585e8..1dbf7d7 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -78,12 +78,38 @@ class ChEBI50GraphPropertiesPartial(ChEBI50GraphProperties, ChEBIOverXPartial): pass +# ---- Augmentation: Variants with graph Node connected to FG nodes only ------------- class ChEBI50_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): """ChEBIOver50 with with FG nodes and FG edges and graph node.""" READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): + """ChEBIOver50 with FG nodes but without FG edges, with graph node.""" + + READER = AtomFGReader_NoFGEdges_WithGraphNode + + +class ChEBI50_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): + """ChEBIOver50 with FG nodes and FG edges, no graph node.""" + + READER = AtomFGReader_WithFGEdges_NoGraphNode + + +class ChEBI50_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): + """ChEBIOver50 with FG nodes but without FG edges or graph node.""" + + READER = AtomsFGReader_NoFGEdges_NoGraphNode + + +class ChEBI50_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): + """ChEBIOver50 with atom-level nodes and graph node only.""" + + READER = AtomReader_WithGraphNodeOnly + + +# ------- Augmentation: Variants with graph Node connected to all others nodes (FG and atoms) -------------- class ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE( AugGraphPropMixIn_WithGraphNode, ChEBIOver50 ): @@ -106,6 +132,7 @@ class ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE( READER = GN_WithAllNodes_FG_WithAtoms_NoFGE +# ------- Augmentation: Variants with graph node connected to atom nodes ONLY ----------- class ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE( AugGraphPropMixIn_WithGraphNode, ChEBIOver50 ): @@ -128,30 +155,7 @@ class ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE( READER = GN_WithAtoms_FG_WithAtoms_NoFGE -class ChEBI50_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes but without FG edges, with graph node.""" - - READER = AtomFGReader_NoFGEdges_WithGraphNode - - -class ChEBI50_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes and FG edges, no graph node.""" - - READER = AtomFGReader_WithFGEdges_NoGraphNode - - -class ChEBI50_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes but without FG edges or graph node.""" - - READER = AtomsFGReader_NoFGEdges_NoGraphNode - - -class ChEBI50_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): - """ChEBIOver50 with atom-level nodes and graph node only.""" - - READER = AtomReader_WithGraphNodeOnly - - +# ---------------------- Ablation: Properties ------------------------------ class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode @@ -186,6 +190,7 @@ class ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly( READER = AtomFGReader_WithFGEdges_WithGraphNode +# ---------- Final Augmentation: Different Thresholds ------------------------------ class ChEBI100_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver100): READER = AtomFGReader_WithFGEdges_WithGraphNode diff --git a/configs/data/augmented/aug-ablation/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/FGN+E+WGN.yml new file mode 100644 index 0000000..e8520f3 --- /dev/null +++ b/configs/data/augmented/aug-ablation/FGN+E+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/FGN+E.yml b/configs/data/augmented/aug-ablation/FGN+E.yml new file mode 100644 index 0000000..c50baac --- /dev/null +++ b/configs/data/augmented/aug-ablation/FGN+E.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/FGN+WGN.yml b/configs/data/augmented/aug-ablation/FGN+WGN.yml new file mode 100644 index 0000000..cae7d3a --- /dev/null +++ b/configs/data/augmented/aug-ablation/FGN+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/FGN.yml b/configs/data/augmented/aug-ablation/FGN.yml new file mode 100644 index 0000000..44b3705 --- /dev/null +++ b/configs/data/augmented/aug-ablation/FGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/WGN.yml b/configs/data/augmented/aug-ablation/WGN.yml new file mode 100644 index 0000000..7b5f253 --- /dev/null +++ b/configs/data/augmented/aug-ablation/WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_Atom_WGNOnly_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity From 14c7ee31f1726c915c5ed89cdbd0ee303fc1916f Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 2 Aug 2026 15:49:58 +0200 Subject: [PATCH 34/34] configs for additional augs --- .../augmented/aug-ablation/gn_wall_fgwa_nfge.yml | 12 ++++++++++++ .../augmented/aug-ablation/gn_wall_fgwa_wfge.yml | 12 ++++++++++++ .../data/augmented/aug-ablation/gnwa_fgwa_nfge.yml | 12 ++++++++++++ .../data/augmented/aug-ablation/gnwa_fgwa_wfge.yml | 12 ++++++++++++ 4 files changed, 48 insertions(+) create mode 100644 configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml create mode 100644 configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml new file mode 100644 index 0000000..c75cbbf --- /dev/null +++ b/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml new file mode 100644 index 0000000..4eda133 --- /dev/null +++ b/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml new file mode 100644 index 0000000..5d6ff23 --- /dev/null +++ b/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml new file mode 100644 index 0000000..648ecfc --- /dev/null +++ b/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity