GraphGAMolecularGenerator

Genetic search using property labels or a scoring oracle.

from torch_molecule import GraphGAMolecularGenerator

Examples · Methods

Constructor

class GraphGAMolecularGenerator(num_task=0, population_size=100, offspring_size=50, mutation_rate=0.0067, n_jobs=1, iteration=5, verbose='none', *, device=None, model_name='GraphGAMolecularGenerator')

This generator implements the Graph Genetic Algorithm for molecular generation.

References

• A Graph-Based Genetic Algorithm and Its Application to the Multiobjective Evolution of Median Molecules. Journal of Chemical Information and Computer Sciences. https://pubs.acs.org/doi/10.1021/ci034290p

• Implementation: https://github.com/wenhao-gao/mol_opt

Parameters:

• num_task (int, default=0) – Number of properties to condition on. Set to 0 for unconditional generation.

• population_size (int, default=100) – Size of the population in each iteration.

• offspring_size (int, default=50) – Number of offspring molecules to generate in each iteration.

• mutation_rate (float, default=0.0067) – Probability of mutation occurring during reproduction.

• n_jobs (int, default=1) – Number of parallel jobs to run. -1 means using all processors.

• iteration (int, default=5) – Number of iterations for each target label (or random sample) to run the genetic algorithm.

• 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="GraphGAMolecularGenerator") – Name identifier for the model.

Example

examples/generation/graphga.py

import os
import numpy as np
import warnings
from tqdm import tqdm
from joblib import delayed

import torch
from torch_molecule.generator.graph_ga.modeling_graph_ga import GraphGAMolecularGenerator
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import Descriptors
from rdkit.Chem import rdMolDescriptors

# Define the minimum value for score adjustment
MINIMUM = 1e-3

def test_graph_ga_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
    
    # Create synthetic properties for conditional generation
    # Let's use molecular weight and logP as example properties
    properties = []
    for smiles in smiles_list:
        mol = Chem.MolFromSmiles(smiles)
        if mol:
            mw = rdMolDescriptors.CalcExactMolWt(mol)
            logp = Descriptors.MolLogP(mol)
            properties.append([mw, logp])
        else:
            properties.append([np.nan, np.nan])
    
    properties = np.array(properties)
    
    print("\n=== Testing GraphGA model initialization ===")
    # 1. Initialize model for unconditional generation
    model_uncond = GraphGAMolecularGenerator(
        num_task=0,  # 0 for unconditional generation
        population_size=50,
        offspring_size=25,
        mutation_rate=0.01,
        n_jobs=1,  # Use 1 for easier debugging
        iteration=3,
        verbose="progress_bar"
    )
    print("GraphGA Model (unconditional) initialized successfully")
    
    # 2. Initialize model for conditional generation
    model_cond = GraphGAMolecularGenerator(
        num_task=2,  # 2 properties: MW and logP
        population_size=50,
        offspring_size=25,
        mutation_rate=0.01,
        n_jobs=1,  # Use 1 for easier debugging
        iteration=3,
        verbose="progress_bar"
    )
    print("GraphGA Model (conditional) initialized successfully")
    
    # 3. Test fitting
    print("\n=== Testing GraphGA model fitting (unconditional) ===")
    model_uncond.fit(smiles_list)
    print("GraphGA Model (unconditional) fitting completed")
    
    print("\n=== Testing GraphGA model fitting (conditional) ===")
    model_cond.fit(smiles_list, properties)
    print("GraphGA Model (conditional) fitting completed")
    
    # 4. Test unconditional generation
    print("\n=== Testing GraphGA model unconditional generation ===")
    generated_smiles_uncond = model_uncond.generate(num_samples=5)
    print(f"Generated {len(generated_smiles_uncond)} molecules unconditionally")
    print("Example generated SMILES:", generated_smiles_uncond[:2])

    
    # 5. Test conditional generation
    print("\n=== Testing GraphGA model conditional generation ===")
    # Create target properties (MW around 400, logP around 3)
    target_properties = np.array([[400, 3.0], [450, 4.0]])
    generated_smiles_cond = model_cond.generate(labels=target_properties)
    print(f"Generated {len(generated_smiles_cond)} molecules conditionally")
    print("Example generated SMILES:", generated_smiles_cond)
    
    # Verify properties of generated molecules
    print("\nProperties of conditionally generated molecules:")
    for i, smiles in enumerate(generated_smiles_cond):
        mol = Chem.MolFromSmiles(smiles)
        if mol:
            mw = Chem.Descriptors.MolWt(mol)
            logp = Chem.Descriptors.MolLogP(mol)
            print(f"Target: MW={target_properties[i][0]}, logP={target_properties[i][1]}")
            print(f"Actual: MW={mw:.2f}, logP={logp:.2f}")

    
    # 6. Test model saving and loading
    print("\n=== Testing GraphGA model saving and loading ===")
    save_path = "graph_ga_test_model.pkl"
    try:
        model_cond.save_to_local(save_path)
        print(f"GraphGA Model saved to {save_path}")
        
        # Load oracles manually with joblib
        import joblib
        saved_oracle = joblib.load(save_path)
        
        # Create new model and pass the loaded oracles to fit
        new_model = GraphGAMolecularGenerator(num_task=2)
        new_model.fit(smiles_list, oracle=saved_oracle)
        print("GraphGA Model loaded successfully")
        
        # Test generation with loaded model
        generated_smiles = new_model.generate(labels=target_properties[:1])
        print("Generated molecules with loaded model:", len(generated_smiles), generated_smiles)
    except Exception as e:
        print(f"Error in saving/loading: {e}")
    
    # Clean up
    if os.path.exists(save_path):
        os.remove(save_path)
        print(f"Cleaned up {save_path}")

