diff --git a/README.md b/README.md index 312f3df..b2a8ec0 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ The list can be found in the `configs/data/chebi50_graph_properties.yml` file. python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml ``` -## Augmented Graphs +## Augmented Graphs _See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_. Graph Neural Networks (GNNs) often fail to explicitly leverage the chemically meaningful substructures present within molecules (i.e. **functional groups (FGs)**). To make this implicit information explicitly accessible to GNNs, we augment molecular graphs with **artificial nodes** that represent these substructures. The resulting graph are referred to as **augmented graphs**. diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 9e20b2d..515a513 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -1,19 +1,25 @@ +from .architectures.gat import GATGraphPred +from .architectures.gine import GINEGraphPred +from .architectures.resgated import ResGatedGraphPred from .augmented import ( - GATAugNodePoolGraphPred, - GATGraphNodeFGNodePoolGraphPred, - ResGatedAugNodePoolGraphPred, - ResGatedGraphNodeFGNodePoolGraphPred, + GATAAPoolGraphPred, + GATAMGPoolGraphPred, + GINEAAPoolGraphPred, + GINEAMGPoolGraphPred, + ResGatedAAPoolGraphPred, + ResGatedAMGPoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred -from .gat import GATGraphPred -from .resgated import ResGatedGraphPred __all__ = [ "ResGatedGraphPred", - "ResGatedAugNodePoolGraphPred", - "ResGatedGraphNodeFGNodePoolGraphPred", - "GATGraphPred", - "GATAugNodePoolGraphPred", - "GATGraphNodeFGNodePoolGraphPred", + "ResGatedAAPoolGraphPred", + "ResGatedAMGPoolGraphPred", "ResGatedDynamicGNIGraphPred", + "GATGraphPred", + "GATAAPoolGraphPred", + "GATAMGPoolGraphPred", + "GINEGraphPred", + "GINEAAPoolGraphPred", + "GINEAMGPoolGraphPred", ] diff --git a/chebai_graph/models/architectures/__init__.py b/chebai_graph/models/architectures/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py new file mode 100644 index 0000000..514e287 --- /dev/null +++ b/chebai_graph/models/architectures/base.py @@ -0,0 +1,185 @@ +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/gat.py b/chebai_graph/models/architectures/gat.py similarity index 100% rename from chebai_graph/models/gat.py rename to chebai_graph/models/architectures/gat.py diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py new file mode 100644 index 0000000..f9e89a3 --- /dev/null +++ b/chebai_graph/models/architectures/gine.py @@ -0,0 +1,146 @@ +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/resgated.py b/chebai_graph/models/architectures/resgated.py similarity index 100% rename from chebai_graph/models/resgated.py rename to chebai_graph/models/architectures/resgated.py diff --git a/chebai_graph/models/augmented.py b/chebai_graph/models/augmented.py index fdb5388..87f1c30 100644 --- a/chebai_graph/models/augmented.py +++ b/chebai_graph/models/augmented.py @@ -1,9 +1,10 @@ -from .base import AugmentedNodePoolingNet, GraphNodeFGNodePoolingNet -from .gat import GATGraphPred -from .resgated import ResGatedGraphPred +from .architectures.gat import GATGraphPred +from .architectures.gine import GINEGraphPred +from .architectures.resgated import ResGatedGraphPred +from .pooling import AAPool, AMGPool -class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): +class ResGatedAAPoolGraphPred(AAPool, ResGatedGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -13,7 +14,7 @@ class ResGatedAugNodePoolGraphPred(AugmentedNodePoolingNet, ResGatedGraphPred): ... -class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred): +class GATAAPoolGraphPred(AAPool, GATGraphPred): """ Combines: - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). @@ -23,9 +24,17 @@ class GATAugNodePoolGraphPred(AugmentedNodePoolingNet, GATGraphPred): ... -class ResGatedGraphNodeFGNodePoolGraphPred( - GraphNodeFGNodePoolingNet, ResGatedGraphPred -): +class GINEAAPoolGraphPred(AAPool, GINEGraphPred): + """ + Combines: + - AugmentedNodePoolingNet: Pools atom and augmented node embeddings (optionally with molecule attributes). + - GINEGraphPred: Graph isomorphism network for final graph prediction. + """ + + ... + + +class ResGatedAMGPoolGraphPred(AMGPool, ResGatedGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -35,7 +44,7 @@ class ResGatedGraphNodeFGNodePoolGraphPred( ... -class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, GATGraphPred): +class GATAMGPoolGraphPred(AMGPool, GATGraphPred): """ Combines: - GraphNodeFGNodePoolingNet: Pools atom, functional group, and graph nodes (optionally with molecule attributes). @@ -43,3 +52,13 @@ class GATGraphNodeFGNodePoolGraphPred(GraphNodeFGNodePoolingNet, 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 deleted file mode 100644 index 3226d21..0000000 --- a/chebai_graph/models/base.py +++ /dev/null @@ -1,707 +0,0 @@ -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 8cb6c7b..4fa1b1f 100644 --- a/chebai_graph/models/dynamic_gni.py +++ b/chebai_graph/models/dynamic_gni.py @@ -28,8 +28,8 @@ from chebai_graph.preprocessing.reader import RandomFeatureInitializationReader -from .base import GraphModelBase, GraphNetWrapper -from .resgated import ResGatedModel +from .architectures.base import GraphModelBase, GraphNetWrapper +from .architectures.resgated import ResGatedModel class ResGatedDynamicGNI(GraphModelBase): diff --git a/chebai_graph/models/gin_net.py b/chebai_graph/models/gin_net.py deleted file mode 100644 index 6fed4c6..0000000 --- a/chebai_graph/models/gin_net.py +++ /dev/null @@ -1,96 +0,0 @@ -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 3be0fdd..69d27c0 100644 --- a/chebai_graph/models/graph.py +++ b/chebai_graph/models/graph.py @@ -10,7 +10,7 @@ from chebai_graph.loss.pretraining import MaskPretrainingLoss -from .base import GraphBaseNet +from .architectures.base import GraphBaseNet logging.getLogger("pysmiles").setLevel(logging.CRITICAL) @@ -88,11 +88,6 @@ def __init__(self, config: typing.Dict, **kwargs): self.n_bond_properties = ( int(config["n_bond_properties"]) if "n_bond_properties" in config else 7 ) - self.n_molecule_properties = ( - int(config["n_molecule_properties"]) - if "n_molecule_properties" in config - else 0 - ) self.activation = F.elu self.dropout = nn.Dropout(self.dropout_rate) @@ -158,7 +153,7 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length + (i == 0) * self.gnn.n_molecule_properties, + self.gnn.hidden_length, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) @@ -196,9 +191,7 @@ def __init__( self.linear_layers = torch.nn.ModuleList( [ torch.nn.Linear( - self.gnn.hidden_length - + (i == 0) * self.gnn.n_molecule_properties - + (i == 0) * self.gnn.hidden_length, + self.gnn.hidden_length + (i == 0) * self.gnn.hidden_length, self.gnn.hidden_length, ) for i in range(n_linear_layers - 1) diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py new file mode 100644 index 0000000..a1ee442 --- /dev/null +++ b/chebai_graph/models/pooling.py @@ -0,0 +1,151 @@ +from abc import ABC + +import torch +from torch_geometric.data import Data as GraphData +from torch_scatter import scatter_add + +from .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/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt index eae0a1d..273cc9b 100644 --- a/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt +++ b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt @@ -157,3 +157,4 @@ RING_71 RING_46 orthoester RING_55 +triiodomethyl diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 8708c28..1dcb2f3 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -14,6 +14,17 @@ ChEBI50GraphProperties, ChEBI100GraphProperties, ) +from .molecule_net_classification import ( + BACE_WFGE_WGN_AsPerNodeType, + BBBP_WFGE_WGN_AsPerNodeType, + ClinTox_WFGE_WGN_AsPerNodeType, + HIV_WFGE_WGN_AsPerNodeType, + MUV_WFGE_WGN_AsPerNodeType, + PCBA_WFGE_WGN_AsPerNodeType, + SIDER_WFGE_WGN_AsPerNodeType, + Tox21_WFGE_WGN_AsPerNodeType, + ToxCast_WFGE_WGN_AsPerNodeType, +) from .pubchem import PubChemGraphProperties __all__ = [ @@ -33,4 +44,13 @@ "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE", + "BACE_WFGE_WGN_AsPerNodeType", + "BBBP_WFGE_WGN_AsPerNodeType", + "ClinTox_WFGE_WGN_AsPerNodeType", + "HIV_WFGE_WGN_AsPerNodeType", + "MUV_WFGE_WGN_AsPerNodeType", + "SIDER_WFGE_WGN_AsPerNodeType", + "Tox21_WFGE_WGN_AsPerNodeType", + "ToxCast_WFGE_WGN_AsPerNodeType", + "PCBA_WFGE_WGN_AsPerNodeType", ] diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py new file mode 100644 index 0000000..9e2bf61 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -0,0 +1,50 @@ +from abc import ABC + +import pandas as pd +from torch_geometric.data.data import Data as GeomData + +from .base import GraphPropertiesMixIn + + +class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): + """Mixin for augmented graph data without additional graph nodes.""" + + READER = None + + def _merge_props_into_base(self, row: pd.Series) -> GeomData: + data = super()._merge_props_into_base(row) + geom_data = row["features"] + assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) + + is_atom_node = geom_data.is_atom_node + assert is_atom_node is not None, "is_atom_node must be set in the geom_data" + data.is_atom_node = is_atom_node + return data + + +class AugGraphPropMixIn_WithGraphNode(AugGraphPropMixIn_NoGraphNode, ABC): + """Mixin for augmented graph data with graph-level nodes.""" + + READER = None + + def _merge_props_into_base(self, row: pd.Series) -> GeomData: + data = super()._merge_props_into_base(row) + return self._add_graph_node_mask(data, row) + + def _add_graph_node_mask(self, data: GeomData, row: pd.Series) -> GeomData: + """ + Add a graph node mask to the GeomData object. + + Args: + data: A GeomData object with features. + row: A dictionary containing 'features' and other metadata. + + Returns: + Modified GeomData with graph node mask added. + """ + geom_data = row["features"] + assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) + is_graph_node = geom_data.is_graph_node + assert is_graph_node is not None, "is_graph_node must be set in the geom_data" + data.is_graph_node = is_graph_node + return data diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py new file mode 100644 index 0000000..13bacbb --- /dev/null +++ b/chebai_graph/preprocessing/datasets/base.py @@ -0,0 +1,657 @@ +import os +from abc import ABC +from collections.abc import Callable +from pprint import pformat +from typing import Optional + +import pandas as pd +import torch +import tqdm +from chebai.preprocessing.datasets.base import XYBaseDataModule +from lightning_utilities.core.rank_zero import rank_zero_info +from rdkit import Chem +from torch_geometric.data.data import Data as GeomData + +from chebai_graph.preprocessing.datasets.utils import resolve_property +from chebai_graph.preprocessing.properties import ( + AllNodeTypeProperty, + AtomNodeTypeProperty, + AtomProperty, + BondProperty, + FGNodeTypeProperty, + MolecularProperty, + MoleculeProperty, +) +from chebai_graph.preprocessing.reader import ( + GraphPropertyReader, + RandomFeatureInitializationReader, +) +from chebai_graph.preprocessing.reader.augmented_reader import _AugmentorReader + + +class DataPropertiesSetter(XYBaseDataModule, ABC): + """Mixin for adding molecular property encodings to graph-based given datasets.""" + + READER = GraphPropertyReader + + def __init__( + self, + properties: list | None = None, + transform: Callable | None = None, + **kwargs, + ): + """ + Initialize GraphPropertiesMixIn. + + Args: + properties: Optional list of MolecularProperty class paths or instances. + transform: Optional transformation applied to each data sample. + """ + super().__init__(**kwargs) + # atom_properties and bond_properties are given as lists containing class_paths + if properties is not None: + properties = [resolve_property(prop) for prop in properties] + properties = self._sort_properties(properties) + else: + properties = [] + self.properties: list[MolecularProperty] = properties + assert isinstance(self.properties, list) and all( + isinstance(p, MolecularProperty) for p in self.properties + ) + self.transform = transform + + def _sort_properties( + self, properties: list[MolecularProperty] + ) -> list[MolecularProperty]: + return sorted(properties, key=lambda prop: self.get_property_path(prop)) + + def _setup_properties(self) -> None: + """ + Process and cache molecular properties to disk. + + Returns: + None + """ + raw_data = [] + os.makedirs(self.processed_properties_dir, exist_ok=True) + + try: + file_names = self.processed_main_file_names + except NotImplementedError: + file_names = self.raw_file_names + + for file in file_names: + # processed_dir_main only exists for ChEBI datasets + path = os.path.join( + ( + self.processed_dir_main + if hasattr(self, "processed_dir_main") + else self.raw_dir + ), + file, + ) + raw_data += list(self._load_dict(path)) + + idents = [row["ident"] for row in raw_data] + features = [row["features"] for row in raw_data] + + # use vectorized version of encode function, apply only if value is present + def enc_if_not_none(encode, value): + return ( + [encode(v) for v in value] + if value is not None and len(value) > 0 + else None + ) + + if any( + not os.path.isfile(self.get_property_path(property)) + for property in self.properties + ): + # augment molecule graph if possible (this would also happen for the properties if needed, but this avoids redundancy) + if isinstance(self.reader, _AugmentorReader): + returned_results = [] + for mol in features: + try: + r = self.reader._create_augmented_graph(mol) + except Exception: + r = None + returned_results.append(r) + mols = [ + augmented_mol[1] if augmented_mol is not None else None + for augmented_mol in returned_results + ] + else: + mols = features + + for property in self.properties: + if not os.path.isfile(self.get_property_path(property)): + rank_zero_info(f"Processing property {property.name}") + # read all property values first, then encode + rank_zero_info(f"\tReading property values of {property.name}...") + property_values = [ + self.reader.read_property(mol, property) + if mol is not None + else None + for mol in tqdm.tqdm(mols) + ] + rank_zero_info(f"\tEncoding property values of {property.name}...") + property.encoder.on_start(property_values=property_values) + encoded_values = [ + enc_if_not_none(property.encoder.encode, value) + for value in tqdm.tqdm(property_values) + ] + assert len(encoded_values) == len(idents) == len(features) + torch.save( + [ + {property.name: torch.cat(feat), "ident": id} + for feat, id in zip(encoded_values, idents) + if feat is not None + ], + self.get_property_path(property), + ) + property.on_finish() + + @property + def processed_properties_dir(self) -> str: + return os.path.join(self.processed_dir, "properties") + + def get_property_path(self, property: MolecularProperty) -> str: + """ + Construct the cache path for a given molecular property. + + Args: + property: Instance of a MolecularProperty. + + Returns: + Path to the cached property file. + """ + return os.path.join( + self.processed_properties_dir, + f"{property.name}_{property.encoder.name}.pt", + ) + + def _after_setup(self, **kwargs) -> None: + """ + Finalize setup after ensuring properties are processed. + + Args: + **kwargs: Additional keyword arguments passed to superclass. + + Returns: + None + """ + self._setup_properties() + super()._after_setup(**kwargs) + + def _preprocess_smiles_for_pred( + self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None + ) -> Optional[dict]: + """Preprocess prediction data.""" + # Add dummy labels because the collate function requires them. + # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, + # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. + result = self.reader.to_data( + {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} + ) + # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object + if isinstance(result["features"], tuple): + result["features"], raw_data = result["features"] + if result is None or result["features"] is None: + return None + for property in self.properties: + property.encoder.eval = True + property_value = self.reader.read_property(raw_data, property) + if property_value is None or len(property_value) == 0: + encoded_value = None + else: + encoded_value = torch.stack( + [property.encoder.encode(v) for v in property_value] + ) + if len(encoded_value.shape) == 3: + encoded_value = encoded_value.squeeze(0) + result[property.name] = encoded_value + + result["features"] = self._prediction_merge_props_into_base_wrapper( + result, model_hparams + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + result["features"] = self.transform(result["features"]) + + return result + + def _prediction_merge_props_into_base_wrapper( + self, row: pd.Series | dict, model_hparams: Optional[dict] = None + ) -> GeomData: + """ + Wrapper to merge properties into base features for prediction. + + Args: + row: A dictionary or pd.Series containing 'features' and encoded properties. + Returns: + A GeomData object with merged features. + """ + return self._merge_props_into_base(row) + + +class GraphPropertiesMixIn(DataPropertiesSetter, ABC): + def __init__( + self, + properties=None, + transform=None, + pad_node_features: int | None = None, + pad_edge_features: int | None = None, + distribution: str = "normal", + **kwargs, + ): + super().__init__(properties, transform, **kwargs) + self.pad_edge_features = int(pad_edge_features) if pad_edge_features else None + self.pad_node_features = int(pad_node_features) if pad_node_features else None + if self.pad_node_features or self.pad_edge_features: + assert ( + distribution is not None + and distribution in RandomFeatureInitializationReader.DISTRIBUTIONS + ), ( + "When using padding for features, a valid distribution must be specified." + ) + self.distribution = distribution + if self.pad_node_features: + print( + f"[Info] Node-level features will be padded with random" + f"{self.pad_node_features} values from {self.distribution} distribution." + ) + if self.pad_edge_features: + print( + f"[Info] Edge-level features will be padded with random" + f"{self.pad_edge_features} values from {self.distribution} distribution." + ) + + if self.properties: + print( + f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}" + ) + + def _merge_props_into_base(self, row: pd.Series | dict) -> GeomData: + """ + Merge encoded molecular properties into the GeomData object. + + Args: + row: A dictionary containing 'features' and encoded properties. + + Returns: + A GeomData object with merged features. + """ + if isinstance(row["features"], tuple): + geom_data, _ = row[ + "features" + ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) + else: + geom_data = row["features"] + assert isinstance(geom_data, GeomData) + edge_attr = geom_data.edge_attr + x = geom_data.x + molecule_attr = torch.empty((1, 0)) + + for property in self.properties: + property_values = row[f"{property.name}"] + if isinstance(property_values, torch.Tensor): + if len(property_values.size()) == 0: + property_values = property_values.unsqueeze(0) + if len(property_values.size()) == 1: + property_values = property_values.unsqueeze(1) + else: + property_values = torch.zeros( + (0, property.encoder.get_encoding_length()) + ) + + if isinstance(property, AtomProperty): + x = torch.cat([x, property_values], dim=1) + elif isinstance(property, BondProperty): + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr = torch.cat( + [edge_attr, torch.cat([property_values, property_values], dim=0)], + dim=1, + ) + elif isinstance(property, MoleculeProperty): + molecule_attr = torch.cat([molecule_attr, property_values], dim=1) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + if self.pad_node_features: + padding_values = torch.empty((x.shape[0], self.pad_node_features)) + RandomFeatureInitializationReader.random_gni( + padding_values, self.distribution + ) + x = torch.cat([x, padding_values], dim=1) + + if self.pad_edge_features: + padding_values = torch.empty((edge_attr.shape[0], self.pad_edge_features)) + RandomFeatureInitializationReader.random_gni( + padding_values, self.distribution + ) + edge_attr = torch.cat([edge_attr, padding_values], dim=1) + + return GeomData( + x=x, + edge_index=geom_data.edge_index, + edge_attr=edge_attr, + molecule_attr=molecule_attr, + ) + + def load_processed_data( + self, kind: Optional[str] = None, filename: Optional[str] = None + ) -> list[dict]: + """ + Load dataset and merge cached properties into base features. + + Args: + filename: The path to the file to load. + + Returns: + List of data entries, each a dictionary. + """ + base_data = super().load_processed_data(kind, filename) + base_df = pd.DataFrame(base_data) + + for property in self.properties: + property_data = torch.load( + self.get_property_path(property), weights_only=False + ) + if len(property_data[0][property.name].shape) > 1: + property.encoder.set_encoding_length( + property_data[0][property.name].shape[1] + ) + + property_df = pd.DataFrame(property_data) + property_df.rename( + columns={property.name: f"{property.name}"}, inplace=True + ) + base_df = base_df.merge(property_df, on="ident", how="left") + + base_df["features"] = base_df.apply( + lambda row: self._merge_props_into_base(row), axis=1 + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + base_df["features"] = base_df["features"].apply(self.transform) + + prop_lengths = [ + (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties + ] + + # -------------------------- Count total node properties + n_node_properties = sum( + p.encoder.get_encoding_length() + for p in self.properties + if isinstance(p, AtomProperty) + ) + + in_channels_str = "" + if self.pad_node_features: + n_node_properties += self.pad_node_features + in_channels_str += f" (with {self.pad_node_features} padded random values from {self.distribution} distribution)" + + in_channels_str = f"in_channels: {n_node_properties}" + in_channels_str + + # -------------------------- Count total edge properties + n_edge_properties = sum( + p.encoder.get_encoding_length() + for p in self.properties + if isinstance(p, BondProperty) + ) + edge_dim_str = "" + if self.pad_edge_features: + n_edge_properties += self.pad_edge_features + edge_dim_str += f" (with {self.pad_edge_features} padded random values from {self.distribution} distribution)" + + edge_dim_str = f"edge_dim: {n_edge_properties}" + edge_dim_str + + rank_zero_info( + f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" + f"Use following values for given parameters for model configuration: \n\t" + f"{in_channels_str} \n\t" + f"{edge_dim_str} \n\t" + ) + + return base_df[base_data[0].keys()].to_dict("records") + + +class GraphPropAsPerNodeType(DataPropertiesSetter, ABC): + def __init__(self, properties=None, transform=None, **kwargs): + super().__init__(properties, transform, **kwargs) + # Sort properties so that AllNodeTypeProperty instances come first, rest of the properties order remain same + first = self._sort_properties( + [prop for prop in self.properties if isinstance(prop, AllNodeTypeProperty)] + ) + rest = self._sort_properties( + [ + prop + for prop in self.properties + if not isinstance(prop, AllNodeTypeProperty) + ] + ) + self.properties = first + rest + print( + "Properties are sorted so that `AllNodeTypeProperty` properties are first in sequence and rest of the order remains same\n", + f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}", + ) + + def load_processed_data( + self, kind: Optional[str] = None, filename: Optional[str] = None + ) -> list[dict]: + """ + Load dataset and merge cached properties into base features. + + Args: + filename: The path to the file to load. + + Returns: + List of data entries, each a dictionary. + """ + base_data = super().load_processed_data(kind, filename) + base_df = pd.DataFrame(base_data) + props_categories = { + "AllNodeTypeProperties": [], + "FGNodeTypeProperties": [], + "AtomNodeTypeProperties": [], + "GraphNodeTypeProperties": [], + "BondProperties": [], + } + n_atom_node_properties, n_fg_node_properties = 0, 0 + n_bond_properties, n_graph_node_properties = 0, 0 + prop_lengths = [] + for prop in self.properties: + prop_length = prop.encoder.get_encoding_length() + prop_name = prop.name + prop_lengths.append((prop_name, prop_length)) + if isinstance(prop, AllNodeTypeProperty): + n_atom_node_properties += prop_length + n_fg_node_properties += prop_length + n_graph_node_properties += prop_length + props_categories["AllNodeTypeProperties"].append(prop_name) + elif isinstance(prop, FGNodeTypeProperty): + n_fg_node_properties += prop_length + props_categories["FGNodeTypeProperties"].append(prop_name) + elif isinstance(prop, AtomNodeTypeProperty): + n_atom_node_properties += prop_length + props_categories["AtomNodeTypeProperties"].append(prop_name) + elif isinstance(prop, BondProperty): + n_bond_properties += prop_length + props_categories["BondProperties"].append(prop_name) + elif isinstance(prop, MoleculeProperty): + # molecule props will be used as graph node props + n_graph_node_properties += prop_length + props_categories["GraphNodeTypeProperties"].append(prop_name) + else: + raise TypeError(f"Unsupported property type: {type(prop).__name__}") + + n_node_properties = max( + n_atom_node_properties, n_fg_node_properties, n_graph_node_properties + ) + rank_zero_info( + f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" + f"Properties Categories:\n{pformat(props_categories)}\n\n" + f"n_atom_node_properties: {n_atom_node_properties}, " + f"n_fg_node_properties: {n_fg_node_properties}, " + f"n_bond_properties: {n_bond_properties}, " + f"n_graph_node_properties: {n_graph_node_properties}\n\n" + f"Use following values for given parameters for model configuration: \n\t" + f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}\n" + ) + + for property in self.properties: + rank_zero_info(f"Loading property {property.name}...") + property_data = torch.load( + self.get_property_path(property), weights_only=False + ) + if len(property_data[0][property.name].shape) > 1: + property.encoder.set_encoding_length( + property_data[0][property.name].shape[1] + ) + + property_df = pd.DataFrame(property_data) + property_df.rename( + columns={property.name: f"{property.name}"}, inplace=True + ) + base_df = base_df.merge(property_df, on="ident", how="left") + + base_df["features"] = base_df.apply( + lambda row: self._merge_props_into_base( + row, + max_len_node_properties=n_node_properties, + ), + axis=1, + ) + + # apply transformation, e.g. masking for pretraining task + if self.transform is not None: + base_df["features"] = base_df["features"].apply(self.transform) + + return base_df[base_data[0].keys()].to_dict("records") + + def _merge_props_into_base( + self, row: pd.Series, max_len_node_properties: int + ) -> GeomData: + """ + Merge encoded molecular properties into the GeomData object. + + Args: + row: A dictionary containing 'features' and encoded properties. + + Returns: + A GeomData object with merged features. + """ + geom_data = row["features"] + if geom_data is None: + return None + if isinstance(geom_data, tuple): + geom_data = geom_data[ + 0 + ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) + assert isinstance(geom_data, GeomData) + + is_atom_node = geom_data.is_atom_node + assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" + is_graph_node = geom_data.is_graph_node + assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" + + is_fg_node = ~is_atom_node & ~is_graph_node + num_nodes = geom_data.x.size(0) + edge_attr = geom_data.edge_attr + + # Initialize node feature matrix + assert max_len_node_properties is not None, ( + "Maximum len of node properties should not be None" + ) + x = torch.zeros((num_nodes, max_len_node_properties)) + + # Track column offsets for each node type + atom_offset, fg_offset, graph_offset = 0, 0, 0 + + for property in self.properties: + property_values = row[f"{property.name}"].to(dtype=torch.float32) + if isinstance(property_values, torch.Tensor): + if len(property_values.size()) == 0: + property_values = property_values.unsqueeze(0) + if len(property_values.size()) == 1: + property_values = property_values.unsqueeze(1) + else: + property_values = torch.zeros( + (0, property.encoder.get_encoding_length()) + ) + + enc_len = property_values.shape[1] + # -------------- Node properties --------------- + if isinstance(property, AllNodeTypeProperty): + x[:, atom_offset : atom_offset + enc_len] = property_values + atom_offset += enc_len + fg_offset += enc_len + graph_offset += enc_len + + elif isinstance(property, AtomNodeTypeProperty): + x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ + is_atom_node + ] + atom_offset += enc_len + + elif isinstance(property, FGNodeTypeProperty): + x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ + is_fg_node + ] + fg_offset += enc_len + + elif isinstance(property, MoleculeProperty): + x[is_graph_node, graph_offset : graph_offset + enc_len] = ( + property_values + ) + graph_offset += enc_len + + # ------------- Bond Properties -------------- + elif isinstance(property, BondProperty): + # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges + edge_attr = torch.cat( + [edge_attr, torch.cat([property_values, property_values], dim=0)], + dim=1, + ) + else: + raise TypeError(f"Unsupported property type: {type(property).__name__}") + + total_used_columns = max(atom_offset, fg_offset, graph_offset) + assert total_used_columns <= max_len_node_properties, ( + f"Used {total_used_columns} columns, but max allowed is {max_len_node_properties}" + ) + + return GeomData( + x=x, + edge_index=geom_data.edge_index, + edge_attr=edge_attr, + molecule_attr=torch.empty((1, 0)), # empty as not used for this class + is_atom_node=is_atom_node, + is_fg_node=is_fg_node, + is_graph_node=is_graph_node, + ) + + def _prediction_merge_props_into_base_wrapper( + self, row: pd.Series | dict, model_hparams: Optional[dict] = None + ) -> GeomData: + """ + Wrapper to merge properties into base features for prediction. + + Args: + row: A dictionary or pd.Series containing 'features' and encoded properties. + Returns: + A GeomData object with merged features. + """ + if ( + model_hparams is None + or "in_channels" not in model_hparams["config"] + or model_hparams["config"]["in_channels"] is None + ): + raise ValueError( + f"model_hparams must be provided for data class: {self.__class__.__name__}" + f" which should contain 'in_channels' key with valid value in 'config' dictionary." + ) + max_len_node_properties = int(model_hparams["config"]["in_channels"]) + return self._merge_props_into_base(row, max_len_node_properties) diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 2d03876..263747e 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -1,13 +1,4 @@ -import os -from abc import ABC -from collections.abc import Callable -from pprint import pformat -from typing import Optional - import pandas as pd -from chebai_graph.preprocessing.reader.augmented_reader import _AugmentorReader -import torch -import tqdm from chebai.preprocessing.datasets.chebi import ( ChEBIOver50, ChEBIOver100, @@ -15,18 +6,7 @@ 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, @@ -37,12 +17,15 @@ GN_WithAllNodes_FG_WithAtoms_NoFGE, GN_WithAtoms_FG_WithAtoms_FGE, GN_WithAtoms_FG_WithAtoms_NoFGE, - GraphPropertyReader, GraphReader, RandomFeatureInitializationReader, ) -from chebai_graph.preprocessing.datasets.utils import resolve_property +from .augmentation_base import ( + AugGraphPropMixIn_NoGraphNode, + AugGraphPropMixIn_WithGraphNode, +) +from .base import DataPropertiesSetter, GraphPropAsPerNodeType, GraphPropertiesMixIn class ChEBI50GraphData(ChEBIOver50): @@ -54,635 +37,6 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -class DataPropertiesSetter(ChEBIOverX, ABC): - """Mixin for adding molecular property encodings to graph-based ChEBI datasets.""" - - READER = GraphPropertyReader - - def __init__( - self, - properties: list | None = None, - transform: Callable | None = None, - **kwargs, - ): - """ - Initialize GraphPropertiesMixIn. - - Args: - properties: Optional list of MolecularProperty class paths or instances. - transform: Optional transformation applied to each data sample. - """ - super().__init__(**kwargs) - # atom_properties and bond_properties are given as lists containing class_paths - if properties is not None: - properties = [resolve_property(prop) for prop in properties] - properties = self._sort_properties(properties) - else: - properties = [] - self.properties: list[MolecularProperty] = properties - assert isinstance(self.properties, list) and all( - isinstance(p, MolecularProperty) for p in self.properties - ) - self.transform = transform - - def _sort_properties( - self, properties: list[MolecularProperty] - ) -> list[MolecularProperty]: - return sorted(properties, key=lambda prop: self.get_property_path(prop)) - - def _setup_properties(self) -> None: - """ - Process and cache molecular properties to disk. - - Returns: - None - """ - raw_data = [] - os.makedirs(self.processed_properties_dir, exist_ok=True) - - try: - file_names = self.processed_main_file_names - except NotImplementedError: - file_names = self.raw_file_names - - for file in file_names: - # processed_dir_main only exists for ChEBI datasets - path = os.path.join( - ( - self.processed_dir_main - if hasattr(self, "processed_dir_main") - else self.raw_dir - ), - file, - ) - raw_data += list(self._load_dict(path)) - - idents = [row["ident"] for row in raw_data] - features = [row["features"] for row in raw_data] - - # use vectorized version of encode function, apply only if value is present - def enc_if_not_none(encode, value): - return ( - [encode(v) for v in value] - if value is not None and len(value) > 0 - else None - ) - - if any( - not os.path.isfile(self.get_property_path(property)) - for property in self.properties - ): - # augment molecule graph if possible (this would also happen for the properties if needed, but this avoids redundancy) - if isinstance(self.reader, _AugmentorReader): - returned_results = [] - for mol in features: - try: - r = self.reader._create_augmented_graph(mol) - except Exception: - r = None - returned_results.append(r) - mols = [ - augmented_mol[1] if augmented_mol is not None else None - for augmented_mol in returned_results - ] - else: - mols = features - - for property in self.properties: - if not os.path.isfile(self.get_property_path(property)): - rank_zero_info(f"Processing property {property.name}") - # read all property values first, then encode - rank_zero_info(f"\tReading property values of {property.name}...") - property_values = [ - self.reader.read_property(mol, property) - if mol is not None - else None - for mol in tqdm.tqdm(mols) - ] - rank_zero_info(f"\tEncoding property values of {property.name}...") - property.encoder.on_start(property_values=property_values) - encoded_values = [ - enc_if_not_none(property.encoder.encode, value) - for value in tqdm.tqdm(property_values) - ] - assert len(encoded_values) == len(idents) == len(features) - torch.save( - [ - {property.name: torch.cat(feat), "ident": id} - for feat, id in zip(encoded_values, idents) - if feat is not None - ], - self.get_property_path(property), - ) - property.on_finish() - - @property - def processed_properties_dir(self) -> str: - return os.path.join(self.processed_dir, "properties") - - def get_property_path(self, property: MolecularProperty) -> str: - """ - Construct the cache path for a given molecular property. - - Args: - property: Instance of a MolecularProperty. - - Returns: - Path to the cached property file. - """ - return os.path.join( - self.processed_properties_dir, - f"{property.name}_{property.encoder.name}.pt", - ) - - def _after_setup(self, **kwargs) -> None: - """ - Finalize setup after ensuring properties are processed. - - Args: - **kwargs: Additional keyword arguments passed to superclass. - - Returns: - None - """ - self._setup_properties() - super()._after_setup(**kwargs) - - def _preprocess_smiles_for_pred( - self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None - ) -> Optional[dict]: - """Preprocess prediction data.""" - # Add dummy labels because the collate function requires them. - # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, - # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. - result = self.reader.to_data( - {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} - ) - # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object - if isinstance(result["features"], tuple): - result["features"], raw_data = result["features"] - if result is None or result["features"] is None: - return None - for property in self.properties: - property.encoder.eval = True - property_value = self.reader.read_property(raw_data, property) - if property_value is None or len(property_value) == 0: - encoded_value = None - else: - encoded_value = torch.stack( - [property.encoder.encode(v) for v in property_value] - ) - if len(encoded_value.shape) == 3: - encoded_value = encoded_value.squeeze(0) - result[property.name] = encoded_value - - result["features"] = self._prediction_merge_props_into_base_wrapper( - result, model_hparams - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - result["features"] = self.transform(result["features"]) - - return result - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - return self._merge_props_into_base(row) - - -class GraphPropertiesMixIn(DataPropertiesSetter, ABC): - def __init__( - self, - properties=None, - transform=None, - pad_node_features: int = None, - pad_edge_features: int = None, - distribution: str = "normal", - **kwargs, - ): - super().__init__(properties, transform, **kwargs) - self.pad_edge_features = int(pad_edge_features) if pad_edge_features else None - self.pad_node_features = int(pad_node_features) if pad_node_features else None - if self.pad_node_features or self.pad_edge_features: - assert ( - distribution is not None - and distribution in RandomFeatureInitializationReader.DISTRIBUTIONS - ), ( - "When using padding for features, a valid distribution must be specified." - ) - self.distribution = distribution - if self.pad_node_features: - print( - f"[Info] Node-level features will be padded with random" - f"{self.pad_node_features} values from {self.distribution} distribution." - ) - if self.pad_edge_features: - print( - f"[Info] Edge-level features will be padded with random" - f"{self.pad_edge_features} values from {self.distribution} distribution." - ) - - if self.properties: - print( - f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}" - ) - - def _merge_props_into_base(self, row: pd.Series | dict) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - if isinstance(row["features"], tuple): - geom_data, _ = row[ - "features" - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - else: - geom_data = row["features"] - assert isinstance(geom_data, GeomData) - edge_attr = geom_data.edge_attr - x = geom_data.x - molecule_attr = torch.empty((1, 0)) - - for property in self.properties: - property_values = row[f"{property.name}"] - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - if isinstance(property, AtomProperty): - x = torch.cat([x, property_values], dim=1) - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - elif isinstance(property, MoleculeProperty): - molecule_attr = torch.cat([molecule_attr, property_values], dim=1) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - if self.pad_node_features: - padding_values = torch.empty((x.shape[0], self.pad_node_features)) - RandomFeatureInitializationReader.random_gni( - padding_values, self.distribution - ) - x = torch.cat([x, padding_values], dim=1) - - if self.pad_edge_features: - padding_values = torch.empty((edge_attr.shape[0], self.pad_edge_features)) - RandomFeatureInitializationReader.random_gni( - padding_values, self.distribution - ) - edge_attr = torch.cat([edge_attr, padding_values], dim=1) - - return GeomData( - x=x, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=molecule_attr, - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - - for property in self.properties: - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base(row), axis=1 - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - prop_lengths = [ - (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties - ] - - # -------------------------- Count total node properties - n_node_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, AtomProperty) - ) - - in_channels_str = "" - if self.pad_node_features: - n_node_properties += self.pad_node_features - in_channels_str += f" (with {self.pad_node_features} padded random values from {self.distribution} distribution)" - - in_channels_str = f"in_channels: {n_node_properties}" + in_channels_str - - # -------------------------- Count total edge properties - n_edge_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, BondProperty) - ) - edge_dim_str = "" - if self.pad_edge_features: - n_edge_properties += self.pad_edge_features - edge_dim_str += f" (with {self.pad_edge_features} padded random values from {self.distribution} distribution)" - - edge_dim_str = f"edge_dim: {n_edge_properties}" + edge_dim_str - - rank_zero_info( - f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" - f"Use following values for given parameters for model configuration: \n\t" - f"{in_channels_str} \n\t" - f"{edge_dim_str} \n\t" - f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}" - ) - - return base_df[base_data[0].keys()].to_dict("records") - - -class GraphPropAsPerNodeType(DataPropertiesSetter, ABC): - def __init__(self, properties=None, transform=None, **kwargs): - super().__init__(properties, transform, **kwargs) - # Sort properties so that AllNodeTypeProperty instances come first, rest of the properties order remain same - first = self._sort_properties( - [prop for prop in self.properties if isinstance(prop, AllNodeTypeProperty)] - ) - rest = self._sort_properties( - [ - prop - for prop in self.properties - if not isinstance(prop, AllNodeTypeProperty) - ] - ) - self.properties = first + rest - print( - "Properties are sorted so that `AllNodeTypeProperty` properties are first in sequence and rest of the order remains same\n", - f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}", - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - props_categories = { - "AllNodeTypeProperties": [], - "FGNodeTypeProperties": [], - "AtomNodeTypeProperties": [], - "GraphNodeTypeProperties": [], - "BondProperties": [], - } - n_atom_node_properties, n_fg_node_properties = 0, 0 - n_bond_properties, n_graph_node_properties = 0, 0 - prop_lengths = [] - for prop in self.properties: - prop_length = prop.encoder.get_encoding_length() - prop_name = prop.name - prop_lengths.append((prop_name, prop_length)) - if isinstance(prop, AllNodeTypeProperty): - n_atom_node_properties += prop_length - n_fg_node_properties += prop_length - n_graph_node_properties += prop_length - props_categories["AllNodeTypeProperties"].append(prop_name) - elif isinstance(prop, FGNodeTypeProperty): - n_fg_node_properties += prop_length - props_categories["FGNodeTypeProperties"].append(prop_name) - elif isinstance(prop, AtomNodeTypeProperty): - n_atom_node_properties += prop_length - props_categories["AtomNodeTypeProperties"].append(prop_name) - elif isinstance(prop, BondProperty): - n_bond_properties += prop_length - props_categories["BondProperties"].append(prop_name) - elif isinstance(prop, MoleculeProperty): - # molecule props will be used as graph node props - n_graph_node_properties += prop_length - props_categories["GraphNodeTypeProperties"].append(prop_name) - else: - raise TypeError(f"Unsupported property type: {type(prop).__name__}") - - n_node_properties = max( - n_atom_node_properties, n_fg_node_properties, n_graph_node_properties - ) - rank_zero_info( - f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" - f"Properties Categories:\n{pformat(props_categories)}\n\n" - f"n_atom_node_properties: {n_atom_node_properties}, " - f"n_fg_node_properties: {n_fg_node_properties}, " - f"n_bond_properties: {n_bond_properties}, " - f"n_graph_node_properties: {n_graph_node_properties}\n\n" - f"Use following values for given parameters for model configuration: \n\t" - f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0\n" - ) - - for property in self.properties: - rank_zero_info(f"Loading property {property.name}...") - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base( - row, - max_len_node_properties=n_node_properties, - ), - axis=1, - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - return base_df[base_data[0].keys()].to_dict("records") - - def _merge_props_into_base( - self, row: pd.Series, max_len_node_properties: int - ) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - geom_data = row["features"] - if geom_data is None: - return None - if isinstance(geom_data, tuple): - geom_data = geom_data[ - 0 - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - assert isinstance(geom_data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" - - is_fg_node = ~is_atom_node & ~is_graph_node - num_nodes = geom_data.x.size(0) - edge_attr = geom_data.edge_attr - - # Initialize node feature matrix - assert max_len_node_properties is not None, ( - "Maximum len of node properties should not be None" - ) - x = torch.zeros((num_nodes, max_len_node_properties)) - - # Track column offsets for each node type - atom_offset, fg_offset, graph_offset = 0, 0, 0 - - for property in self.properties: - property_values = row[f"{property.name}"].to(dtype=torch.float32) - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - enc_len = property_values.shape[1] - # -------------- Node properties --------------- - if isinstance(property, AllNodeTypeProperty): - x[:, atom_offset : atom_offset + enc_len] = property_values - atom_offset += enc_len - fg_offset += enc_len - graph_offset += enc_len - - elif isinstance(property, AtomNodeTypeProperty): - x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ - is_atom_node - ] - atom_offset += enc_len - - elif isinstance(property, FGNodeTypeProperty): - x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ - is_fg_node - ] - fg_offset += enc_len - - elif isinstance(property, MoleculeProperty): - x[is_graph_node, graph_offset : graph_offset + enc_len] = ( - property_values - ) - graph_offset += enc_len - - # ------------- Bond Properties -------------- - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - total_used_columns = max(atom_offset, fg_offset, graph_offset) - assert total_used_columns <= max_len_node_properties, ( - f"Used {total_used_columns} columns, but max allowed is {max_len_node_properties}" - ) - - return GeomData( - x=x, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=torch.empty((1, 0)), # empty as not used for this class - is_atom_node=is_atom_node, - is_fg_node=is_fg_node, - is_graph_node=is_graph_node, - ) - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - if ( - model_hparams is None - or "in_channels" not in model_hparams["config"] - or model_hparams["config"]["in_channels"] is None - ): - raise ValueError( - f"model_hparams must be provided for data class: {self.__class__.__name__}" - f" which should contain 'in_channels' key with valid value in 'config' dictionary." - ) - max_len_node_properties = int(model_hparams["config"]["in_channels"]) - return self._merge_props_into_base(row, max_len_node_properties) - - class ChEBI50_StaticGNI(DataPropertiesSetter, ChEBIOver50): READER = RandomFeatureInitializationReader @@ -696,7 +50,6 @@ def load_processed_data_from_file(self, filename): f"Use following values for given parameters for model configuration: \n\t" f"in_channels: {self.reader.num_node_properties} , " f"edge_dim: {self.reader.num_bond_properties}, " - f"n_molecule_properties: {self.reader.num_molecule_properties}" ) return base_df[base_data[0].keys()].to_dict("records") @@ -719,50 +72,6 @@ class ChEBI50GraphPropertiesPartial(ChEBI50GraphProperties, ChEBIOverXPartial): pass -class AugGraphPropMixIn_NoGraphNode(GraphPropertiesMixIn, ABC): - """Mixin for augmented graph data without additional graph nodes.""" - - READER = None - - def _merge_props_into_base(self, row: pd.Series) -> GeomData: - data = super()._merge_props_into_base(row) - geom_data = row["features"] - assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "is_atom_node must be set in the geom_data" - data.is_atom_node = is_atom_node - return data - - -class AugGraphPropMixIn_WithGraphNode(AugGraphPropMixIn_NoGraphNode, ABC): - """Mixin for augmented graph data with graph-level nodes.""" - - READER = None - - def _merge_props_into_base(self, row: pd.Series) -> GeomData: - data = super()._merge_props_into_base(row) - return self._add_graph_node_mask(data, row) - - def _add_graph_node_mask(self, data: GeomData, row: pd.Series) -> GeomData: - """ - Add a graph node mask to the GeomData object. - - Args: - data: A GeomData object with features. - row: A dictionary containing 'features' and other metadata. - - Returns: - Modified GeomData with graph node mask added. - """ - geom_data = row["features"] - assert isinstance(geom_data, GeomData) and isinstance(data, GeomData) - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "is_graph_node must be set in the geom_data" - data.is_graph_node = is_graph_node - return data - - class ChEBI50_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50): """ChEBIOver50 with with FG nodes and FG edges and graph node.""" diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py new file mode 100644 index 0000000..9afaf71 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -0,0 +1,60 @@ +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/configs/data/augmented/BACE_final_augmented.yml b/configs/data/augmented/BACE_final_augmented.yml new file mode 100644 index 0000000..7ea67f0 --- /dev/null +++ b/configs/data/augmented/BACE_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BACE_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/BBBP_final_augmented.yml b/configs/data/augmented/BBBP_final_augmented.yml new file mode 100644 index 0000000..9a7ee90 --- /dev/null +++ b/configs/data/augmented/BBBP_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/ClinTox_final_augmented.yml b/configs/data/augmented/ClinTox_final_augmented.yml new file mode 100644 index 0000000..a4023b3 --- /dev/null +++ b/configs/data/augmented/ClinTox_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ClinTox_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/HIV_final_augmented.yml b/configs/data/augmented/HIV_final_augmented.yml new file mode 100644 index 0000000..596fc19 --- /dev/null +++ b/configs/data/augmented/HIV_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.HIV_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/MUV_final_augmented.yml b/configs/data/augmented/MUV_final_augmented.yml new file mode 100644 index 0000000..56a3e1e --- /dev/null +++ b/configs/data/augmented/MUV_final_augmented.yml @@ -0,0 +1,24 @@ +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/PCBA_final_augmented.yml b/configs/data/augmented/PCBA_final_augmented.yml new file mode 100644 index 0000000..40688f4 --- /dev/null +++ b/configs/data/augmented/PCBA_final_augmented.yml @@ -0,0 +1,24 @@ +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/SIDER_final_augmented.yml b/configs/data/augmented/SIDER_final_augmented.yml new file mode 100644 index 0000000..5eef128 --- /dev/null +++ b/configs/data/augmented/SIDER_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.SIDER_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/Tox21_final_augmented.yml b/configs/data/augmented/Tox21_final_augmented.yml new file mode 100644 index 0000000..1c88112 --- /dev/null +++ b/configs/data/augmented/Tox21_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/ToxCast_final_augmented.yml b/configs/data/augmented/ToxCast_final_augmented.yml new file mode 100644 index 0000000..4831ae5 --- /dev/null +++ b/configs/data/augmented/ToxCast_final_augmented.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/chebi50_aug_prop_as_per_node.yml b/configs/data/augmented/chebi50_final_augmented.yml similarity index 100% rename from configs/data/chebi50_aug_prop_as_per_node.yml rename to configs/data/augmented/chebi50_final_augmented.yml diff --git a/configs/data/chebi50_baseline.yml b/configs/data/chebi50_baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/chebi50_baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +init_args: + properties: + - chebai_graph.preprocessing.properties.AtomType + - chebai_graph.preprocessing.properties.NumAtomBonds + - chebai_graph.preprocessing.properties.AtomCharge + - chebai_graph.preprocessing.properties.AtomAromaticity + - chebai_graph.preprocessing.properties.AtomHybridization + - chebai_graph.preprocessing.properties.AtomNumHs + - chebai_graph.preprocessing.properties.BondType + - chebai_graph.preprocessing.properties.BondInRing + - chebai_graph.preprocessing.properties.BondAromaticity diff --git a/configs/model/gat_aug_amgpool.yml b/configs/model/augmented/pooling/aa_pool/gat.yml similarity index 84% rename from configs/model/gat_aug_amgpool.yml rename to configs/model/augmented/pooling/aa_pool/gat.yml index e596487..36750e1 100644 --- a/configs/model/gat_aug_amgpool.yml +++ b/configs/model/augmented/pooling/aa_pool/gat.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.models.GATGraphNodeFGNodePoolGraphPred +class_path: chebai_graph.models.GATAAPoolGraphPred init_args: optimizer_kwargs: lr: 1e-3 @@ -11,5 +11,4 @@ init_args: heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/aa_pool/gine.yml b/configs/model/augmented/pooling/aa_pool/gine.yml new file mode 100644 index 0000000..7c60a73 --- /dev/null +++ b/configs/model/augmented/pooling/aa_pool/gine.yml @@ -0,0 +1,13 @@ +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 new file mode 100644 index 0000000..4f6e94f --- /dev/null +++ b/configs/model/augmented/pooling/aa_pool/rggcn.yml @@ -0,0 +1,12 @@ +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/gat.yml b/configs/model/augmented/pooling/amg_pool/gat.yml new file mode 100644 index 0000000..f88159d --- /dev/null +++ b/configs/model/augmented/pooling/amg_pool/gat.yml @@ -0,0 +1,14 @@ +class_path: chebai_graph.models.GATAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 12 # number of bond properties + heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) + v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0 + n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml new file mode 100644 index 0000000..773088a --- /dev/null +++ b/configs/model/augmented/pooling/amg_pool/gine.yml @@ -0,0 +1,13 @@ +class_path: chebai_graph.models.GINEAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 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 new file mode 100644 index 0000000..b2657f2 --- /dev/null +++ b/configs/model/augmented/pooling/amg_pool/rggcn.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.models.ResGatedAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 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/gat.yml b/configs/model/baselines/gat.yml similarity index 76% rename from configs/model/gat.yml rename to configs/model/baselines/gat.yml index dda84dc..da84e40 100644 --- a/configs/model/gat.yml +++ b/configs/model/baselines/gat.yml @@ -3,12 +3,12 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 158 # number of node/atom properties + in_channels: 161 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 7 # number of bond properties + edge_dim: 8 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` - n_molecule_properties: 0 + dropout: 0 n_linear_layers: 1 diff --git a/configs/model/baselines/gine.yml b/configs/model/baselines/gine.yml new file mode 100644 index 0000000..9d56063 --- /dev/null +++ b/configs/model/baselines/gine.yml @@ -0,0 +1,13 @@ +class_path: chebai_graph.models.GINEGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 161 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 8 # number of bond properties + dropout: 0 + train_eps: true + n_linear_layers: 1 diff --git a/configs/model/resgated.yml b/configs/model/baselines/rggcn.yml similarity index 62% rename from configs/model/resgated.yml rename to configs/model/baselines/rggcn.yml index ccc6615..d1847cc 100644 --- a/configs/model/resgated.yml +++ b/configs/model/baselines/rggcn.yml @@ -3,11 +3,10 @@ init_args: optimizer_kwargs: lr: 1e-3 config: - in_channels: 158 # number of node/atom properties + in_channels: 161 # number of node/atom properties hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 7 # number of bond properties + edge_dim: 8 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/gat_aug_aapool.yml b/configs/model/gat_aug_aapool.yml index fae47c3..38a326c 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` - n_molecule_properties: 0 + dropout: 0 n_linear_layers: 1 diff --git a/configs/model/gnn_res_gated.yml b/configs/model/gnn_res_gated.yml index 27d1e78..b8cfe54 100644 --- a/configs/model/gnn_res_gated.yml +++ b/configs/model/gnn_res_gated.yml @@ -10,4 +10,3 @@ init_args: n_linear_layers: 3 n_atom_properties: 158 n_bond_properties: 7 - n_molecule_properties: 200 diff --git a/configs/model/gnn_resgated_pretrain.yml b/configs/model/gnn_resgated_pretrain.yml index fad8c27..370fc9e 100644 --- a/configs/model/gnn_resgated_pretrain.yml +++ b/configs/model/gnn_resgated_pretrain.yml @@ -13,4 +13,3 @@ init_args: n_linear_layers: 3 n_atom_properties: 151 n_bond_properties: 7 - n_molecule_properties: 200 diff --git a/configs/model/res_aug_aapool.yml b/configs/model/res_aug_aapool.yml index de28d1c..f5c2e83 100644 --- a/configs/model/res_aug_aapool.yml +++ b/configs/model/res_aug_aapool.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 12 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/res_aug_amgpool.yml b/configs/model/res_aug_amgpool.yml index 9194cd7..9a59240 100644 --- a/configs/model/res_aug_amgpool.yml +++ b/configs/model/res_aug_amgpool.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 12 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/configs/model/resgated_dynamic_gni.yml b/configs/model/resgated_dynamic_gni.yml index 4749795..ad55d88 100644 --- a/configs/model/resgated_dynamic_gni.yml +++ b/configs/model/resgated_dynamic_gni.yml @@ -9,5 +9,4 @@ init_args: num_layers: 4 edge_dim: 7 # number of bond properties dropout: 0 - n_molecule_properties: 0 n_linear_layers: 1 diff --git a/pyproject.toml b/pyproject.toml index a517cac..c524311 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,20 +13,23 @@ dependencies = [ # torch-geometric # torch_scatter ] -requires-python = ">=3.8" +requires-python = ">=3.10" [project.optional-dependencies] dev = [ "tox", + "omegaconf", + "chebi_utils", ] linters = [ - "isort", + "ruff", "pre-commit", - "black", ] +wandb = ["wandb"] + [build-system] build-backend = "flit_core.buildapi" requires = ["flit_core >=3.2,<4"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0b7cd32 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,120 @@ +aiohappyeyeballs==2.7.1 +aiohttp==3.14.1 +aiosignal==1.4.0 +annotated-doc==0.0.4 +annotated-types==0.7.0 +antlr4-python3-runtime==4.9.3 +anyio==4.14.2 +async-timeout==5.0.1 +attrs==26.1.0 +cachetools==7.1.4 +certifi==2026.6.17 +cfgv==3.5.0 +chardet==7.4.3 +charset-normalizer==3.4.9 +chebai==1.3.0 +chebi-utils==0.2.1 +chembl-structure-pipeline==1.2.4 +click==8.4.2 +colorama==0.4.6 +descriptastorus==2.8.0 +distlib==0.4.3 +docstring-parser==0.18.0 +exceptiongroup==1.3.1 +fastobo==0.14.1 +filelock==3.30.2 +frozenlist==1.8.0 +fsspec==2025.12.0 +h11==0.16.0 +hf-xet==1.5.2 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.23.0 +identify==2.6.19 +idna==3.18 +importlib-resources==7.1.0 +iterative-stratification==0.1.9 +jinja2==3.1.6 +joblib==1.5.3 +jsonargparse==4.49.0 +lightning==2.5.1 +lightning-utilities==0.15.3 +markdown-it-py==4.2.0 +markupsafe==3.0.3 +mdurl==0.1.2 +mpmath==1.3.0 +multidict==6.7.1 +networkx==3.4.2 +nodeenv==1.10.0 +numpy==2.2.6 +nvidia-cublas-cu12==12.4.5.8 +nvidia-cuda-cupti-cu12==12.4.127 +nvidia-cuda-nvrtc-cu12==12.4.127 +nvidia-cuda-runtime-cu12==12.4.127 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cufft-cu12==11.2.1.3 +nvidia-curand-cu12==10.3.5.147 +nvidia-cusolver-cu12==11.6.1.9 +nvidia-cusparse-cu12==12.3.1.170 +nvidia-cusparselt-cu12==0.6.2 +nvidia-nccl-cu12==2.21.5 +nvidia-nvjitlink-cu12==12.4.127 +nvidia-nvtx-cu12==12.4.127 +omegaconf==2.3.1 +packaging==24.2 +pandas==2.3.3 +pandas-flavor==0.8.1 +pbr==7.0.3 +pillow==12.3.0 +platformdirs==4.10.0 +pluggy==1.6.0 +pre-commit==4.6.0 +propcache==0.5.2 +protobuf==7.35.1 +psutil==7.2.2 +pydantic==2.13.4 +pydantic-core==2.46.4 +pygments==2.20.0 +pyparsing==3.3.2 +pyproject-api==1.9.0 +pysmiles==1.1.2 +python-dateutil==2.9.0.post0 +python-discovery==1.4.4 +pytorch-lightning==2.6.5 +pytz==2026.2 +pyyaml==6.0.3 +rdkit==2024.3.6 +regex==2026.7.10 +requests==2.34.2 +rich==15.0.0 +safetensors==0.8.0 +scikit-learn==1.7.2 +scipy==1.15.3 +sentry-sdk==2.66.0 +setuptools==83.0.0 +shellingham==1.5.4 +six==1.17.0 +sympy==1.13.1 +threadpoolctl==3.6.0 +tokenizers==0.22.2 +tomli==2.4.1 +torch==2.6.0 +torch-geometric==2.8.0 +torch-scatter==2.1.2+pt26cu124 +torch-sparse==0.6.18+pt26cu124 +torchmetrics==1.9.0 +tox==4.27.0 +tqdm==4.68.4 +transformers==5.14.1 +triton==3.2.0 +typer==0.27.0 +typeshed-client==2.12.0 +typing-extensions==4.16.0 +typing-inspection==0.4.2 +tzdata==2026.3 +urllib3==2.7.0 +virtualenv==21.6.1 +wandb==0.28.1 +xarray==2025.6.1 +xxhash==3.8.1 +yarl==1.24.2