GRINMolecularPredictor¶

Repetition-invariant graph representations for polymers.

from torch_molecule import GRINMolecularPredictor

Examples · Methods

Constructor¶

class GRINMolecularPredictor(repetition_augmentation=False, l1_penalty=0.001, epochs_to_penalize=100, 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='GRINMolecularPredictor')[source]¶

This predictor implements GRIN for Max Spanning Tree algorithm aligned GNN.

The full name of GRIN is Graph Invariant Representation Learning.

References

Learning Repetition-Invariant Representations for Polymer Informatics. NeurIPS 2025 https://arxiv.org/pdf/2505.10726

Parameters:

repetition_augmentation (bool, default=False) – Whether to enable polymer augmentation for training.
l1_penalty (float, default=1e-3) – Weight for the L1 penalty.
epochs_to_penalize (int, default=100) – Number of epochs to train before starting L1 penalty.
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.
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.
learning_rate (float, default=0.001) – Learning rate for optimizer.
grad_clip_value (float, optional) – Maximum norm of gradients for gradient clipping.
weight_decay (float, default=0.0) – L2 regularization strength.
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.
verbose (str, default="none") – Whether to display progress info. Options are: “none”, “progress_bar”, “print_statement”. If any other, “none” is automatically chosen.
device (torch.device or str, optional) – Device to use for computation.
model_name (str, default="GRINMolecularPredictor") – Name of the model.

Example¶

examples/prediction/grin.py¶

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

def test_grin_predictor():
    # Test data
    smiles_list = [
        'FC(F)(F)C(C1=CC2=C(C=C1)C(=O)N(C2=O)C1=CC=C(CC2=CC=C(*)C=C2)C=C1)(C1=CC=C2C(=O)N(*)C(=O)C2=C1)C(F)(F)F',
        'CCCCCCCCCCCCC1=C(*)SC(*)=C1',
        'CC(C)C1=C(O*)C=CC(=C1)C(C1C=CC=CC1C(O)=O)C1=CC(C(C)C)=C(OC2=CC=C(C=C2)C(=O)C2=CC=C(*)C=C2)C=C1',
        '*OC1=CC=C(C=C1)C1(OC(=O)C2=C1C=CC=C2)C1=CC=C(OC2=CC=C(C=C2)S(=O)(=O)C2=CC=C(*)C=C2)C=C1' 
    ]
    properties = np.array([19, 88.2, 22.8, 5.74])  # Regression
    print('smiles_list', len(smiles_list))
    print('properties', len(properties))
    # 1. Basic initialization test
    print("\n=== Testing model initialization ===")
    model = GRINMolecularPredictor(
        num_task=1,
        task_type="regression",
        num_layer=3,
        hidden_size=128,
        batch_size=4,
        epochs=5,  # Small number for testing
        verbose="progress_bar",
    )
    print("Model initialized successfully")

    # 1.2. Basic fitting test
    print("\n=== Testing model fitting ===")
    model.fit(smiles_list, properties)
    print("Model fitting completed")

    # 1.3. Prediction test
    print("\n=== Testing model prediction ===")
    predictions = model.predict(smiles_list)
    print(f"Prediction shape: {predictions['prediction'].shape}")
    print(f"Prediction for new polymer repeat times 1: {predictions['prediction']}")

    # 2. Initialize with polymer train augmentation
    print("\n=== Testing model initialization with polymer train augmentation ===")
    model = GRINMolecularPredictor(
        num_task=1,
        task_type="regression",
        repetition_augmentation=3,
        num_layer=3,
        hidden_size=128,
        batch_size=4,
        epochs=5,  # Small number for testing
        verbose="progress_bar",
    )
    print("Model initialized with polymer train augmentation successfully")

    # 2.1. Fitting test
    print("\n=== Testing model fitting with polymer train augmentation ===")
    model.fit(smiles_list, properties)
    print("Model fitting completed")

    # 2.2. Prediction test on different repeat times
    print("\n=== Testing model prediction with polymer train augmentation ===")
    predictions = model.predict(smiles_list, test_time_augmentation=3)
    print(f"Prediction shape: {predictions['prediction'].shape}")
    print(f"Prediction for new polymer repeat times 3: {predictions['prediction']}")
    predictions = model.predict(smiles_list, test_time_augmentation=5)
    print(f"Prediction for new polymer repeat times 5: {predictions['prediction']}")

    # 3. Auto-fitting test with custom parameters
    print("\n=== Testing model auto-fitting ===")
    search_parameters = {
        'num_layer': (2, 4),
        'hidden_size': (64, 256),
        'learning_rate': (1e-4, 1e-2),
        'l1_penalty': (1e-6, 1e-3),
        'epochs_to_penalize': (0, 100)
    }
    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)
        ),
        'l1_penalty': ParameterSpec(
            param_type=ParameterType.LOG_FLOAT,
            value_range=(1e-6, 1e-3)
        ),
        'epochs_to_penalize': ParameterSpec(
            param_type=ParameterType.INTEGER,
            value_range=(0, 100)
        )
    }
    model_auto = GRINMolecularPredictor(
        num_task=1,
        task_type="regression",
        epochs=3,  # Small number for testing
        verbose="none"
    )
    
    model_auto.autofit(
        smiles_list,
        properties,
        search_parameters=search_parameters,
        n_trials=2  # Small number for testing
    )
    print("Model auto-fitting completed")

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

    new_model = GRINMolecularPredictor(
        num_task=1,
        task_type="regression"
    )
    new_model.load_from_local(save_path)
    print("Model loaded successfully")

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

if __name__ == "__main__":
    test_grin_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¶

GRINMolecularPredictor.fit(X_train, y_train, X_val=None, y_val=None, X_unlbl=None)[source]¶

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:

GRINMolecularPredictor

autofit¶

GRINMolecularPredictor.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¶

GRINMolecularPredictor.predict(X, test_time_augmentation=False)[source]¶

Make predictions using the fitted model.

Parameters:

X (List[str]) – List of SMILES strings to make predictions for
test_time_augmentation (bool, default=False) – Whether to enable polymer augmentation for making predictions.

Returns:

Dictionary containing:

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

Return type:

Dict[str, np.ndarray]