DIRMolecularPredictor¶
Learning invariant graph rationales.
from torch_molecule import DIRMolecularPredictor
Constructor¶
- class DIRMolecularPredictor(causal_ratio=0.8, lw_invariant=0.0001, 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='DIRMolecularPredictor')[source]¶
This predictor implements the DIR for molecular property prediction tasks.
The full name of DIR is Discovering Invariant Rationales.
References
Discovering Invariant Rationales for Graph Neural Networks. https://openreview.net/forum?id=hGXij5rfiHw
- Parameters:
causal_ratio (float, default=0.8) – The ratio of causal edges to keep during training. A higher ratio means more edges are considered causal/important for the prediction. This controls the sparsity of the learned rationales.
lw_invariant (float, default=1e-4) – The weight of the invariance loss term. This loss encourages the model to learn rationales that are invariant across different environments/perturbations. A higher value puts more emphasis on learning invariant features.
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 (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="DIRMolecularPredictor") – Name of the model.
Example¶
import numpy as np
from torch_molecule import DIRMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec
import os
from pathlib import Path
import tempfile
def test_dir_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'
]
properties = np.array([0, 0, 1, 1]) # Binary classification
# 1. Basic initialization test
print("\n=== Testing DIR model initialization ===")
model = DIRMolecularPredictor(
num_task=1,
task_type="classification",
causal_ratio=0.8,
lw_invariant=1e-4,
num_layer=3,
hidden_size=128,
batch_size=4,
epochs=5,
verbose="progress_bar"
)
print("DIR model initialized successfully")
# 2. Basic fitting test
print("\n=== Testing DIR model fitting ===")
model.fit(smiles_list[:3], properties[:3])
print("DIR model fitting completed")
# 3. Prediction test
print("\n=== Testing DIR 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 DIR 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)
),
'causal_ratio': ParameterSpec(
param_type=ParameterType.FLOAT,
value_range=(0.2, 0.9)
),
'lw_invariant': ParameterSpec(
param_type=ParameterType.LOG_FLOAT,
value_range=(1e-5, 1e-3)
),
'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 = DIRMolecularPredictor(
num_task=1,
task_type="classification",
epochs=3,
verbose="progress_bar"
)
model_auto.autofit(
smiles_list,
properties,
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'],
'causal_ratio': search_parameters['causal_ratio']
}
model_partial = DIRMolecularPredictor(
num_task=1,
task_type="classification",
epochs=3,
verbose="progress_bar"
)
model_partial.autofit(
smiles_list,
properties,
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 = DIRMolecularPredictor(
num_task=1,
task_type="classification",
epochs=3,
verbose="progress_bar"
)
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 DIR model saving and loading ===")
save_path = "test_dir_model.pt"
model.save_to_local(save_path)
print(f"DIR model saved to {save_path}")
new_model = DIRMolecularPredictor(
num_task=1,
task_type="classification"
)
new_model.load_from_local(save_path)
print("DIR 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 = DIRMolecularPredictor(
num_task=1,
task_type="classification",
verbose="progress_bar"
)
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}")
def test_dir_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 = DIRMolecularPredictor(
num_task=1,
task_type="regression",
causal_ratio=0.8,
lw_invariant=1e-4,
num_layer=2, # Small model for quick testing
hidden_size=64, # Small embedding dimension
batch_size=4,
epochs=2, # Few epochs for quick testing
verbose="none",
model_name='DIR_Invariant'
)
# Fit the model with sample data
model_for_upload.autofit(smiles_list[:3], 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 DIR model for invariant rationale discovery",
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-dir-") as download_dir:
downloaded_model = DIRMolecularPredictor()
downloaded_model.load_from_hf(
repo_id=repo_id, local_cache=str(Path(download_dir) / "DIR_Invariant.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_dir_predictor()
Methods¶
Fit the model to the training data with optional validation set. |
|
Automatically find the best hyperparameters using Optuna optimization. |
|
Make predictions using the fitted model. |
For shared checkpoint and parameter operations, see Base Modules.
fit¶
- DIRMolecularPredictor.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:
autofit¶
- DIRMolecularPredictor.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: