GREAMolecularPredictor

Graph rationales and environment-based data augmentation.

from torch_molecule import GREAMolecularPredictor

Examples · Methods

Constructor

class GREAMolecularPredictor(gamma=0.4, num_task=1, task_type='regression', num_layer=5, hidden_size=300, gnn_type='gin-virtual', drop_ratio=0.5, norm_layer='batch_norm', graph_pooling='sum', augmented_feature=None, batch_size=128, epochs=500, learning_rate=0.001, weight_decay=0.0, grad_clip_value=None, patience=50, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, loss_criterion=None, evaluate_criterion=None, evaluate_higher_better=None, verbose='none', device=None, model_name='GREAMolecularPredictor')

This predictor implements GREA model from the paper “Graph Rationalization with Environment-based Augmentations”.

During model training, it learns the rationales (explainable subgraphs). Use them for data augmentation. During prediction, the model uses the rationales to make predictions and the rationales themselves can also explain the predictions.

References

• Graph Rationalization with Environment-based Augmentations. https://dl.acm.org/doi/10.1145/3534678.3539347

• Code: https://github.com/liugangcode/GREA

Parameters:

• gamma (float, default=0.4) – GREA-specific parameter that penalize the size of the rationales (ratio between the number of nodes in the rationales and the number of nodes in the original graph).

• num_task (int, default=1) – Number of prediction tasks.

• task_type (str, default="regression") – Type of prediction task, either “regression” or “classification”.

• num_layer (int, default=5) – Number of GNN layers.

• hidden_size (int, default=300) – Dimension of hidden node features.

• gnn_type (str, default="gin-virtual") – Type of GNN architecture to use. One of [“gin-virtual”, “gcn-virtual”, “gin”, “gcn”].

• drop_ratio (float, default=0.5) – Dropout probability.

• norm_layer (str, default="batch_norm") – Type of normalization layer to use. One of [“batch_norm”, “layer_norm”, “instance_norm”, “graph_norm”, “size_norm”, “pair_norm”].

• graph_pooling (str, default="sum") – Method for aggregating node features to graph-level representations. One of [“sum”, “mean”, “max”].

• augmented_feature (list or None, default=None) – Additional molecular fingerprints to use as features. It will be concatenated with the graph representation after pooling. Examples like [“morgan”, “maccs”] or None.

• batch_size (int, default=128) – Number of samples per batch for training.

• epochs (int, default=500) – Maximum number of training epochs.

• learning_rate (float, default=0.001) – Learning rate for optimizer.

• weight_decay (float, default=0.0) – L2 regularization strength.

• grad_clip_value (float, optional) – Maximum norm of gradients for gradient clipping.

• patience (int, default=50) – Number of epochs to wait for improvement before early stopping.

• use_lr_scheduler (bool, default=False) – Whether to use learning rate scheduler.

• scheduler_factor (float, default=0.5) – Factor by which to reduce learning rate when plateau is reached.

• scheduler_patience (int, default=5) – Number of epochs with no improvement after which learning rate will be reduced.

• loss_criterion (callable, optional) – Loss function for training.

• evaluate_criterion (str or callable, optional) – Metric for model evaluation.

• evaluate_higher_better (bool, optional) – Whether higher values of the evaluation metric are better.

• verbose (bool, default=False) – Whether to print progress information during training.

• device (torch.device or str, optional) – Device to use for computation.

• model_name (str, default="GREAMolecularPredictor") – Name of the model.

Example

examples/prediction/grea.py

import numpy as np
from torch_molecule import GREAMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec
import os
from pathlib import Path
import tempfile

EPOCHS = 5

