DeFoGMolecularGenerator¶
Discrete graph flow matching with optional property labels.
from torch_molecule import DeFoGMolecularGenerator
Constructor¶
- class DeFoGMolecularGenerator(num_layer=6, hidden_mlp_dims=None, hidden_dims=None, transition='marginal', time_distortion='polydec', lambda_train=None, extra_features_type='rrwp', rrwp_steps=16, sample_steps=500, rdb='general', rdb_crit='p_x1_g_xt', eta=0, omega=0.1, guidance_weight=0.2, batch_size=128, epochs=1000, learning_rate=0.0002, grad_clip_value=1.0, weight_decay=0.0, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=10, task_type=None, verbose='none', *, device=None, model_name='DeFoGMolecularGenerator')[source]¶
This generator implements the Discrete Flow Graph Model (DeFoG) for (multi-conditional and unconditional) molecular generation. It uses a flow matching approach with a graph transformer backbone.
References
DeFoG: A Discrete Flow Model for Graph Generation. ICML 2025. (https://openreview.net/forum?id=KPRIwWhqAZ)
Implementation based on: https://github.com/manuelmlmadeira/DeFoG
- Parameters:
num_layer (int, default=6) – Number of transformer layers
hidden_mlp_dims (Dict[str, int], default={'X': 256, 'E': 128, 'y': 128} if None) – Hidden dimensions for MLP layers in X (node dim), E (edge dim), and y (property dim) components
hidden_dims (Dict[str, Any], default={'dx': 256, 'de': 64, 'dy': 64, 'n_head': 8, 'dim_ffX': 256, 'dim_ffE': 128, 'dim_ffy': 128} if None) – Hidden dimensions for transformer components including attention heads and feed-forward layers Keys: ‘dx’ (node dim), ‘de’ (edge dim), ‘dy’ (property dim), ‘n_head’ (number of attention heads), ‘dim_ffX’ (feed-forward dim for node features), ‘dim_ffE’ (feed-forward dim for edge features), ‘dim_ffy’ (feed-forward dim for property features)
transition (str, default='marginal') – Transition type for flow matching. Options: ‘marginal’, ‘absorbing’, ‘uniform’, ‘absorbfirst’, ‘argmax’, ‘edge_marginal’, ‘node_marginal’
time_distortion (str, default="polydec") – Time distortion schedule used during training/sampling. Options: ‘identity’, ‘cos’, ‘revcos’, ‘polyinc’, ‘polydec’
lambda_train (List[float], default=[5.0, 1.0] if None) – Loss weights: [edge_loss_weight, property_loss_weight]
extra_features_type (str, default='rrwp') – Extra feature type. Options: ‘rrwp’, ‘rrwp_double’, ‘rrwp_only’, ‘rrwp_comp’, ‘cycles’, ‘eigenvalues’, ‘all’
rrwp_steps (int, default=16) – Number of steps for (R)RWP features
sample_steps (int, default=500) – Number of sampling steps during generation
rdb (str, default="general") – Rate matrix design (R^db) method. Options: ‘general’, ‘marginal’, ‘column’, ‘entry’
rdb_crit (str, default="p_x1_g_xt") – Criterion used by ‘column’/’entry’ designs. Options: ‘max_marginal’, ‘x_t’, ‘abs_state’, ‘p_x1_g_xt’, ‘x_1’, ‘p_xt_g_x1’, ‘xhat_t’, ‘first’
eta (float, default=0) – Stochasticity scaling for R^db (higher = more stochastic)
omega (float, default=0.1) – Target-guidance scaling for R^tg (higher = stronger guidance)
guidance_weight (float, default=0.2) – Classifier-free guidance weight during sampling
batch_size (int, default=128) – Batch size for training
epochs (int, default=1000) – Number of training epochs
learning_rate (float, default=0.0002) – Learning rate for optimization
grad_clip_value (Optional[float], default=1.0) – Value for gradient clipping (None = no clipping)
weight_decay (float, default=0.0) – Weight decay for optimization
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 on plateau
scheduler_patience (int, default=10) – Number of epochs with no improvement after which learning rate will be reduced
task_type (List[str], default=[]) – List specifying type of each task (‘regression’ or ‘classification’)
verbose (str, default="none") – Whether to display progress info. Options are: “none”, “progress_bar”, “print_statement”. If any other, “none” is automatically chosen.
device (Optional[Union[torch.device, str]], default=None) – Device to run the model on (CPU or GPU)
model_name (str, default="DeFoGMolecularGenerator") – Name identifier for the model
Example¶
import os
import numpy as np
import torch
from torch_molecule import DeFoGMolecularGenerator
EPOCHS = 10
BATCH_SIZE = 32
def test_defog_generator():
"""
Test suite for the DeFoGMolecularGenerator, covering conditional and unconditional
initialization, fitting, generation, saving, and loading.
"""
# 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'
]
smiles_list = smiles_list * 25 # Create 100 molecules for training
# Multi-dimensional properties: each row is a molecule, each column is a property
# Properties range from 0 to 1
np.random.seed(42) # For reproducible results
properties = np.random.rand(100, 3) # 100 molecules, 3 properties each
# 1. Multi-Conditional Model Testing
print("\n=== Testing Multi-Conditional DeFoG Model ===")
conditional_model = DeFoGMolecularGenerator(
task_type=['regression', 'regression', 'regression'], # 3 regression tasks
epochs=EPOCHS,
batch_size=BATCH_SIZE,
learning_rate=5e-4,
sample_steps=10, # Fewer steps for faster testing
guidance_weight=0.2,
verbose="progress_bar",
)
print("Multi-Conditional DeFoG Model initialized successfully.")
print(f"Input dim y: {conditional_model.input_dim_y}")
print("\n--- Fitting multi-conditional model ---")
conditional_model.fit(smiles_list, properties)
print("Multi-Conditional DeFoG Model fitting completed.")
print("\n--- Testing multi-conditional generation ---")
# Generate molecules with specific multi-dimensional properties
target_properties = [
[0.1, 0.2, 0.3], # Low values for all properties
[0.4, 0.5, 0.6], # Medium values for all properties
[0.7, 0.8, 0.9], # High values for all properties
[0.9, 0.1, 0.5] # Mixed values
]
generated_smiles = conditional_model.generate(labels=target_properties)
print(f"Multi-conditionally generated {len(generated_smiles)} molecules.")
assert len(generated_smiles) == len(target_properties)
print("Example SMILES:", generated_smiles[:2])
print("Target properties for first molecule:", target_properties[0])
print("\n--- Testing model saving and loading ---")
save_path = "conditional_defog_test_model.pt"
conditional_model.save_to_local(save_path)
print(f"Model saved to {save_path}")
new_conditional_model = DeFoGMolecularGenerator()
new_conditional_model.load_from_local(save_path)
print("Model loaded successfully.")
loaded_generated_smiles = new_conditional_model.generate(labels=target_properties)
print(f"Generated {len(loaded_generated_smiles)} molecules with loaded model.")
assert len(loaded_generated_smiles) > 0
print("\n--- Testing generation with specified node counts ---")
num_nodes = np.array([[20], [25], [30], [35]])
generated_smiles_nodes = conditional_model.generate(labels=target_properties, num_nodes=num_nodes)
print(f"Generated {len(generated_smiles_nodes)} molecules with specified node counts.")
assert len(generated_smiles_nodes) == len(num_nodes)
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
# 2. Single-property conditional testing (backwards compatibility)
print("\n=== Testing Single-Property Conditional DeFoG Model ===")
single_properties = properties[:, 0:1] # Use only first property
single_conditional_model = DeFoGMolecularGenerator(
task_type=['regression'], # Single regression task
epochs=EPOCHS,
batch_size=BATCH_SIZE,
learning_rate=5e-4,
sample_steps=10,
guidance_weight=0.2,
verbose="progress_bar",
)
print("Single-Property Conditional DeFoG Model initialized successfully.")
single_conditional_model.fit(smiles_list, single_properties)
print("Single-Property DeFoG Model fitting completed.")
single_target_properties = [[0.2], [0.5], [0.8], [0.1]]
single_generated_smiles = single_conditional_model.generate(labels=single_target_properties)
print(f"Single-conditionally generated {len(single_generated_smiles)} molecules.")
assert len(single_generated_smiles) == len(single_target_properties)
# 3. Unconditional Model Testing
print("\n=== Testing Unconditional DeFoG Model ===")
unconditional_model = DeFoGMolecularGenerator(
task_type=[], # Empty task_type for unconditional generation
epochs=EPOCHS,
batch_size=BATCH_SIZE,
learning_rate=5e-4,
sample_steps=10, # Fewer steps for faster testing
guidance_weight=0.0, # No guidance for unconditional generation
verbose="progress_bar",
)
print("Unconditional DeFoG Model initialized successfully.")
print(f"Input dim y: {unconditional_model.input_dim_y}")
print("\n--- Fitting unconditional model ---")
# For unconditional model, we don't need properties
unconditional_model.fit(smiles_list)
print("Unconditional DeFoG Model fitting completed.")
print("\n--- Testing unconditional generation (no parameters) ---")
unconditional_smiles = unconditional_model.generate()
print(f"Unconditionally generated {len(unconditional_smiles)} molecules (no parameters).")
assert len(unconditional_smiles) > 0
print("Example SMILES:", unconditional_smiles[:2])
print("\n--- Testing unconditional generation with batch_size ---")
unconditional_smiles_batch = unconditional_model.generate(batch_size=5)
print(f"Unconditionally generated {len(unconditional_smiles_batch)} molecules (batch_size=5).")
assert len(unconditional_smiles_batch) == 5
print("Example SMILES:", unconditional_smiles_batch[:2])
print("\n--- Testing unconditional generation with specified node counts ---")
num_nodes_uncond = np.array([[15], [20], [25]])
unconditional_smiles_nodes = unconditional_model.generate(num_nodes=num_nodes_uncond)
print(f"Unconditionally generated {len(unconditional_smiles_nodes)} molecules with specified node counts.")
assert len(unconditional_smiles_nodes) == len(num_nodes_uncond)
print("\n--- Testing unconditional model saving and loading ---")
save_path_uncond = "unconditional_defog_test_model.pt"
unconditional_model.save_to_local(save_path_uncond)
print(f"Unconditional model saved to {save_path_uncond}")
new_unconditional_model = DeFoGMolecularGenerator()
new_unconditional_model.load_from_local(save_path_uncond)
print("Unconditional model loaded successfully.")
loaded_unconditional_smiles = new_unconditional_model.generate(batch_size=3)
print(f"Generated {len(loaded_unconditional_smiles)} molecules with loaded unconditional model.")
assert len(loaded_unconditional_smiles) == 3
if os.path.exists(save_path_uncond):
os.remove(save_path_uncond)
print(f"Cleaned up {save_path_uncond}")
print("\n=== All DeFoG tests completed successfully! ===")
if __name__ == "__main__":
test_defog_generator()
Methods¶
For shared checkpoint and parameter operations, see Base Modules.
fit¶
generate¶
- DeFoGMolecularGenerator.generate(labels=None, num_nodes=None, batch_size=32)[source]¶
Generate molecules with specified properties and optional node counts.
- Parameters:
labels (Optional[Union[List[List], np.ndarray, torch.Tensor]], default=None) – Target properties for the generated molecules. Can be provided as: - A list of lists for multiple properties - A numpy array of shape (batch_size, n_properties) - A torch tensor of shape (batch_size, n_properties) For single label (properties values), can also be provided as 1D array/tensor. If None, generates unconditional samples.
num_nodes (Optional[Union[List[List], np.ndarray, torch.Tensor]], default=None) – Number of nodes for each molecule in the batch. If None, samples from the training distribution. Can be provided as: - A list of lists - A numpy array of shape (batch_size, 1) - A torch tensor of shape (batch_size, 1)
batch_size (int, default=32) – Number of molecules to generate. Only used if labels is None.
- Returns:
List of generated molecules in SMILES format.
- Return type:
List[str]