DigressMolecularGenerator

Unconditional discrete graph diffusion.

from torch_molecule import DigressMolecularGenerator

Examples · Methods

Constructor

class DigressMolecularGenerator(hidden_size_X=256, hidden_size_E=128, num_layer=5, n_head=8, dropout=0.1, timesteps=500, batch_size=512, epochs=1000, learning_rate=0.0002, grad_clip_value=None, weight_decay=1e-12, lw_X=1, lw_E=5, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, verbose='none', device=None, model_name='DigressMolecularGenerator')[source]

This generator implements the DiGress for unconditional molecular generation.

References

Parameters:

  • hidden_size_X (int, optional) – Hidden dimension size for node features, defaults to 256

  • hidden_size_E (int, optional) – Hidden dimension size for edge features, defaults to 128

  • num_layer (int, optional) – Number of transformer layers, defaults to 5

  • n_head (int, optional) – Number of attention heads, defaults to 8

  • dropout (float, optional) – Dropout rate for transformer layers, defaults to 0.1

  • timesteps (int, optional) – Number of diffusion timesteps, defaults to 500

  • batch_size (int, optional) – Batch size for training, defaults to 128

  • epochs (int, optional) – Number of training epochs, defaults to 10000

  • learning_rate (float, optional) – Learning rate for optimization, defaults to 0.0002

  • grad_clip_value (Optional[float], optional) – Value for gradient clipping (None = no clipping), defaults to None

  • weight_decay (float, optional) – Weight decay for optimization, defaults to 0.0

  • lw_X (float, optional) – Loss weight for node reconstruction, defaults to 1

  • lw_E (float, optional) – Loss weight for edge reconstruction, defaults to 10

  • use_lr_scheduler (bool, optional) – Whether to use learning rate scheduler, defaults to False

  • scheduler_factor (float, optional) – Factor for learning rate scheduler (if use_lr_scheduler is True), defaults to 0.5

  • scheduler_patience (int, optional) – Patience for learning rate scheduler (if use_lr_scheduler is True), defaults to 5

  • 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]], optional) – Device to use for computation (cuda/cpu)

  • model_name (str, optional) – Name of the model, defaults to “DigressMolecularGenerator”

Example

examples/generation/digress.py
import os
import numpy as np
from tqdm import tqdm

import torch
from torch_molecule import DigressMolecularGenerator
from torch_molecule.utils.search import ParameterType, ParameterSpec

EPOCHS = 100
BATCH_SIZE = 32

def test_digress_generator():
    # 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

    # 1. Basic initialization test
    print("\n=== Testing DiGress model initialization ===")
    model = DigressMolecularGenerator(
        hidden_size_X=256,
        hidden_size_E=128,
        num_layer=5,
        dropout=0.1,
        timesteps=500,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        verbose="progress_bar",
        lw_X=1,
        lw_E=5
    )
    print("DiGress Model initialized successfully")

    # 2. Basic fitting test
    print("\n=== Testing DiGress model fitting ===")
    model.fit(smiles_list)
    print("DiGress Model fitting completed")

    # 3. Generation test
    print("\n=== Testing DiGress model generation ===")
    generated_smiles = model.generate(batch_size=BATCH_SIZE)
    print(f"Generated {len(generated_smiles)} molecules")
    print("Example generated SMILES:", generated_smiles[:5])

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

    new_model = DigressMolecularGenerator()
    new_model.load_from_local(save_path)
    print("DiGress Model loaded successfully")

    # Test generation with loaded model
    generated_smiles = new_model.generate(batch_size=BATCH_SIZE)
    print(f"Generated molecules with loaded model: {len(generated_smiles)}")
    print("Example generated SMILES:", generated_smiles[:5])

    # 5. Test generation with specific number of nodes
    print("\n=== Testing generation with specific node counts ===")
    num_nodes = np.array([[20], [25], [30], [35]])  # Specify different node counts
    generated_smiles = model.generate(num_nodes=num_nodes)
    print(f"Generated molecules with specific node counts: {len(generated_smiles)}")
    print("Example generated SMILES:", generated_smiles)

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

if __name__ == "__main__":
    test_digress_generator()

Methods

fit

generate

Randomly generate molecules with specified node counts.

For shared checkpoint and parameter operations, see Base Modules.

fit

DigressMolecularGenerator.fit(X_train)[source]
Parameters:

X_train (List[str])

Return type:

DigressMolecularGenerator

generate

DigressMolecularGenerator.generate(num_nodes=None, batch_size=32)[source]

Randomly generate molecules with specified node counts.

Parameters:

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

Returns:

List of generated molecules in SMILES format.

Return type:

List[str]