IRMMolecularPredictor

Invariant risk minimization with graph neural networks.

from torch_molecule import IRMMolecularPredictor

Examples · Methods

Constructor

class IRMMolecularPredictor(IRM_environment='random', scale=1.0, penalty_weight=1.0, penalty_anneal_iters=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='IRMMolecularPredictor')[source]

This predictor implements a Invariant Risk Minimization model with the GNN.

The full name of IRM is Invariant Risk Minimization.

References

Parameters:

  • IRM_environment (Union[torch.Tensor, np.ndarray, List, str], default="random") – Environment assignments for IRM. Can be a list of integers (one per sample), or “random” to assign environments randomly.

  • scale (float, default=1.0) – Scaling factor for the IRM penalty term.

  • penalty_weight (float, default=1.0) – Weight of the IRM penalty in the loss function.

  • penalty_anneal_iters (int, default=100) – Number of iterations for annealing the penalty weight.

  • 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 run the model on.

  • model_name (str)

Example

examples/prediction/irm_quickstart.py
import numpy as np

from torch_molecule import IRMMolecularPredictor

smiles = ["CCO", "CCCO", "c1ccccc1", "Cc1ccccc1"]
targets = np.array([[0.0], [0.3], [1.6], [1.8]])  # Illustrative property values.

model = IRMMolecularPredictor(
    task_type="regression",
    IRM_environment="random",
    num_layer=2,
    hidden_size=32,
    batch_size=4,
    epochs=5,
    device="cpu",
)
model.fit(smiles, targets)
predictions = model.predict(["CCCO", "CCC"])["prediction"]
print(predictions.shape)  # (2, 1)

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

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

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

IRMMolecularPredictor.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])

Return type:

Dict[str, np.ndarray]