GDSSMolecularGenerator¶

Unconditional score-based graph diffusion.

from torch_molecule import GDSSMolecularGenerator

Examples · Methods

Constructor¶

class GDSSMolecularGenerator(*, num_layer=3, hidden_size_adj=8, hidden_size=16, attention_dim=16, num_head=4, sde_type_x='VE', sde_beta_min_x=0.1, sde_beta_max_x=1, sde_num_scales_x=1000, sde_type_adj='VE', sde_beta_min_adj=0.1, sde_beta_max_adj=1, sde_num_scales_adj=1000, batch_size=128, epochs=500, learning_rate=0.005, grad_clip_value=1.0, weight_decay=0.0001, use_loss_reduce_mean=False, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, sampler_predictor='Reverse', sampler_corrector='Langevin', sampler_snr=0.2, sampler_scale_eps=0.7, sampler_n_steps=1, sampler_probability_flow=False, sampler_noise_removal=True, verbose='none', device=None, model_name='GDSSMolecularGenerator')[source]¶

This generator implements “Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations”

References

Paper: https://arxiv.org/abs/2202.02514
Official Implementation: https://github.com/harryjo97/GDSS

Parameters:

num_layer (int, default=3) – Number of layers in the score networks.
hidden_size_adj (float, default=8) – Hidden dimension size for the adjacency in the adjacency score network.
hidden_size (int, default=16) – Hidden dimension size latent representation.
attention_dim (int, default=16) – Dimension of attention layers.
num_head (int, default=4) – Number of attention heads.
sde_type_x (str, default='VE') – SDE type for node features. One of ‘VP’, ‘VE’, ‘subVP’.
sde_beta_min_x (float, default=0.1) – Minimum noise level for node features.
sde_beta_max_x (float, default=1) – Maximum noise level for node features.
sde_num_scales_x (int, default=1000) – Number of noise scales for node features.
sde_type_adj (str, default='VE') – SDE type for adjacency matrix. One of ‘VP’, ‘VE’, ‘subVP’.
sde_beta_min_adj (float, default=0.1) – Minimum noise level for adjacency matrix.
sde_beta_max_adj (float, default=1) – Maximum noise level for adjacency matrix.
sde_num_scales_adj (int, default=1000) – Number of noise scales for adjacency matrix.
batch_size (int, default=128) – Batch size for training.
epochs (int, default=500) – Number of training epochs.
learning_rate (float, default=0.005) – Learning rate for optimizer.
grad_clip_value (Optional[float], default=1) – Value for gradient clipping. None means no clipping.
weight_decay (float, default=1e-4) – Weight decay for optimizer.
use_loss_reduce_mean (bool, default=False) – Whether to use mean reduction for loss calculation.
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 using scheduler (only used if use_lr_scheduler is True).
scheduler_patience (int, default=5) – Number of epochs with no improvement after which learning rate will be reduced (only used if use_lr_scheduler is True).
sampler_predictor (str, default='Reverse') – Predictor method for sampling. One of ‘Euler’, ‘Reverse’.
sampler_corrector (str, default='Langevin') – Corrector method for sampling. One of ‘Langevin’, ‘None’.
sampler_snr (float, default=0.2) – Signal-to-noise ratio for corrector.
sampler_scale_eps (float, default=0.7) – Scale factor for noise level in corrector.
sampler_n_steps (int, default=1) – Number of corrector steps per predictor step.
sampler_probability_flow (bool, default=False) – Whether to use probability flow ODE for sampling.
sampler_noise_removal (bool, default=True) – Whether to remove noise in the final step of sampling.
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 “GDSSMolecularGenerator”

Example¶

examples/generation/gdss.py¶

import os
import numpy as np
from tqdm import tqdm

import torch
from torch_molecule import GDSSMolecularGenerator

EPOCHS = 5
BATCH_SIZE = 16

def test_gdss_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 GDSS model initialization ===")
    model = GDSSMolecularGenerator(
        num_layer=3,
        hidden_size_adj=8,
        hidden_size=16,
        attention_dim=16,
        num_head=4,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        learning_rate=0.005,
        verbose="print_statement"
    )
    print("GDSS Model initialized successfully")

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

    # 3. Generation test
    print("\n=== Testing GDSS 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 GDSS model saving and loading ===")
    save_path = "gdss_test_model.pt"
    model.save_to_local(save_path)
    print(f"GDSS Model saved to {save_path}")

    new_model = GDSSMolecularGenerator()
    new_model.load_from_local(save_path)
    print("GDSS 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_gdss_generator()

Methods¶

`fit`	Fit the model to the training data.
`generate`	Randomly generate molecules with specified node counts.

For shared checkpoint and parameter operations, see Base Modules.

fit¶

GDSSMolecularGenerator.fit(X_train)[source]¶

Fit the model to the training data.

Parameters:: X_train (List[str]) – List of training data in SMILES format.
Returns:: self – The fitted model.
Return type:: GDSSMolecularGenerator

generate¶

GDSSMolecularGenerator.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]