if __name__ == "__main__":
    test_graph_ga_generator()

Methods

fit	Fit the model to the training data.
generate	Generate molecules using genetic algorithm optimization.

For shared checkpoint and parameter operations, see Base Modules.

fit

GraphGAMolecularGenerator.fit(X_train, y_train=None, oracle=None)

Fit the model to the training data.

Parameters:

• X_train (List[str]) – Training data, which will be used as the initial population.

• y_train (Optional[Union[List, np.ndarray]]) – Training labels for conditional generation (num_task is not 0).

• oracle (Optional[Callable]) –

  Oracle used to score the generated molecules. If not provided, default oracles based on sklearn.ensemble.RandomForestRegressor are trained on the X_train and y_train.

  For a customized oracle, it should be a Callable object, i.e., oracle(X, y). Please properly wrap your oracle to take two inputs:

  • a list of rdkit.Chem.rdchem.Mol objects and

  • a (1, num_task) numpy array of target values that all the molecules in the list target to achieve. Take care of NaN values if any.

  Scores for different tasks should be aggregated, i.e., mean or sum. The return should be a list of scores (float). Smaller scores mean closer to the target goal.

  Oracles are not needed for unconditional generation.

Returns:

self – Fitted model.

Return type:

GraphGAMolecularGenerator

generate

GraphGAMolecularGenerator.generate(labels=None, num_samples=32)

Generate molecules using genetic algorithm optimization.

Parameters:

• labels (List[List] | ndarray | None)

• num_samples (int)

Return type:

List[str]

Scoring oracles

class Oracle(models=None, num_task=1)

The default Oracle class for scoring molecules in GraphGA.

This class wraps predictive models (like RandomForestRegressor) to score molecules based on their properties. It handles conversion of SMILES to fingerprints.

Parameters:

• models (List[Any], optional) – List of trained models that implement a predict method. If None, RandomForestRegressors will be created when fit is called.

• num_task (int, default=1) – Number of properties to predict.

fit(X_train, y_train)

Fit the underlying models with training data.

Parameters:

• X_train (List[str] or List[RDKit.Mol]) – Training molecules as SMILES strings or RDKit Mol objects.

• y_train (np.ndarray) – Training labels with shape (n_samples, num_task).

Returns:

self – Fitted oracle.

Return type:

Oracle