Initial commit
This commit is contained in:
Zhongwei Li
404
skills/pytdc/scripts/molecular_generation.py
Normal file
404
skills/pytdc/scripts/molecular_generation.py
Normal file
@@ -0,0 +1,404 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
TDC Molecular Generation with Oracles Template
|
||||
|
||||
This script demonstrates how to use TDC oracles for molecular generation
|
||||
tasks including goal-directed generation and distribution learning.
|
||||
|
||||
Usage:
|
||||
python molecular_generation.py
|
||||
"""
|
||||
|
||||
from tdc.generation import MolGen
|
||||
from tdc import Oracle
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_generation_dataset():
|
||||
"""
|
||||
Load molecular generation dataset
|
||||
"""
|
||||
print("=" * 60)
|
||||
print("Loading Molecular Generation Dataset")
|
||||
print("=" * 60)
|
||||
|
||||
# Load ChEMBL dataset
|
||||
data = MolGen(name='ChEMBL_V29')
|
||||
|
||||
# Get training molecules
|
||||
split = data.get_split()
|
||||
train_smiles = split['train']['Drug'].tolist()
|
||||
|
||||
print(f"\nDataset: ChEMBL_V29")
|
||||
print(f"Training molecules: {len(train_smiles)}")
|
||||
|
||||
# Display sample molecules
|
||||
print("\nSample SMILES:")
|
||||
for i, smiles in enumerate(train_smiles[:5], 1):
|
||||
print(f" {i}. {smiles}")
|
||||
|
||||
return train_smiles
|
||||
|
||||
|
||||
def single_oracle_example():
|
||||
"""
|
||||
Example: Using a single oracle for molecular evaluation
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 1: Single Oracle Evaluation")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize oracle for GSK3B target
|
||||
oracle = Oracle(name='GSK3B')
|
||||
|
||||
# Test molecules
|
||||
test_molecules = [
|
||||
'CC(C)Cc1ccc(cc1)C(C)C(O)=O', # Ibuprofen
|
||||
'CC(=O)Oc1ccccc1C(=O)O', # Aspirin
|
||||
'Cn1c(=O)c2c(ncn2C)n(C)c1=O', # Caffeine
|
||||
'CN1C=NC2=C1C(=O)N(C(=O)N2C)C' # Theophylline
|
||||
]
|
||||
|
||||
print("\nEvaluating molecules with GSK3B oracle:")
|
||||
print("-" * 60)
|
||||
|
||||
for smiles in test_molecules:
|
||||
score = oracle(smiles)
|
||||
print(f"SMILES: {smiles}")
|
||||
print(f"GSK3B score: {score:.4f}\n")
|
||||
|
||||
|
||||
def multiple_oracles_example():
|
||||
"""
|
||||
Example: Using multiple oracles for multi-objective optimization
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 2: Multiple Oracles (Multi-Objective)")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize multiple oracles
|
||||
oracles = {
|
||||
'QED': Oracle(name='QED'), # Drug-likeness
|
||||
'SA': Oracle(name='SA'), # Synthetic accessibility
|
||||
'GSK3B': Oracle(name='GSK3B'), # Target binding
|
||||
'LogP': Oracle(name='LogP') # Lipophilicity
|
||||
}
|
||||
|
||||
# Test molecule
|
||||
test_smiles = 'CC(C)Cc1ccc(cc1)C(C)C(O)=O'
|
||||
|
||||
print(f"\nEvaluating: {test_smiles}")
|
||||
print("-" * 60)
|
||||
|
||||
scores = {}
|
||||
for name, oracle in oracles.items():
|
||||
score = oracle(test_smiles)
|
||||
scores[name] = score
|
||||
print(f"{name:10s}: {score:.4f}")
|
||||
|
||||
# Multi-objective score (weighted combination)
|
||||
print("\n--- Multi-Objective Scoring ---")
|
||||
|
||||
# Invert SA (lower is better, so we invert for maximization)
|
||||
sa_score = 1.0 / (1.0 + scores['SA'])
|
||||
|
||||
# Weighted combination
|
||||
weights = {'QED': 0.3, 'SA': 0.2, 'GSK3B': 0.4, 'LogP': 0.1}
|
||||
multi_score = (
|
||||
weights['QED'] * scores['QED'] +
|
||||
weights['SA'] * sa_score +
|
||||
weights['GSK3B'] * scores['GSK3B'] +
|
||||
weights['LogP'] * (scores['LogP'] / 5.0) # Normalize LogP
|
||||
)
|
||||
|
||||
print(f"Multi-objective score: {multi_score:.4f}")
|
||||
print(f"Weights: {weights}")
|
||||
|
||||
|
||||
def batch_evaluation_example():
|
||||
"""
|
||||
Example: Batch evaluation of multiple molecules
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 3: Batch Evaluation")
|
||||
print("=" * 60)
|
||||
|
||||
# Generate sample molecules
|
||||
molecules = [
|
||||
'CC(C)Cc1ccc(cc1)C(C)C(O)=O',
|
||||
'CC(=O)Oc1ccccc1C(=O)O',
|
||||
'Cn1c(=O)c2c(ncn2C)n(C)c1=O',
|
||||
'CN1C=NC2=C1C(=O)N(C(=O)N2C)C',
|
||||
'CC(C)NCC(COc1ccc(cc1)COCCOC(C)C)O'
|
||||
]
|
||||
|
||||
# Initialize oracle
|
||||
oracle = Oracle(name='DRD2')
|
||||
|
||||
print(f"\nBatch evaluating {len(molecules)} molecules with DRD2 oracle...")
|
||||
|
||||
# Batch evaluation (more efficient than individual calls)
|
||||
scores = oracle(molecules)
|
||||
|
||||
print("\nResults:")
|
||||
print("-" * 60)
|
||||
for smiles, score in zip(molecules, scores):
|
||||
print(f"{smiles[:40]:40s}... Score: {score:.4f}")
|
||||
|
||||
# Statistics
|
||||
print(f"\nStatistics:")
|
||||
print(f" Mean score: {np.mean(scores):.4f}")
|
||||
print(f" Std score: {np.std(scores):.4f}")
|
||||
print(f" Min score: {np.min(scores):.4f}")
|
||||
print(f" Max score: {np.max(scores):.4f}")
|
||||
|
||||
|
||||
def goal_directed_generation_template():
|
||||
"""
|
||||
Template for goal-directed molecular generation
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 4: Goal-Directed Generation Template")
|
||||
print("=" * 60)
|
||||
|
||||
template = '''
|
||||
# Template for goal-directed molecular generation
|
||||
|
||||
from tdc.generation import MolGen
|
||||
from tdc import Oracle
|
||||
import numpy as np
|
||||
|
||||
# 1. Load training data
|
||||
data = MolGen(name='ChEMBL_V29')
|
||||
train_smiles = data.get_split()['train']['Drug'].tolist()
|
||||
|
||||
# 2. Initialize oracle(s)
|
||||
oracle = Oracle(name='GSK3B')
|
||||
|
||||
# 3. Initialize your generative model
|
||||
# model = YourGenerativeModel()
|
||||
# model.fit(train_smiles)
|
||||
|
||||
# 4. Generation loop
|
||||
num_iterations = 100
|
||||
num_molecules_per_iter = 100
|
||||
best_molecules = []
|
||||
|
||||
for iteration in range(num_iterations):
|
||||
# Generate candidate molecules
|
||||
# candidates = model.generate(num_molecules_per_iter)
|
||||
|
||||
# Evaluate with oracle
|
||||
scores = oracle(candidates)
|
||||
|
||||
# Select top molecules
|
||||
top_indices = np.argsort(scores)[-10:]
|
||||
top_molecules = [candidates[i] for i in top_indices]
|
||||
top_scores = [scores[i] for i in top_indices]
|
||||
|
||||
# Store best molecules
|
||||
best_molecules.extend(zip(top_molecules, top_scores))
|
||||
|
||||
# Optional: Fine-tune model on top molecules
|
||||
# model.fine_tune(top_molecules)
|
||||
|
||||
# Print progress
|
||||
print(f"Iteration {iteration}: Best score = {max(scores):.4f}")
|
||||
|
||||
# Sort and display top molecules
|
||||
best_molecules.sort(key=lambda x: x[1], reverse=True)
|
||||
print("\\nTop 10 molecules:")
|
||||
for smiles, score in best_molecules[:10]:
|
||||
print(f"{smiles}: {score:.4f}")
|
||||
'''
|
||||
|
||||
print("\nGoal-Directed Generation Template:")
|
||||
print("=" * 60)
|
||||
print(template)
|
||||
|
||||
|
||||
def distribution_learning_example(train_smiles):
|
||||
"""
|
||||
Example: Distribution learning evaluation
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 5: Distribution Learning")
|
||||
print("=" * 60)
|
||||
|
||||
# Use subset for demonstration
|
||||
train_subset = train_smiles[:1000]
|
||||
|
||||
# Initialize oracle
|
||||
oracle = Oracle(name='QED')
|
||||
|
||||
print("\nEvaluating property distribution...")
|
||||
|
||||
# Evaluate training set
|
||||
print("Computing training set distribution...")
|
||||
train_scores = oracle(train_subset)
|
||||
|
||||
# Simulate generated molecules (in practice, use your generative model)
|
||||
# For demo: add noise to training molecules
|
||||
print("Computing generated set distribution...")
|
||||
generated_scores = train_scores + np.random.normal(0, 0.1, len(train_scores))
|
||||
generated_scores = np.clip(generated_scores, 0, 1) # QED is [0, 1]
|
||||
|
||||
# Compare distributions
|
||||
print("\n--- Distribution Statistics ---")
|
||||
print(f"Training set (n={len(train_subset)}):")
|
||||
print(f" Mean: {np.mean(train_scores):.4f}")
|
||||
print(f" Std: {np.std(train_scores):.4f}")
|
||||
print(f" Median: {np.median(train_scores):.4f}")
|
||||
|
||||
print(f"\nGenerated set (n={len(generated_scores)}):")
|
||||
print(f" Mean: {np.mean(generated_scores):.4f}")
|
||||
print(f" Std: {np.std(generated_scores):.4f}")
|
||||
print(f" Median: {np.median(generated_scores):.4f}")
|
||||
|
||||
# Distribution similarity metrics
|
||||
from scipy.stats import ks_2samp
|
||||
ks_statistic, p_value = ks_2samp(train_scores, generated_scores)
|
||||
|
||||
print(f"\nKolmogorov-Smirnov Test:")
|
||||
print(f" KS statistic: {ks_statistic:.4f}")
|
||||
print(f" P-value: {p_value:.4f}")
|
||||
|
||||
if p_value > 0.05:
|
||||
print(" → Distributions are similar (p > 0.05)")
|
||||
else:
|
||||
print(" → Distributions are significantly different (p < 0.05)")
|
||||
|
||||
|
||||
def available_oracles_info():
|
||||
"""
|
||||
Display information about available oracles
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 6: Available Oracles")
|
||||
print("=" * 60)
|
||||
|
||||
oracle_info = {
|
||||
'Biochemical Targets': [
|
||||
'DRD2', 'GSK3B', 'JNK3', '5HT2A', 'ACE',
|
||||
'MAPK', 'CDK', 'P38', 'PARP1', 'PIK3CA'
|
||||
],
|
||||
'Physicochemical Properties': [
|
||||
'QED', 'SA', 'LogP', 'MW', 'Lipinski'
|
||||
],
|
||||
'Composite Metrics': [
|
||||
'Isomer_Meta', 'Median1', 'Median2',
|
||||
'Rediscovery', 'Similarity', 'Uniqueness', 'Novelty'
|
||||
],
|
||||
'Specialized': [
|
||||
'ASKCOS', 'Docking', 'Vina'
|
||||
]
|
||||
}
|
||||
|
||||
print("\nAvailable Oracle Categories:")
|
||||
print("-" * 60)
|
||||
|
||||
for category, oracles in oracle_info.items():
|
||||
print(f"\n{category}:")
|
||||
for oracle_name in oracles:
|
||||
print(f" - {oracle_name}")
|
||||
|
||||
print("\nFor detailed oracle documentation, see:")
|
||||
print(" references/oracles.md")
|
||||
|
||||
|
||||
def constraint_satisfaction_example():
|
||||
"""
|
||||
Example: Molecular generation with constraints
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("Example 7: Constraint Satisfaction")
|
||||
print("=" * 60)
|
||||
|
||||
# Define constraints
|
||||
constraints = {
|
||||
'QED': (0.5, 1.0), # Drug-likeness >= 0.5
|
||||
'SA': (1.0, 5.0), # Easy to synthesize
|
||||
'MW': (200, 500), # Molecular weight 200-500 Da
|
||||
'LogP': (0, 3) # Lipophilicity 0-3
|
||||
}
|
||||
|
||||
# Initialize oracles
|
||||
oracles = {name: Oracle(name=name) for name in constraints.keys()}
|
||||
|
||||
# Test molecules
|
||||
test_molecules = [
|
||||
'CC(C)Cc1ccc(cc1)C(C)C(O)=O',
|
||||
'CC(=O)Oc1ccccc1C(=O)O',
|
||||
'Cn1c(=O)c2c(ncn2C)n(C)c1=O'
|
||||
]
|
||||
|
||||
print("\nConstraints:")
|
||||
for prop, (min_val, max_val) in constraints.items():
|
||||
print(f" {prop}: [{min_val}, {max_val}]")
|
||||
|
||||
print("\n" + "-" * 60)
|
||||
print("Evaluating molecules against constraints:")
|
||||
print("-" * 60)
|
||||
|
||||
for smiles in test_molecules:
|
||||
print(f"\nSMILES: {smiles}")
|
||||
|
||||
satisfies_all = True
|
||||
for prop, (min_val, max_val) in constraints.items():
|
||||
score = oracles[prop](smiles)
|
||||
satisfies = min_val <= score <= max_val
|
||||
|
||||
status = "✓" if satisfies else "✗"
|
||||
print(f" {prop:10s}: {score:7.2f} [{min_val:5.1f}, {max_val:5.1f}] {status}")
|
||||
|
||||
satisfies_all = satisfies_all and satisfies
|
||||
|
||||
result = "PASS" if satisfies_all else "FAIL"
|
||||
print(f" Overall: {result}")
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Main function to run all molecular generation examples
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print("TDC Molecular Generation with Oracles Examples")
|
||||
print("=" * 60)
|
||||
|
||||
# Load generation dataset
|
||||
train_smiles = load_generation_dataset()
|
||||
|
||||
# Example 1: Single oracle
|
||||
single_oracle_example()
|
||||
|
||||
# Example 2: Multiple oracles
|
||||
multiple_oracles_example()
|
||||
|
||||
# Example 3: Batch evaluation
|
||||
batch_evaluation_example()
|
||||
|
||||
# Example 4: Goal-directed generation template
|
||||
goal_directed_generation_template()
|
||||
|
||||
# Example 5: Distribution learning
|
||||
distribution_learning_example(train_smiles)
|
||||
|
||||
# Example 6: Available oracles
|
||||
available_oracles_info()
|
||||
|
||||
# Example 7: Constraint satisfaction
|
||||
constraint_satisfaction_example()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Molecular generation examples completed!")
|
||||
print("=" * 60)
|
||||
print("\nNext steps:")
|
||||
print("1. Implement your generative model")
|
||||
print("2. Use oracles to guide generation")
|
||||
print("3. Evaluate generated molecules")
|
||||
print("4. Iterate and optimize")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user