SupervisedMolecularEncoder

Molecular property labels or computed molecular descriptors.

from torch_molecule import SupervisedMolecularEncoder

Examples · Methods

Constructor

class SupervisedMolecularEncoder(num_task=None, predefined_task=None, encoder_type='gin-virtual', readout='sum', num_layer=5, hidden_size=300, drop_ratio=0.5, norm_layer='batch_norm', batch_size=128, epochs=500, learning_rate=0.001, grad_clip_value=None, weight_decay=0.0, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, device=None, model_name='SupervisedMolecularEncoder', verbose='none')

This encoder implements a GNN model for supervised molecular representation learning with user-defined or predefined fingerprint/calculated property tasks.

Parameters:

• num_task (int, optional) – Number of user-defined tasks for supervised pretraining. If it is specified, user must provide y_train in the fit function.

• predefined_task (List[str], optional) – List of predefined tasks to use. Must be from the supported task list [“morgan”, “maccs”, “logP”]. If None and num_task is None, all predefined tasks will be used.

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

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

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

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

• 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”].

• 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.

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

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

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

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

• 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. Inherited from BaseMolecularEncoder.

• model_name (str, default="SupervisedMolecularEncoder") – Name of the model. Inherited from BaseMolecularEncoder.

Example

examples/encoder/supervised.py

import numpy as np
from torch_molecule import SupervisedMolecularEncoder

def test_supervised_encoder():
    # Test molecules (simple examples)
    molecules = [
        "CC(=O)O",  # Acetic acid
        "CCO",      # Ethanol
        "CCCC",     # Butane
        "c1ccccc1", # Benzene
        "CCN",      # Ethylamine
    ]
    
    # Test case 1: Using only predefined tasks
    print("\nTest 1: Using predefined tasks (morgan and maccs, logp)")
    encoder1 = SupervisedMolecularEncoder(
        predefined_task=["morgan", "maccs", "logP"],
        epochs=2,  # Small number for testing
        verbose='print_statement'
    )
    encoder1.fit(molecules, y_train=None)
    encodings1 = encoder1.encode(molecules)
    print(f"Encoding shape: {encodings1.shape}")
    
    # Test case 2: Using custom tasks
    print("\nTest 2: Using custom tasks")
    # Generate some dummy regression and classification targets
    n_samples = len(molecules)
    y_custom = np.zeros((n_samples, 2))
    # First column: regression task (random values between 0 and 1)
    y_custom[:, 0] = np.random.rand(n_samples)
    # Second column: binary classification task (0 or 1)
    y_custom[:, 1] = np.random.randint(0, 2, n_samples)
    
    encoder2 = SupervisedMolecularEncoder(
        num_task=2,  # Two custom tasks
        epochs=2,    # Small number for testing
        verbose='progress_bar'
    )
    encoder2.fit(molecules, y_train=y_custom)
    encodings2 = encoder2.encode(molecules)
    print(f"Encoding shape: {encodings2.shape}")
    
    # Test case 3: Combining predefined and custom tasks
    print("\nTest 3: Combining predefined and custom tasks")
    encoder3 = SupervisedMolecularEncoder(
        predefined_task=["morgan"],
        num_task=2,  # 1024 (morgan) + 2 (custom)
        epochs=2,    # Small number for testing
        verbose='progress_bar'
    )
    encoder3.fit(molecules, y_train=y_custom)
    encodings3 = encoder3.encode(molecules)
    print(f"Encoding shape: {encodings3.shape}")

if __name__ == "__main__":
    test_supervised_encoder()

Methods

fit	Fit the model to the training data with optional validation set.
encode	Encode molecules into vector representations.

For shared checkpoint and parameter operations, see Base Modules.

fit

SupervisedMolecularEncoder.fit(X_train, y_train=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 representation learning

Returns:

self – Fitted estimator

Return type:

SupervisedMolecularEncoder

encode

SupervisedMolecularEncoder.encode(X, return_type='pt')

Encode molecules into vector representations.

Parameters:

• X (List[str]) – List of SMILES strings

• return_type (Literal["np", "pt"], default="pt") – Return type of the representations

Returns:

representations – Molecular representations

Return type:

ndarray or torch.Tensor