def test_grea_predictor():
    # Test data
    smiles_list = [
        'CNC[C@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@@H]1C',
        'CNC[C@@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@H]1C',
        'C[C@H]1CN([C@@H](C)CO)C(=O)CCCn2cc(nn2)CO[C@@H]1CN(C)C(=O)CCC(F)(F)F',
        'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F'
    ] * 10
    properties = np.array([0, 0, 1, 1] * 10)  # Binary classification

    # 1. Basic initialization test
    print("\n=== Testing GREA model initialization ===")
    model = GREAMolecularPredictor(
        num_task=1,
        task_type="classification",
        gamma=0.8,
        num_layer=3,
        hidden_size=128,
        batch_size=4,
        epochs=EPOCHS,
        verbose="print_statement"
    )
    print("GREA model initialized successfully")

    # 2. Basic fitting test
    print("\n=== Testing GREA model fitting ===")
    model.fit(smiles_list, properties, X_val=smiles_list[3:], y_val=properties[3:])
    print("GREA model fitting completed")

    # 3. Prediction test
    print("\n=== Testing GREA model prediction ===")
    predictions = model.predict(smiles_list[3:])
    print(f"Prediction shape: {predictions['prediction'].shape}")
    print(f"Prediction for new molecule: {predictions['prediction']}")

    # 4. Auto-fitting test with search space parameters
    print("\n=== Testing GREA model auto-fitting ===")
    
    # Define search parameters using ParameterSpec
    search_parameters = {
        'num_layer': ParameterSpec(
            param_type=ParameterType.INTEGER,
            value_range=(2, 4)
        ),
        'hidden_size': ParameterSpec(
            param_type=ParameterType.INTEGER,
            value_range=(64, 256)
        ),
        'learning_rate': ParameterSpec(
            param_type=ParameterType.LOG_FLOAT,
            value_range=(1e-4, 1e-2)
        ),
        'drop_ratio': ParameterSpec(
            param_type=ParameterType.FLOAT,
            value_range=(0.1, 0.5)
        ),
        'gnn_type': ParameterSpec(
            param_type=ParameterType.CATEGORICAL,
            value_range=['gin-virtual', 'gcn']
        )
    }
    
    # Test different search configurations
    print("\n--- Testing with full search space ---")
    model_auto = GREAMolecularPredictor(
        num_task=1,
        task_type="classification",
        epochs=EPOCHS,
        verbose="print_statement"
    )
    
    model_auto.autofit(
        smiles_list,
        properties,
        X_val=smiles_list[3:],
        y_val=properties[3:],
        search_parameters=search_parameters,
        n_trials=2
    )
    print("Full search space auto-fitting completed")
    
    # Test with subset of parameters
    print("\n--- Testing with partial search space ---")
    partial_search = {
        'num_layer': search_parameters['num_layer'],
        'learning_rate': search_parameters['learning_rate']
    }
    
    model_partial = GREAMolecularPredictor(
        num_task=1,
        task_type="classification",
        epochs=EPOCHS,
        verbose="print_statement"
    )
    
    model_partial.autofit(
        smiles_list,
        properties,
        X_val=smiles_list[3:],
        y_val=properties[3:],
        search_parameters=partial_search,
        n_trials=2
    )
    print("Partial search space auto-fitting completed")
    
    # Test with default search space
    print("\n--- Testing with default search space ---")
    model_default = GREAMolecularPredictor(
        num_task=1,
        task_type="classification",
        epochs=EPOCHS,
        verbose="print_statement"
    )
    
    model_default.autofit(
        smiles_list,
        properties,
        n_trials=2
    )
    print("Default search space auto-fitting completed")

    # 5. Model saving and loading test
    print("\n=== Testing GREA model saving and loading ===")
    save_path = "test_grea_model.pt"
    model.save_to_local(save_path)
    print(f"GREA model saved to {save_path}")

    new_model = GREAMolecularPredictor(
        num_task=1,
        task_type="classification"
    )
    new_model.load_from_local(save_path)
    print("GREA model loaded successfully")

    # 6. Test invalid SMILES handling
    print("\n=== Testing invalid SMILES handling ===")
    try:
        invalid_smiles = ['CC', 'INVALID_SMILES']
        predictions = model.predict(invalid_smiles)
        print("Invalid SMILES handling test completed")
    except Exception as e:
        print(f"Invalid SMILES handled with error: {str(e)}")

    # 7. Test error handling for invalid search parameters
    print("\n=== Testing invalid search parameters handling ===")
    try:
        invalid_search = {
            'invalid_param': ParameterSpec(
                param_type=ParameterType.FLOAT,
                value_range=(0, 1)
            )
        }
        model_invalid = GREAMolecularPredictor(
            num_task=1,
            task_type="classification",
            verbose="print_statement"
        )
        model_invalid.autofit(
            smiles_list,
            properties,
            search_parameters=invalid_search,
            n_trials=1
        )
    except ValueError as e:
        print(f"Invalid parameter handled correctly with error: {str(e)}")

    # Clean up
    if os.path.exists(save_path):
        os.remove(save_path)
        print(f"Cleaned up {save_path}")

