From 15780755bd7a02c9aad9c136f58f2e03346e8ab8 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Fri, 28 Aug 2026 17:31:29 +0200 Subject: [PATCH 1/3] bump version --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d27f1a7..8c13c8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,12 +2,12 @@ name = "chebai-graph" description = "GNNs for ChEB-AI" readme = "README.md" -version = "1.1.0" +version = "1.2.0" authors = [ { name = "Martin Glauer", email = "martin.glauer@ovgu.de" } ] dependencies = [ - "chebai", + "chebai>=1.4.0", "chebi_utils>=0.4", "descriptastorus", # below packages need to manually installed as mentioned in readme From f83e84e37dab25abee351b66bfcdee4a702a37c1 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 1 Sep 2026 11:37:11 +0200 Subject: [PATCH 2/3] fix method name --- chebai_graph/preprocessing/datasets/chebi.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 80171a6..14489bc 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -214,7 +214,7 @@ def _after_setup(self, **kwargs) -> None: self._setup_properties() super()._after_setup(**kwargs) - def _preprocess_smiles_for_pred( + def _preprocess_molecule_for_pred( self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None ) -> Optional[dict]: """Preprocess prediction data.""" diff --git a/pyproject.toml b/pyproject.toml index 8c13c8d..b2295db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "chebai-graph" description = "GNNs for ChEB-AI" readme = "README.md" -version = "1.2.0" +version = "1.2.3" authors = [ { name = "Martin Glauer", email = "martin.glauer@ovgu.de" } ] From 91345a89d2b58ee1eaaa9215a003eec0b6e2de27 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Tue, 1 Sep 2026 11:41:04 +0200 Subject: [PATCH 3/3] Revert "Merge branch 'dev' into fix/method-names-this-time-for-real-for-real" This reverts commit 97ca5ca1517a8dcd6157dff5f66e4e7c980dff69, reversing changes made to f83e84e37dab25abee351b66bfcdee4a702a37c1. --- chebai_graph/models/__init__.py | 28 +- chebai_graph/models/architectures/__init__.py | 0 chebai_graph/models/architectures/base.py | 185 ----- chebai_graph/models/architectures/gine.py | 146 ---- chebai_graph/models/augmented.py | 37 +- chebai_graph/models/base.py | 707 ++++++++++++++++ chebai_graph/models/dynamic_gni.py | 4 +- .../models/{architectures => }/gat.py | 0 chebai_graph/models/gin_net.py | 96 +++ chebai_graph/models/graph.py | 13 +- chebai_graph/models/pooling.py | 151 ---- .../models/{architectures => }/resgated.py | 0 .../AtomCharge/indices_one_hot.txt | 0 .../AtomFunctionalGroup/indices_one_hot.txt | 0 .../AtomHybridization/indices_one_hot.txt | 0 .../AtomNodeLevel/indices_one_hot.txt | 0 .../{chebi => }/AtomNumHs/indices_one_hot.txt | 0 .../{chebi => }/AtomType/indices_one_hot.txt | 0 .../{chebi => }/BondLevel/indices_one_hot.txt | 0 .../{chebi => }/BondType/indices_one_hot.txt | 0 .../NumAtomBonds/indices_one_hot.txt | 0 .../preprocessing/datasets/__init__.py | 34 - .../datasets/augmentation_base.py | 179 ---- chebai_graph/preprocessing/datasets/base.py | 781 ------------------ chebai_graph/preprocessing/datasets/chebi.py | 130 +-- .../datasets/molecule_net_classification.py | 60 -- chebai_graph/preprocessing/datasets/utils.py | 10 +- .../properties/augmented_properties.py | 40 +- chebai_graph/preprocessing/properties/base.py | 16 +- .../preprocessing/properties/properties.py | 72 +- .../preprocessing/property_encoder.py | 16 +- .../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 - .../aug-ablation/gn_wall_fgwa_nfge.yml | 12 - .../aug-ablation/gn_wall_fgwa_wfge.yml | 12 - .../augmented/aug-ablation/gnwa_fgwa_nfge.yml | 12 - .../augmented/aug-ablation/gnwa_fgwa_wfge.yml | 12 - .../augmented/final/BACE_final_augmented.yml | 24 - .../augmented/final/BBBP_final_augmented.yml | 24 - .../final/ClinTox_final_augmented.yml | 24 - .../augmented/final/HIV_final_augmented.yml | 24 - .../augmented/final/MUV_final_augmented.yml | 24 - .../augmented/final/PCBA_final_augmented.yml | 24 - .../augmented/final/SIDER_final_augmented.yml | 24 - .../augmented/final/Tox21_final_augmented.yml | 24 - .../final/ToxCast_final_augmented.yml | 24 - .../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 - .../chebi50_node_level_prop_only.yml | 24 - ...d.yml => chebi50_aug_prop_as_per_node.yml} | 0 configs/data/chebi50_baseline.yml | 12 - .../model/augmented/pooling/aa_pool/gat.yml | 14 - .../model/augmented/pooling/aa_pool/gine.yml | 13 - .../model/augmented/pooling/aa_pool/rggcn.yml | 12 - .../model/augmented/pooling/amg_pool/gine.yml | 13 - .../augmented/pooling/amg_pool/rggcn.yml | 12 - .../augmented/pooling/no_pooling/gat.yml | 14 - configs/model/baselines/gine.yml | 13 - configs/model/{baselines => }/gat.yml | 1 + configs/model/gat_aug_aapool.yml | 2 +- .../amg_pool/gat.yml => gat_aug_amgpool.yml} | 3 +- 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 + .../{baselines/rggcn.yml => resgated.yml} | 1 + configs/model/resgated_dynamic_gni.yml | 1 + pyproject.toml | 9 +- requirements.txt | 120 --- 76 files changed, 971 insertions(+), 2470 deletions(-) delete mode 100644 chebai_graph/models/architectures/__init__.py delete mode 100644 chebai_graph/models/architectures/base.py delete mode 100644 chebai_graph/models/architectures/gine.py create mode 100644 chebai_graph/models/base.py rename chebai_graph/models/{architectures => }/gat.py (100%) create mode 100644 chebai_graph/models/gin_net.py delete mode 100644 chebai_graph/models/pooling.py rename chebai_graph/models/{architectures => }/resgated.py (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomCharge/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomFunctionalGroup/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomHybridization/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomNodeLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomNumHs/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/BondLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/BondType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/NumAtomBonds/indices_one_hot.txt (100%) delete mode 100644 chebai_graph/preprocessing/datasets/augmentation_base.py delete mode 100644 chebai_graph/preprocessing/datasets/base.py delete mode 100644 chebai_graph/preprocessing/datasets/molecule_net_classification.py delete mode 100644 configs/data/augmented/aug-ablation/FGN+E+WGN.yml delete mode 100644 configs/data/augmented/aug-ablation/FGN+E.yml delete mode 100644 configs/data/augmented/aug-ablation/FGN+WGN.yml delete mode 100644 configs/data/augmented/aug-ablation/FGN.yml delete mode 100644 configs/data/augmented/aug-ablation/WGN.yml delete mode 100644 configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml delete mode 100644 configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml delete mode 100644 configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml delete mode 100644 configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml delete mode 100644 configs/data/augmented/final/BACE_final_augmented.yml delete mode 100644 configs/data/augmented/final/BBBP_final_augmented.yml delete mode 100644 configs/data/augmented/final/ClinTox_final_augmented.yml delete mode 100644 configs/data/augmented/final/HIV_final_augmented.yml delete mode 100644 configs/data/augmented/final/MUV_final_augmented.yml delete mode 100644 configs/data/augmented/final/PCBA_final_augmented.yml delete mode 100644 configs/data/augmented/final/SIDER_final_augmented.yml delete mode 100644 configs/data/augmented/final/Tox21_final_augmented.yml delete mode 100644 configs/data/augmented/final/ToxCast_final_augmented.yml delete mode 100644 configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_atom_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_fg_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_graph_node_prop_only.yml delete mode 100644 configs/data/augmented/properties/chebi50_node_level_prop_only.yml rename configs/data/{augmented/final/chebi50_final_augmented.yml => chebi50_aug_prop_as_per_node.yml} (100%) delete mode 100644 configs/data/chebi50_baseline.yml delete mode 100644 configs/model/augmented/pooling/aa_pool/gat.yml delete mode 100644 configs/model/augmented/pooling/aa_pool/gine.yml delete mode 100644 configs/model/augmented/pooling/aa_pool/rggcn.yml delete mode 100644 configs/model/augmented/pooling/amg_pool/gine.yml delete mode 100644 configs/model/augmented/pooling/amg_pool/rggcn.yml delete mode 100644 configs/model/augmented/pooling/no_pooling/gat.yml delete mode 100644 configs/model/baselines/gine.yml rename configs/model/{baselines => }/gat.yml (95%) rename configs/model/{augmented/pooling/amg_pool/gat.yml => gat_aug_amgpool.yml} (84%) rename configs/model/{baselines/rggcn.yml => resgated.yml} (91%) delete mode 100644 requirements.txt diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 515a513..9e20b2d 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -1,25 +1,19 @@ -from .architectures.gat import GATGraphPred -from .architectures.gine import GINEGraphPred -from .architectures.resgated import ResGatedGraphPred from .augmented import ( - GATAAPoolGraphPred, - GATAMGPoolGraphPred, - GINEAAPoolGraphPred, - GINEAMGPoolGraphPred, - ResGatedAAPoolGraphPred, - ResGatedAMGPoolGraphPred, + GATAugNodePoolGraphPred, + GATGraphNodeFGNodePoolGraphPred, + ResGatedAugNodePoolGraphPred, + ResGatedGraphNodeFGNodePoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred +from .gat import GATGraphPred +from .resgated import ResGatedGraphPred __all__ = [ "ResGatedGraphPred", - "ResGatedAAPoolGraphPred", - "ResGatedAMGPoolGraphPred", - "ResGatedDynamicGNIGraphPred", + "ResGatedAugNodePoolGraphPred", + "ResGatedGraphNodeFGNodePoolGraphPred", "GATGraphPred", - "GATAAPoolGraphPred", - "GATAMGPoolGraphPred", - "GINEGraphPred", - "GINEAAPoolGraphPred", - "GINEAMGPoolGraphPred", + "GATAugNodePoolGraphPred", + "GATGraphNodeFGNodePoolGraphPred", + "ResGatedDynamicGNIGraphPred", ] diff --git a/chebai_graph/models/architectures/__init__.py b/chebai_graph/models/architectures/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py deleted file mode 100644 index 514e287..0000000 --- a/chebai_graph/models/architectures/base.py +++ /dev/null @@ -1,185 +0,0 @@ -from abc import ABC, abstractmethod - -import torch -from chebai.models.base import ChebaiBaseNet -from chebai.preprocessing.structures import XYData -from torch_geometric.data import Data as GraphData -from torch_scatter import scatter_add - - -class GraphBaseNet(ChebaiBaseNet, ABC): - """ - Base class for graph-based prediction networks. - """ - - def _get_prediction_and_labels( - self, data: XYData, labels: torch.Tensor, output: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Apply sigmoid activation to outputs and return processed labels. - - Args: - data (XYData): Input batch data. - labels (torch.Tensor): Ground-truth labels. - output (torch.Tensor): Raw model output. - - Returns: - tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels). - """ - return torch.sigmoid(output), labels.int() - - def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None: - """ - Process labels from XYData batch. - - Returns: - torch.Tensor | None: Processed labels if present, else None. - """ - return batch.y.float() if batch.y is not None else None - - -class GraphModelBase(torch.nn.Module, ABC): - """ - Abstract base class for graph models with configurable architecture. - """ - - def __init__(self, config: dict, **kwargs) -> None: - """ - Initialize model hyperparameters from configuration. - - Args: - config (dict): Configuration dictionary with keys: - - 'num_layers' - - 'in_channels' - - 'hidden_channels' - - 'out_channels' - - 'edge_dim' - - 'dropout' - **kwargs: Additional keyword arguments for torch.nn.Module. - """ - super().__init__(**kwargs) - self.num_layers = int(config["num_layers"]) - assert self.num_layers > 1, "Need atleast two convolution layers" - self.in_channels = int(config["in_channels"]) # number of node/atom properties - self.hidden_channels = int(config["hidden_channels"]) - self.out_channels = int(config["out_channels"]) - self.edge_dim = int(config["edge_dim"]) # number of bond properties - self.dropout = float(config["dropout"]) - - -class GraphNetWrapper(GraphBaseNet, ABC): - """ - Base wrapper class for GNNs with linear layers for graph classification - with standard pooling . - """ - - def __init__( - self, - config: dict, - n_linear_layers: int, - use_batch_norm: bool = False, - **kwargs, - ): - """ - Args: - config (dict): Model configuration. - n_linear_layers (int): Number of linear layers. - **kwargs: Additional arguments. - """ - super().__init__(**kwargs) - self.gnn = self._get_gnn(config) - gnn_out_dim = int(config["out_channels"]) - self.activation = torch.nn.ELU - self.lin_input_dim = self._get_lin_seq_input_dim( - gnn_out_dim=gnn_out_dim, - ) - self.use_batch_norm = use_batch_norm - if self.use_batch_norm: - self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim) - - lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim) - self.lin_sequential: torch.nn.Sequential = self._get_linear_module_list( - n_linear_layers=n_linear_layers, - in_dim=self.lin_input_dim, - hidden_dim=lin_hidden_dim, - out_dim=self.out_dim, - ) - - @abstractmethod - def _get_gnn(self, config: dict) -> torch.nn.Module: - """ - Create the graph neural network. - - Args: - config (dict): Configuration dictionary. - - Returns: - torch.nn.Module: Instantiated GNN module. - """ - pass - - 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. - - Returns: - int: Total input dimension. - """ - return gnn_out_dim - - def _get_linear_module_list( - self, - n_linear_layers: int, - in_dim: int, - hidden_dim: int, - out_dim: int, - ) -> torch.nn.Sequential: - """ - Construct a sequential module of linear layers. - - Args: - n_linear_layers (int): Number of linear layers. - in_dim (int): Input dimension. - hidden_dim (int): Hidden dimension. - out_dim (int): Output dimension. - - Returns: - torch.nn.Sequential: Linear layers with activations. - """ - if n_linear_layers < 1: - raise ValueError("n_linear_layers must be at least 1") - - layers = [] - if n_linear_layers == 1: - layers.append(torch.nn.Linear(in_dim, out_dim)) - else: - layers.append(torch.nn.Linear(in_dim, hidden_dim)) - layers.append(self.activation()) - for _ in range(n_linear_layers - 2): - layers.append(torch.nn.Linear(hidden_dim, hidden_dim)) - layers.append(self.activation()) - layers.append(torch.nn.Linear(hidden_dim, out_dim)) - - return torch.nn.Sequential(*layers) - - def forward(self, batch: dict) -> torch.Tensor: - """ - Forward pass through GNN, pooling and linear layers. - - Args: - batch (dict): Input batch with graph features. - - Returns: - torch.Tensor: Predicted output. - """ - graph_data = batch["features"][0] - graph_data.to(self.device) - assert isinstance(graph_data, GraphData) - a = self.gnn(batch) - a = scatter_add(a, graph_data.batch, dim=0) - if self.use_batch_norm: - a = self.batch_norm(a) - return self.lin_sequential(a) diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py deleted file mode 100644 index f9e89a3..0000000 --- a/chebai_graph/models/architectures/gine.py +++ /dev/null @@ -1,146 +0,0 @@ -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 - - https://github.com/pyg-team/pytorch_geometric/issues/1311 - - 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/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index 87f1c30..fdb5388 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,10 +1,9 @@ -from .architectures.gat import GATGraphPred -from .architectures.gine import GINEGraphPred -from .architectures.resgated import ResGatedGraphPred -from .pooling import AAPool, AMGPool +from .base import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet +from .gat import GATGraphPred +from .resgated import ResGatedGraphPred -class ResGatedAAPoolGraphPred(AAPool, ResGatedGraphPred): +class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -14,7 +13,7 @@ class ResGatedAAPoolGraphPred(AAPool, ResGatedGraphPred): ... -class GATAAPoolGraphPred(AAPool, GATGraphPred): +class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -24,17 +23,9 @@ 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): +class ResGatedGraphNodeFGNodePoolGraphPred( + GraphNodeFGNodePoolingNet, ResGatedGraphPred +): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -44,7 +35,7 @@ class ResGatedAMGPoolGraphPred(AMGPool, ResGatedGraphPred): ... -class GATAMGPoolGraphPred(AMGPool, GATGraphPred): +class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, GATGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -52,13 +43,3 @@ class GATAMGPoolGraphPred(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/chebai_graph/models/base.py b/chebai_graph/models/base.py new file mode 100644 index 0000000..3226d21 --- /dev/null +++ b/chebai_graph/models/base.py @@ -0,0 +1,707 @@ +from abc import ABC, abstractmethod +from typing import Optional + +import torch +from chebai.models.base import ChebaiBaseNet +from chebai.preprocessing.structures import XYData +from torch_geometric.data import Data as GraphData +from torch_scatter import scatter_add + + +class GraphBaseNet(ChebaiBaseNet, ABC): + """ + Base class for graph-based prediction networks. + """ + + def _get_prediction_and_labels( + self, data: XYData, labels: torch.Tensor, output: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Apply sigmoid activation to outputs and return processed labels. + + Args: + data (XYData): Input batch data. + labels (torch.Tensor): Ground-truth labels. + output (torch.Tensor): Raw model output. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels). + """ + return torch.sigmoid(output), labels.int() + + def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None: + """ + Process labels from XYData batch. + + Returns: + torch.Tensor | None: Processed labels if present, else None. + """ + return batch.y.float() if batch.y is not None else None + + +class GraphModelBase(torch.nn.Module, ABC): + """ + Abstract base class for graph models with configurable architecture. + """ + + def __init__(self, config: dict, **kwargs) -> None: + """ + Initialize model hyperparameters from configuration. + + Args: + config (dict): Configuration dictionary with keys: + - 'num_layers' + - 'in_channels' + - 'hidden_channels' + - 'out_channels' + - 'edge_dim' + - 'dropout' + **kwargs: Additional keyword arguments for torch.nn.Module. + """ + super().__init__(**kwargs) + self.num_layers = int(config["num_layers"]) + assert self.num_layers > 1, "Need atleast two convolution layers" + self.in_channels = int(config["in_channels"]) # number of node/atom properties + self.hidden_channels = int(config["hidden_channels"]) + self.out_channels = int(config["out_channels"]) + self.edge_dim = int(config["edge_dim"]) # number of bond properties + self.dropout = float(config["dropout"]) + + +class GraphNetWrapper(GraphBaseNet, ABC): + """ + Base wrapper class for GNNs with linear layers for property prediction. + """ + + def __init__( + self, + config: dict, + n_linear_layers: int, + n_molecule_properties: Optional[int] = 0, + use_batch_norm: bool = False, + **kwargs, + ): + """ + 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) + self.gnn = self._get_gnn(config) + gnn_out_dim = int(config["out_channels"]) + 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: + self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim) + + lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim) + self.lin_sequential: torch.nn.Sequential = self._get_linear_module_list( + n_linear_layers=n_linear_layers, + in_dim=self.lin_input_dim, + hidden_dim=lin_hidden_dim, + out_dim=self.out_dim, + ) + + @abstractmethod + def _get_gnn(self, config: dict) -> torch.nn.Module: + """ + Create the graph neural network. + + Args: + config (dict): Configuration dictionary. + + Returns: + torch.nn.Module: Instantiated GNN module. + """ + pass + + def _get_lin_seq_input_dim( + self, gnn_out_dim: int, n_molecule_properties: 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 + + def _get_linear_module_list( + self, + n_linear_layers: int, + in_dim: int, + hidden_dim: int, + out_dim: int, + ) -> torch.nn.Sequential: + """ + Construct a sequential module of linear layers. + + Args: + n_linear_layers (int): Number of linear layers. + in_dim (int): Input dimension. + hidden_dim (int): Hidden dimension. + out_dim (int): Output dimension. + + Returns: + torch.nn.Sequential: Linear layers with activations. + """ + if n_linear_layers < 1: + raise ValueError("n_linear_layers must be at least 1") + + layers = [] + if n_linear_layers == 1: + layers.append(torch.nn.Linear(in_dim, out_dim)) + else: + layers.append(torch.nn.Linear(in_dim, hidden_dim)) + layers.append(self.activation()) + for _ in range(n_linear_layers - 2): + layers.append(torch.nn.Linear(hidden_dim, hidden_dim)) + layers.append(self.activation()) + layers.append(torch.nn.Linear(hidden_dim, out_dim)) + + return torch.nn.Sequential(*layers) + + def forward(self, batch: dict) -> torch.Tensor: + """ + Forward pass through GNN, pooling and linear layers. + + Args: + batch (dict): Input batch with graph features. + + Returns: + torch.Tensor: Predicted output. + """ + graph_data = batch["features"][0] + graph_data.to(self.device) + 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) + + +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. + """ + + def _get_lin_seq_input_dim( + self, gnn_out_dim: int, n_molecule_properties: 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 + + 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 + - Molecular attributes + - 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, graph_data.molecule_attr, aug_nodes_vec], dim=1 + ) + + 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, n_molecule_properties: 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 + + 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, graph_data.molecule_attr, 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 + + 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: + """ + 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 + + 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 + - Molecular attributes + - 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, graph_data.molecule_attr, 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, n_molecule_properties: 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 + n_molecule_properties + 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_data.molecule_attr, 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, n_molecule_properties: 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 + + 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, graph_data.molecule_attr, 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, n_molecule_properties: 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 + + 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_data.molecule_attr, graph_node_vec], dim=1 + ) + + return self.lin_sequential(graph_vector) + + +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. + + 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, graph_data.molecule_attr], dim=1) + + return self.lin_sequential(graph_vector) + + +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. + + 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, graph_data.molecule_attr], dim=1) + + return self.lin_sequential(graph_vector) + + +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. + + 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, graph_data.molecule_attr], dim=1) + + return self.lin_sequential(graph_vector) diff --git a/chebai_graph/models/dynamic_gni.py b/chebai_graph/models/dynamic_gni.py index 4fa1b1f..8cb6c7b 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 .architectures.base import GraphModelBase, GraphNetWrapper -from .architectures.resgated import ResGatedModel +from .base import GraphModelBase, GraphNetWrapper +from .resgated import ResGatedModel class ResGatedDynamicGNI(GraphModelBase): diff --git a/chebai_graph/models/architectures/gat.py b/chebai_graph/models/gat.py similarity index 100% rename from chebai_graph/models/architectures/gat.py rename to chebai_graph/models/gat.py diff --git a/chebai_graph/models/gin_net.py b/chebai_graph/models/gin_net.py new file mode 100644 index 0000000..6fed4c6 --- /dev/null +++ b/chebai_graph/models/gin_net.py @@ -0,0 +1,96 @@ +import typing + +import torch +import torch.nn.functional as F +import torch_geometric +from torch_scatter import scatter_add + +from chebai_graph.models.graph import GraphBaseNet + + +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(GraphBaseNet): + """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.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.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 + self.convs.append( + torch_geometric.nn.GINEConv( + AggregateMLP(in_length, out_length, self.hidden_size), + edge_dim=self.n_bond_properties, + ) + ) + # 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) + 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(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) + + 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 diff --git a/chebai_graph/models/graph.py b/chebai_graph/models/graph.py index 69d27c0..3be0fdd 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 .architectures.base import GraphBaseNet +from .base import GraphBaseNet logging.getLogger("pysmiles").setLevel(logging.CRITICAL) @@ -88,6 +88,11 @@ 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) @@ -153,7 +158,7 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length, + self.gnn.hidden_length + (i == 0) * self.gnn.n_molecule_properties, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) @@ -191,7 +196,9 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length + (i == 0) * self.gnn.hidden_length, + self.gnn.hidden_length + + (i == 0) * self.gnn.n_molecule_properties + + (i == 0) * self.gnn.hidden_length, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py deleted file mode 100644 index a1ee442..0000000 --- a/chebai_graph/models/pooling.py +++ /dev/null @@ -1,151 +0,0 @@ -from abc import ABC - -import torch -from torch_geometric.data import Data as GraphData -from torch_scatter import scatter_add - -from .architectures.base import GraphNetWrapper - - -class AAPool(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 AMGPool(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/architectures/resgated.py b/chebai_graph/models/resgated.py similarity index 100% rename from chebai_graph/models/architectures/resgated.py rename to chebai_graph/models/resgated.py diff --git a/chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 9ee9856..950a738 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -9,30 +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_ForNodeLevelPropOnlyForAllNodes, ChEBI50_WFGE_WGN_GraphProp, ChEBI50GraphData, ChEBI25GraphProperties, 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__ = [ @@ -53,20 +35,4 @@ "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", - "ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes", - "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/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py deleted file mode 100644 index 9bf6cdc..0000000 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ /dev/null @@ -1,179 +0,0 @@ -from abc import ABC - -import pandas as pd -import torch -from torch_geometric.data.data import Data as GeomData - -from .base import GraphPropAsPerNodeType, 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) - if isinstance(row["features"], tuple): - geom_data, _ = row["features"] - else: - 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. - """ - if isinstance(row["features"], tuple): - geom_data, _ = row["features"] - else: - 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 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: - 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: - 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: - 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: - return node_tensor - - -class GraphPropForGraphNodeOnly(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: - 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: - 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: - return node_tensor - - -class GraphPropForAtomLevelAndGraphNodeOnly(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: - return node_tensor - - -class GraphPropForFGLevelAndGraphNodeOnly(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: - return node_tensor - - -class GraphPropNodeLevelPropOnlyForAllNodes(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: - 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: - 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: - return node_tensor diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py deleted file mode 100644 index d01b53b..0000000 --- a/chebai_graph/preprocessing/datasets/base.py +++ /dev/null @@ -1,781 +0,0 @@ -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, self.data_type) 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: property.encoder.compress( - 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_molecule_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, - ) - - 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) - base_df["ident"] = base_df["ident"].astype(str) - - for property in self.properties: - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - 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 - ) - property_df["ident"] = property_df["ident"].astype(str) - 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" - ) - - 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) - base_df["ident"] = base_df["ident"].astype(str) - 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" - ) - - 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 - ) - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - 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["ident"] = property_df["ident"].astype(str) - 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 - 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, ( - "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()) - ) - - 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, - 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"] - - 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, ( - 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, - is_atom_node=is_atom_node, - is_fg_node=is_fg_node, - is_graph_node=is_graph_node, - ) - - def _build_node_property_tensor( - self, - property: MolecularProperty, - 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 = 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 = 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 = 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 = 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 - - return { - "node_tensor": node_tensor, - "atom_offset": atom_offset, - "fg_offset": fg_offset, - "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: - # 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] - ) - 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: - # 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 - ) - return node_tensor - - def _build_edge_property_tensor( - self, - edge_attr_tensor: torch.Tensor, - property_values: torch.Tensor, - ) -> torch.Tensor: - # 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( - 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 ac58a84..14489bc 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -1,4 +1,13 @@ +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 ( ChEBIOver25, ChEBIOver50, @@ -7,7 +16,18 @@ ChEBIOverXPartial, ) 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, @@ -18,22 +38,12 @@ GN_WithAllNodes_FG_WithAtoms_NoFGE, GN_WithAtoms_FG_WithAtoms_FGE, GN_WithAtoms_FG_WithAtoms_NoFGE, + GraphPropertyReader, GraphReader, RandomFeatureInitializationReader, ) -from .augmentation_base import ( - AugGraphPropMixIn_NoGraphNode, - AugGraphPropMixIn_WithGraphNode, - GraphPropForAtomAndFGLevelOnly, - GraphPropForAtomLevelAndGraphNodeOnly, - GraphPropForAtomLevelOnly, - GraphPropForFGLevelAndGraphNodeOnly, - GraphPropForFGLevelOnly, - GraphPropForGraphNodeOnly, - GraphPropNodeLevelPropOnlyForAllNodes, -) -from .base import DataPropertiesSetter, GraphPropAsPerNodeType, GraphPropertiesMixIn +from chebai_graph.preprocessing.datasets.utils import resolve_property class ChEBI50GraphData(ChEBIOver50): @@ -700,6 +710,7 @@ 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") @@ -728,38 +739,56 @@ 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.""" +class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): + """Mixin for augmented graph data without additional graph nodes.""" - READER = AtomFGReader_WithFGEdges_WithGraphNode + 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) -class ChEBI50_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes but without FG edges, with graph node.""" + 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 - READER = AtomFGReader_NoFGEdges_WithGraphNode +class AugGraphPropMixIn_WithGraphNode(AugGraphPropMixIn_NoGraphNode, ABC): + """Mixin for augmented graph data with graph-level nodes.""" -class ChEBI50_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes and FG edges, no graph node.""" + READER = None - READER = AtomFGReader_WithFGEdges_NoGraphNode + 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. -class ChEBI50_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): - """ChEBIOver50 with FG nodes but without FG edges or graph node.""" + Args: + data: A GeomData object with features. + row: A dictionary containing 'features' and other metadata. - READER = AtomsFGReader_NoFGEdges_NoGraphNode + 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_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): - """ChEBIOver50 with atom-level nodes and graph node only.""" +class ChEBI50_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): + """ChEBIOver50 with with FG nodes and FG edges and graph node.""" - READER = AtomReader_WithGraphNodeOnly + READER = AtomFGReader_WithFGEdges_WithGraphNode -# ------- Augmentation: Variants with graph Node connected to all others nodes (FG and atoms) -------------- class ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE( AugGraphPropMixIn_WithGraphNode, ChEBIOver50 ): @@ -782,7 +811,6 @@ 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 ): @@ -805,48 +833,34 @@ class ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE( READER = GN_WithAtoms_FG_WithAtoms_NoFGE -# ---------------------- Ablation: Properties ------------------------------ -class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): - READER = AtomFGReader_WithFGEdges_WithGraphNode - - -class ChEBI50_WFGE_WGN_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): - 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_WGN_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): - READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ChEBIOver50): + """ChEBIOver50 with FG nodes and FG edges, no graph node.""" -class ChEBI50_WFGE_WGN_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): - READER = AtomFGReader_WithFGEdges_WithGraphNode + READER = AtomFGReader_WithFGEdges_NoGraphNode -class ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly( - GraphPropForAtomAndFGLevelOnly, ChEBIOver50 -): - READER = AtomFGReader_WithFGEdges_WithGraphNode +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_WFGE_WGN_ForAtomLevelAndGraphNodeOnly( - GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 -): - READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): + """ChEBIOver50 with atom-level nodes and graph node only.""" -class ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly( - GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 -): - READER = AtomFGReader_WithFGEdges_WithGraphNode + READER = AtomReader_WithGraphNodeOnly -class ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes( - GraphPropNodeLevelPropOnlyForAllNodes, ChEBIOver50 -): +class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode -# ---------- Final Augmentation: Different Thresholds ------------------------------ class ChEBI100_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver100): READER = AtomFGReader_WithFGEdges_WithGraphNode diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py deleted file mode 100644 index 9afaf71..0000000 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ /dev/null @@ -1,60 +0,0 @@ -from chebai.preprocessing.datasets.molecule_net_classification import ( - BACE, - BBBP, - HIV, - MUV, - PCBA, - SIDER, - ClinTox, - 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/chebai_graph/preprocessing/datasets/utils.py b/chebai_graph/preprocessing/datasets/utils.py index 79f2fb3..3ec6515 100644 --- a/chebai_graph/preprocessing/datasets/utils.py +++ b/chebai_graph/preprocessing/datasets/utils.py @@ -4,9 +4,7 @@ from chebai_graph.preprocessing.properties import MolecularProperty -def resolve_property( - property: str | MolecularProperty, data_type: str -) -> MolecularProperty: +def resolve_property(property: str | MolecularProperty) -> MolecularProperty: """ Resolves a molecular property specification (either as a class instance or class path string) into a MolecularProperty instance. @@ -21,8 +19,6 @@ def resolve_property( property (str | MolecularProperty): The property to resolve. Can be a class instance, a fully qualified class name (e.g. "module.ClassName"), or a class name assumed to be in `chebai_graph.preprocessing.properties`. - data_type (str): The data type associated with the property. This used to determine or set - tokens file path for the property if applicable. Returns: MolecularProperty: An instance of the resolved MolecularProperty. @@ -41,7 +37,7 @@ def resolve_property( module_name = property[:last_dot] class_name = property[last_dot + 1 :] module = importlib.import_module(module_name) - return getattr(module, class_name)(data_type=data_type) + return getattr(module, class_name)() except ValueError: # if only a class name is given, assume the module is chebai_graph.processing.properties - return getattr(graph_properties, property)(data_type=data_type) + return getattr(graph_properties, property)() diff --git a/chebai_graph/preprocessing/properties/augmented_properties.py b/chebai_graph/preprocessing/properties/augmented_properties.py index 6f112ea..f5f7b1d 100644 --- a/chebai_graph/preprocessing/properties/augmented_properties.py +++ b/chebai_graph/preprocessing/properties/augmented_properties.py @@ -23,17 +23,14 @@ class AtomNodeLevel(AllNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomNodeLevel with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -49,17 +46,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomFunctionalGroup(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomFunctionalGroup with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -75,17 +69,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomRingSize(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomRingSize with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -122,14 +113,14 @@ def _check_modify_atom_prop_value( class IsHydrogenBondDonorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize IsHydrogenBondDonorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L26-L31 self._hydrogen_bond_donor: set[str] = { @@ -155,14 +146,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsHydrogenBondAcceptorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize IsHydrogenBondAcceptorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L33-L39 self._hydrogen_bond_acceptor: set[str] = { @@ -189,13 +180,13 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsFGAlkyl(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Args: encoder (PropertyEncoder | None): Optional encoder to use for this property. Defaults to BoolEncoder if not provided. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -332,15 +323,12 @@ class AugAtomAromaticity(AugNodeValueDefaulter, pr.AtomAromaticity): class BondLevel(AugmentedBondProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Args: encoder (PropertyEncoder | None): Optional encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond | dict) -> str: """ diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index f148df8..ea2725a 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -15,7 +15,7 @@ 3, 7, ), "This code requires Python 3.7 or higher." -# Order preservation is necessary to to create `prop_list` in Augmented properties +# Order preservation is necessary to to create `prop_list`in Augmented properties class MolecularProperty(ABC): @@ -31,16 +31,10 @@ class MolecularProperty(ABC): Defaults to IndexEncoder if not provided. """ - def __init__( - self, - data_type: str, - encoder: PropertyEncoder | None = None, - ) -> None: - assert data_type is not None, "data_type must be provided for MolecularProperty" + def __init__(self, encoder: PropertyEncoder | None = None) -> None: if encoder is None: - encoder = IndexEncoder(self, data_type=data_type) + encoder = IndexEncoder(self) self.encoder: PropertyEncoder = encoder - self._data_type = data_type @property def name(self) -> str: @@ -180,8 +174,8 @@ class AugAtomType(FrozenPropertyAlias, AtomType): ... ValueError: If new tokens are added to the frozen encoder during processing. """ - def __init__(self, encoder: PropertyEncoder, **kwargs) -> None: - super().__init__(encoder=encoder, **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder) # Lock the encoder's cache to prevent adding new tokens if hasattr(self.encoder, "cache") and isinstance(self.encoder.cache, dict): self.encoder.cache = MappingProxyType(self.encoder.cache) diff --git a/chebai_graph/preprocessing/properties/properties.py b/chebai_graph/preprocessing/properties/properties.py index 87112da..2154f9c 100644 --- a/chebai_graph/preprocessing/properties/properties.py +++ b/chebai_graph/preprocessing/properties/properties.py @@ -19,11 +19,8 @@ class AtomType(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -45,11 +42,8 @@ class NumAtomBonds(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -71,11 +65,8 @@ class AtomCharge(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -97,11 +88,8 @@ class AtomChirality(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.ChiralType: """ @@ -123,11 +111,8 @@ class AtomHybridization(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.HybridizationType: """ @@ -149,11 +134,8 @@ class AtomNumHs(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -175,8 +157,8 @@ class AtomAromaticity(AtomProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> bool: """ @@ -198,8 +180,8 @@ class BondAromaticity(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -221,11 +203,8 @@ class BondType(BondProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> Chem.rdchem.BondType: """ @@ -247,8 +226,8 @@ class BondInRing(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -270,11 +249,8 @@ class MoleculeNumRings(MoleculeProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_property_value(self, mol: Chem.rdchem.Mol) -> list[int]: """ @@ -296,8 +272,8 @@ class RDKit2DNormalized(MoleculeProperty): Uses an identity encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or AsIsEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or AsIsEncoder(self)) self.generator_normalized = rdNormalizedDescriptors.RDKit2DNormalized() # Create a dummy molecule (e.g., methane) to extract the length of descriptor vector dummy_mol = Chem.MolFromSmiles("C") diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 5daeb4c..38cb279 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -99,15 +99,12 @@ class IndexEncoder(PropertyEncoder): **kwargs: Additional keyword arguments. """ - def __init__( - self, property, data_type: str, indices_dir: str | None = None, **kwargs - ) -> None: + def __init__(self, property, indices_dir: str | None = None, **kwargs) -> None: super().__init__(property, **kwargs) if indices_dir is None: indices_dir = os.path.dirname(inspect.getfile(self.__class__)) self.dirname = indices_dir # load already existing cache - self._data_type = data_type with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { token.strip(): idx for idx, token in enumerate(pk) @@ -129,15 +126,12 @@ def index_path(self) -> str: Returns: Path to index file. """ - assert self._data_type is not None, "data_type must be set for IndexEncoder" index_path = os.path.join( - self.dirname, - "bin", - self._data_type, - self.property.name, - f"indices_{self.name}.txt", + self.dirname, "bin", self.property.name, f"indices_{self.name}.txt" + ) + os.makedirs( + os.path.join(self.dirname, "bin", self.property.name), exist_ok=True ) - os.makedirs(os.path.dirname(index_path), exist_ok=True) if not os.path.exists(index_path): with open(index_path, "x"): pass diff --git a/configs/data/augmented/aug-ablation/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/FGN+E+WGN.yml deleted file mode 100644 index e8520f3..0000000 --- a/configs/data/augmented/aug-ablation/FGN+E+WGN.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index c50baac..0000000 --- a/configs/data/augmented/aug-ablation/FGN+E.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index cae7d3a..0000000 --- a/configs/data/augmented/aug-ablation/FGN+WGN.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 44b3705..0000000 --- a/configs/data/augmented/aug-ablation/FGN.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 7b5f253..0000000 --- a/configs/data/augmented/aug-ablation/WGN.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml deleted file mode 100644 index c75cbbf..0000000 --- a/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 4eda133..0000000 --- a/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 5d6ff23..0000000 --- a/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index 648ecfc..0000000 --- a/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 diff --git a/configs/data/augmented/final/BACE_final_augmented.yml b/configs/data/augmented/final/BACE_final_augmented.yml deleted file mode 100644 index 7ea67f0..0000000 --- a/configs/data/augmented/final/BACE_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/BBBP_final_augmented.yml b/configs/data/augmented/final/BBBP_final_augmented.yml deleted file mode 100644 index 9a7ee90..0000000 --- a/configs/data/augmented/final/BBBP_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/ClinTox_final_augmented.yml b/configs/data/augmented/final/ClinTox_final_augmented.yml deleted file mode 100644 index a4023b3..0000000 --- a/configs/data/augmented/final/ClinTox_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/HIV_final_augmented.yml b/configs/data/augmented/final/HIV_final_augmented.yml deleted file mode 100644 index 596fc19..0000000 --- a/configs/data/augmented/final/HIV_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/MUV_final_augmented.yml b/configs/data/augmented/final/MUV_final_augmented.yml deleted file mode 100644 index 56a3e1e..0000000 --- a/configs/data/augmented/final/MUV_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -class_path: chebai_graph.preprocessing.datasets.MUV_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/final/PCBA_final_augmented.yml b/configs/data/augmented/final/PCBA_final_augmented.yml deleted file mode 100644 index 40688f4..0000000 --- a/configs/data/augmented/final/PCBA_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -class_path: chebai_graph.preprocessing.datasets.PCBA_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/final/SIDER_final_augmented.yml b/configs/data/augmented/final/SIDER_final_augmented.yml deleted file mode 100644 index 5eef128..0000000 --- a/configs/data/augmented/final/SIDER_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/Tox21_final_augmented.yml b/configs/data/augmented/final/Tox21_final_augmented.yml deleted file mode 100644 index 1c88112..0000000 --- a/configs/data/augmented/final/Tox21_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/final/ToxCast_final_augmented.yml b/configs/data/augmented/final/ToxCast_final_augmented.yml deleted file mode 100644 index 4831ae5..0000000 --- a/configs/data/augmented/final/ToxCast_final_augmented.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 diff --git a/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml deleted file mode 100644 index a06a948..0000000 --- a/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 deleted file mode 100644 index 26ec990..0000000 --- a/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 deleted file mode 100644 index 31821f1..0000000 --- a/configs/data/augmented/properties/chebi50_atom_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 deleted file mode 100644 index 785f798..0000000 --- a/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 deleted file mode 100644 index 10ed37d..0000000 --- a/configs/data/augmented/properties/chebi50_fg_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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 deleted file mode 100644 index 9827723..0000000 --- a/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -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/data/augmented/properties/chebi50_node_level_prop_only.yml b/configs/data/augmented/properties/chebi50_node_level_prop_only.yml deleted file mode 100644 index d1df497..0000000 --- a/configs/data/augmented/properties/chebi50_node_level_prop_only.yml +++ /dev/null @@ -1,24 +0,0 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes -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/final/chebi50_final_augmented.yml b/configs/data/chebi50_aug_prop_as_per_node.yml similarity index 100% rename from configs/data/augmented/final/chebi50_final_augmented.yml rename to configs/data/chebi50_aug_prop_as_per_node.yml diff --git a/configs/data/chebi50_baseline.yml b/configs/data/chebi50_baseline.yml deleted file mode 100644 index ce9dfa0..0000000 --- a/configs/data/chebi50_baseline.yml +++ /dev/null @@ -1,12 +0,0 @@ -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 diff --git a/configs/model/augmented/pooling/aa_pool/gat.yml b/configs/model/augmented/pooling/aa_pool/gat.yml deleted file mode 100644 index 36750e1..0000000 --- a/configs/model/augmented/pooling/aa_pool/gat.yml +++ /dev/null @@ -1,14 +0,0 @@ -class_path: chebai_graph.models.GATAAPoolGraphPred -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/augmented/pooling/aa_pool/gine.yml b/configs/model/augmented/pooling/aa_pool/gine.yml deleted file mode 100644 index 7c60a73..0000000 --- a/configs/model/augmented/pooling/aa_pool/gine.yml +++ /dev/null @@ -1,13 +0,0 @@ -class_path: chebai_graph.models.GINEAAPoolGraphPred -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/pooling/aa_pool/rggcn.yml b/configs/model/augmented/pooling/aa_pool/rggcn.yml deleted file mode 100644 index 4f6e94f..0000000 --- a/configs/model/augmented/pooling/aa_pool/rggcn.yml +++ /dev/null @@ -1,12 +0,0 @@ -class_path: chebai_graph.models.ResGatedAAPoolGraphPred -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/augmented/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml deleted file mode 100644 index 773088a..0000000 --- a/configs/model/augmented/pooling/amg_pool/gine.yml +++ /dev/null @@ -1,13 +0,0 @@ -class_path: chebai_graph.models.GINEAMGPoolGraphPred -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 - 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 deleted file mode 100644 index b2657f2..0000000 --- a/configs/model/augmented/pooling/amg_pool/rggcn.yml +++ /dev/null @@ -1,12 +0,0 @@ -class_path: chebai_graph.models.ResGatedAMGPoolGraphPred -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 - dropout: 0 - n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/no_pooling/gat.yml b/configs/model/augmented/pooling/no_pooling/gat.yml deleted file mode 100644 index f9d3575..0000000 --- a/configs/model/augmented/pooling/no_pooling/gat.yml +++ /dev/null @@ -1,14 +0,0 @@ -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 diff --git a/configs/model/baselines/gine.yml b/configs/model/baselines/gine.yml deleted file mode 100644 index 9d56063..0000000 --- a/configs/model/baselines/gine.yml +++ /dev/null @@ -1,13 +0,0 @@ -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/baselines/gat.yml b/configs/model/gat.yml similarity index 95% rename from configs/model/baselines/gat.yml rename to configs/model/gat.yml index da84e40..1e96d57 100644 --- a/configs/model/baselines/gat.yml +++ b/configs/model/gat.yml @@ -11,4 +11,5 @@ 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_aapool.yml b/configs/model/gat_aug_aapool.yml index 38a326c..fae47c3 100644 --- a/configs/model/gat_aug_aapool.yml +++ b/configs/model/gat_aug_aapool.yml @@ -10,5 +10,5 @@ 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 diff --git a/configs/model/augmented/pooling/amg_pool/gat.yml b/configs/model/gat_aug_amgpool.yml similarity index 84% rename from configs/model/augmented/pooling/amg_pool/gat.yml rename to configs/model/gat_aug_amgpool.yml index f88159d..e596487 100644 --- a/configs/model/augmented/pooling/amg_pool/gat.yml +++ b/configs/model/gat_aug_amgpool.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.GATAMGPoolGraphPred +class_path: chebai_graph.models.GATGraphNodeFGNodePoolGraphPred init_args: optimizer_kwargs: lr: 1e-3 @@ -11,4 +11,5 @@ 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 b8cfe54..27d1e78 100644 --- a/configs/model/gnn_res_gated.yml +++ b/configs/model/gnn_res_gated.yml @@ -10,3 +10,4 @@ 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 370fc9e..fad8c27 100644 --- a/configs/model/gnn_resgated_pretrain.yml +++ b/configs/model/gnn_resgated_pretrain.yml @@ -13,3 +13,4 @@ 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 f5c2e83..de28d1c 100644 --- a/configs/model/res_aug_aapool.yml +++ b/configs/model/res_aug_aapool.yml @@ -9,4 +9,5 @@ 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 9a59240..9194cd7 100644 --- a/configs/model/res_aug_amgpool.yml +++ b/configs/model/res_aug_amgpool.yml @@ -9,4 +9,5 @@ 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/baselines/rggcn.yml b/configs/model/resgated.yml similarity index 91% rename from configs/model/baselines/rggcn.yml rename to configs/model/resgated.yml index d1847cc..863890d 100644 --- a/configs/model/baselines/rggcn.yml +++ b/configs/model/resgated.yml @@ -9,4 +9,5 @@ init_args: num_layers: 4 edge_dim: 8 # 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 ad55d88..4749795 100644 --- a/configs/model/resgated_dynamic_gni.yml +++ b/configs/model/resgated_dynamic_gni.yml @@ -9,4 +9,5 @@ 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/pyproject.toml b/pyproject.toml index 67e3cc5..b2295db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,23 +14,20 @@ dependencies = [ # torch-geometric # torch_scatter ] -requires-python = ">=3.10" +requires-python = ">=3.8" [project.optional-dependencies] dev = [ "tox", - "omegaconf", - "chebi_utils", ] linters = [ - "ruff", + "isort", "pre-commit", + "black", ] -wandb = ["wandb"] - [build-system] build-backend = "flit_core.buildapi" requires = ["flit_core >=3.2,<4"] diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 0b7cd32..0000000 --- a/requirements.txt +++ /dev/null @@ -1,120 +0,0 @@ -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