diff --git a/.gitignore b/.gitignore index d01eecf..cc9e754 100644 --- a/.gitignore +++ b/.gitignore @@ -174,3 +174,9 @@ electra_pretrained.ckpt *.out *.sh *.ckpt +*.csv + +!configs/data/augmented/aug-ablation/*.sh +!.vscode/launch.json +!configs/data/augmented/final/*.sh +!configs/data/augmented/pretraining/*.sh diff --git a/README.md b/README.md index b2a8ec0..cc30202 100644 --- a/README.md +++ b/README.md @@ -76,17 +76,156 @@ 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 ``` +## Experiment Scripts + +The scripts below reproduce the augmented-graph experiments introduced in this +work. They expect both repositories to be checked out as described in +[Recommended Folder Structure](#recommended-folder-structure), and they run +`python -m chebai` from the `python-chebai` checkout. Update the paths, +environment activation, and logging configuration in the scripts for a +different machine. + +### Augmentation ablations + +Each ablation script evaluates the same ten graph-construction variants: +`baseline`, `FGN`, `FGN+E`, `FGN+WGN`, `FGN+E+WGN`, `WGN`, +`gn_wall_fgwa_nfge`, `gn_wall_fgwa_wfge`, `gnwa_fgwa_nfge`, and +`gnwa_fgwa_wfge`. The experiments use seeds `0`, `42`, and `12345`, so each +dataset runs 30 configurations in total. The variants compare whether +functional-group nodes (FGN), functional-group edges (FGE), and the graph node +(WGN) are included, and whether their connections use the corresponding +weighted or unweighted schemes. + +- [`aug-ablation-chebi.sh`](configs/data/augmented/aug-ablation/aug-ablation-chebi.sh) + is a Slurm job-array runner for ChEBI50. It schedules the 10 configurations + across the three seeds, trains the baseline GAT, and uses the ChEBI v252 + split. +- [`aug-ablation-tox21.sh`](configs/data/augmented/aug-ablation/aug-ablation-tox21.sh) + runs the same 30-experiment matrix for Tox21 serially on one local GPU. It + reports micro/macro F1 and ROC-AUC using the Tox21 split. +- [`aug-ablation-toxcast.sh`](configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh) + runs the corresponding ToxCast matrix serially on one local GPU. It uses + ToxCast-specific input dimensions and learning rates and reports micro/macro + F1 and ROC-AUC. + +Run a local ablation script from this repository with, for example: + +```bash +bash configs/data/augmented/aug-ablation/aug-ablation-tox21.sh +``` + +The ChEBI script is intended for Slurm and should be submitted with +`sbatch`. The Tox21 and ToxCast scripts run every seed and configuration in +sequence, rather than requiring Slurm. + +### MoleculeNet datasets + +[MoleculeNet](https://moleculenet.org/) is a collection of standard molecular +machine-learning benchmarks. These datasets cover binary classification tasks +with different label counts, class imbalance, and chemical domains, making +them useful for testing whether an architecture generalizes beyond ChEBI50. +The augmented-graph runner currently evaluates six MoleculeNet datasets: +**BBBP**, **BACE**, **SIDER**, **ClinTox**, **Tox21**, and **ToxCast**. Their +dataset configurations are in +`configs/data/augmented/final/_final_augmented.yml`. + +[`molnet.sh`](configs/data/augmented/final/molnet.sh) runs each of the six +datasets for seeds `0`, `42`, and `12345` on one local GPU. For each dataset +and seed it: + +1. trains the augmented atom-motif-graph pooling GAT; +2. selects the best validation checkpoint using ROC-AUC; +3. evaluates that checkpoint on the test split; and +4. writes per-run metrics and checkpoint paths to + `logs/molnet_test_results.tsv`, followed by mean and sample standard + deviation summaries. + +The script selects binary or multilabel trainers and metrics per dataset and +uses dataset-specific dropout and learning rates. Run it with: + +```bash +bash configs/data/augmented/final/molnet.sh +``` + +For comparisons with published MoleculeNet results, use scaffold splits for +BBBP and BACE and random splits for Tox21, ToxCast, SIDER, and ClinTox. The +reported benchmark metric is the mean test ROC-AUC over three independent +runs; the runner also records the F1 metrics selected for each dataset. The +MoleculeNet experiments can be run with or without initializing the model from +the ZINC15 MSP checkpoint, allowing the contribution of pretraining to be +measured separately. + +### Final ChEBI50 training + +[`final-chebi.sh`](configs/data/augmented/final/final-chebi.sh) is a Slurm +job-array script that trains the final augmented-graph GAT on ChEBI50 v252 for +the three seeds. It uses the `gnwa_fgwa_wfge` configuration, the +atom-motif-graph pooling architecture, weighted binary cross-entropy, and the +standard micro/macro F1 metrics. Submit it with: + +```bash +sbatch configs/data/augmented/final/final-chebi.sh +``` + +### ZINC15 pretraining + +[`zinc_msp_pretraining.sh`](configs/data/augmented/pretraining/zinc_msp_pretraining.sh) +trains the augmented GAT with the masked-subgraph prediction (MSP) +pretraining objective on ZINC15. It uses the dedicated ZINC pretrainer, +`gat_msp.yml`, and the augmented ZINC15 data configuration. Validation batches +are disabled because this is a pretraining run; the resulting checkpoint can +then be used to initialize downstream augmented-graph experiments. The script +requests three GPUs on Slurm and is submitted with: + +```bash +sbatch configs/data/augmented/pretraining/zinc_msp_pretraining.sh +``` + ## 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**. > Note: Rings are also treated as functional groups in our work. -In these augmented graphs, each functional group node is connected to the atoms that constitute the group. Additionally, two functional group nodes are connected if any atom belonging to one group shares a bond with an atom from the other group. We further introduce a **graph node**, an extra node connected to all functional group nodes. - -Among all the connection schemes we evaluated, this configuration delivered the strongest performance. We denote it using the abbreviation **WFG_WFGE_WGN** in our work and is shown in below figure. - -mol_to_aug_mol +In these augmented graphs, each functional group node is connected to the atoms +that constitute the group. Additionally, two functional group nodes are +connected if any atom belonging to one group shares a bond with an atom from +the other group. We also introduce a **graph node**, which is connected to all +atom nodes. Rings are treated as functional groups. + +Functional-group nodes encode both the presence and the identity of the +detected group, rather than representing an unlabeled fragment. The graph node +acts as a global or master node: it provides a shortcut for information +exchange between distant parts of a molecule and collects a whole-molecule +representation. + +### Best configuration: `GNWA_FGWA_WFGE` + +Among the connection schemes evaluated, **GNWA_FGWA_WFGE** delivered the +strongest performance. The name describes the complete augmented topology: + +- **GNWA** (*Graph Node With Atom connections*): the graph node is connected + to every atom node. +- **FGWA** (*Functional Group Nodes With Atom connections*): every functional + group node is connected to the atoms that belong to that group. +- **WFGE** (*With Functional Group Edges*): functional group nodes are + connected when their member atoms are connected by a molecular bond. + +This configuration is used by the final ChEBI50 and MoleculeNet experiments. +Its ChEBI50 data definition is +[`chebi50_final_augmented.yml`](configs/data/augmented/final/chebi50_final_augmented.yml), +and its atom-motif-graph pooling model is +[`amg_pool/gat.yml`](configs/model/augmented/pooling/amg_pool/gat.yml). +The topology is illustrated below. + +The final augmented representation initializes atoms, bonds, functional-group +nodes, and the graph node with chemically meaningful, node-type-specific +properties. This is different from the isolated topology ablations, where +artificial nodes and added edges use zero initialization so that the effect of +connectivity can be measured independently of feature initialization and +type-specific pooling. + +mol graph drawio

@@ -94,50 +233,86 @@ Among all the connection schemes we evaluated, this configuration delivered the Below is the command for the model and data configuration that achieved the best classification performance using augmented graphs. ```bash -python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/wandb_logger.yml --model=../python-chebai-graph/configs/model/gat_aug_amgpool.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_aug_prop_as_per_node.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 --trainer.logger.init_args.name=gatv2_amg_s0 +python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/wandb_logger.yml --model=../python-chebai-graph/configs/model/augmented/pooling/amg_pool/gat.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/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml --data.init_args.batch_size=64 --trainer.accumulate_grad_batches=1 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=252 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_unweighted.yml --trainer.logger.init_args.name=gnwa_fgwa_wfge_final_42 --model.init_args.optimizer_kwargs.lr=0.002 --seed_everything=42 --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv --trainer.logger.init_args.tags=["augmented_paper","final:gnwa_fgwa_wfge"] ``` ### Model Hyperparameters -#### **GAT Architecture** - -To use a GAT-based model, choose **one** of the following configs: +The model paths below are relative to the `python-chebai` working directory +used by the training commands. In this repository they correspond to the +following configuration files: -- **Standard Pooling**: `--model=../python-chebai-graph/configs/model/gat.yml` - > Standard pooling sums the learned representations from all the nodes to produce a single representation which is used for classification. +#### **GAT Architecture** -- **Atom-Augmented Node Pooling**: `--model=../python-chebai-graph/configs/model/gat_aug_aagpool.yml` - > With this pooling stratergy, the learned representations are first separated into **two distinct sets**: those from atom nodes and those from all artificial nodes (both functional groups and the graph node). The representations within each set are aggregated separately (using summation) to yield two distinct single vectors. These two resulting vectors are then concatenated before being passed to the classification layer. +Choose one of the following GAT configurations: + +- **Standard pooling**: + [`configs/model/baselines/gat.yml`](configs/model/baselines/gat.yml) + (`--model=../python-chebai-graph/configs/model/baselines/gat.yml`). + `GATGraphPred` pools all node representations into one graph + representation. This configuration is used for the ablation baselines. +- **Atom-augmented pooling**: + [`configs/model/augmented/pooling/aa_pool/gat.yml`](configs/model/augmented/pooling/aa_pool/gat.yml). + `GATAAPoolGraphPred` aggregates atom-node representations separately from + artificial-node representations, then concatenates the two pooled vectors. +- **Atom-motif-graph pooling**: + [`configs/model/augmented/pooling/amg_pool/gat.yml`](configs/model/augmented/pooling/amg_pool/gat.yml). + `GATAMGPoolGraphPred` separately aggregates atom nodes, functional-group + nodes, and the graph node before classification. This is the configuration + used by the final augmented experiments. +- **Augmented graphs without separate pooling**: + [`configs/model/augmented/pooling/no_pooling/gat.yml`](configs/model/augmented/pooling/no_pooling/gat.yml). + This uses `GATGraphPred` with augmented node and edge features while pooling + all nodes together. + +The augmented GAT configurations use 203 node features and 12 edge features; +the baseline configuration uses 161 node features and 8 edge features. All +current GAT configurations default to four message-passing layers, 256 hidden +channels, 512 output channels, eight attention heads, and GATv2 +(`v2: True`). The number of heads must divide the output or hidden channel +size. These values can be overridden, for example: -- **Atom–Motif–Graph Node Pooling**: `--model=../python-chebai-graph/configs/model/gat_aug_amgpool.yml` - > This approach employs a finer granularity of separation, distinguishing learned representations into **three distinct sets**: atom nodes, Functional Group (FG) nodes, and the single graph node. Summation is performed separately on the atom node set and the FG node set, yielding two vectors. These two vectors are then concatenated along with the single vector corresponding to the graph node before the final linear layer. +```bash +--model.config.num_layers=5 +--model.config.heads=4 +--model.config.v2=False +--model.config.dropout=0.1 +``` -#### GAT-specific hyperparameters +Set `v2=False` to use the original `GATConv`; the default `v2=True` uses +`GATv2Conv`, which provides dynamic rather than static attention. For details, +see the [GATv2 paper](https://arxiv.org/abs/2105.14491). -- **Number of message-passing layers**: `--model.config.num_layers=5` (default: 4) -- **Attention heads**: `--model.config.heads=4` (default: 8) - > **Note**: The number of heads should be divisible by the output channels (or hidden channels if output channels are not specified). +The augmentation and property initialization do not add trainable parameters +to the core GNN. Type-specific pooling can increase the size of the +classification head, but the main performance gain is intended to come from +the augmented molecular representation and its chemically meaningful +features. -- **To Use different GAT versions**: - - **Use GAT**: `--model.config.v2=False` +#### **ResGated Architecture** - - **Use GATv2**: `--model.config.v2=True` (__default__) - > **Note**: GATv2 addresses the limitation of static attention in GAT by introducing a dynamic attention mechanism. For further details, please refer to the [original GATv2 paper](https://arxiv.org/abs/2105.14491). +Choose one of the following ResGated configurations: -#### **ResGated Architecture** +- **Atom-motif-graph pooling**: + [`configs/model/res_aug_amgpool.yml`](configs/model/res_aug_amgpool.yml). +- **Atom-augmented pooling**: + [`configs/model/res_aug_aapool.yml`](configs/model/res_aug_aapool.yml). -To use a ResGated GNN model, choose **one** of the following configs: +Both augmented ResGated configurations use 203 node features, 12 edge +features, four message-passing layers, 256 hidden channels, and 512 output +channels. The ResGated convolution does not expose GAT attention-head or +`v2` settings. -- **Atom–Motif–Graph Node Pooling**: `--model=../python-chebai-graph/configs/model/res_aug_amgpool.yml` -- **Atom-Augmented Node Pooling**: `--model=../python-chebai-graph/configs/model/res_aug_aagpool.yml` -- **Standard Pooling**: `--model=../python-chebai-graph/configs/model/resgated.yml` +#### **Shared hyperparameters** -#### **Common Hyperparameters** +These settings are available for both GAT and ResGated model configurations: -These can be used for both GAT and ResGated architectures: +- **Dropout**: `--model.config.dropout=0.1` (default: `0`) +- **Number of final linear layers**: `--model.n_linear_layers=2` + (default: `1`) -- **Dropout**: `--model.config.dropout=0.1` (default: 0) -- **Number of final linear layers**: `--model.n_linear_layers=2` (default: 1) +The complete defaults for each architecture are defined directly in the +linked YAML files above. ## Random Node Initialization diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py new file mode 100644 index 0000000..7973698 --- /dev/null +++ b/chebai_graph/loss/msp.py @@ -0,0 +1,300 @@ +""" +Multi-level Self-supervised Pre-training from +Zang, X., Zhao, X. & Tang, B. Hierarchical Molecular Graph Self-Supervised Learning for property prediction. Commun Chem 6, 34 (2023). https://doi.org/10.1038/s42004-023-00825-5 + +This file contains modified code from: + - https://github.com/ZangXuan/HiMol/blob/main/pretrain.py +""" + +import torch +from sklearn.metrics import average_precision_score, roc_auc_score +from torch_geometric.data import Data + +from chebai_graph.preprocessing.properties.augmented_properties import ( + AugAtomType, + AugBondType, +) +from chebai_graph.preprocessing.property_encoder import IndexEncoder + + +class MSPPretrainingLoss(torch.nn.Module): + """HiMol-style molecular structure pretraining loss. + + The model output must be ``(graph_data, node_embeddings)``. The graph data + must contain ``is_atom_node`` and a PyG ``batch`` vector. If present, + ``is_graph_node`` identifies the molecule-level representation. + """ + + def __init__( + self, + hidden_size: int, + dropout: float = 0.2, + tokens_directory: str = "_shared", + ) -> None: + super().__init__() + self.tokens_directory = tokens_directory + self.atom_type_encoder = IndexEncoder( + property=AugAtomType(tokens_directory=self.tokens_directory), + use_one_hot_cache=True, + eval=True, + tokens_directory=self.tokens_directory, + ) + self.bond_type_encoder = IndexEncoder( + property=AugBondType(tokens_directory=self.tokens_directory), + use_one_hot_cache=True, + eval=True, + tokens_directory=self.tokens_directory, + ) + + # Bond existence projection + self.bond_if_proj = torch.nn.Sequential( + torch.nn.Linear(hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, hidden_size), + ) + # Bond existence classifier + self.bond_if_s = torch.nn.Sequential( + torch.nn.Linear(2 * hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, 1), + ) + + # Bond type classifier + self.bond_type_s = torch.nn.Sequential( + torch.nn.Linear(2 * hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, len(self.bond_type_encoder.cache)), + ) + # Atom type classifier + self.atom_type_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, len(self.atom_type_encoder.cache)), + ) + + count_hidden = max(hidden_size // 4, 1) + # Atom count regressor + self.atom_num_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, count_hidden), + torch.nn.Softplus(), + torch.nn.Linear(count_hidden, 1), + ) + # Bond count regressor + self.bond_num_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, count_hidden), + torch.nn.Softplus(), + torch.nn.Linear(count_hidden, 1), + ) + self.feat_drop = ( + torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity() + ) + self.bond_if_loss = torch.nn.BCEWithLogitsLoss() + self.bond_type_loss = torch.nn.CrossEntropyLoss() + self.atom_type_loss = torch.nn.CrossEntropyLoss() + self.count_loss = torch.nn.SmoothL1Loss() + + self.loss_weight = torch.nn.Parameter(torch.rand(5), requires_grad=True) + + @staticmethod + def _bond_count(edge_index: torch.Tensor) -> int: + if edge_index.numel() == 0: + return 0 + return int(edge_index.t().sort(dim=1).values.unique(dim=0).shape[0]) + + def _super( + self, data: Data, embeddings: torch.Tensor, graphs: list[torch.Tensor] + ) -> torch.Tensor: + if not hasattr(data, "is_graph_node") or data.is_graph_node is None: + raise ValueError( + "MSP pretraining requires a boolean is_graph_node attribute in the graph data" + ) + + graph_mask = data.is_graph_node.bool() + graph_representations = [] + for graph_index in range(len(graphs)): + representations = embeddings[graph_mask & (data.batch == graph_index)] + if representations.numel() == 0: + raise ValueError( + f"Graph {graph_index} has no node marked is_graph_node" + ) + if len(representations) > 1: + raise ValueError( + f"Graph {graph_index} has multiple nodes marked is_graph_node" + ) + graph_representations.append(representations[0]) + return torch.stack(graph_representations) + + @staticmethod + def _molecule_metadata(custom_metadata, graph_index: int) -> dict: + metadata = custom_metadata[graph_index] + if isinstance(metadata, list): + metadata = next( + ( + item + for item in metadata + if isinstance(item, dict) and "nodes" in item + ), + None, + ) + if ( + not isinstance(metadata, dict) + or "nodes" not in metadata + or "edges" not in metadata + ): + raise ValueError( + f"Missing molecular metadata for graph {graph_index}; expected nodes and edges" + ) + return metadata + + def _metadata_targets(self, metadata: dict, device: torch.device): + # atom_vocab = list(range(1, 119)) # Atomic numbers from 1 to 118 + # bond_vocab = ["SINGLE", "DOUBLE", "AROMATIC", "TRIPLE", "DATIVE", "UNSPECIFIED"] + atom_mol = metadata["nodes"]["atom_nodes"] + + atom_type_targets = torch.tensor( + [ + self.atom_type_encoder.encode(str(atom.GetAtomicNum())) + for atom in atom_mol.GetAtoms() + ], + dtype=torch.long, + device=device, + ) + within_atoms = metadata["edges"]["within_atoms_lvl"] + bond_indices = [] + bond_type_targets = [] + for bond in within_atoms.GetBonds(): + begin = bond.GetBeginAtomIdx() + end = bond.GetEndAtomIdx() + bond_class = self.bond_type_encoder.encode(str(bond.GetBondType())) + bond_indices.extend(((begin, end), (end, begin))) + bond_type_targets.extend((bond_class, bond_class)) + edge_index = ( + torch.tensor(bond_indices, dtype=torch.long, device=device).t() + if bond_indices + else torch.empty((2, 0), dtype=torch.long, device=device) + ) + return ( + atom_type_targets, + torch.tensor(bond_type_targets, dtype=torch.long, device=device), + edge_index, + ) + + def forward( + self, output, target, features, custom_metadata, **loss_kwargs + ) -> torch.Tensor: + node_embeddings = output + graph_data = features[0] + if not isinstance(graph_data, Data): + raise TypeError("features[0] must be a torch_geometric.data.Data object") + if graph_data.batch is None or not hasattr(graph_data, "is_atom_node"): + raise ValueError("MSP pretraining requires batch and is_atom_node") + is_atom_node = graph_data.is_atom_node.bool() + batch_index = graph_data.batch.long() + graph_nodes = [ + torch.where(batch_index == graph_index)[0] + for graph_index in range(int(batch_index.max().item()) + 1) + ] + graph_node_embeddings = self._super(graph_data, node_embeddings, graph_nodes) + + # Number of atoms prediction loss + atom_counts = torch.stack( + [ + (is_atom_node & (batch_index == graph_index)).sum() + for graph_index in range(len(graph_nodes)) + ] + ).to(node_embeddings.dtype) + atom_count_pred = self.atom_num_s(graph_node_embeddings).squeeze(-1) + atom_count_loss = self.count_loss(atom_count_pred, atom_counts) + + bond_if_loss = node_embeddings.new_zeros(()) + bond_type_loss = node_embeddings.new_zeros(()) + atom_type_loss = node_embeddings.new_zeros(()) + bond_counts = [] + atom_correct = bond_correct = 0 + atom_total = bond_total = 0 + auc_values, ap_values = [], [] + for graph_index, nodes in enumerate(graph_nodes): + metadata = self._molecule_metadata(custom_metadata, graph_index) + mol = metadata["nodes"]["atom_nodes"] + atom_type_targets, bond_type_targets, local_edges = self._metadata_targets( + metadata, node_embeddings.device + ) + + # Bond existence prediction loss + atoms = nodes[is_atom_node[nodes]] + if atoms.numel() != atom_type_targets.numel(): + raise ValueError( + f"Graph {graph_index} has {atoms.numel()} atom nodes but metadata has " + f"{atom_type_targets.numel()} atoms" + ) + atom_embeddings = self.feat_drop(self.bond_if_proj(node_embeddings[atoms])) + bond_counts.append(mol.GetNumBonds()) + n_atoms = atoms.numel() + pairs = torch.cat( + [ + atom_embeddings[:, None, :].expand(-1, n_atoms, -1), + atom_embeddings[None, :, :].expand(n_atoms, -1, -1), + ], + dim=-1, + ) + bond_pred = self.bond_if_s(pairs).squeeze(-1) + bond_target = torch.zeros_like(bond_pred) + if local_edges.numel(): + bond_target[local_edges[0], local_edges[1]] = 1 + bond_if_loss = bond_if_loss + self.bond_if_loss(bond_pred, bond_target) + + if loss_kwargs.get("compute_metrics", True) and local_edges.numel(): + target_cpu = bond_target.flatten().detach().cpu() + pred_cpu = bond_pred.flatten().detach().sigmoid().cpu() + if target_cpu.unique().numel() > 1: + auc_values.append(roc_auc_score(target_cpu, pred_cpu)) + ap_values.append(average_precision_score(target_cpu, pred_cpu)) + + # Atom type prediction loss + atom_pred = self.atom_type_s(node_embeddings[atoms]) + atom_type_loss = atom_type_loss + self.atom_type_loss( + atom_pred, atom_type_targets + ) + atom_correct += int((atom_pred.argmax(1) == atom_type_targets).sum()) + atom_total += atom_type_targets.numel() + + # Bond type prediction loss + if local_edges.numel(): + edge_rep = torch.cat( + [atom_embeddings[local_edges[0]], atom_embeddings[local_edges[1]]], + dim=1, + ) + bond_pred_type = self.bond_type_s(edge_rep) + bond_type_loss = bond_type_loss + self.bond_type_loss( + bond_pred_type, bond_type_targets + ) + bond_correct += int( + (bond_pred_type.argmax(1) == bond_type_targets).sum() + ) + bond_total += bond_type_targets.numel() + + # Number of bonds prediction loss + bond_counts = torch.tensor( + bond_counts, + dtype=graph_node_embeddings.dtype, + device=graph_node_embeddings.device, + ) + bond_count_pred = self.bond_num_s(graph_node_embeddings).squeeze(-1) + bond_count_loss = self.count_loss(bond_count_pred, bond_counts) + + divisor = max(len(graph_nodes), 1) + losses = torch.stack( + [ + bond_if_loss / divisor, + bond_type_loss / divisor, + atom_type_loss / divisor, + atom_count_loss, + bond_count_loss, + ] + ) + if self.loss_weight.numel() != losses.numel(): + raise RuntimeError( + f"Expected {losses.numel()} loss weights, got {self.loss_weight.numel()}" + ) + loss = (losses * torch.softmax(self.loss_weight, dim=0)).sum() + return loss diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 515a513..27a4b64 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -10,6 +10,7 @@ ResGatedAMGPoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred +from .pretraining import GATMSPretraining __all__ = [ "ResGatedGraphPred", @@ -22,4 +23,5 @@ "GINEGraphPred", "GINEAAPoolGraphPred", "GINEAMGPoolGraphPred", + "GATMSPretraining", ] diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 514e287..15f3775 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -26,7 +26,20 @@ def _get_prediction_and_labels( Returns: tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels). """ - return torch.sigmoid(output), labels.int() + valid_label_mask = data["loss_kwargs"]["valid_label_mask"] + predictions = torch.sigmoid(output) + labels = labels.int() + + if valid_label_mask is not None: + labels[~valid_label_mask] = -1 # Mark invalid labels as -1 + # https://lightning.ai/docs/torchmetrics/stable/classification/auroc#multilabelauroc + # -1 as we torchmetrics ignores -1 labels in multilabel metrics + # metric = MultilabelAUROC( + # num_labels=labels.shape[1], + # ignore_index=-1, + # ) + + return predictions, labels def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None: """ @@ -64,7 +77,10 @@ def __init__(self, config: dict, **kwargs) -> None: 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"]) + self.dropout = float(config.get("dropout", 0.0)) + self.norm = config.get("norm", None) + print("----------Dropout set to: ", self.dropout) + print("----------GNN Norm set to: ", self.norm) class GraphNetWrapper(GraphBaseNet, ABC): @@ -78,6 +94,7 @@ def __init__( config: dict, n_linear_layers: int, use_batch_norm: bool = False, + pretrained_checkpoint: str | None = None, **kwargs, ): """ @@ -86,8 +103,11 @@ def __init__( n_linear_layers (int): Number of linear layers. **kwargs: Additional arguments. """ + lin_hidden_dim = kwargs.pop("lin_hidden_dim", None) super().__init__(**kwargs) self.gnn = self._get_gnn(config) + if pretrained_checkpoint is not None: + self._load_pretrained_gnn(pretrained_checkpoint) gnn_out_dim = int(config["out_channels"]) self.activation = torch.nn.ELU self.lin_input_dim = self._get_lin_seq_input_dim( @@ -96,8 +116,9 @@ def __init__( self.use_batch_norm = use_batch_norm if self.use_batch_norm: self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim) + if lin_hidden_dim is None: + lin_hidden_dim = gnn_out_dim // 2 - 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, @@ -105,6 +126,68 @@ def __init__( out_dim=self.out_dim, ) + def _load_pretrained_gnn(self, checkpoint_path: str) -> None: + """Load encoder weights while leaving the task-specific head untouched.""" + checkpoint = torch.load( + checkpoint_path, weights_only=False, map_location=self.device + ) + state_dict = checkpoint.get("state_dict", checkpoint) + + if not isinstance(state_dict, dict): + raise ValueError( + f"Checkpoint {checkpoint_path!r} does not contain a state dictionary." + ) + + gnn_state = self.gnn.state_dict() + gnn_keys = set(gnn_state) + matching_keys = set(state_dict) & gnn_keys + if not matching_keys: + for prefix in ("gnn.", "model.gnn."): + prefixed_keys = { + key.removeprefix(prefix): value + for key, value in state_dict.items() + if key.startswith(prefix) + } + if set(prefixed_keys) & gnn_keys: + state_dict = prefixed_keys + matching_keys = set(state_dict) & gnn_keys + break + + if not matching_keys: + raise ValueError( + f"Checkpoint {checkpoint_path!r} contains no weights matching the GNN." + ) + + pretrained_gnn_state = { + key: value for key, value in state_dict.items() if key in gnn_keys + } + skipped_edge_weights = [] + # for key in list(pretrained_gnn_state): + # if ( + # key.endswith(".lin_edge.weight") + # and pretrained_gnn_state[key].shape != gnn_state[key].shape + # ): + # skipped_edge_weights.append(key) + # del pretrained_gnn_state[key] + + load_result = self.gnn.load_state_dict(pretrained_gnn_state, strict=False) + allowed_missing = set(skipped_edge_weights) + unexpected_keys = set(load_result.unexpected_keys) + missing_keys = set(load_result.missing_keys) - allowed_missing + if unexpected_keys or missing_keys: + raise RuntimeError( + "Pretrained GNN checkpoint is incompatible: " + f"missing={sorted(missing_keys)}, " + f"unexpected={sorted(unexpected_keys)}" + ) + + if skipped_edge_weights: + print( + "Skipped incompatible pretrained edge weights: " + + ", ".join(sorted(skipped_edge_weights)) + ) + print(f"Loaded pretrained GNN weights from {checkpoint_path!r} successfully.") + @abstractmethod def _get_gnn(self, config: dict) -> torch.nn.Module: """ @@ -183,3 +266,25 @@ def forward(self, batch: dict) -> torch.Tensor: if self.use_batch_norm: a = self.batch_norm(a) return self.lin_sequential(a) + + def configure_optimizers(self, **kwargs) -> torch.optim.Optimizer: + optimizer_kwargs = dict(self.optimizer_kwargs) + + default_lr = optimizer_kwargs.pop("lr", None) + gnn_lr = optimizer_kwargs.pop("lr_gnn", default_lr) + linear_lr = optimizer_kwargs.pop("lr_linear", default_lr) + + if gnn_lr is None or linear_lr is None: + raise ValueError("Set lr or both lr_gnn and lr_linear") + + linear_params = list(self.lin_sequential.parameters()) + if self.use_batch_norm: + linear_params.extend(self.batch_norm.parameters()) + + return torch.optim.Adamax( + [ + {"params": self.gnn.parameters(), "lr": gnn_lr}, + {"params": linear_params, "lr": linear_lr}, + ], + **optimizer_kwargs, + ) diff --git a/chebai_graph/models/architectures/gat.py b/chebai_graph/models/architectures/gat.py index 13dcabc..f3dae9f 100644 --- a/chebai_graph/models/architectures/gat.py +++ b/chebai_graph/models/architectures/gat.py @@ -42,6 +42,7 @@ def __init__(self, config: dict, **kwargs): heads=self.heads, v2=self.v2, act=self.activation, + norm=self.norm, **local_kwargs, ) diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py index f9e89a3..b7087d7 100644 --- a/chebai_graph/models/architectures/gine.py +++ b/chebai_graph/models/architectures/gine.py @@ -1,7 +1,7 @@ from typing import Any, Final from torch import Tensor -from torch.nn import ELU +from torch.nn import ELU, ReLU from torch_geometric import nn as tgnn from torch_geometric.data import Data as GraphData from torch_geometric.nn.conv import MessagePassing @@ -65,6 +65,55 @@ def init_conv( return tgnn.GINEConv(mlp, **kwargs) +class HiMolGINEModel(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: + + 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, 2 * 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. @@ -85,20 +134,35 @@ class GINEConvNetBase(GraphModelBase): """ def __init__(self, config: dict[str, Any], **kwargs: Any): + use_himol_model = bool(config.pop("use_himol_model", False)) + 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, - ) + if use_himol_model: + self.gine: BasicGNN = HiMolGINEModel( + 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, + norm=self.norm, + act=ReLU(), + ) + else: + 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, + norm=self.norm, + ) def forward(self, batch: dict[str, Any]) -> Tensor: """ diff --git a/chebai_graph/models/architectures/resgated.py b/chebai_graph/models/architectures/resgated.py index 521a853..a9e4663 100644 --- a/chebai_graph/models/architectures/resgated.py +++ b/chebai_graph/models/architectures/resgated.py @@ -65,6 +65,7 @@ def __init__(self, config: dict[str, Any], **kwargs: Any): num_layers=self.num_layers, edge_dim=self.edge_dim, act=self.activation, + norm=self.norm, ) def forward(self, batch: dict[str, Any]) -> Tensor: diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py index a1ee442..df5d558 100644 --- a/chebai_graph/models/pooling.py +++ b/chebai_graph/models/pooling.py @@ -103,24 +103,22 @@ def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ return gnn_out_dim + gnn_out_dim + gnn_out_dim - def forward(self, batch: dict) -> torch.Tensor: + def pool_representations( + self, batch: dict + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - Forward pass for pooling node embeddings. + Return atom-, functional-group-, and graph-level pooled 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. + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Pooled embeddings + in atom, functional-group, and graph order. """ graph_data = batch["features"][0] assert isinstance(graph_data, GraphData) @@ -145,7 +143,12 @@ def forward(self, batch: dict) -> torch.Tensor: atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - # Concatenate all - graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) + return atoms_vec, fg_nodes_vec, graph_node_vec + + def forward(self, batch: dict) -> torch.Tensor: + """Pool the three representation levels and apply the classifier.""" + representations = self.pool_representations(batch) + # Keep this order aligned with the classifier input layout. + graph_vector = torch.cat(representations, dim=1) return self.lin_sequential(graph_vector) diff --git a/chebai_graph/models/pretraining/__init__.py b/chebai_graph/models/pretraining/__init__.py new file mode 100644 index 0000000..5e7e638 --- /dev/null +++ b/chebai_graph/models/pretraining/__init__.py @@ -0,0 +1,3 @@ +from .msp import GATMSPretraining + +__all__ = ["GATMSPretraining"] diff --git a/chebai_graph/models/pretraining/msp.py b/chebai_graph/models/pretraining/msp.py new file mode 100644 index 0000000..1f252d6 --- /dev/null +++ b/chebai_graph/models/pretraining/msp.py @@ -0,0 +1,75 @@ +from typing import Any + +import torch +from chebai.preprocessing.structures import XYData +from torch_geometric.data import Data as GraphData + +from chebai_graph.models.architectures.base import GraphBaseNet +from chebai_graph.models.architectures.gat import GATGraphConvNetBase + + +class GATMSPretraining(GraphBaseNet): + def __init__( + self, + config: dict, + **kwargs, + ): + super().__init__(**kwargs) + self.gnn = GATGraphConvNetBase(config=config) + + def _process_batch(self, batch: XYData, batch_idx: int) -> dict[str, Any]: + """ + Processes the batch data. + + Args: + batch (XYData): The input batch of data. + batch_idx (int): The index of the current batch. + + Returns: + Dict[str, Any]: Processed batch data. + """ + return dict( + features=batch.x, + labels=self._process_labels_in_batch(batch), + model_kwargs=batch.additional_fields["model_kwargs"], + loss_kwargs=batch.additional_fields["loss_kwargs"], + idents=batch.additional_fields["idents"], + custom_metadata=batch.additional_fields["custom_metadata"], + ) + + 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) + return self.gnn(batch) + + def _process_for_loss( + self, + model_output: torch.Tensor, + labels: torch.Tensor, + data: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """ + Processes the data for loss computation. + + Args: + model_output (torch.Tensor): The model output. + labels (torch.Tensor): The true labels. + data (dict[str, Any]): The processed batch data. + + Returns: + Tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: Model output, labels, and data. + """ + data["valid_label_mask"] = data.get("loss_kwargs", dict()).get( + "valid_label_mask", None + ) + return model_output, labels, data diff --git a/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..b0f2214 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt @@ -0,0 +1,13 @@ +-1 +1 +0 +2 +3 +-4 +-2 +4 +-3 +5 +6 +7 +-5 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..273cc9b --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,160 @@ +NO_FG +graph_fg +alkyl +4_ammonium_ion +fluoro +chloro +ether +RING_6 +nitrile +primary_amine +secondary_amine +tertiary_carbon +RING_5 +alkene +RING_7 +tertiary_amine +ketone +hydroxyl +sulfide +sulfonyl +amide +ester +carboxyl +quaternary_carbon +secondary_ketimine +alkyne +alkene_carbon +nitro +trifluoromethyl +RING_4 +secondary_aldimine +primary_ketimine +acetal +carbonate_ester +azo +carbamate +bromo +RING_3 +phosphono +RING_40 +imide +RING_19 +nitroso +RING_8 +sulfhydryl +aldehyde +ketoxime +borono +difluoromethyl +thiolester +aldoxime +carboxylate +thiocyanate +isothiocyanate +phosphate +phosphodiester +iodo +sulfonic_acid +carbodithio +sulfinyl +azide +thioketone +RING_23 +trimethylsilyl +RING_21 +sulfonate_ester +RING_16 +RING_15 +RING_9 +RING_18 +RING_20 +RING_32 +RING_14 +RING_22 +nitrate +primary_aldimine +disulfide +RING_10 +hemiacetal +dichloromethyl +trichloromethyl +haloformyl +hydroperoxy +hemiketal +isonitrile +RING_24 +RING_11 +RING_12 +nitrosooxy +RING_33 +phosphino +sulfino +RING_17 +cyanate +RING_13 +carbothioic_O-acid +RING_36 +RING_72 +hydrazone +RING_31 +RING_25 +RING_28 +RING_35 +RING_34 +RING_42 +RING_38 +RING_29 +RING_26 +RING_51 +RING_75 +RING_27 +RING_45 +RING_60 +RING_30 +RING_43 +RING_41 +RING_56 +RING_81 +RING_39 +RING_37 +RING_47 +RING_49 +RING_44 +phosphoryl +RING_54 +RING_53 +RING_63 +carboxylic_anhydride +carbothioic_S-acid +silyl_ether +borinate +carbodithioic_acid +peroxy +tribromomethyl +chlorobromomethyl +amidine +dibromomethyl +isocyanate +bromodichloromethyl +ketal +boronate +borino +chloroiodomethyl +bromoiodomethyl +thionoester +fluorochloromethyl +RING_68 +bromodifluoromethyl +RING_58 +difluorochloromethyl +thial +RING_52 +RING_48 +RING_62 +RING_164 +RING_71 +RING_46 +orthoester +RING_55 +triiodomethyl diff --git a/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..dfc8e85 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +UNSPECIFIED +SP3 +S +SP3D2 +SP3D +SP +SP2 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..5ad8538 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,8 @@ +0 +3 +2 +4 +1 +5 +6 +8 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..a4766ec --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt @@ -0,0 +1,119 @@ +0 +13 +9 +11 +7 +47 +6 +4 +19 +50 +51 +17 +12 +8 +20 +27 +38 +81 +3 +30 +29 +39 +55 +25 +56 +70 +88 +40 +37 +26 +82 +78 +80 +16 +33 +1 +35 +15 +5 +44 +53 +85 +79 +24 +14 +87 +32 +2 +48 +63 +76 +92 +23 +64 +57 +34 +22 +28 +74 +42 +46 +52 +83 +62 +49 +58 +71 +65 +67 +77 +31 +59 +75 +45 +36 +54 +18 +10 +86 +43 +84 +118 +90 +73 +113 +114 +115 +116 +117 +112 +72 +104 +41 +105 +106 +107 +108 +109 +110 +111 +21 +103 +60 +61 +66 +68 +69 +89 +91 +93 +94 +95 +96 +97 +98 +99 +100 +101 +102 diff --git a/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..b389485 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +atom_fg_lvl +to_graphNode_lvl +within_atoms_lvl +within_fg_lvl diff --git a/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f36bdf7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt @@ -0,0 +1,6 @@ +DATIVE +SINGLE +AROMATIC +TRIPLE +DOUBLE +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..22da8da --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,13 @@ +0 +1 +2 +4 +5 +3 +6 +8 +7 +10 +12 +11 +9 diff --git a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..b621c2a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,44 @@ +RING_6 +difluoromethyl +RING_5 +ether +alkyne +alkyl +primary_amine +fluoro +amide +graph_fg +NO_FG +alkene +tertiary_carbon +alkene_carbon +nitrile +trifluoromethyl +ketone +chloro +quaternary_carbon +sulfonyl +hydroxyl +bromo +secondary_amine +ester +aldoxime +sulfide +iodo +tertiary_amine +RING_4 +secondary_ketimine +ketoxime +RING_3 +carboxylate +acetal +sulfinyl +RING_7 +RING_16 +RING_15 +nitro +primary_ketimine +sulfonate_ester +RING_8 +RING_14 +RING_13 diff --git a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..c4b216a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +3 +2 diff --git a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..6da4980 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt @@ -0,0 +1,8 @@ +6 +7 +8 +9 +17 +16 +35 +53 diff --git a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f1a2ce5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +TRIPLE +DOUBLE diff --git a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..cf286d6 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +2 +4 +1 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7b6a00a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +-1 +2 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..d2ee211 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,80 @@ +alkyl +ester +chloro +RING_6 +tertiary_amine +graph_fg +tertiary_carbon +primary_amine +dichloromethyl +sulfonyl +ketone +secondary_amine +hydroxyl +ether +quaternary_carbon +amide +NO_FG +carboxyl +nitro +alkene +alkene_carbon +secondary_ketimine +secondary_aldimine +carboxylate +iodo +aldoxime +bromo +alkyne +fluoro +carbamate +trifluoromethyl +sulfide +trichloromethyl +hemiacetal +primary_ketimine +azide +4_ammonium_ion +sulfonate_ester +ketoxime +nitroso +sulfhydryl +difluoromethyl +fluorochloromethyl +chlorobromomethyl +phosphono +carbodithio +disulfide +difluorochloromethyl +carbonate_ester +phosphate +RING_5 +thiolester +RING_7 +nitrile +phosphoryl +aldehyde +sulfinyl +thioketone +acetal +RING_4 +sulfonic_acid +imide +RING_14 +RING_24 +amidine +RING_3 +RING_9 +RING_19 +RING_16 +RING_8 +isonitrile +RING_23 +RING_12 +RING_21 +RING_17 +RING_30 +RING_33 +RING_36 +azo +RING_15 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..0e4a9aa --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,4 @@ +SP3 +SP2 +S +SP diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..35f5032 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +0 +1 +3 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9500570 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt @@ -0,0 +1,13 @@ +6 +17 +8 +7 +16 +11 +53 +35 +9 +1 +15 +5 +20 diff --git a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt new file mode 100644 index 0000000..a134437 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +DOUBLE +SINGLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..5d07eee --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,5 @@ +2 +4 +1 +3 +0 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..374de68 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..36bc343 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,84 @@ +alkyl +chloro +RING_6 +ether +secondary_amine +fluoro +graph_fg +tertiary_carbon +hydroxyl +ketone +sulfonyl +carboxylate +RING_5 +NO_FG +amide +sulfide +RING_4 +quaternary_carbon +alkyne +alkene +primary_amine +carboxyl +trifluoromethyl +ester +alkene_carbon +RING_20 +sulfhydryl +nitrile +RING_3 +RING_29 +RING_7 +4_ammonium_ion +acetal +RING_9 +iodo +RING_38 +sulfinyl +difluoromethyl +imide +RING_16 +tertiary_amine +nitro +RING_8 +secondary_ketimine +nitrate +sulfonic_acid +carbamate +secondary_aldimine +bromo +thioketone +RING_14 +azo +RING_33 +trichloromethyl +RING_15 +ketoxime +phosphodiester +RING_24 +RING_22 +aldehyde +RING_18 +RING_21 +nitroso +dichloromethyl +carbonate_ester +silyl_ether +thiolester +borono +aldoxime +sulfonate_ester +RING_23 +RING_19 +RING_25 +carbodithio +disulfide +RING_36 +RING_31 +fluorochloromethyl +RING_12 +phosphate +azide +chlorobromomethyl +bromodifluoromethyl +phosphino diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..a96b9b3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +UNSPECIFIED +SP3D diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..1127527 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +0 +2 +1 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..2c8a324 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt @@ -0,0 +1,30 @@ +6 +17 +7 +8 +9 +16 +15 +53 +80 +1 +13 +35 +79 +0 +27 +22 +78 +33 +34 +30 +14 +5 +29 +24 +81 +43 +83 +25 +20 +26 diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt new file mode 100644 index 0000000..bdc90d7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..3c2db51 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +3 +2 +4 +0 +6 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..22aacd7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,9 @@ +0 +1 +-1 +3 +-4 +-3 +-2 +2 +-5 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..e5b33d4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,122 @@ +alkyl +RING_6 +chloro +NO_FG +secondary_amine +graph_fg +carboxyl +primary_amine +sulfonic_acid +nitro +ketone +hydroxyl +sulfonyl +ester +nitroso +sulfhydryl +tertiary_carbon +tertiary_amine +dichloromethyl +nitrile +primary_ketimine +alkyne +alkene +ether +iodo +sulfide +secondary_aldimine +phosphoryl +thioketone +alkene_carbon +secondary_ketimine +nitrate +4_ammonium_ion +ketoxime +amide +amidine +trichloromethyl +fluoro +quaternary_carbon +hydrazone +disulfide +thiocyanate +carbamate +trifluoromethyl +phosphate +sulfonate_ester +sulfinyl +azo +bromo +imide +carbodithioic_acid +carbodithio +sulfino +silyl_ether +aldehyde +phosphono +azide +aldoxime +phosphino +trimethylsilyl +thionoester +thiolester +isonitrile +acetal +difluorochloromethyl +dibromomethyl +ketal +diiodomethyl +borono +isothiocyanate +difluoromethyl +tribromomethyl +carbonate_ester +hemiacetal +phosphodiester +carboxylate +orthoester +isocyanate +primary_aldimine +hemiketal +borinate +RING_5 +cyanate +RING_4 +RING_24 +RING_3 +thial +RING_7 +RING_12 +RING_9 +RING_16 +hydroperoxy +RING_15 +RING_18 +RING_8 +RING_35 +RING_10 +RING_19 +RING_17 +RING_29 +RING_23 +RING_27 +RING_11 +RING_14 +RING_20 +RING_33 +RING_13 +RING_31 +RING_21 +RING_28 +RING_22 +RING_30 +RING_34 +RING_54 +RING_38 +RING_36 +RING_26 +RING_32 +carboxylic_anhydride +RING_25 +carbothioic_S-acid +RING_40 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..95c9cba --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt @@ -0,0 +1,54 @@ +6 +17 +8 +7 +15 +16 +33 +53 +9 +35 +14 +11 +34 +79 +50 +80 +32 +51 +82 +5 +19 +3 +83 +29 +67 +44 +65 +74 +27 +55 +25 +26 +75 +42 +78 +47 +24 +46 +89 +20 +31 +12 +30 +28 +77 +1 +40 +81 +45 +22 +52 +64 +92 +23 diff --git a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..58c0a5c --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..88ddb48 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,11 @@ +1 +2 +3 +4 +0 +5 +6 +9 +7 +8 +10 diff --git a/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..c55dc92 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,68 @@ +primary_amine +RING_6 +ketone +carboxyl +alkyl +tertiary_carbon +secondary_amine +graph_fg +hydroxyl +alkene +thioketone +bromo +fluoro +ester +ether +tertiary_amine +NO_FG +chloro +sulfonyl +difluoromethyl +quaternary_carbon +nitro +amide +alkyne +trifluoromethyl +aldoxime +alkene_carbon +nitrile +sulfide +secondary_aldimine +ketoxime +carbodithio +secondary_ketimine +thionoester +thiolester +sulfinyl +carbamate +thiocyanate +aldehyde +sulfonate_ester +amidine +trichloromethyl +sulfhydryl +sulfonic_acid +RING_5 +nitroso +acetal +carbonate_ester +imide +azo +primary_ketimine +dichloromethyl +RING_7 +RING_3 +RING_8 +hydrazone +RING_4 +RING_15 +RING_12 +azide +RING_13 +RING_18 +difluorochloromethyl +RING_14 +RING_9 +RING_11 +disulfide +4_ammonium_ion diff --git a/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..52db5ae --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +1 +0 +3 diff --git a/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..c36ddc7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt @@ -0,0 +1,7 @@ +7 +6 +8 +35 +9 +16 +17 diff --git a/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..94ebaf9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +1 +2 +3 +4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7897486 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt @@ -0,0 +1,8 @@ +0 +1 +-1 +3 +2 +4 +5 +6 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..89e00f3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,127 @@ +alkyl +RING_6 +sulfide +RING_5 +ketone +trifluoromethyl +secondary_amine +graph_fg +NO_FG +amide +sulfonyl +ether +quaternary_carbon +tertiary_carbon +imide +carboxyl +nitrile +ester +primary_amine +alkyne +fluoro +tertiary_amine +nitro +RING_7 +chloro +hydroxyl +RING_3 +alkene +secondary_aldimine +bromo +thioketone +sulfinyl +RING_4 +secondary_ketimine +RING_12 +primary_ketimine +alkene_carbon +carbamate +aldehyde +phosphoryl +trimethylsilyl +carboxylate +ketoxime +RING_23 +difluoromethyl +sulfonate_ester +RING_14 +iodo +thiocyanate +RING_8 +disulfide +carbonate_ester +RING_19 +sulfhydryl +nitroso +acetal +aldoxime +borono +RING_16 +phosphono +RING_9 +RING_15 +hydrazone +RING_11 +4_ammonium_ion +thiolester +trichloromethyl +carbodithio +azo +azide +sulfonic_acid +dichloromethyl +phosphodiester +borinate +RING_13 +RING_18 +thial +RING_21 +RING_10 +isocyanate +isonitrile +RING_22 +RING_36 +RING_20 +phosphate +amidine +dibromomethyl +RING_17 +peroxy +RING_40 +silyl_ether +thionoester +nitrate +haloformyl +hemiacetal +RING_25 +isothiocyanate +difluorochloromethyl +RING_24 +nitrosooxy +RING_32 +carbodithioic_acid +hydroperoxy +sulfino +RING_26 +tribromomethyl +RING_30 +RING_29 +diiodomethyl +phosphino +chlorobromomethyl +RING_33 +RING_38 +RING_34 +bromodifluoromethyl +carboxylic_anhydride +bromodichloromethyl +RING_27 +dichlorosilane +RING_42 +ketal +RING_28 +RING_56 +RING_31 +RING_35 +fluorochloromethyl +RING_54 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..fb3232b --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +SP3D +S +UNSPECIFIED +SP3D2 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..cad2399 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +1 +0 +2 +6 +4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..166ffe2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt @@ -0,0 +1,44 @@ +6 +16 +7 +8 +9 +17 +35 +15 +14 +53 +5 +29 +79 +38 +50 +30 +33 +34 +48 +51 +1 +26 +82 +27 +78 +80 +3 +28 +22 +56 +44 +11 +13 +46 +25 +83 +47 +60 +24 +64 +32 +49 +74 +31 diff --git a/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..2efbda4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +5 diff --git a/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..6f89a5d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..6fbc2c2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,93 @@ +alkyl +RING_6 +carboxyl +alkene +fluoro +sulfinyl +RING_5 +graph_fg +tertiary_carbon +secondary_amine +ketone +primary_amine +primary_ketimine +NO_FG +sulfide +hydroxyl +sulfonic_acid +bromo +4_ammonium_ion +chloro +quaternary_carbon +ether +carboxylate +tertiary_amine +amide +RING_3 +trifluoromethyl +ester +nitro +sulfonyl +iodo +secondary_ketimine +RING_7 +aldehyde +phosphono +alkyne +nitrile +carbamate +alkene_carbon +carbonate_ester +RING_4 +RING_18 +RING_24 +RING_14 +imide +sulfonate_ester +RING_81 +RING_20 +RING_8 +nitroso +RING_29 +RING_16 +RING_12 +RING_21 +sulfhydryl +thiolester +difluoromethyl +thioketone +RING_25 +RING_78 +RING_68 +RING_77 +nitrate +azo +RING_36 +phosphate +ketoxime +secondary_aldimine +dichloromethyl +carboxylic_anhydride +RING_33 +borono +RING_15 +diiodomethyl +RING_23 +carbodithio +disulfide +acetal +RING_40 +phosphodiester +fluorochloromethyl +RING_22 +RING_10 +RING_9 +azide +phosphino +RING_31 +RING_28 +RING_53 +RING_56 +RING_41 +RING_43 +RING_35 diff --git a/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..e9958eb --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,6 @@ +SP3 +SP2 +S +SP3D +SP +SP3D2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..384eb37 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt @@ -0,0 +1,40 @@ +6 +8 +9 +16 +7 +35 +17 +11 +64 +15 +24 +53 +47 +49 +1 +25 +29 +20 +62 +3 +19 +78 +39 +43 +31 +26 +5 +12 +38 +79 +33 +30 +56 +98 +34 +32 +27 +88 +81 +57 diff --git a/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..dfb5c6d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..a8bd8c3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..5aa0016 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,115 @@ +alkyl +hemiacetal +ether +graph_fg +RING_18 +NO_FG +RING_6 +ester +tertiary_carbon +secondary_amine +ketone +carboxyl +RING_5 +amide +primary_amine +RING_4 +carboxylate +alkene +nitro +imide +4_ammonium_ion +bromo +chloro +hydroxyl +sulfonyl +fluoro +quaternary_carbon +alkene_carbon +RING_3 +alkyne +secondary_ketimine +nitrile +isocyanate +azo +tertiary_amine +trifluoromethyl +aldehyde +thiolester +sulfide +nitroso +phosphate +phosphoryl +RING_7 +iodo +acetal +carbodithio +carbonate_ester +RING_21 +difluoromethyl +disulfide +RING_25 +nitrate +carbamate +secondary_aldimine +sulfonic_acid +haloformyl +amidine +sulfhydryl +phosphono +dibromomethyl +RING_23 +trichloromethyl +sulfonate_ester +RING_40 +RING_24 +RING_15 +primary_ketimine +RING_29 +thioketone +aldoxime +RING_16 +nitrosooxy +sulfino +RING_36 +peroxy +RING_14 +ketoxime +thiocyanate +RING_8 +RING_17 +sulfinyl +dichloromethyl +primary_aldimine +tribromomethyl +RING_19 +RING_9 +silyl_ether +phosphodiester +RING_33 +azide +RING_27 +RING_26 +RING_20 +phosphino +RING_12 +trimethylsilyl +carboxylic_anhydride +orthoester +RING_10 +thionoester +isothiocyanate +borono +chlorobromomethyl +dichlorosilane +hydroperoxy +RING_35 +RING_30 +RING_28 +fluorochloromethyl +RING_13 +bromodichloromethyl +ketal +difluorochloromethyl +bromodifluoromethyl +RING_11 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..5a694e8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +2 +1 +0 +4 +6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..64e481b --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt @@ -0,0 +1,50 @@ +6 +8 +7 +16 +15 +35 +17 +9 +56 +30 +26 +80 +53 +28 +5 +51 +14 +47 +11 +78 +40 +79 +4 +29 +20 +34 +81 +50 +82 +23 +32 +38 +48 +46 +60 +27 +83 +64 +25 +19 +33 +22 +3 +12 +42 +66 +24 +1 +49 +70 diff --git a/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..21b19a8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +1 +2 +3 +4 +0 +5 +6 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..c857ab5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..500ebf5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,117 @@ +ester +RING_6 +sulfonyl +alkyl +NO_FG +4_ammonium_ion +graph_fg +sulfide +primary_amine +carboxyl +chloro +quaternary_carbon +hydroxyl +secondary_amine +fluoro +ketone +tertiary_carbon +ether +alkene +nitro +tertiary_amine +aldehyde +RING_5 +alkyne +nitrile +RING_7 +alkene_carbon +carbonate_ester +trifluoromethyl +imide +sulfonic_acid +primary_ketimine +disulfide +sulfonate_ester +amide +nitroso +iodo +RING_3 +carboxylate +RING_8 +RING_20 +acetal +bromo +carbamate +RING_12 +ketoxime +RING_18 +silyl_ether +RING_4 +dichloromethyl +sulfinyl +aldoxime +thioketone +secondary_ketimine +isothiocyanate +isocyanate +RING_14 +sulfhydryl +phosphono +carbodithioic_acid +secondary_aldimine +haloformyl +phosphoryl +azo +nitrate +RING_35 +thiocyanate +hemiacetal +thiolester +difluorochloromethyl +dibromomethyl +difluoromethyl +phosphino +RING_29 +chlorobromomethyl +RING_9 +tribromomethyl +trichloromethyl +RING_28 +RING_16 +nitrosooxy +bromodichloromethyl +carbodithio +phosphodiester +RING_23 +phosphate +carboxylic_anhydride +RING_13 +RING_24 +RING_30 +RING_11 +RING_17 +azide +RING_40 +primary_aldimine +RING_15 +peroxy +RING_33 +trimethylsilyl +sulfino +ketal +RING_10 +thionoester +borono +amidine +orthoester +RING_27 +RING_26 +RING_38 +hydroperoxy +RING_21 +RING_19 +dichlorosilane +diiodomethyl +fluorochloromethyl +bromodifluoromethyl +RING_36 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..5d039f5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP2 +SP3 +SP +S +UNSPECIFIED +SP3D2 +SP3D diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..ed1daa5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +0 +1 +3 +2 +4 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9c59d72 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt @@ -0,0 +1,53 @@ +8 +6 +7 +16 +17 +9 +50 +19 +22 +14 +11 +53 +35 +33 +49 +78 +15 +51 +12 +3 +83 +28 +26 +13 +20 +47 +29 +5 +1 +30 +38 +80 +27 +34 +79 +56 +81 +24 +23 +21 +64 +70 +60 +46 +32 +4 +82 +66 +25 +48 +63 +40 +42 diff --git a/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..a61fff5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +2 +1 +3 +4 +0 +6 +5 diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..89a8b18 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +-1 +1 diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..bc1b07e --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,72 @@ +alkyl +secondary_amine +trifluoromethyl +quaternary_carbon +amide +sulfide +graph_fg +RING_6 +RING_5 +ketone +RING_3 +tertiary_carbon +fluoro +ether +imide +primary_amine +NO_FG +RING_4 +alkene_carbon +alkene +carboxyl +hydroxyl +sulfonyl +ester +tertiary_amine +alkyne +nitrile +secondary_aldimine +chloro +RING_7 +nitro +difluoromethyl +thioketone +secondary_ketimine +bromo +sulfinyl +primary_ketimine +nitroso +phosphate +RING_8 +aldehyde +sulfonate_ester +trimethylsilyl +carbamate +sulfhydryl +acetal +sulfonic_acid +RING_10 +thiolester +phosphoryl +iodo +carbonate_ester +nitrate +4_ammonium_ion +azo +amidine +bromodifluoromethyl +aldoxime +RING_9 +ketal +ketoxime +difluorochloromethyl +trichloromethyl +hydrazone +phosphono +RING_12 +carbodithio +RING_11 +thionoester +azide +RING_16 +phosphino diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..a4a9930 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP3 +SP2 +SP diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..3ed60bd --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +1 +0 +2 diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..81eeb68 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt @@ -0,0 +1,11 @@ +6 +7 +9 +8 +16 +17 +35 +15 +14 +5 +53 diff --git a/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..207d1d2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +1 +2 +4 +3 diff --git a/chebai_graph/preprocessing/collate.py b/chebai_graph/preprocessing/collate.py index 53bdde5..1f5dcca 100644 --- a/chebai_graph/preprocessing/collate.py +++ b/chebai_graph/preprocessing/collate.py @@ -15,22 +15,46 @@ def __call__(self, data): # Unpack labels and optional identifiers y, idents = zip(*((d["labels"], d.get("ident")) for d in data)) - # Replace labels with `y` inside graph features and collect them + valid_label_mask = self._get_valid_label_mask(y) + + # Keep labels in XYGraphData.y instead of graph attributes. Mixed labelled + # and unlabelled batches otherwise make PyG collate a partial `y` key. merged_data = [] for row in data: - row["features"].y = row["labels"] - merged_data.append(row["features"]) + # Dataset samples may be cached and reused across epochs. Collation + # removes temporary attributes below, so never mutate the source graph. + if isinstance(row["features"], tuple): + # If tuple, then metadata is in second element. Move it to the first element for consistency. + row["features"][0]["custom_metadata"] = row["features"][1] + features = row["features"][0].clone() + else: + features = row["features"].clone() + if "y" in features: + del features.y + merged_data.append(features) + custom_metadata = [] # Add empty edge_attr for graphs with no edges to prevent PyG errors for mdata in merged_data: + graph_metadata = [] for store in mdata.stores: if "edge_attr" not in store: store["edge_attr"] = torch.tensor([]) + if "custom_metadata" in store: + graph_metadata.append(store.pop("custom_metadata")) + else: + pass + custom_metadata.append(graph_metadata) # Ensure all attributes are float tensors to prevent torch.cat dtype issues for attr in merged_data[0].keys(): for data in merged_data: for store in data.stores: + if attr not in store: + continue + if store[attr] is None: + del store[attr] + continue # Im not sure why the following conversion is needed, but it solves this error: # packages/torch_geometric/data/collate.py", line 177, in _collate # value = torch.cat(values, dim=cat_dim or 0, out=out) @@ -69,10 +93,12 @@ def __call__(self, data): x[0].x = x[0].x.to(dtype=torch.int64) # x is a Tuple[BaseData, Mapping, Mapping] + loss_kwargs["valid_label_mask"] = valid_label_mask return XYGraphData( x, y, idents=idents, model_kwargs={}, loss_kwargs=loss_kwargs, + custom_metadata=custom_metadata, ) diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 9ee9856..d0b7b72 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -1,14 +1,25 @@ from .chebi import ( + ChEBI25GraphProperties, ChEBI50_Atom_WGNOnly_GraphProp, ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE, ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE, ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE, ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE, + ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType, + ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel, + ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes, ChEBI50_NFGE_NGN_GraphProp, ChEBI50_NFGE_WGN_GraphProp, ChEBI50_StaticGNI, ChEBI50_WFGE_NGN_GraphProp, ChEBI50_WFGE_WGN_AsPerNodeType, + ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel, ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly, ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly, ChEBI50_WFGE_WGN_ForAtomLevelOnly, @@ -18,41 +29,66 @@ ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes, ChEBI50_WFGE_WGN_GraphProp, ChEBI50GraphData, - ChEBI25GraphProperties, ChEBI50GraphProperties, ChEBI100GraphProperties, ) from .molecule_net_classification import ( + BACE_GNWA_FGWA_WFGE_AsPerNodeType, BACE_WFGE_WGN_AsPerNodeType, + BBBP_GNWA_FGWA_WFGE_AsPerNodeType, BBBP_WFGE_WGN_AsPerNodeType, + ClinTox_GNWA_FGWA_WFGE_AsPerNodeType, ClinTox_WFGE_WGN_AsPerNodeType, + HIV_GNWA_FGWA_WFGE_AsPerNodeType, HIV_WFGE_WGN_AsPerNodeType, + MUV_GNWA_FGWA_WFGE_AsPerNodeType, MUV_WFGE_WGN_AsPerNodeType, + PCBA_GNWA_FGWA_WFGE_AsPerNodeType, PCBA_WFGE_WGN_AsPerNodeType, + SIDER_GNWA_FGWA_WFGE_AsPerNodeType, SIDER_WFGE_WGN_AsPerNodeType, + Tox21_Atom_WGNOnly_GraphProp, + Tox21_GN_WithAllNodes_FG_WithAtoms_FGE, + Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE, + Tox21_GN_WithAtoms_FG_WithAtoms_FGE, + Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE, + Tox21_GNWA_FGWA_WFGE_AsPerNodeType, + Tox21_NFGE_NGN_GraphProp, + Tox21_NFGE_WGN_GraphProp, + Tox21_WFGE_NGN_GraphProp, Tox21_WFGE_WGN_AsPerNodeType, + Tox21_WFGE_WGN_GraphProp, + Tox21GraphProperties, + ToxCast_Atom_WGNOnly_GraphProp, + ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE, + ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE, + ToxCast_GN_WithAtoms_FG_WithAtoms_FGE, + ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE, + ToxCast_GNWA_FGWA_WFGE_AsPerNodeType, + ToxCast_NFGE_NGN_GraphProp, + ToxCast_NFGE_WGN_GraphProp, + ToxCast_WFGE_NGN_GraphProp, ToxCast_WFGE_WGN_AsPerNodeType, + ToxCast_WFGE_WGN_GraphProp, + ToxCastGraphProperties, ) +from .pretrain_zinc15 import Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType from .pubchem import PubChemGraphProperties __all__ = [ - "ChEBI50GraphFGAugmentorReader", - "ChEBI25GraphProperties", - "ChEBI50GraphProperties", - "ChEBI100GraphProperties", - "ChEBI50GraphData", - "PubChemGraphProperties", - "ChEBI50_Atom_WGNOnly_GraphProp", + # ---------- CHEBI Augmentation ablation classes -------- + "ChEBI50GraphProperties", # -- baseline class for chebi "ChEBI50_NFGE_NGN_GraphProp", "ChEBI50_NFGE_WGN_GraphProp", "ChEBI50_WFGE_NGN_GraphProp", "ChEBI50_WFGE_WGN_GraphProp", - "ChEBI50_StaticGNI", "ChEBI50_WFGE_WGN_AsPerNodeType", "ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE", "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE", "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE", + "ChEBI50_Atom_WGNOnly_GraphProp", + # ----- CHEBI FGN+E+FGE Properties ablation classes ----- "ChEBI50_WFGE_WGN_ForAtomLevelOnly", "ChEBI50_WFGE_WGN_ForFGLevelOnly", "ChEBI50_WFGE_WGN_ForGraphNodeOnly", @@ -60,6 +96,8 @@ "ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly", "ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly", "ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes", + "ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel", + # ----- MoleculeNet Dataset FGN+E+FGE Final classes ----- "BACE_WFGE_WGN_AsPerNodeType", "BBBP_WFGE_WGN_AsPerNodeType", "ClinTox_WFGE_WGN_AsPerNodeType", @@ -69,4 +107,53 @@ "Tox21_WFGE_WGN_AsPerNodeType", "ToxCast_WFGE_WGN_AsPerNodeType", "PCBA_WFGE_WGN_AsPerNodeType", + # ----- MoleculeNet Dataset GNWA+FGWA+WFGE Final classes ----- + "Tox21_GNWA_FGWA_WFGE_AsPerNodeType", + "ToxCast_GNWA_FGWA_WFGE_AsPerNodeType", + "BACE_GNWA_FGWA_WFGE_AsPerNodeType", + "BBBP_GNWA_FGWA_WFGE_AsPerNodeType", + "ClinTox_GNWA_FGWA_WFGE_AsPerNodeType", + "HIV_GNWA_FGWA_WFGE_AsPerNodeType", + "MUV_GNWA_FGWA_WFGE_AsPerNodeType", + "PCBA_GNWA_FGWA_WFGE_AsPerNodeType", + "SIDER_GNWA_FGWA_WFGE_AsPerNodeType", + # ----- CHEBI GNWA_FGWA_WFGE Properties ablation classes ----- + "ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType", + "ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes", + # ----- Tox21 Augmentation Ablation classes ----- + "Tox21_Atom_WGNOnly_GraphProp", + "Tox21_GN_WithAllNodes_FG_WithAtoms_FGE", + "Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE", + "Tox21_GN_WithAtoms_FG_WithAtoms_FGE", + "Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE", + "Tox21_NFGE_NGN_GraphProp", + "Tox21_NFGE_WGN_GraphProp", + "Tox21_WFGE_NGN_GraphProp", + "Tox21_WFGE_WGN_GraphProp", + "Tox21GraphProperties", # -- baseline class for Tox21 + # ----- ToxCast Augmentation Ablation classes ----- + "ToxCast_Atom_WGNOnly_GraphProp", + "ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE", + "ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE", + "ToxCast_GN_WithAtoms_FG_WithAtoms_FGE", + "ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE", + "ToxCast_NFGE_NGN_GraphProp", + "ToxCast_NFGE_WGN_GraphProp", + "ToxCast_WFGE_NGN_GraphProp", + "ToxCast_WFGE_WGN_GraphProp", + "ToxCastGraphProperties", # -- baseline class for ToxCast + # ------ Others ----- + "ChEBI50_StaticGNI", + "ChEBI25GraphProperties", + "ChEBI100GraphProperties", + "ChEBI50GraphData", + "PubChemGraphProperties", + "Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType", ] diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py index 9bf6cdc..20ce564 100644 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -177,3 +177,10 @@ def _fill_node_tensor_with_molecule_type_property( is_graph_node: torch.Tensor, ) -> torch.Tensor: return node_tensor + + +class GraphPropAllPropsExceptNodeLevel(GraphPropAsPerNodeType): + def _fill_node_tensor_with_all_node_type_property( + self, node_tensor: torch.Tensor, property_values: torch.Tensor, offset: int + ) -> torch.Tensor: + return node_tensor diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index d01b53b..70d62ca 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -50,7 +50,9 @@ def __init__( super().__init__(**kwargs) # atom_properties and bond_properties are given as lists containing class_paths if properties is not None: - properties = [resolve_property(prop, self.data_type) for prop in properties] + properties = [ + resolve_property(prop, self.tokens_directory) for prop in properties + ] properties = self._sort_properties(properties) else: properties = [] @@ -462,6 +464,7 @@ def load_processed_data( base_data = super().load_processed_data(kind, filename) base_df = pd.DataFrame(base_data) base_df["ident"] = base_df["ident"].astype(str) + len_base_df_before_merge = len(base_df) props_categories = { "AllNodeTypeProperties": [], "FGNodeTypeProperties": [], @@ -524,11 +527,21 @@ def load_processed_data( ) property_df = pd.DataFrame(property_data) + assert len(property_df) == len(property_df["ident"].unique()), ( + f"Duplicate entries found in property {property.name} data. " + f"Unique entries {len(property_df['ident'].unique())}. " + f"Total entries {len(property_df)}. " + ) property_df["ident"] = property_df["ident"].astype(str) property_df.rename( columns={property.name: f"{property.name}"}, inplace=True ) + base_df = base_df.merge(property_df, on="ident", how="left") + assert len(base_df) == len_base_df_before_merge, ( + f"Length mismatch after merging property {property.name}. " + f"Expected {len_base_df_before_merge}, got {len(base_df)}" + ) base_df["features"] = base_df.apply( lambda row: self._merge_props_into_base( @@ -560,9 +573,10 @@ def _merge_props_into_base( 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) + # include custom metadata returned data from _read_data (e.g. augmented molecule dict) + custom_metadata = geom_data[1] + geom_data = geom_data[0] + geom_data["custom_metadata"] = custom_metadata assert isinstance(geom_data, GeomData) is_atom_node = geom_data.is_atom_node @@ -632,6 +646,7 @@ def _merge_props_into_base( is_atom_node=is_atom_node, is_fg_node=is_fg_node, is_graph_node=is_graph_node, + custom_metadata=getattr(geom_data, "custom_metadata", None), ) def _build_node_property_tensor( diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index ac58a84..7a68dc6 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -25,6 +25,7 @@ from .augmentation_base import ( AugGraphPropMixIn_NoGraphNode, AugGraphPropMixIn_WithGraphNode, + GraphPropAllPropsExceptNodeLevel, GraphPropForAtomAndFGLevelOnly, GraphPropForAtomLevelAndGraphNodeOnly, GraphPropForAtomLevelOnly, @@ -45,648 +46,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: property.encoder.compress( - torch.cat(feat) - ), - "ident": id, - } - for feat, id in zip(encoded_values, idents) - if feat is not None - ], - self.get_property_path(property), - ) - property.on_finish() - - @property - def processed_properties_dir(self) -> str: - return os.path.join(self.processed_dir, "properties") - - def get_property_path(self, property: MolecularProperty) -> str: - """ - Construct the cache path for a given molecular property. - - Args: - property: Instance of a MolecularProperty. - - Returns: - Path to the cached property file. - """ - return os.path.join( - self.processed_properties_dir, - f"{property.name}_{property.encoder.name}.pt", - ) - - def _after_setup(self, **kwargs) -> None: - """ - Finalize setup after ensuring properties are processed. - - Args: - **kwargs: Additional keyword arguments passed to superclass. - - Returns: - None - """ - self._setup_properties() - super()._after_setup(**kwargs) - - def _preprocess_molecule_for_pred( - self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None - ) -> Optional[dict]: - """Preprocess prediction data.""" - # Add dummy labels because the collate function requires them. - # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, - # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. - result = self.reader.to_data( - {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} - ) - # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object - if isinstance(result["features"], tuple): - result["features"], raw_data = result["features"] - if result is None or result["features"] is None: - return None - for property in self.properties: - property.encoder.eval = True - property_value = self.reader.read_property(raw_data, property) - if property_value is None or len(property_value) == 0: - encoded_value = None - else: - encoded_value = torch.stack( - [property.encoder.encode(v) for v in property_value] - ) - if len(encoded_value.shape) == 3: - encoded_value = encoded_value.squeeze(0) - result[property.name] = encoded_value - - result["features"] = self._prediction_merge_props_into_base_wrapper( - result, model_hparams - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - result["features"] = self.transform(result["features"]) - - return result - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - return self._merge_props_into_base(row) - - -class GraphPropertiesMixIn(DataPropertiesSetter, ABC): - def __init__( - self, - properties=None, - transform=None, - pad_node_features: int = None, - 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) - base_df["ident"] = base_df["ident"].astype(str) - - for property in self.properties: - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - property_df["ident"] = property_df["ident"].astype(str) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base(row), axis=1 - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - prop_lengths = [ - (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties - ] - - # -------------------------- Count total node properties - n_node_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, AtomProperty) - ) - - in_channels_str = "" - if self.pad_node_features: - n_node_properties += self.pad_node_features - in_channels_str += f" (with {self.pad_node_features} padded random values from {self.distribution} distribution)" - - in_channels_str = f"in_channels: {n_node_properties}" + in_channels_str - - # -------------------------- Count total edge properties - n_edge_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, BondProperty) - ) - edge_dim_str = "" - if self.pad_edge_features: - n_edge_properties += self.pad_edge_features - edge_dim_str += f" (with {self.pad_edge_features} padded random values from {self.distribution} distribution)" - - edge_dim_str = f"edge_dim: {n_edge_properties}" + edge_dim_str - - rank_zero_info( - f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" - f"Use following values for given parameters for model configuration: \n\t" - f"{in_channels_str} \n\t" - f"{edge_dim_str} \n\t" - 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) - base_df["ident"] = base_df["ident"].astype(str) - props_categories = { - "AllNodeTypeProperties": [], - "FGNodeTypeProperties": [], - "AtomNodeTypeProperties": [], - "GraphNodeTypeProperties": [], - "BondProperties": [], - } - n_atom_node_properties, n_fg_node_properties = 0, 0 - n_bond_properties, n_graph_node_properties = 0, 0 - prop_lengths = [] - for prop in self.properties: - prop_length = prop.encoder.get_encoding_length() - prop_name = prop.name - prop_lengths.append((prop_name, prop_length)) - if isinstance(prop, AllNodeTypeProperty): - n_atom_node_properties += prop_length - n_fg_node_properties += prop_length - n_graph_node_properties += prop_length - props_categories["AllNodeTypeProperties"].append(prop_name) - elif isinstance(prop, FGNodeTypeProperty): - n_fg_node_properties += prop_length - props_categories["FGNodeTypeProperties"].append(prop_name) - elif isinstance(prop, AtomNodeTypeProperty): - n_atom_node_properties += prop_length - props_categories["AtomNodeTypeProperties"].append(prop_name) - elif isinstance(prop, BondProperty): - n_bond_properties += prop_length - props_categories["BondProperties"].append(prop_name) - elif isinstance(prop, MoleculeProperty): - # molecule props will be used as graph node props - n_graph_node_properties += prop_length - props_categories["GraphNodeTypeProperties"].append(prop_name) - else: - raise TypeError(f"Unsupported property type: {type(prop).__name__}") - - n_node_properties = max( - n_atom_node_properties, n_fg_node_properties, n_graph_node_properties - ) - rank_zero_info( - f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" - f"Properties Categories:\n{pformat(props_categories)}\n\n" - f"n_atom_node_properties: {n_atom_node_properties}, " - f"n_fg_node_properties: {n_fg_node_properties}, " - f"n_bond_properties: {n_bond_properties}, " - f"n_graph_node_properties: {n_graph_node_properties}\n\n" - f"Use following values for given parameters for model configuration: \n\t" - f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_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 - ) - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df["ident"] = property_df["ident"].astype(str) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base( - row, - max_len_node_properties=n_node_properties, - ), - axis=1, - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - return base_df[base_data[0].keys()].to_dict("records") - - def _merge_props_into_base( - self, row: pd.Series, max_len_node_properties: int - ) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - geom_data = row["features"] - if geom_data is None: - return None - if isinstance(geom_data, tuple): - geom_data = geom_data[ - 0 - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - assert isinstance(geom_data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" - - is_fg_node = ~is_atom_node & ~is_graph_node - num_nodes = geom_data.x.size(0) - edge_attr = geom_data.edge_attr - - # 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 @@ -810,42 +169,94 @@ class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly( GraphPropForAtomAndFGLevelOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly( + GraphPropForAtomAndFGLevelOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly( GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly( + GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly( GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly( + GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes( GraphPropNodeLevelPropOnlyForAllNodes, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes( + GraphPropNodeLevelPropOnlyForAllNodes, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +class ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel( + GraphPropAllPropsExceptNodeLevel, ChEBIOver50 +): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel( + GraphPropAllPropsExceptNodeLevel, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + # ---------- Final Augmentation: Different Thresholds ------------------------------ class ChEBI100_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver100): READER = AtomFGReader_WithFGEdges_WithGraphNode diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py index 9afaf71..2132965 100644 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -12,9 +12,23 @@ from chebai_graph.preprocessing.datasets.base import ( GraphPropAsPerNodeType, + GraphPropertiesMixIn, ) -from chebai_graph.preprocessing.reader.augmented_reader import ( +from chebai_graph.preprocessing.reader import ( + AtomFGReader_NoFGEdges_WithGraphNode, + AtomFGReader_WithFGEdges_NoGraphNode, AtomFGReader_WithFGEdges_WithGraphNode, + AtomReader_WithGraphNodeOnly, + AtomsFGReader_NoFGEdges_NoGraphNode, + GN_WithAllNodes_FG_WithAtoms_FGE, + GN_WithAllNodes_FG_WithAtoms_NoFGE, + GN_WithAtoms_FG_WithAtoms_FGE, + GN_WithAtoms_FG_WithAtoms_NoFGE, +) + +from .augmentation_base import ( + AugGraphPropMixIn_NoGraphNode, + AugGraphPropMixIn_WithGraphNode, ) @@ -22,38 +36,226 @@ class PCBA_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, PCBA): READER = AtomFGReader_WithFGEdges_WithGraphNode +class PCBA_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, PCBA): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class BACE_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BACE): READER = AtomFGReader_WithFGEdges_WithGraphNode +class BACE_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, BACE): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class BBBP_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BBBP): READER = AtomFGReader_WithFGEdges_WithGraphNode +class BBBP_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, BBBP): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ClinTox_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ClinTox): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ClinTox_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ClinTox): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class HIV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIV): READER = AtomFGReader_WithFGEdges_WithGraphNode +class HIV_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, HIV): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class SIDER_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SIDER): READER = AtomFGReader_WithFGEdges_WithGraphNode +class SIDER_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, SIDER): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class MUV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, MUV): READER = AtomFGReader_WithFGEdges_WithGraphNode +class MUV_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, MUV): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class Tox21_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Tox21): READER = AtomFGReader_WithFGEdges_WithGraphNode +class Tox21_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, Tox21): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ToxCast_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ToxCast_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +# ---- Augmentation: Variants with graph Node connected to FG nodes only ------------- +class Tox21_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with with FG nodes and FG edges and graph node.""" + + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ToxCast_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with with FG nodes and FG edges and graph node.""" + + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Tox21_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with FG nodes but without FG edges, with graph node.""" + + READER = AtomFGReader_NoFGEdges_WithGraphNode + + +class ToxCast_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with FG nodes but without FG edges, with graph node.""" + + READER = AtomFGReader_NoFGEdges_WithGraphNode + + +class Tox21_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, Tox21): + """Tox21 with FG nodes and FG edges, no graph node.""" + + READER = AtomFGReader_WithFGEdges_NoGraphNode + + +class ToxCast_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ToxCast): + """ToxCast with FG nodes and FG edges, no graph node.""" + + READER = AtomFGReader_WithFGEdges_NoGraphNode + + +class Tox21_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, Tox21): + """Tox21 with FG nodes but without FG edges or graph node.""" + + READER = AtomsFGReader_NoFGEdges_NoGraphNode + + +class ToxCast_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ToxCast): + """ToxCast with FG nodes but without FG edges or graph node.""" + + READER = AtomsFGReader_NoFGEdges_NoGraphNode + + +class Tox21_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with atom-level nodes and graph node only.""" + + READER = AtomReader_WithGraphNodeOnly + + +class ToxCast_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with atom-level nodes and graph node only.""" + + READER = AtomReader_WithGraphNodeOnly + + +# ------- Augmentation: Variants with graph Node connected to all others nodes (FG and atoms) -------------- +class Tox21_GN_WithAllNodes_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 with FG nodes (connected to their respective atom nodes) with functional group + edges, and adds a graph-level node connected to all nodes (fg + atoms). + """ + + READER = GN_WithAllNodes_FG_WithAtoms_FGE + + +class ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE( + AugGraphPropMixIn_WithGraphNode, ToxCast +): + """ + ToxCast with FG nodes (connected to their respective atom nodes) with functional group + edges, and adds a graph-level node connected to all nodes (fg + atoms). + """ + + READER = GN_WithAllNodes_FG_WithAtoms_FGE + + +class Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 with FG nodes (connected to their respective atom nodes) without functional group + edges, and adds a graph-level node connected to all nodes (fg + atoms). + """ + + READER = GN_WithAllNodes_FG_WithAtoms_NoFGE + + +class ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE( + AugGraphPropMixIn_WithGraphNode, ToxCast +): + """ + ToxCast with FG nodes (connected to their respective atom nodes) without functional group + edges, and adds a graph-level node connected to all nodes (fg + atoms). + """ + + READER = GN_WithAllNodes_FG_WithAtoms_NoFGE + + +# ------- Augmentation: Variants with graph node connected to atom nodes ONLY ----------- +class Tox21_GN_WithAtoms_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 with FG nodes (connected to their respective atom nodes) with functional group + edges, and adds a graph-level node connected to all atom nodes. + """ + + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +class ToxCast_GN_WithAtoms_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ + ToxCast with FG nodes (connected to their respective atom nodes) with functional group + edges, and adds a graph-level node connected to all atom nodes. + """ + + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +class Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 with FG nodes (connected to their respective atom nodes) without functional group + edges, and adds a graph-level node connected to all atom nodes. + """ + + READER = GN_WithAtoms_FG_WithAtoms_NoFGE + + +class ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ + ToxCast with FG nodes (connected to their respective atom nodes) without functional group + edges, and adds a graph-level node connected to all atom nodes. + """ + + READER = GN_WithAtoms_FG_WithAtoms_NoFGE + + +# ---------------------------Baselines classes for Tox21 and ToxCast datasets +class Tox21GraphProperties(GraphPropertiesMixIn, Tox21): + """Tox21 dataset with molecular property encodings.""" + + pass + + +class ToxCastGraphProperties(GraphPropertiesMixIn, ToxCast): + """ToxCast dataset with molecular property encodings.""" + + pass + + if __name__ == "__main__": dataset = BACE_WFGE_WGN_AsPerNodeType() dataset.prepare_data() diff --git a/chebai_graph/preprocessing/datasets/pretrain_zinc15.py b/chebai_graph/preprocessing/datasets/pretrain_zinc15.py new file mode 100644 index 0000000..3082899 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/pretrain_zinc15.py @@ -0,0 +1,18 @@ +from chebai.preprocessing.datasets.pretrain_zinc15 import Zinc15Pretrain + +from chebai_graph.preprocessing.datasets.base import GraphPropAsPerNodeType +from chebai_graph.preprocessing.reader import GN_WithAtoms_FG_WithAtoms_FGE + + +class Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType( + GraphPropAsPerNodeType, Zinc15Pretrain +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +if __name__ == "__main__": + dataset = Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType( + root="data/Zinc15", name="Zinc15" + ) + dataset.prepare_data() + dataset.setup() diff --git a/chebai_graph/preprocessing/datasets/utils.py b/chebai_graph/preprocessing/datasets/utils.py index 79f2fb3..fa42078 100644 --- a/chebai_graph/preprocessing/datasets/utils.py +++ b/chebai_graph/preprocessing/datasets/utils.py @@ -5,7 +5,7 @@ def resolve_property( - property: str | MolecularProperty, data_type: str + property: str | MolecularProperty, tokens_directory: str ) -> MolecularProperty: """ Resolves a molecular property specification (either as a class instance or class path string) @@ -21,8 +21,8 @@ def resolve_property( property (str | MolecularProperty): The property to resolve. Can be a class instance, a fully qualified class name (e.g. "module.ClassName"), or a class name assumed to be in `chebai_graph.preprocessing.properties`. - data_type (str): The data type associated with the property. This used to determine or set - tokens file path for the property if applicable. + tokens_directory (str): The directory where the tokens for the property are stored. This is used to determine or set + the tokens file path for the property if applicable. Returns: MolecularProperty: An instance of the resolved MolecularProperty. @@ -41,7 +41,7 @@ def resolve_property( module_name = property[:last_dot] class_name = property[last_dot + 1 :] module = importlib.import_module(module_name) - return getattr(module, class_name)(data_type=data_type) + return getattr(module, class_name)(tokens_directory=tokens_directory) except ValueError: # if only a class name is given, assume the module is chebai_graph.processing.properties - return getattr(graph_properties, property)(data_type=data_type) + return getattr(graph_properties, property)(tokens_directory=tokens_directory) diff --git a/chebai_graph/preprocessing/properties/augmented_properties.py b/chebai_graph/preprocessing/properties/augmented_properties.py index 6f112ea..07d514d 100644 --- a/chebai_graph/preprocessing/properties/augmented_properties.py +++ b/chebai_graph/preprocessing/properties/augmented_properties.py @@ -30,9 +30,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: @@ -56,9 +57,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: @@ -82,9 +84,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: @@ -337,9 +340,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Optional encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_bond_value(self, bond: Chem.rdchem.Bond | dict) -> str: diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index f148df8..bf70665 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -33,14 +33,16 @@ class MolecularProperty(ABC): def __init__( self, - data_type: str, + tokens_directory: str, encoder: PropertyEncoder | None = None, ) -> None: - assert data_type is not None, "data_type must be provided for MolecularProperty" + assert tokens_directory is not None, ( + "tokens_directory must be provided for MolecularProperty" + ) if encoder is None: - encoder = IndexEncoder(self, data_type=data_type) + encoder = IndexEncoder(self, tokens_directory=tokens_directory) self.encoder: PropertyEncoder = encoder - self._data_type = data_type + self.tokens_directory = tokens_directory @property def name(self) -> str: @@ -209,8 +211,8 @@ def on_finish(self) -> None: hasattr(self.encoder, "cache") and len(self.encoder.cache) > self.encoder.index_length_start ): - raise ValueError( - f"{self.__class__.__name__} attempted to add new tokens " + print( + f"{self.__class__.__name__} added new tokens " f"to a frozen encoder at {self.encoder.index_path}" ) super().on_finish() diff --git a/chebai_graph/preprocessing/properties/properties.py b/chebai_graph/preprocessing/properties/properties.py index 87112da..0e91d57 100644 --- a/chebai_graph/preprocessing/properties/properties.py +++ b/chebai_graph/preprocessing/properties/properties.py @@ -20,9 +20,10 @@ class AtomType(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -46,9 +47,10 @@ class NumAtomBonds(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -72,9 +74,10 @@ class AtomCharge(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -98,9 +101,10 @@ class AtomChirality(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.ChiralType: @@ -124,9 +128,10 @@ class AtomHybridization(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.HybridizationType: @@ -150,9 +155,10 @@ class AtomNumHs(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -222,9 +228,10 @@ class BondType(BondProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_bond_value(self, bond: Chem.rdchem.Bond) -> Chem.rdchem.BondType: @@ -271,9 +278,10 @@ class MoleculeNumRings(MoleculeProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_property_value(self, mol: Chem.rdchem.Mol) -> list[int]: diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 5daeb4c..1173590 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -100,14 +100,15 @@ class IndexEncoder(PropertyEncoder): """ def __init__( - self, property, data_type: str, indices_dir: str | None = None, **kwargs + self, property, tokens_directory: str, indices_dir: str | None = None, **kwargs ) -> None: super().__init__(property, **kwargs) if indices_dir is None: indices_dir = os.path.dirname(inspect.getfile(self.__class__)) self.dirname = indices_dir + self._use_one_hot_cache = kwargs.get("use_one_hot_cache", False) # load already existing cache - self._data_type = data_type + self.tokens_directory = tokens_directory with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { token.strip(): idx for idx, token in enumerate(pk) @@ -120,6 +121,8 @@ def __init__( @property def name(self) -> str: """Name of this encoder.""" + if self._use_one_hot_cache: + return "one_hot" return "index" @property @@ -129,11 +132,13 @@ def index_path(self) -> str: Returns: Path to index file. """ - assert self._data_type is not None, "data_type must be set for IndexEncoder" + assert self.tokens_directory is not None, ( + "tokens_directory must be set for IndexEncoder" + ) index_path = os.path.join( self.dirname, "bin", - self._data_type, + self.tokens_directory, self.property.name, f"indices_{self.name}.txt", ) diff --git a/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh new file mode 100644 index 0000000..eb82031 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh @@ -0,0 +1,191 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 + +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --constraint="A100|H100.80gb" + +# ============================================================ +# Job array +# +# 10 data configurations × 3 seeds = 30 experiments +# Maximum 10 experiments running simultaneously +# ============================================================ + +#SBATCH --array=0-29%10 + +#SBATCH --job-name=aug-ablation + +# Separate output/error file for every array task +#SBATCH --output=aug-ablation_%A_%a.out +#SBATCH --error=aug-ablation_%A_%a.err + + +# ============================================================ +# Shell settings +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Seeds +# ============================================================ + +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations +# +# Add/remove configs here. +# +# IMPORTANT: +# The number of array tasks must be: +# +# number of data configs × number of seeds +# +# ============================================================ + +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/tox21" +) +DATA_CONFIGS=( + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" +) + +# ============================================================ +# Determine data config and seed from array task ID +# ============================================================ + +# Array task ordering: +# All data configurations are run for seed 0 first, +# followed by all data configurations for seed 42, +# and finally all data configurations for seed 12345. +# This is intentionally because the data is generated in first run for each data configuration. +# Then the following runs for the same data configuration can re-use the same generated data. + +NUM_SEEDS=${#SEEDS[@]} +NUM_DATA_CONFIGS=${#DATA_CONFIGS[@]} + +SEED_INDEX=$((SLURM_ARRAY_TASK_ID / NUM_DATA_CONFIGS)) +DATA_INDEX=$((SLURM_ARRAY_TASK_ID % NUM_DATA_CONFIGS)) + +DATA_CONFIG=${DATA_CONFIGS[$DATA_INDEX]} +SEED=${SEEDS[$SEED_INDEX]} + +# ============================================================ +# Create experiment name +# ============================================================ + +DATA_NAME=$(basename "$DATA_CONFIG" .yml) +RUN_NAME="${DATA_NAME}_s${SEED}" + +# ============================================================ +# Print job information +# ============================================================ + +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Job ID: $SLURM_ARRAY_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data index: $DATA_INDEX" +echo "Seed index: $SEED_INDEX" +echo "Data config: $DATA_CONFIG" +echo "Data name: $DATA_NAME" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Check selected configuration +# ============================================================ + +if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist:" + echo "$DATA_CONFIG" + exit 1 +fi + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model=../python-chebai-graph/configs/model/baselines/gat.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="$DATA_CONFIG" \ + --data.init_args.batch_size=64 \ + --trainer.accumulate_grad_batches=1 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --data.init_args.chebi_version=252 \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --model.init_args.optimizer_kwargs.lr=0.002 \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ + --seed_everything="$SEED" + + +# ============================================================ +# Finished +# ============================================================ + +echo "============================================================" +echo "Training finished successfully" +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data config: $DATA_CONFIG" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh new file mode 100644 index 0000000..bf2c8c2 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh @@ -0,0 +1,129 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# Run this script using: `bash configs/data/augmented/aug-ablation/aug-ablation-tox21.sh` +# +# The original is a Slurm job-array script (10 data configs x 3 +# seeds = 30 experiments, 10 concurrent). This version runs the +# SAME experiments serially on a single local GPU, without Slurm. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial +# loop over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations (all resolved to absolute local paths) +# ============================================================ +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/tox21" +) +DATA_CONFIGS=( + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" +) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +for SEED in "${SEEDS[@]}"; do + for DATA_CONFIG in "${DATA_CONFIGS[@]}"; do + + if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist: $DATA_CONFIG" + exit 1 + fi + + DATA_NAME=$(basename "$DATA_CONFIG" .yml) + RUN_NAME="${DATA_NAME}_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Seed: $SEED" + echo "Config: $DATA_CONFIG" + echo "============================================================" + + python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "tox21"]' \ + --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ + --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ + --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --trainer.min_epochs=100 \ + --trainer.max_epochs=100 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path=data/Tox21:MNClassification/Tox21/processed/splits.csv \ + --model.config.edge_dim=6 \ + --model.config.in_channels=77 \ + --model.config.dropout=0.7 \ + --model.optimizer_kwargs.lr_gnn=0.0005 \ + --model.optimizer_kwargs.lr_linear=0.001 \ + --seed_everything="$SEED" + + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh new file mode 100644 index 0000000..a4de723 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh @@ -0,0 +1,129 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# # Run this script using: `bash configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh` +# +# The original is a Slurm job-array script (10 data configs x 3 +# seeds = 30 experiments, 10 concurrent). This version runs the +# SAME experiments serially on a single local GPU, without Slurm. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial +# loop over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations (all resolved to absolute local paths) +# ============================================================ +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/toxcast" +) +DATA_CONFIGS=( + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" +) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +for SEED in "${SEEDS[@]}"; do + for DATA_CONFIG in "${DATA_CONFIGS[@]}"; do + + if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist: $DATA_CONFIG" + exit 1 + fi + + DATA_NAME=$(basename "$DATA_CONFIG" .yml) + RUN_NAME="${DATA_NAME}_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Seed: $SEED" + echo "Config: $DATA_CONFIG" + echo "============================================================" + + python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "toxcast"]' \ + --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ + --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ + --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --trainer.min_epochs=100 \ + --trainer.max_epochs=100 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path=data/ToxCast:MNClassification/ToxCast/processed/splits.csv \ + --model.config.edge_dim=6 \ + --model.config.in_channels=79 \ + --model.config.dropout=0.7 \ + --model.optimizer_kwargs.lr_gnn=0.001 \ + --model.optimizer_kwargs.lr_linear=0.001 \ + --seed_everything="$SEED" + + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN+E+WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+E+WGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+E+WGN.yml diff --git a/configs/data/augmented/aug-ablation/FGN+E.yml b/configs/data/augmented/aug-ablation/chebi/FGN+E.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+E.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+E.yml diff --git a/configs/data/augmented/aug-ablation/FGN+WGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN+WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+WGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+WGN.yml diff --git a/configs/data/augmented/aug-ablation/FGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN.yml diff --git a/configs/data/augmented/aug-ablation/WGN.yml b/configs/data/augmented/aug-ablation/chebi/WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/WGN.yml rename to configs/data/augmented/aug-ablation/chebi/WGN.yml diff --git a/configs/data/augmented/aug-ablation/chebi/baseline.yml b/configs/data/augmented/aug-ablation/chebi/baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/chebi/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/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_nfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml rename to configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_nfge.yml diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_wfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml rename to configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_wfge.yml diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_nfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml rename to configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_nfge.yml diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_wfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml rename to configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_wfge.yml diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml new file mode 100644 index 0000000..9b1f918 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml new file mode 100644 index 0000000..193eec0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml new file mode 100644 index 0000000..3515dfb --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_NFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN.yml new file mode 100644 index 0000000..50feae1 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_NFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/WGN.yml b/configs/data/augmented/aug-ablation/tox21/WGN.yml new file mode 100644 index 0000000..bdb347f --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_Atom_WGNOnly_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/baseline.yml b/configs/data/augmented/aug-ablation/tox21/baseline.yml new file mode 100644 index 0000000..c0aeeda --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21GraphProperties +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/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml new file mode 100644 index 0000000..88ddd67 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml new file mode 100644 index 0000000..a227629 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAllNodes_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml new file mode 100644 index 0000000..5102fdf --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml new file mode 100644 index 0000000..d79a0ea --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAtoms_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml new file mode 100644 index 0000000..662da11 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml new file mode 100644 index 0000000..1c5b321 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml new file mode 100644 index 0000000..875a57d --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_NFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN.yml new file mode 100644 index 0000000..d046445 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_NFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/WGN.yml b/configs/data/augmented/aug-ablation/toxcast/WGN.yml new file mode 100644 index 0000000..1464be0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_Atom_WGNOnly_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/baseline.yml b/configs/data/augmented/aug-ablation/toxcast/baseline.yml new file mode 100644 index 0000000..5a7e18c --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCastGraphProperties +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/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml new file mode 100644 index 0000000..c38ced5 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml new file mode 100644 index 0000000..8d28ad9 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml new file mode 100644 index 0000000..bae368f --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml new file mode 100644 index 0000000..d754e25 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAtoms_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/final/BACE_final_augmented.yml b/configs/data/augmented/final/BACE_final_augmented.yml index 7ea67f0..fc95a55 100644 --- a/configs/data/augmented/final/BACE_final_augmented.yml +++ b/configs/data/augmented/final/BACE_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BACE_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.BACE_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/BBBP_final_augmented.yml b/configs/data/augmented/final/BBBP_final_augmented.yml index 9a7ee90..fdef5ca 100644 --- a/configs/data/augmented/final/BBBP_final_augmented.yml +++ b/configs/data/augmented/final/BBBP_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.BBBP_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/ClinTox_final_augmented.yml b/configs/data/augmented/final/ClinTox_final_augmented.yml index a4023b3..774103a 100644 --- a/configs/data/augmented/final/ClinTox_final_augmented.yml +++ b/configs/data/augmented/final/ClinTox_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ClinTox_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ClinTox_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/HIV_final_augmented.yml b/configs/data/augmented/final/HIV_final_augmented.yml index 596fc19..126a919 100644 --- a/configs/data/augmented/final/HIV_final_augmented.yml +++ b/configs/data/augmented/final/HIV_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.HIV_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.HIV_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/MUV_final_augmented.yml b/configs/data/augmented/final/MUV_final_augmented.yml index 56a3e1e..a3ade9c 100644 --- a/configs/data/augmented/final/MUV_final_augmented.yml +++ b/configs/data/augmented/final/MUV_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.MUV_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.MUV_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/PCBA_final_augmented.yml b/configs/data/augmented/final/PCBA_final_augmented.yml index 40688f4..dbacc76 100644 --- a/configs/data/augmented/final/PCBA_final_augmented.yml +++ b/configs/data/augmented/final/PCBA_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.PCBA_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.PCBA_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/SIDER_final_augmented.yml b/configs/data/augmented/final/SIDER_final_augmented.yml index 5eef128..bcb071c 100644 --- a/configs/data/augmented/final/SIDER_final_augmented.yml +++ b/configs/data/augmented/final/SIDER_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.SIDER_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.SIDER_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/Tox21_final_augmented.yml b/configs/data/augmented/final/Tox21_final_augmented.yml index 1c88112..cebd0dc 100644 --- a/configs/data/augmented/final/Tox21_final_augmented.yml +++ b/configs/data/augmented/final/Tox21_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.Tox21_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/ToxCast_final_augmented.yml b/configs/data/augmented/final/ToxCast_final_augmented.yml index 4831ae5..88f1a17 100644 --- a/configs/data/augmented/final/ToxCast_final_augmented.yml +++ b/configs/data/augmented/final/ToxCast_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ToxCast_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/chebi50_final_augmented.yml b/configs/data/augmented/final/chebi50_final_augmented.yml index 576cf75..3ff3232 100644 --- a/configs/data/augmented/final/chebi50_final_augmented.yml +++ b/configs/data/augmented/final/chebi50_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/final-chebi.sh b/configs/data/augmented/final/final-chebi.sh new file mode 100644 index 0000000..4b2cb84 --- /dev/null +++ b/configs/data/augmented/final/final-chebi.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 + +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --constraint="A100|H100.80gb" + +# 3 experiments, maximum 3 running at the same time +#SBATCH --array=0-2%3 + +#SBATCH --job-name=chebi_gatv2_final + +# Separate log files for every array task +#SBATCH --output=chebi_gatv2_final_%A_%a.out +#SBATCH --error=chebi_gatv2_final_%A_%a.err + + +# ============================================================ +# Shell settings +# ============================================================ +set -x +set -euo pipefail +# ============================================================ +# Define seeds +# ============================================================ + +SEEDS=(0 42 12345) + +if [[ -z "${SEEDS[$SLURM_ARRAY_TASK_ID]}" ]]; then + echo "ERROR: Invalid array task ID: $SLURM_ARRAY_TASK_ID" + exit 1 +fi + +SEED=${SEEDS[$SLURM_ARRAY_TASK_ID]} +RUN_NAME="chebi_gatv2_final_${SEED}" + +# ============================================================ +# Print job information +# ============================================================ + +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Job ID: $SLURM_ARRAY_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model=../python-chebai-graph/configs/model/augmented/pooling/amg_pool/gat.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/augmented/final/chebi50_final_augmented.yml \ + --data.init_args.batch_size=64 \ + --data.init_args.num_workers=10 \ + --data.init_args.chebi_version=252 \ + --model.pass_loss_kwargs=false \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_weighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --seed_everything="$SEED" \ + --trainer.logger.init_args.tags='["augmented_paper", "chebi", "final:gnwa_fgwa_wfge"]' \ + +echo "============================================================" +echo "Training finished" +echo "Seed: $SEED" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" diff --git a/configs/data/augmented/final/molnet.sh b/configs/data/augmented/final/molnet.sh new file mode 100644 index 0000000..f7c40ed --- /dev/null +++ b/configs/data/augmented/final/molnet.sh @@ -0,0 +1,270 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# Run this script using: `bash configs/data/augmented/final/molnet.sh` +# +# The original is a Slurm job-array script. This version runs the +# six MolNet data configurations x three seeds serially on one GPU. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial loop +# over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -Eeuo pipefail +trap 'status=$?; echo "ERROR: command failed with status ${status} at line ${BASH_LINENO[0]}: ${BASH_COMMAND}" >&2; exit "$status"' ERR + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" +LOG_DIR="$CHEBAI_DIR/logs" +LOG_FILE="$LOG_DIR/molnet_$(date +%Y%m%d_%H%M%S).txt" +mkdir -p "$LOG_DIR" + +# Keep output visible in the terminal while recording stdout, stderr, and +# Bash trace output in a timestamped text file. +exec > >(tee -a "$LOG_FILE") 2>&1 +set -x +echo "Writing MolNet run log to: $LOG_FILE" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +DATASETS=(BBBP BACE SIDER ClinTox Tox21 ToxCast) +RESULTS_FILE="$CHEBAI_DIR/logs/molnet_test_results.tsv" +mkdir -p "$(dirname "$RESULTS_FILE")" +printf 'dataset\tseed\tmetric\tvalue\tcheckpoint\n' > "$RESULTS_FILE" + +for SEED in "${SEEDS[@]}"; do + for DATASET in "${DATASETS[@]}"; do + case "$DATASET" in + BBBP) + TRAINER="configs/training/binary_trainer.yml" + METRICS="configs/metrics/binary-f1-roc-auc.yml" + CALLBACKS="configs/training/binary_callbacks.yml" + CHECKPOINT_PREFIX="best_roc-auc_" + METRICS_TO_REPORT=(test_roc-auc test_f1) + DROPOUT="0.5" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + BACE) + TRAINER="configs/training/binary_trainer.yml" + METRICS="configs/metrics/binary-f1-roc-auc.yml" + CALLBACKS="configs/training/binary_callbacks.yml" + CHECKPOINT_PREFIX="best_roc-auc_" + METRICS_TO_REPORT=(test_roc-auc test_f1) + DROPOUT="0.5" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + SIDER) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) + DROPOUT="0.5" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + ClinTox) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) + DROPOUT="0.5" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + Tox21) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) + DROPOUT="0.7" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + ToxCast) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) + DROPOUT="0.7" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + *) + echo "Unsupported dataset: $DATASET" >&2 + exit 1 + ;; + esac + + DATASET_LOWER=$(printf '%s' "$DATASET" | tr '[:upper:]' '[:lower:]') + RUN_NAME="${DATASET_LOWER}_gat_amg_augmented_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Dataset: $DATASET" + echo "Seed: $SEED" + echo "============================================================" + + RUN_MARKER=$(mktemp) + CALLBACK_ARGS=(--trainer.callbacks="$CALLBACKS") + + python -m chebai fit \ + --trainer="$TRAINER" \ + --trainer.logger=configs/training/wandb_logger.yml \ + "--trainer.logger.init_args.tags=[\"augmented_paper\",\"$DATASET_LOWER\"]" \ + "${CALLBACK_ARGS[@]}" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ + --model.train_metrics="$METRICS" \ + --model.test_metrics="$METRICS" \ + --model.val_metrics="$METRICS" \ + --data="$CHEBAI_GRAPH_DIR/configs/data/augmented/final/${DATASET}_final_augmented.yml" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path="data/${DATASET}:MNClassification/${DATASET}/processed/splits.csv" \ + --model.config.dropout="$DROPOUT" \ + --model.optimizer_kwargs.lr_gnn="$LR_GNN" \ + --model.optimizer_kwargs.lr_linear="$LR_LINEAR" \ + --seed_everything="$SEED" \ + --data.init_args.tokens_directory=_shared \ + # --model.pretrained_checkpoint=/home/aditya/python-chebai-graph/msp_best_epoch=92_train_loss=0.0009.ckpt \ + + + + CHECKPOINT_PATH=$( + find "$CHEBAI_DIR/logs" -type f \ + -name "${CHECKPOINT_PREFIX}*.ckpt" \ + -newer "$RUN_MARKER" -print | sort | tail -n 1 + ) + rm -f "$RUN_MARKER" + if [[ -z "$CHECKPOINT_PATH" ]]; then + echo "Could not find best checkpoint for $RUN_NAME" >&2 + exit 1 + fi + + TEST_NAME="${RUN_NAME}_test" + TEST_MARKER=$(mktemp) + python -m chebai test \ + --trainer="$TRAINER" \ + --trainer.logger=configs/training/csv_logger.yml \ + --trainer.logger.init_args.name="$TEST_NAME" \ + --trainer.devices=1 \ + --trainer.num_nodes=1 \ + --ckpt_path="$CHECKPOINT_PATH" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ + --model.test_metrics="$METRICS" \ + --data="$CHEBAI_GRAPH_DIR/configs/data/augmented/final/${DATASET}_final_augmented.yml" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --data.init_args.splits_file_path="data/${DATASET}:MNClassification/${DATASET}/processed/splits.csv" \ + --data.init_args.tokens_directory=_shared \ + + TEST_METRICS_FILE=$( + find "$CHEBAI_DIR/logs/$TEST_NAME" -type f \ + -name metrics.csv -newer "$TEST_MARKER" -print | sort | tail -n 1 + ) + rm -f "$TEST_MARKER" + if [[ ! -f "$TEST_METRICS_FILE" ]]; then + echo "Could not find fresh test metrics for $TEST_NAME" >&2 + exit 1 + fi + METRIC_VALUES=$(python - "$TEST_METRICS_FILE" "${METRICS_TO_REPORT[@]}" <<'PY' +import csv +import sys + +metrics_path, *metric_names = sys.argv[1:] +values = {metric_name: None for metric_name in metric_names} +with open(metrics_path, newline="") as metrics_file: + for row in csv.DictReader(metrics_file): + for metric_name in metric_names: + if row.get(metric_name): + values[metric_name] = row[metric_name] +for metric_name in metric_names: + value = values[metric_name] + if value is None: + raise SystemExit(f"Metric {metric_name!r} not found in {metrics_path}") + print(f"{metric_name}\t{value}") +PY +) + while IFS=$'\t' read -r TEST_METRIC TEST_VALUE; do + printf '%s\t%s\t%s\t%s\t%s\n' \ + "$DATASET" "$SEED" "$TEST_METRIC" "$TEST_VALUE" "$CHECKPOINT_PATH" \ + >> "$RESULTS_FILE" + echo "Test $TEST_METRIC: $TEST_VALUE" + done <<< "$METRIC_VALUES" + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" + +python - "$RESULTS_FILE" <<'PY' +import csv +import math +import statistics +import sys +from collections import defaultdict + +results = defaultdict(list) +with open(sys.argv[1], newline="") as results_file: + for row in csv.DictReader(results_file, delimiter="\t"): + results[(row["dataset"], row["metric"])].append(float(row["value"])) + +print("\nDataset test performance (best validation checkpoint per seed)") +print("dataset\tmetric\tmean\tsample_sd\tn") +for (dataset, metric), values in sorted(results.items()): + mean = statistics.mean(values) + sample_sd = statistics.stdev(values) if len(values) > 1 else math.nan + print(f"{dataset}\t{metric}\t{mean:.6f}\t{sample_sd:.6f}\t{len(values)}") +PY diff --git a/configs/data/augmented/pretraining/zinc15.yml b/configs/data/augmented/pretraining/zinc15.yml new file mode 100644 index 0000000..dd8ab65 --- /dev/null +++ b/configs/data/augmented/pretraining/zinc15.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.Zinc15Pretrain_GNWA_FGWA_WFGE_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/pretraining/zinc_msp_pretraining.sh b/configs/data/augmented/pretraining/zinc_msp_pretraining.sh new file mode 100644 index 0000000..0c3c2e9 --- /dev/null +++ b/configs/data/augmented/pretraining/zinc_msp_pretraining.sh @@ -0,0 +1,65 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:3 +#SBATCH --constraint="A100|H100.80gb" +#SBATCH --job-name=zinc +#SBATCH --output=zinc.out +#SBATCH --error=zinc.err + + +# ============================================================ +# Shell settings +# ============================================================ +set -x +# ============================================================ +# Define seeds +# ============================================================ + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/zinc_pretrainer.yml \ + --trainer.limit_val_batches=0 \ + --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","zinc:pretraining","final:gnwa_fgwa_wfge"]' \ + --trainer.logger.init_args.name=zinc15pretrain_gat_amg_augmented \ + --model=../python-chebai-graph/configs/model/pretraining/gat_msp.yml \ + --data=../python-chebai-graph/configs/data/augmented/pretraining/zinc15.yml \ + --data.init_args.batch_size=32 \ + --model.pass_loss_kwargs=true \ + --data.init_args.num_workers=10 \ + --trainer.logger.init_args.name=zinc15pretrain_gat_amg_augmented \ + --data.init_args.splits_file_path=data/Zinc15:Pretrain/Zinc15/processed/splits.csv \ diff --git a/configs/data/augmented/properties/FGN+E+WGN/chebi50_all_props_except_node_lvl.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_all_props_except_node_lvl.yml new file mode 100644 index 0000000..a03da69 --- /dev/null +++ b/configs/data/augmented/properties/FGN+E+WGN/chebi50_all_props_except_node_lvl.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_fg_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_fg_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_atom_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_fg_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_fg_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_node_level_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_node_level_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_node_level_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_node_level_prop_only.yml diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml new file mode 100644 index 0000000..3ff3232 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml new file mode 100644 index 0000000..fd33150 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml new file mode 100644 index 0000000..d9053ae --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml new file mode 100644 index 0000000..7ab3ecb --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml new file mode 100644 index 0000000..50d18ba --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml new file mode 100644 index 0000000..c56edff --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml new file mode 100644 index 0000000..9b42371 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml new file mode 100644 index 0000000..f1921ff --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml new file mode 100644 index 0000000..db61465 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType diff --git a/configs/model/augmented/pooling/amg_pool/gat.yml b/configs/model/augmented/pooling/amg_pool/gat.yml index f88159d..186fbc8 100644 --- a/configs/model/augmented/pooling/amg_pool/gat.yml +++ b/configs/model/augmented/pooling/amg_pool/gat.yml @@ -10,5 +10,4 @@ init_args: edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` - dropout: 0 n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml index 773088a..a40a4bd 100644 --- a/configs/model/augmented/pooling/amg_pool/gine.yml +++ b/configs/model/augmented/pooling/amg_pool/gine.yml @@ -8,6 +8,5 @@ init_args: 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/himol_gine.yml b/configs/model/augmented/pooling/amg_pool/himol_gine.yml new file mode 100644 index 0000000..10559ac --- /dev/null +++ b/configs/model/augmented/pooling/amg_pool/himol_gine.yml @@ -0,0 +1,15 @@ +class_path: chebai_graph.models.GINEAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 512 + out_channels: 512 + num_layers: 5 + edge_dim: 12 # number of bond properties + train_eps: true + use_himol_model: true + norm: batch_norm + n_linear_layers: 2 + lin_hidden_dim: 256 diff --git a/configs/model/augmented/pooling/no_pooling/gat.yml b/configs/model/augmented/pooling/no_pooling/gat.yml index f9d3575..55b9cfc 100644 --- a/configs/model/augmented/pooling/no_pooling/gat.yml +++ b/configs/model/augmented/pooling/no_pooling/gat.yml @@ -10,5 +10,4 @@ init_args: edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` - dropout: 0 n_linear_layers: 1 diff --git a/configs/model/pretraining/gat_msp.yml b/configs/model/pretraining/gat_msp.yml new file mode 100644 index 0000000..76c5a60 --- /dev/null +++ b/configs/model/pretraining/gat_msp.yml @@ -0,0 +1,17 @@ +class_path: chebai_graph.models.GATMSPretraining +init_args: + criterion: + class_path: chebai_graph.loss.msp.MSPPretrainingLoss + init_args: + hidden_size: 512 + 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.5 diff --git a/results/combine_json_file.py b/results/combine_json_file.py new file mode 100644 index 0000000..41da467 --- /dev/null +++ b/results/combine_json_file.py @@ -0,0 +1,37 @@ +import json + +import pandas as pd + +file1 = "/home/aditya/python-chebai/data/chebi_v252/ChEBI50/processed/gatv2_baseline_classes_test.json" +file2 = "/home/aditya/python-chebai/data/chebi_v252/ChEBI50/processed/gatv2_aug_classes_test.json" + +# Output CSV file +output_file = "combined_f1_scores.csv" + +# Load JSON data +with open(file1, "r") as f: + data1 = json.load(f) + +with open(file2, "r") as f: + data2 = json.load(f) + +# Create a combined set of labels from both files +labels = sorted(set(data1.keys()) | set(data2.keys())) + +# Create output rows +rows = [] +for label in labels: + rows.append( + { + "label": label, + file1: data1.get(label, {}).get("f1", None), + file2: data2.get(label, {}).get("f1", None), + } + ) + +# Create DataFrame and save to CSV +df = pd.DataFrame(rows) +df.to_csv(output_file, index=False) + +print(f"CSV file created: {output_file}") +print(df) diff --git a/results/molecule_distance_statistics.py b/results/molecule_distance_statistics.py new file mode 100644 index 0000000..e1b7901 --- /dev/null +++ b/results/molecule_distance_statistics.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +"""Compute molecular graph-diameter statistics from a saved PyTorch dataset.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import torch +from rdkit import Chem, RDLogger +from rdkit.Chem import Draw, Mol, rdmolops + +# Cached molecules may have been serialized with an older RDKit release. +RDLogger.DisableLog("rdApp.warning") + + +def extract_molecule(sample: Any) -> Mol: + """Extract the atom-level RDKit molecule from one saved dataset sample.""" + try: + molecule = sample["features"][1]["nodes"]["atom_nodes"] + except (IndexError, KeyError, TypeError) as error: + raise ValueError( + "Sample does not contain features[1]['nodes']['atom_nodes']" + ) from error + + if not isinstance(molecule, Mol): + raise TypeError( + "features[1]['nodes']['atom_nodes'] is not an RDKit molecule: " + f"{type(molecule).__name__}" + ) + return molecule + + +def molecule_diameter(molecule: Mol) -> int: + """Return the largest finite shortest-path distance in the molecule.""" + if molecule.GetNumAtoms() == 0: + raise ValueError("Cannot calculate a diameter for a molecule with no atoms") + + distances = rdmolops.GetDistanceMatrix(molecule) + components = rdmolops.GetMolFrags(molecule, asMols=False) + return max( + int(distances[atom_i][atom_j]) + for component in components + for atom_i in component + for atom_j in component + ) + + +def draw_molecule(molecule: Mol, output_path: Path) -> None: + """Draw an RDKit molecule and save it as an image.""" + output_path.parent.mkdir(parents=True, exist_ok=True) + Draw.MolToFile(molecule, str(output_path), size=(800, 600)) + + +def iter_samples(dataset: Any) -> Iterable[tuple[int, Any]]: + """Yield dataset samples while preserving their instance indices.""" + if isinstance(dataset, dict): + yield from dataset.items() + return + + try: + yield from enumerate(dataset) + except TypeError as error: + raise TypeError( + "The loaded .pt file must contain an iterable dataset" + ) from error + + +def calculate_statistics( + dataset: Any, +) -> tuple[list[dict[str, Any]], dict[str, Any], Mol]: + """Calculate per-instance molecular diameters and aggregate statistics.""" + distances: list[dict[str, Any]] = [] + molecules: list[Mol] = [] + for index, sample in iter_samples(dataset): + molecule = extract_molecule(sample) + molecules.append(molecule) + distances.append( + { + "index": int(index), + "num_atoms": molecule.GetNumAtoms(), + "largest_distance": molecule_diameter(molecule), + "smiles": Chem.MolToSmiles(molecule), + } + ) + + if not distances: + raise ValueError("The loaded dataset contains no samples") + + values = [item["largest_distance"] for item in distances] + largest_index = values.index(max(values)) + largest_molecule = distances[largest_index] + largest_rdkit_molecule = molecules[largest_index] + summary: dict[str, Any] = { + "num_molecules": len(values), + "min": min(values), + "max": max(values), + "mean": statistics.mean(values), + "median": statistics.median(values), + "stdev": statistics.stdev(values) if len(values) > 1 else 0.0, + "largest_distance_index": largest_molecule["index"], + "largest_distance_smiles": largest_molecule["smiles"], + } + return distances, summary, largest_rdkit_molecule + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Compute the largest topological (bond-path) distance between any " + "two atoms for every molecule in a saved .pt dataset." + ) + ) + parser.add_argument("data_path", type=Path, help="Path to the saved data.pt file") + parser.add_argument( + "--output", + type=Path, + help="Optional JSON file for per-molecule distances and aggregate statistics", + ) + parser.add_argument( + "--image-output", + type=Path, + default=Path("largest_distance_molecule.png"), + help=( + "PNG file for the molecule with the largest shortest path " + "(default: largest_distance_molecule.png)" + ), + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + dataset = torch.load(args.data_path, weights_only=False) + distances, summary, largest_molecule = calculate_statistics(dataset) + + print(f"Molecules: {summary['num_molecules']}") + print(f"Largest distance (min): {summary['min']}") + print(f"Largest distance (max): {summary['max']}") + print(f"Largest distance (mean): {summary['mean']:.3f}") + print(f"Largest distance (median): {summary['median']:.3f}") + print(f"Largest distance (stdev): {summary['stdev']:.3f}") + print( + "Molecule with largest shortest path " + f"(distance={summary['max']}, index={summary['largest_distance_index']}): " + f"{summary['largest_distance_smiles']}" + ) + draw_molecule(largest_molecule, args.image_output) + print(f"Molecule drawing written to {args.image_output}") + + if args.output is not None: + result = {"summary": summary, "molecules": distances} + args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(f"Results written to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/results/visualize_augmented_molecule.py b/results/visualize_augmented_molecule.py index fcc406b..76bd49d 100644 --- a/results/visualize_augmented_molecule.py +++ b/results/visualize_augmented_molecule.py @@ -1,9 +1,19 @@ +try: + import matplotlib + + assert matplotlib.__version__ >= "3.10.9", "Matplotlib >= 3.10.9 required" +except (ImportError, AssertionError) as e: + raise ImportError( + "Matplotlib is required for 2D plotting. Install it with `pip install matplotlib`." + ) from e + import io -import matplotlib +matplotlib.use("Agg") import matplotlib.pyplot as plt import networkx as nx import numpy as np +from chebi_utils.read_molecule import smiles_or_inchi_to_mol from jsonargparse import CLI from PIL import Image from rdkit.Chem import AllChem, BondType, Mol, rdDepictor @@ -23,8 +33,6 @@ GN_WithAtoms_FG_WithAtoms_NoFGE, ) -matplotlib.use("TkAgg") - EDGE_COLOR_MAP = { k.WITHIN_ATOMS_EDGE: "#1f77b4", k.ATOM_FG_EDGE: "#9467bd", @@ -241,13 +249,16 @@ def _draw_simple(G: nx.Graph) -> None: plt.show() -def _draw_3d(G: nx.Graph, mol: Mol) -> None: +def _draw_3d( + G: nx.Graph, mol: Mol, output_html: str = "augmented_molecule_3d.html" +) -> None: """ Visualize the graph in 3D using Plotly. Args: G (nx.Graph): Augmented molecular graph. mol (Chem.Mol): RDKit molecule object for 3D coordinates. + output_html (str): Path where the interactive Plotly figure is saved. Raises: ImportError: If Plotly is not installed. @@ -312,7 +323,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: ] ) graph_pos = { - graph_node: (graph_pos_arr[:, 0].mean(), graph_pos_arr[:, 1].mean(), 4) + graph_node: (graph_pos_arr[:, 0].mean(), graph_pos_arr[:, 1].mean(), -2) } pos = {**atom_pos, **fg_pos, **graph_pos} @@ -343,6 +354,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: z=zs, mode="lines", line=dict(color=EDGE_COLOR_MAP[edge_type], width=4), + opacity=0.5, name=edge_type, hoverinfo="none", ) @@ -360,15 +372,25 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: x=pos_x, y=pos_y, z=pos_z, - mode="markers+text", + mode="markers", marker=dict(size=8, color=node_colors, opacity=0.9), - text=node_names, - textposition="top center", hovertext=node_ids, hoverinfo="text", ) - fig = go.Figure(data=edge_traces + [node_trace]) + label_trace = go.Scatter3d( + x=pos_x, + y=pos_y, + z=pos_z, + mode="text", + text=node_names, + textposition="top center", + textfont=dict(color="black", size=12), + hoverinfo="none", + showlegend=False, + ) + + fig = go.Figure(data=edge_traces + [node_trace, label_trace]) fig.update_layout( title="3D Augmented Molecule Graph", showlegend=True, @@ -379,6 +401,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: ), margin=dict(l=0, r=0, b=0, t=40), ) + # fig.write_html(output_html) fig.show() @@ -387,6 +410,7 @@ def plot_augmented_graph( augmented_molecule: dict, mol: Mol, plot_type: str, + output_html: str = "augmented_molecule_3d.html", ) -> None: """ Main plotting function to visualize the augmented graph. @@ -396,6 +420,7 @@ def plot_augmented_graph( augmented_molecule (dict): Augmented Molecule. mol (Chem.Mol): RDKit molecule object. plot_type (str): One of ["simple", "h", "3d"]. + output_html (str): Path where the 3D Plotly figure is saved. """ G = _create_graph( edge_index, augmented_molecule["nodes"], augmented_molecule["edges"] @@ -406,7 +431,7 @@ def plot_augmented_graph( elif plot_type == "simple": _draw_simple(G) elif plot_type == "3d": - _draw_3d(G, mol) + _draw_3d(G, mol, output_html) else: raise ValueError(f"Unknown plot type: {plot_type}") @@ -471,6 +496,7 @@ def plot( smiles: str = "OC(=O)c1ccccc1O", plot_type: str = "simple", reader: str = "w_fge_w_gn", + output_html: str = "augmented_molecule_3d.html", ) -> None: """ Plot an augmented molecular graph from SMILES. @@ -491,9 +517,10 @@ def plot( - 'gnwa_fgwa_wfge': Graph node connected to atoms and FG nodes connected to atoms, with FG edges. - 'gn_wall_fgwa_nfge': Graph node connected to all atoms and FG nodes, no FG edges. - 'gn_wall_fgwa_wfge': Graph node connected to all atoms and FG nodes, with FG edges. + output_html (str): Output path for the interactive 3D HTML figure. """ fg_reader = READER[reader]() - mol = fg_reader._smiles_to_mol(smiles) + mol = smiles_or_inchi_to_mol(smiles) if mol is None: raise ValueError(f"Invalid SMILES: {smiles}") @@ -502,7 +529,9 @@ def plot( return edge_index, augmented_molecule = fg_reader._create_augmented_graph(mol) - plot_augmented_graph(edge_index, augmented_molecule, mol, plot_type) + plot_augmented_graph( + edge_index, augmented_molecule, mol, plot_type, output_html + ) if __name__ == "__main__": diff --git a/scripts/analyze_amgpool.py b/scripts/analyze_amgpool.py new file mode 100644 index 0000000..05f3b21 --- /dev/null +++ b/scripts/analyze_amgpool.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""Measure representation scales and classifier contributions in AMGPool models. + +Example (preferred; uses hyperparameters stored in the checkpoint): + python scripts/analyze_amgpool.py \ + --checkpoint path/to/model.ckpt \ + --output-dir analysis/bbbp + +Explicit model and data configs can also be supplied together when the +checkpoint does not contain Lightning hyperparameters. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import importlib +import json +from pathlib import Path +from typing import Any + +import torch +import yaml +from chebai.models.base import ChebaiBaseNet +from chebai.preprocessing.datasets.base import XYBaseDataModule +from lightning.pytorch.cli import instantiate_module + +LEVELS = ("atom", "motif", "graph") + + +def _import_class(path: str) -> type: + module_name, class_name = path.rsplit(".", 1) + return getattr(importlib.import_module(module_name), class_name) + + +def _instantiate(config: dict[str, Any]) -> Any: + class_path = config.get("class_path") + if not class_path: + raise ValueError("Expected a config with a 'class_path' entry.") + init_args = config.get("init_args", {}) + return _import_class(class_path)(**init_args) + + +def _load_yaml(path: Path) -> dict[str, Any]: + with path.open() as handle: + config = yaml.safe_load(handle) + if not isinstance(config, dict): + raise ValueError(f"{path} must contain a YAML mapping.") + return config + + +def _load_checkpoint(model: torch.nn.Module, path: Path) -> None: + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + state_dict = checkpoint.get("state_dict", checkpoint) + if not isinstance(state_dict, dict): + raise ValueError(f"{path} does not contain a state dictionary.") + + model_state = model.state_dict() + candidates = [] + for prefix in ("", "model.", "module."): + candidates.append( + { + key.removeprefix(prefix): value + for key, value in state_dict.items() + if key.startswith(prefix) + } + ) + matching = max( + (candidate for candidate in candidates if candidate), + key=lambda candidate: len(set(candidate) & set(model_state)), + default={}, + ) + filtered = {key: value for key, value in matching.items() if key in model_state} + if not filtered: + raise ValueError(f"{path} contains no weights matching the configured model.") + result = model.load_state_dict(filtered, strict=False) + if result.missing_keys: + raise RuntimeError( + f"Checkpoint {path} is missing model weights: " + + ", ".join(sorted(result.missing_keys)) + ) + + +def _load_experiment( + checkpoint_path: Path, +) -> tuple[ChebaiBaseNet, XYBaseDataModule]: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + try: + data_hparams = copy.deepcopy(checkpoint["datamodule_hyper_parameters"]) + model_hparams = copy.deepcopy(checkpoint["hyper_parameters"]) + except KeyError as error: + raise ValueError( + "Checkpoint must contain Lightning model and datamodule hyperparameters " + "when --model-config and --data-config are omitted." + ) from error + + data_hparams.pop("_instantiator", None) + model_hparams.pop("_instantiator", None) + data = instantiate_module(XYBaseDataModule, data_hparams) + + load_kwargs = ( + {"pretrained_checkpoint": None} + if "pretrained_checkpoint" in model_hparams + else {} + ) + model = ChebaiBaseNet.load_from_checkpoint( + checkpoint_path, + map_location="cpu", + **load_kwargs, + ) + return model, data + + +def _prepare_model_batch(model: ChebaiBaseNet, batch: Any, batch_index: int) -> dict: + """Apply the repository's standard XYData-to-model batch conversion.""" + if not hasattr(batch, "to") or not hasattr(batch, "x"): + raise TypeError( + "The test dataloader must return an XYData batch with an x attribute." + ) + return model._process_batch(batch.to(model.device), batch_index) + + +def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + if not rows: + return + with path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def _plot_summary(summary: dict[str, Any], output_dir: Path) -> None: + try: + import matplotlib.pyplot as plt + except ImportError: + return + + metrics = { + "embedding_rms": [summary["embedding"][level]["rms"] for level in LEVELS], + "weight_fraction": [summary["classifier"][level]["alpha"] for level in LEVELS], + "logit_fraction": [summary["contribution"][level]["beta"] for level in LEVELS], + } + fig, axes = plt.subplots(1, 3, figsize=(12, 4)) + for axis, (title, values) in zip(axes, metrics.items()): + axis.bar(LEVELS, values) + axis.set_title(title.replace("_", " ")) + axis.tick_params(axis="x", rotation=25) + fig.tight_layout() + fig.savefig(output_dir / "representation_summary.png", dpi=180) + plt.close(fig) + + +def analyse(args: argparse.Namespace) -> dict[str, Any]: + model, data = _load_experiment( + Path(args.checkpoint), + ) + model.to(args.device).eval() + + linear_layers = [ + layer for layer in model.lin_sequential if isinstance(layer, torch.nn.Linear) + ] + if len(linear_layers) != 1: + raise ValueError( + "This analysis requires a single-layer linear classifier; " + f"found {len(linear_layers)} linear layers." + ) + classifier = linear_layers[0] + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + loader = data.test_dataloader() + if isinstance(loader, (list, tuple)): + if len(loader) != 1: + raise ValueError("Multiple dataloaders are not supported by this script.") + loader = loader[0] + + sums = {level: {"std": 0.0, "rms": 0.0, "l2": 0.0} for level in LEVELS} + contribution_sum = {level: 0.0 for level in LEVELS} + class_contribution = torch.zeros((classifier.out_features, 3), dtype=torch.float64) + ablation_sum = torch.zeros(3, dtype=torch.float64) + sample_count = 0 + batch_limit = args.max_batches + + with torch.inference_mode(): + for batch_index, raw_batch in enumerate(loader): + if batch_limit is not None and batch_index >= batch_limit: + break + batch = _prepare_model_batch(model, raw_batch, batch_index) + representations = model.pool_representations(batch) + batch_size = representations[0].shape[0] + sample_count += batch_size + blocks = [] + for index, (level, representation) in enumerate( + zip(LEVELS, representations) + ): + representation_cpu = representation.detach().cpu() + sums[level]["std"] += ( + float(representation_cpu.std().item()) * batch_size + ) + sums[level]["rms"] += ( + float(representation_cpu.square().mean().sqrt().item()) * batch_size + ) + sums[level]["l2"] += ( + float(representation_cpu.norm(dim=1).mean().item()) * batch_size + ) + start = index * representation.shape[1] + weight = classifier.weight[:, start : start + representation.shape[1]] + contribution = representation @ weight.t() + contribution_sum[level] += ( + float(contribution.norm(dim=1).mean().item()) * batch_size + ) + class_contribution[:, index] += ( + contribution.abs().sum(dim=0).detach().cpu().double() + ) + blocks.append(representation) + full_logits = classifier(torch.cat(blocks, dim=1)) + for index in range(3): + ablated = list(blocks) + ablated[index] = torch.zeros_like(ablated[index]) + ablation_sum[index] += ( + (full_logits - classifier(torch.cat(ablated, dim=1))) + .abs() + .sum() + .detach() + .cpu() + .double() + ) + + if sample_count == 0: + raise RuntimeError("The selected split produced no samples.") + widths = [classifier.in_features // len(LEVELS)] * len(LEVELS) + if classifier.in_features % len(LEVELS): + raise ValueError( + "The classifier input dimension is not divisible into three AMGPool blocks." + ) + weight_norms = [] + for index, level in enumerate(LEVELS): + width = widths[index] + weight_norms.append( + float( + classifier.weight[:, index * width : (index + 1) * width].norm().item() + ) + ) + weight_total = sum(weight_norms) + contribution_total = sum(contribution_sum.values()) + if weight_total == 0 or contribution_total == 0: + raise RuntimeError( + "The checkpoint has zero total classifier weight or logit contribution; " + "normalized shares are undefined." + ) + summary = { + "checkpoint": str(Path(args.checkpoint).resolve()), + "split": "test", + "samples": sample_count, + "embedding": { + level: { + metric: value / sample_count for metric, value in sums[level].items() + } + for level in LEVELS + }, + "classifier": { + level: { + "frobenius_norm": weight_norms[index], + "alpha": weight_norms[index] / weight_total, + } + for index, level in enumerate(LEVELS) + }, + "contribution": { + level: { + "mean_logit_l2": contribution_sum[level] / sample_count, + "beta": contribution_sum[level] / contribution_total, + } + for level in LEVELS + }, + "ablation_logit_change_mean": { + level: float( + ablation_sum[index].item() / sample_count / classifier.out_features + ) + for index, level in enumerate(LEVELS) + }, + } + class_rows = [ + { + "class_index": index, + **{ + f"{level}_mean_abs_logit_contribution": float( + class_contribution[index, level_index].item() / sample_count + ) + for level_index, level in enumerate(LEVELS) + }, + } + for index in range(classifier.out_features) + ] + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + _write_csv(out_dir / "class_contributions.csv", class_rows) + _plot_summary(summary, out_dir) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--output-dir", type=Path, default=Path("amgpool_analysis")) + parser.add_argument("--device", default="cpu") + parser.add_argument("--max-batches", type=int, default=None) + analyse(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_collate.py b/tests/unit/test_collate.py new file mode 100644 index 0000000..8c81979 --- /dev/null +++ b/tests/unit/test_collate.py @@ -0,0 +1,63 @@ +import unittest + +import torch +from torch_geometric.data import Data as GeomData + +from chebai_graph.preprocessing.collate import GraphCollator + + +class TestGraphCollator(unittest.TestCase): + @staticmethod + def _graph(): + return GeomData( + x=torch.tensor([[6]]), + edge_index=torch.empty((2, 0), dtype=torch.long), + ) + + def test_mixed_missing_labels_do_not_enter_graph_collate(self): + batch = GraphCollator()( + [ + { + "features": self._graph(), + "labels": [1, None, 0, 0, 1], + "ident": "a", + }, + {"features": self._graph(), "labels": None, "ident": "b"}, + { + "features": self._graph(), + "labels": [0, 1, None, 0, 1], + "ident": "c", + }, + ] + ) + + expected_labels = torch.tensor( + [ + [True, False, False, False, True], + [False, True, False, False, True], + ] + ) + self.assertTrue(torch.equal(batch.y, expected_labels)) + self.assertEqual( + batch.additional_fields["loss_kwargs"]["non_null_labels"], [0, 2] + ) + self.assertNotIn("y", batch.x[0]) + + def test_custom_metadata_survives_repeated_collation(self): + metadata = {"nodes": {"atom_nodes": [6]}, "edges": {}} + graph = self._graph() + graph.custom_metadata = metadata + row = {"features": graph, "labels": None, "ident": "a"} + + first_batch = GraphCollator()([row]) + second_batch = GraphCollator()([row]) + + self.assertEqual(first_batch.additional_fields["custom_metadata"], [[metadata]]) + self.assertEqual( + second_batch.additional_fields["custom_metadata"], [[metadata]] + ) + self.assertIn("custom_metadata", graph) + + +if __name__ == "__main__": + unittest.main() diff --git a/tutorials/classwise_stats.ipynb b/tutorials/classwise_stats.ipynb index 072a128..a64d8e0 100644 --- a/tutorials/classwise_stats.ipynb +++ b/tutorials/classwise_stats.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "id": "07f2eb38", "metadata": {}, "outputs": [], @@ -12,172 +12,46 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": null, "id": "eb64bfee", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
labelres_baseline_propsres_wfgn_wfge_wgn_wprops_apoolgatv2_wfgn_wfge_wgn_wprops_apoolgatv2 - baselineres - baseline
01461800.00.00000.00000.00000.0000
1600040.00.11760.00000.00000.1176
2167330.00.23530.00000.00000.2353
3222990.00.28570.00000.00000.2857
4178150.00.12500.11110.11110.1250
\n", - "
" - ], - "text/plain": [ - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0 0.0000 \n", - "1 60004 0.0 0.1176 \n", - "2 16733 0.0 0.2353 \n", - "3 22299 0.0 0.2857 \n", - "4 17815 0.0 0.1250 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \n", - "0 0.0000 0.0000 0.0000 \n", - "1 0.0000 0.0000 0.1176 \n", - "2 0.0000 0.0000 0.2353 \n", - "3 0.0000 0.0000 0.2857 \n", - "4 0.1111 0.1111 0.1250 " - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ - "df = pd.read_excel(\"G:\\github-aditya0by0\\python-chebai-graph\\classwise_f1_score.xlsx\")\n", + "df = pd.read_csv(\"/home/aditya/python-chebai-graph/combined_f1_scores.csv\")\n", + "df[\"aug-baseline\"] = df[\"augmented\"] - df[\"baseline\"]\n", "df.head()" ] }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "id": "14c10198", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(1528, 6)" - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "df.shape" ] }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "32062eb6", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "Index(['label', 'res_baseline_props', 'res_wfgn_wfge_wgn_wprops_apool',\n", - " 'gatv2_wfgn_wfge_wgn_wprops_apool', 'gatv2 - baseline ',\n", - " 'res - baseline'],\n", - " dtype='object')" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "df.columns" ] }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, "id": "5cd9265a", "metadata": {}, "outputs": [], "source": [ - "BASELINE = \"res_baseline_props\"\n", - "RES = \"res_wfgn_wfge_wgn_wprops_apool\"\n", - "GATv2 = \"gatv2_wfgn_wfge_wgn_wprops_apool\"\n", - "GATv2_BASELINE_DIFF = \"gatv2 - baseline \"\n", - "RES_BASELINE_DIFF = \"res - baseline\"" + "BASELINE = \"baseline\"\n", + "AUGMENTED = \"augmented\"\n", + "DIFF_COL = \"aug-baseline\"" ] }, { @@ -192,45 +66,24 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "id": "b38fa321", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean improvement over baseline:\n", - " res - baseline 0.054853\n", - "gatv2 - baseline 0.068136\n", - "dtype: float64\n", - "\n", - "Median improvement over baseline:\n", - " res - baseline 0.03330\n", - "gatv2 - baseline 0.04535\n", - "dtype: float64\n", - "\n", - "Variation in improvement:\n", - " res - baseline 0.144831\n", - "gatv2 - baseline 0.139299\n", - "dtype: float64\n" - ] - } - ], + "outputs": [], "source": [ "# mean improvement\n", - "mean_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].mean()\n", + "mean_diff = df[[DIFF_COL]].mean()\n", "print(\"Mean improvement over baseline:\\n\", mean_diff)\n", "print()\n", "\n", "# median improvement\n", - "median_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].median()\n", + "median_diff = df[[DIFF_COL]].median()\n", "print(\"Median improvement over baseline:\\n\", median_diff)\n", "print()\n", "\n", "\n", "# standard deviation of improvements\n", - "std_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].std()\n", + "std_diff = df[[DIFF_COL]].std()\n", "print(\"Variation in improvement:\\n\", std_diff)" ] }, @@ -244,24 +97,13 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "id": "d046db01", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Labels improved in ResGated Final model over baseline: 1010\n", - "Labels improved in GATv2 Final model over baseline: 1091\n" - ] - } - ], + "outputs": [], "source": [ - "res_improvement = (df[RES_BASELINE_DIFF] > 0).sum()\n", - "gat_v2_improvement = (df[GATv2_BASELINE_DIFF] > 0).sum()\n", + "gat_v2_improvement = (df[DIFF_COL] > 0).sum()\n", "\n", - "print(f\"Labels improved in ResGated Final model over baseline: {res_improvement}\")\n", "print(f\"Labels improved in GATv2 Final model over baseline: {gat_v2_improvement}\")" ] }, @@ -275,77 +117,15 @@ }, { "cell_type": "code", - "execution_count": 27, - "id": "3c3972a1", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "best_model\n", - "gatv2_wfgn_wfge_wgn_wprops_apool 678\n", - "res_wfgn_wfge_wgn_wprops_apool 524\n", - "res_baseline_props 326\n", - "Name: count, dtype: int64\n" - ] - } - ], - "source": [ - "df[\"best_model\"] = df[[BASELINE, RES, GATv2]].idxmax(axis=1)\n", - "print(df[\"best_model\"].value_counts())" - ] - }, - { - "cell_type": "code", - "execution_count": 32, + "execution_count": null, "id": "02ecb78e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Top 10 classes with highest improvement for res - baseline:\n", - " 25 0.8889\n", - "24 0.8889\n", - "37 0.7692\n", - "34 0.7143\n", - "35 0.7143\n", - "29 0.7143\n", - "77 0.6889\n", - "41 0.6487\n", - "36 0.6400\n", - "23 0.6364\n", - "Name: res - baseline, dtype: float64\n", - "\n", - "Top 10 classes with highest improvement for gatv2 - baseline :\n", - " 25 0.8889\n", - "24 0.8889\n", - "37 0.7692\n", - "34 0.7143\n", - "35 0.7143\n", - "29 0.7143\n", - "77 0.6889\n", - "41 0.6487\n", - "36 0.6400\n", - "23 0.6364\n", - "Name: res - baseline, dtype: float64\n" - ] - } - ], + "outputs": [], "source": [ "# Classes with highest difference\n", - "top_diff_res = df[RES_BASELINE_DIFF].sort_values(ascending=False).head(10)\n", + "top_diff_res = df[DIFF_COL].sort_values(ascending=False).head(10)\n", "print(\n", - " f\"\\nTop 10 classes with highest improvement for {RES_BASELINE_DIFF}:\\n\",\n", - " top_diff_res,\n", - ")\n", - "\n", - "top_diff_gat = df[GATv2_BASELINE_DIFF].sort_values(ascending=False).head(10)\n", - "print(\n", - " f\"\\nTop 10 classes with highest improvement for {GATv2_BASELINE_DIFF}:\\n\",\n", + " f\"\\nTop 10 classes with highest improvement for {DIFF_COL}:\\n\",\n", " top_diff_res,\n", ")" ] @@ -355,397 +135,155 @@ "execution_count": null, "id": "524a01d8", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
labelres_baseline_propsres_wfgn_wfge_wgn_wprops_apoolgatv2_wfgn_wfge_wgn_wprops_apoolgatv2 - baselineres - baselinebest_model
01461800.00.00.00000.00000.0res_baseline_props
9257110.00.00.25000.25000.0gatv2_wfgn_wfge_wgn_wprops_apool
10390930.00.00.25000.25000.0gatv2_wfgn_wfge_wgn_wprops_apool
1841940.00.00.46150.46150.0gatv2_wfgn_wfge_wgn_wprops_apool
19176080.00.00.46150.46150.0gatv2_wfgn_wfge_wgn_wprops_apool
........................
1523766291.01.01.00000.00000.0res_baseline_props
1524787761.01.01.00000.00000.0res_baseline_props
1525792041.01.01.00000.00000.0res_baseline_props
15261327421.01.01.00000.00000.0res_baseline_props
15271366371.01.01.00000.00000.0res_baseline_props
\n", - "

87 rows × 7 columns

\n", - "
" - ], - "text/plain": [ - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0 0.0 \n", - "9 25711 0.0 0.0 \n", - "10 39093 0.0 0.0 \n", - "18 4194 0.0 0.0 \n", - "19 17608 0.0 0.0 \n", - "... ... ... ... \n", - "1523 76629 1.0 1.0 \n", - "1524 78776 1.0 1.0 \n", - "1525 79204 1.0 1.0 \n", - "1526 132742 1.0 1.0 \n", - "1527 136637 1.0 1.0 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "0 0.0000 0.0000 0.0 \n", - "9 0.2500 0.2500 0.0 \n", - "10 0.2500 0.2500 0.0 \n", - "18 0.4615 0.4615 0.0 \n", - "19 0.4615 0.4615 0.0 \n", - "... ... ... ... \n", - "1523 1.0000 0.0000 0.0 \n", - "1524 1.0000 0.0000 0.0 \n", - "1525 1.0000 0.0000 0.0 \n", - "1526 1.0000 0.0000 0.0 \n", - "1527 1.0000 0.0000 0.0 \n", - "\n", - " best_model \n", - "0 res_baseline_props \n", - "9 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "10 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "18 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "19 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "... ... \n", - "1523 res_baseline_props \n", - "1524 res_baseline_props \n", - "1525 res_baseline_props \n", - "1526 res_baseline_props \n", - "1527 res_baseline_props \n", - "\n", - "[87 rows x 7 columns]" - ] - }, - "execution_count": 42, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Classes with minimal or no difference\n", "minimal_diff_threshold = 0.01 # you can adjust\n", - "minimal_diff = df[df[RES_BASELINE_DIFF].abs() <= minimal_diff_threshold]\n", + "minimal_diff = df[df[DIFF_COL].abs() <= minimal_diff_threshold]\n", "minimal_diff" ] }, { "cell_type": "code", - "execution_count": 37, + "execution_count": null, "id": "5f15e872", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where baseline performs better than our model:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "45 59999 0.1176 0.0541 \n", - "47 59266 0.1176 0.1143 \n", - "50 25095 0.1333 0.0000 \n", - "70 48927 0.2000 0.0000 \n", - "72 83943 0.2000 0.0000 \n", - "... ... ... ... \n", - "1508 24846 1.0000 0.9333 \n", - "1509 26187 1.0000 0.9333 \n", - "1510 78231 1.0000 0.9333 \n", - "1511 26377 1.0000 0.9565 \n", - "1512 38834 1.0000 0.9677 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "45 0.1250 0.0074 -0.0635 \n", - "47 0.5385 0.4209 -0.0033 \n", - "50 0.2667 0.1334 -0.1333 \n", - "70 0.0000 -0.2000 -0.2000 \n", - "72 0.2000 0.0000 -0.2000 \n", - "... ... ... ... \n", - "1508 1.0000 0.0000 -0.0667 \n", - "1509 1.0000 0.0000 -0.0667 \n", - "1510 1.0000 0.0000 -0.0667 \n", - "1511 1.0000 0.0000 -0.0435 \n", - "1512 1.0000 0.0000 -0.0323 \n", - "\n", - " best_model \n", - "45 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "47 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "50 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "70 res_baseline_props \n", - "72 res_baseline_props \n", - "... ... \n", - "1508 res_baseline_props \n", - "1509 res_baseline_props \n", - "1510 res_baseline_props \n", - "1511 res_baseline_props \n", - "1512 res_baseline_props \n", - "\n", - "[431 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where baseline performs better\n", - "baseline_better = df[df[RES_BASELINE_DIFF] < 0]\n", + "baseline_better = df[df[DIFF_COL] < 0]\n", "print(\"\\nClasses where baseline performs better than our model:\\n\", baseline_better)" ] }, { "cell_type": "code", - "execution_count": 38, + "execution_count": null, "id": "0c4292bf", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where both baseline and res_wfgn_wfge_wgn_wprops_apool have F1 < 0.5:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0000 0.0000 \n", - "1 60004 0.0000 0.1176 \n", - "2 16733 0.0000 0.2353 \n", - "3 22299 0.0000 0.2857 \n", - "4 17815 0.0000 0.1250 \n", - ".. ... ... ... \n", - "265 32877 0.4615 0.4706 \n", - "271 139358 0.4667 0.4000 \n", - "283 35346 0.4762 0.4706 \n", - "284 35507 0.4762 0.4737 \n", - "295 24586 0.4878 0.4762 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "0 0.0000 0.0000 0.0000 \n", - "1 0.0000 0.0000 0.1176 \n", - "2 0.0000 0.0000 0.2353 \n", - "3 0.0000 0.0000 0.2857 \n", - "4 0.1111 0.1111 0.1250 \n", - ".. ... ... ... \n", - "265 0.5333 0.0718 0.0091 \n", - "271 0.5517 0.0850 -0.0667 \n", - "283 0.3529 -0.1233 -0.0056 \n", - "284 0.3871 -0.0891 -0.0025 \n", - "295 0.6667 0.1789 -0.0116 \n", - "\n", - " best_model \n", - "0 res_baseline_props \n", - "1 res_wfgn_wfge_wgn_wprops_apool \n", - "2 res_wfgn_wfge_wgn_wprops_apool \n", - "3 res_wfgn_wfge_wgn_wprops_apool \n", - "4 res_wfgn_wfge_wgn_wprops_apool \n", - ".. ... \n", - "265 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "271 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "283 res_baseline_props \n", - "284 res_baseline_props \n", - "295 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "\n", - "[148 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where both baseline and our model perform worse\n", "# Assuming \"worse\" means below a threshold, e.g., f1 < 0.5\n", "threshold = 0.5\n", - "both_worse = df[(df[BASELINE] < threshold) & (df[RES] < threshold)]\n", - "print(f\"\\nClasses where both baseline and {RES} have F1 < {threshold}:\\n\", both_worse)" + "both_worse = df[(df[BASELINE] < threshold) & (df[AUGMENTED] < threshold)]\n", + "print(\n", + " f\"\\nClasses where both baseline and Augmented have F1 < {threshold}:\\n\", both_worse\n", + ")" ] }, { "cell_type": "code", - "execution_count": 40, + "execution_count": null, "id": "8948b7bb", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where baseline performs better than our model:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "45 59999 0.1176 0.0541 \n", - "47 59266 0.1176 0.1143 \n", - "50 25095 0.1333 0.0000 \n", - "70 48927 0.2000 0.0000 \n", - "72 83943 0.2000 0.0000 \n", - "... ... ... ... \n", - "1508 24846 1.0000 0.9333 \n", - "1509 26187 1.0000 0.9333 \n", - "1510 78231 1.0000 0.9333 \n", - "1511 26377 1.0000 0.9565 \n", - "1512 38834 1.0000 0.9677 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "45 0.1250 0.0074 -0.0635 \n", - "47 0.5385 0.4209 -0.0033 \n", - "50 0.2667 0.1334 -0.1333 \n", - "70 0.0000 -0.2000 -0.2000 \n", - "72 0.2000 0.0000 -0.2000 \n", - "... ... ... ... \n", - "1508 1.0000 0.0000 -0.0667 \n", - "1509 1.0000 0.0000 -0.0667 \n", - "1510 1.0000 0.0000 -0.0667 \n", - "1511 1.0000 0.0000 -0.0435 \n", - "1512 1.0000 0.0000 -0.0323 \n", - "\n", - " best_model \n", - "45 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "47 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "50 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "70 res_baseline_props \n", - "72 res_baseline_props \n", - "... ... \n", - "1508 res_baseline_props \n", - "1509 res_baseline_props \n", - "1510 res_baseline_props \n", - "1511 res_baseline_props \n", - "1512 res_baseline_props \n", - "\n", - "[431 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where baseline performs better\n", - "baseline_better = df[df[RES_BASELINE_DIFF] < 0]\n", + "baseline_better = df[df[DIFF_COL] < 0]\n", "print(\"\\nClasses where baseline performs better than our model:\\n\", baseline_better)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0bc0f4c1", + "metadata": {}, + "outputs": [], + "source": [ + "len(df)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61ab61c7", + "metadata": {}, + "outputs": [], + "source": [ + "print(len(df[df[DIFF_COL] >= 0.1]))\n", + "print(len(df[df[DIFF_COL] >= 0.01]))\n", + "\n", + "print(len(df[abs(df[DIFF_COL]) < 0.01]))\n", + "\n", + "print(len(df[df[DIFF_COL] <= -0.1]))\n", + "print(len(df[df[DIFF_COL] <= -0.01]))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "828c727f", + "metadata": {}, + "outputs": [], + "source": [ + "1005 + 221 + 283" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b5476b8", + "metadata": {}, + "outputs": [], + "source": [ + "len(df[df[DIFF_COL] == 0.0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2caab4d5", + "metadata": {}, + "outputs": [], + "source": [ + "len(df[(df[DIFF_COL] == 0.0) & ~((df[BASELINE] == 1.0) & (df[AUGMENTED] == 1.0))])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb50e15a", + "metadata": {}, + "outputs": [], + "source": [ + "for label in [57880, 28802, 49172, 140949, 59268, 59412, 25627]:\n", + " print(label, label in df[df[DIFF_COL] >= 0.1][\"label\"].tolist())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bba6bb1a", + "metadata": {}, + "outputs": [], + "source": [ + "df[df[DIFF_COL] >= 0.1]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20e2ea95", + "metadata": {}, + "outputs": [], + "source": [ + "df[(df[BASELINE] <= 0.1) & (df[AUGMENTED] >= 0.8)]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9e043cf", + "metadata": {}, + "outputs": [], + "source": [ + "df.iloc[df[DIFF_COL].sort_values(ascending=False).head(10).index]" + ] } ], "metadata": { "kernelspec": { - "display_name": "gnn3", + "display_name": "python-chebai-graph (3.10.20)", "language": "python", "name": "python3" }, @@ -759,7 +297,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.16" + "version": "3.10.20" } }, "nbformat": 4,