DOI: 10.5281/zenodo.16927285
title: Code and Data for "Molecular Contrastive Learning with Graph Attention Network (MoCL-GAT) for Enhanced Molecular Representation"

description: Description

This archive contains the source code, datasets, pre-trained models, and environment specification for the paper "Molecular  Contrastive Learning with Graph Attention Network (MoCL-GAT) for Enhanced Molecular Representation". The contents of this archive allow for the full reproduction of the results presented in our study on MoCL-GAT.

Workflow Overview

The MoCL-GAT framework consists of two main stages:

1.  Pre-training (Self-Supervised): A GAT model is pre-trained on a large-scale unlabeled chemical database (e.g., ChEMBL). 2.  Fine-tuning (Supervised): The pre-trained model is then fine-tuned on a smaller, labeled dataset for a specific downstream task (e.g., aqueous solubility (ESOL), predicting Blood-Brain Barrier Penetration (BBBP)).

Repository Structure

.├── training_files/ # Directory for all dataset files│ ├── chembl/ # Pre-training dataset│ │ ├── smiles.tsv│ │ ├── comp_targ_rdkit2.tsv│ │ └── data/folds/│ │ ├── train_fold_setting1.txt # Contains indices for all 5 validation folds│ │ └── test_fold_setting1.txt # Contains indices for the held-out test set│ ├── esol/ # Example downstream dataset│ │ └── esol.csv│ └── ... # Other datasets (clintox, bbbp, etc.)│├── trained_models/ # Directory for saved model weights│ └── chembl/ # Saved weights from pre-training on ChEMBL│ ├── best_state_dict_loss1.pth│ └── ...│├── MoCL-GAT_env.yaml Conda environment file├── data_processing.py # Data loading and pre-processing utilities├── evaluation_metrics.py # Functions for calculating performance metrics├── models.py # GAT and ContrastiveGAT model definitions├── train_FNN_DTI.py # Core training, validation, and testing loops└── main_training.py # Main script to run experiments with argument parsing

Installation

The recommended way to set up the environment is by using the provided Conda environment file.

1.  Download and Unpack the Archive:    Download the `.zip` file and extract it to a location of your choice.

2.  Create and activate the Conda environment:    Navigate to the extracted directory in your terminal and run the following commands:    ```bash    conda env create -f MoCL-GAT_env.yml    conda activate MoCL-GAT_env    ```    This will install all the necessary dependencies with the correct versions used in our study.

Data Preparation

The `training_files` directory included in this archive is already structured correctly for all experiments.

Downstream Task Datasets (e.g., from MoleculeNet):Each dataset is located in its own subdirectory (e.g., `training_files/bbbp/`) and contains a single `.csv` file.*   Example format for `bbbp.csv`:    ```csv    smiles,p_np    [Cl].CC(C)NCC(O)COc1cccc2ccccc12,1    C(=O)(OC(C)(C)C)CCCc1ccc(cc1)N(CCCl)CCCl,1    ```

Pre-training Dataset (ChEMBL):The `training_files/chembl/` directory contains all necessary files for the self-supervised pre-training, including SMILES, RDKit descriptors, and the 5-fold cross-validation splits.

How to Run the Experiments

The `main_training.py` script is the main entry point for all experiments. The key arguments are:*   `--setting`: `1` for 5-fold CV pre-training, `2` for train/val/test split on downstream tasks.*   `--td`: Target dataset name (e.g., `bbbp`, `esol`, `chembl`).*   `--selfsuper`: `1` to enable the self-supervised pre-training task.*   `--tlf`: `1` to enable transfer learning (fine-tuning) from a pre-trained model.*   `--epoch`: Number of training epochs.*   `--rflag`: to enable the regression task.*   `--nc`: number of classes for the multiclass classification task.

Scenario 1: Training from Scratch (Baseline)

This trains a model on a downstream task (e.g., `ESOL`) without any pre-training.

```bashpython main_training.py --setting 2 --epoch 500 --tlf 0 --td esol --rflag 1```

Scenario 2: Self-Supervised Pre-training

This runs the MoCL-GAT self-supervised task on the included `ChEMBL` dataset. The resulting models will be saved in `trained_models/chembl/`.Warning: This pre-training process is computationally intensive. Running the full 100 epochs on the large ChEMBL dataset may take several days to complete, depending on the specifications of your GPU.```bashpython main_training.py --setting 1 --epoch 100 --td chembl --selfsuper 1 --rflag 1```

Scenario 3: Transfer Learning (Fine-tuning)

This loads the pre-trained model from Scenario 2 and fine-tunes it on the `ESOL` dataset.```bashpython main_training.py --setting 2 --epoch 500 --tlf 1 --td esol --rflag 1```

Pre-trained Models

The pre-trained models generated from the `ChEMBL` dataset are included in this archive in the `trained_models/chembl/` directory. These five models (`best_state_dict_loss1.pth` through `best_state_dict_loss5.pth`) correspond to the five folds of the cross-validation pre-training. The fine-tuning script (`--tlf 1`) will automatically load these models for downstream tasks.