import json
import pandas as pd
import numpy as np
from rdkit import Chem
from rdkit.Chem import QED, AllChem
from rdkit.Chem import Descriptors
from rdkit.Chem import RDConfig
import os
import sys
sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
import sascorer
from fcd_torch import FCD # You'll need to install fcd_torch
sys.path.append(os.path.dirname(os.path.dirname(__file__))) # Add parent directory to path
from src.util.utils import (fraction_valid, fraction_unique, novelty,
internal_diversity,obey_lipinski, obey_veber, load_pains_filters,
is_pains, FragMetric, ScafMetric)
import torch
class MoleculeEvaluator:
def __init__(self, gen_smiles, ref_smiles_1, ref_smiles_2=None, n_jobs=8):
"""
Initialize evaluator with generated and reference SMILES
ref_smiles_2 is optional
"""
self.gen_smiles = gen_smiles
self.ref_smiles_1 = ref_smiles_1
self.ref_smiles_2 = ref_smiles_2
self.n_jobs = n_jobs
# Convert SMILES to RDKit molecules and filter out invalid ones
self.gen_mols = [mol for s in gen_smiles if s and (mol := Chem.MolFromSmiles(s)) is not None]
self.ref_mols_1 = [mol for s in ref_smiles_1 if s and (mol := Chem.MolFromSmiles(s)) is not None]
self.ref_mols_2 = [mol for s in ref_smiles_2 if s and (mol := Chem.MolFromSmiles(s)) is not None] if ref_smiles_2 else None
# Initialize metrics that need setup
self.fcd = FCD(device='cuda' if torch.cuda.is_available() else 'cpu')
self.pains_catalog = load_pains_filters()
self.frag_metric = FragMetric(n_jobs=1)
self.scaf_metric = ScafMetric(n_jobs=1)
def calculate_basic_metrics(self):
"""Calculate validity, uniqueness, novelty, and internal diversity"""
# Generate Morgan fingerprints for internal diversity calculation
fps = np.array([AllChem.GetMorganFingerprintAsBitVect(mol, 2, 1024) for mol in self.gen_mols if mol is not None])
internal_diversity_mean, internal_diversity_std = internal_diversity(fps)
results = {
'validity': fraction_valid(self.gen_smiles, n_jobs=self.n_jobs),
'uniqueness': fraction_unique(self.gen_smiles, n_jobs=self.n_jobs),
'novelty_ref1': novelty(self.gen_smiles, self.ref_smiles_1, n_jobs=self.n_jobs),
'internal_diversity': internal_diversity_mean,
'internal_diversity_std': internal_diversity_std
}
if self.ref_smiles_2:
results['novelty_ref2'] = novelty(self.gen_smiles, self.ref_smiles_2, n_jobs=self.n_jobs)
return results
def calculate_property_metrics(self):
"""Calculate QED and SA scores"""
qed_scores = [QED.qed(mol) for mol in self.gen_mols if mol is not None]
sa_scores = [sascorer.calculateScore(mol) for mol in self.gen_mols if mol is not None]
return {
'qed_mean': np.mean(qed_scores),
'qed_std': np.std(qed_scores),
'sa_mean': np.mean(sa_scores),
'sa_std': np.std(sa_scores)
}
def calculate_fcd(self):
"""Calculate FCD score against both reference sets"""
# Filter out None values and convert mols back to SMILES
gen_valid_smiles = [Chem.MolToSmiles(mol) for mol in self.gen_mols if mol is not None]
ref1_valid_smiles = [Chem.MolToSmiles(mol) for mol in self.ref_mols_1 if mol is not None]
results = {
'fcd_ref1': self.fcd(gen_valid_smiles, ref1_valid_smiles)
}
if self.ref_mols_2:
ref2_valid_smiles = [Chem.MolToSmiles(mol) for mol in self.ref_mols_2 if mol is not None]
results['fcd_ref2'] = self.fcd(gen_valid_smiles, ref2_valid_smiles)
return results
def calculate_similarity_metrics(self):
"""Calculate fragment and scaffold similarity"""
results = {
'frag_sim_ref1': self.frag_metric(gen=self.gen_mols, ref=self.ref_mols_1),
'scaf_sim_ref1': self.scaf_metric(gen=self.gen_mols, ref=self.ref_mols_1)
}
if self.ref_mols_2:
results.update({
'frag_sim_ref2': self.frag_metric(gen=self.gen_mols, ref=self.ref_mols_2),
'scaf_sim_ref2': self.scaf_metric(gen=self.gen_mols, ref=self.ref_mols_2)
})
return results
def calculate_drug_likeness(self):
"""Calculate Lipinski, Veber and PAINS filtering results"""
lipinski_scores = [obey_lipinski(mol) for mol in self.gen_mols if mol is not None]
veber_scores = [obey_veber(mol) for mol in self.gen_mols if mol is not None]
pains_results = [not is_pains(mol, self.pains_catalog) for mol in self.gen_mols if mol is not None]
return {
'lipinski_mean': np.mean(lipinski_scores),
'lipinski_std': np.std(lipinski_scores),
'veber_mean': np.mean(veber_scores),
'veber_std': np.std(veber_scores),
'pains_pass_rate': np.mean(pains_results)
}
def evaluate_all(self):
"""Run all evaluations and combine results"""
results = {}
print("\nCalculating basic metrics...")
basic_metrics = self.calculate_basic_metrics()
print("Basic metrics:", {k: f"{v:.3f}" if isinstance(v, float) else v for k, v in basic_metrics.items()})
results.update(basic_metrics)
print("\nCalculating property metrics...")
property_metrics = self.calculate_property_metrics()
print("Property metrics:", {k: f"{v:.3f}" if isinstance(v, float) else v for k, v in property_metrics.items()})
results.update(property_metrics)
print("\nCalculating FCD scores...")
fcd_metrics = self.calculate_fcd()
print("FCD metrics:", {k: f"{v:.3f}" if isinstance(v, float) else v for k, v in fcd_metrics.items()})
results.update(fcd_metrics)
print("\nCalculating similarity metrics...")
similarity_metrics = self.calculate_similarity_metrics()
print("Similarity metrics:", {k: f"{v:.3f}" if isinstance(v, float) else v for k, v in similarity_metrics.items()})
results.update(similarity_metrics)
print("\nCalculating drug-likeness metrics...")
drug_likeness = self.calculate_drug_likeness()
print("Drug-likeness metrics:", {k: f"{v:.3f}" if isinstance(v, float) else v for k, v in drug_likeness.items()})
results.update(drug_likeness)
return results
def read_smi_file(file_path):
"""
Read SMILES strings from a .smi file
Args:
file_path (str): Path to .smi file
Returns:
list: List of SMILES strings
"""
smiles_list = []
try:
with open(file_path, 'r') as f:
for line in f:
# Handle different .smi formats:
# Some .smi files have just SMILES, others have SMILES followed by an identifier
parts = line.strip().split()
if parts: # Skip empty lines
smiles = parts[0] # First part is always the SMILES
smiles_list.append(smiles)
except FileNotFoundError:
raise FileNotFoundError(f"Could not find SMI file: {file_path}")
except Exception as e:
raise ValueError(f"Error reading SMI file {file_path}: {str(e)}")
return smiles_list
def evaluate_molecules_from_files(gen_path, ref_path_1, ref_path_2=None, smiles_col='SMILES', output_prefix="results", n_jobs=8):
"""
Main function to evaluate generated molecules against reference sets
Args:
gen_path (str): Path to CSV file containing generated SMILES
ref_path_1 (str): Path to .smi file containing first reference set SMILES
ref_path_2 (str, optional): Path to .smi file containing second reference set SMILES
smiles_col (str): Name of column containing SMILES strings in the CSV file
output_prefix (str): Prefix for output files
n_jobs (int): Number of parallel jobs
"""
# Read generated SMILES from CSV file
try:
gen_df = pd.read_csv(gen_path)
if smiles_col not in gen_df.columns:
raise ValueError(f"SMILES column '{smiles_col}' not found in generated dataset")
gen_smiles = gen_df[smiles_col].dropna().tolist()
except FileNotFoundError:
raise FileNotFoundError(f"Could not find generated CSV file: {gen_path}")
except pd.errors.EmptyDataError:
raise ValueError("Generated CSV file is empty")
# Read reference SMILES from .smi files
ref_smiles_1 = read_smi_file(ref_path_1)
ref_smiles_2 = read_smi_file(ref_path_2) if ref_path_2 else None
# Validate that we have SMILES to process
if not gen_smiles:
raise ValueError("No valid SMILES found in generated set")
if not ref_smiles_1:
raise ValueError("No valid SMILES found in reference set 1")
if ref_path_2 and not ref_smiles_2:
raise ValueError("No valid SMILES found in reference set 2")
print(f"\nProcessing datasets:")
print(f"Generated molecules: {len(gen_smiles)}")
print(f"Reference set 1: {len(ref_smiles_1)}")
if ref_smiles_2:
print(f"Reference set 2: {len(ref_smiles_2)}")
# Run evaluation
evaluator = MoleculeEvaluator(gen_smiles, ref_smiles_1, ref_smiles_2, n_jobs=n_jobs)
results = evaluator.evaluate_all()
print("\nSaving results...")
# Add dataset sizes to results
results.update({
'n_generated': len(gen_smiles),
'n_reference_1': len(ref_smiles_1),
'n_reference_2': len(ref_smiles_2) if ref_smiles_2 is not None else 0
})
# Save results
# Format float values to 3 decimal places for JSON
formatted_results = {k: round(v, 3) if isinstance(v, float) else v for k, v in results.items()}
with open(f"{output_prefix}.json", 'w') as f:
json.dump(formatted_results, f, indent=4)
# Create DataFrame with formatted results
df = pd.DataFrame([formatted_results])
df.to_csv(f"{output_prefix}.csv", index=False)
return results
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description='Evaluate generated molecules against reference sets')
parser.add_argument('--gen', required=True, help='Path to CSV file with generated SMILES')
parser.add_argument('--ref1', required=True, help='Path to .smi file with first reference set SMILES')
parser.add_argument('--ref2', help='Path to .smi file with second reference set SMILES (optional)')
parser.add_argument('--smiles-col', default='SMILES', help='Name of SMILES column in generated CSV file')
parser.add_argument('--output', default='results', help='Prefix for output files')
parser.add_argument('--n-jobs', type=int, default=8, help='Number of parallel jobs')
args = parser.parse_args()
try:
results = evaluate_molecules_from_files(
args.gen,
args.ref1,
args.ref2,
smiles_col=args.smiles_col,
output_prefix=args.output,
n_jobs=args.n_jobs
)
print(f"Evaluation complete. Results saved to {args.output}.json and {args.output}.csv")
except Exception as e:
print(f"Error during evaluation: {str(e)}")
Datasets
Models