import numpy as np
from torch_molecule import GREAMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec
import os

def test_grea_upload(repo_id):
    # Test data
    smiles_list = [
        'CNC[C@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@@H]1C',
        'CNC[C@@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@H]1C',
        'C[C@H]1CN([C@@H](C)CO)C(=O)CCCn2cc(nn2)CO[C@@H]1CN(C)C(=O)CCC(F)(F)F',
        'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F'
    ]
    properties = np.array([0, 0, 1, 1])  # Binary classification

    # 8. Test Hugging Face upload
    print("\n=== Testing Hugging Face Hub upload ===")
    # Initialize and fit a small model for testing upload
    model_for_upload = GREAMolecularPredictor(
        num_task=1,
        task_type="regression",
        gamma=0.8,
        num_layer=2,  # Small model for quick testing
        hidden_size=64,   # Small embedding dimension
        batch_size=4,
        epochs=EPOCHS,     # Few epochs for quick testing
        verbose="none",
        model_name='GREA_O2'
    )
    
    # Fit the model with sample data
    model_for_upload.autofit(smiles_list, properties, X_val=smiles_list[3:], y_val=properties[3:])
    
    # Push to Hugging Face Hub
    # Note: HF_TOKEN should be set in environment variables
    model_for_upload.save_to_hf(
        repo_id=repo_id,
        commit_message="Upload GREA model for gas separation tasks",
        private=False
    )
    print("Successfully pushed model to Hugging Face Hub")
    
    # Test downloading and loading from Hub
    print("\n=== Testing model loading from Hugging Face Hub ===")        
    with tempfile.TemporaryDirectory(prefix="torch-molecule-grea-") as download_dir:
        downloaded_model = GREAMolecularPredictor()
        downloaded_model.load_from_hf(
            repo_id=repo_id, local_cache=str(Path(download_dir) / "GREA_O2.pt")
        )
        test_pred = downloaded_model.predict(smiles_list[3:])
        print("Successfully loaded and tested model from Hugging Face Hub")
        print(f"Test prediction shape: {test_pred['prediction'].shape}")

if __name__ == "__main__":
    test_grea_predictor()

Methods

fit	Fit the model to the training data with optional validation set.
autofit	Automatically find the best hyperparameters using Optuna optimization.
predict	Make predictions using the fitted model.

For shared checkpoint and parameter operations, see Base Modules.

fit

GREAMolecularPredictor.fit(X_train, y_train, X_val=None, y_val=None, X_unlbl=None)

Fit the model to the training data with optional validation set.

Parameters:

• X_train (List[str]) – Training set input molecular structures as SMILES strings

• y_train (Union[List, np.ndarray]) – Training set target values for property prediction

• X_val (List[str], optional) – Validation set input molecular structures as SMILES strings. If None, training data will be used for validation

• y_val (Union[List, np.ndarray], optional) – Validation set target values. Required if X_val is provided

• X_unlbl (List[str], optional) – Unlabeled set input molecular structures as SMILES strings.

Returns:

self – Fitted estimator

Return type:

GNNMolecularPredictor

autofit

GREAMolecularPredictor.autofit(X_train, y_train, X_val=None, y_val=None, X_unlbl=None, search_parameters=None, n_trials=10)

Automatically find the best hyperparameters using Optuna optimization.

Parameters:

• X_train (List[str])

• y_train (List | ndarray | None)

• X_val (List[str] | None)

• y_val (List | ndarray | None)

• X_unlbl (List[str] | None)

• search_parameters (Dict[str, ParameterSpec] | None)

• n_trials (int)

Return type:

GNNMolecularPredictor

predict

GREAMolecularPredictor.predict(X)

Make predictions using the fitted model.

Parameters:

X (List[str]) – List of SMILES strings to make predictions for

Returns:

Dictionary containing:

• ’prediction’: Model predictions (shape: [n_samples, n_tasks])

• ’variance’: Prediction variances (shape: [n_samples, n_tasks])

• ’node_importance’: A nested list where the outer list has length n_samples and each inner list has length n_nodes for that molecule

Return type:

Dict[str, np.ndarray]