import evaluate import datasets import pandas as pd import numpy as np import scipy.sparse from scipy.spatial.distance import cosine as cos_distance from scipy.stats import wasserstein_distance import torch import warnings from multiprocessing import Pool from functools import partial from fcd_torch import FCD from collections import Counter from tdc import Oracle from rdkit.Chem.Crippen import MolLogP from rdkit import Chem from rdkit.Chem import MACCSkeys from rdkit.Chem import AllChem from rdkit.Chem.AllChem import GetMorganFingerprintAsBitVect as Morgan from rdkit.Chem.QED import qed from rdkit.Contrib.SA_Score import sascorer from rdkit.Chem.Scaffolds import MurckoScaffold from syba.syba import SybaClassifier from myscscore.SCScore import SCScorer def get_mol(smiles_or_mol): """ Converts a SMILES string or RDKit molecule object to an RDKit molecule object. If the input is already an RDKit molecule object, it returns it directly. For a SMILES string, it attempts to create an RDKit molecule object. Parameters: - smiles_or_mol (str or Mol): The SMILES string of the molecule or an RDKit molecule object. Returns: - Mol or None: The RDKit molecule object or None if conversion fails. """ if isinstance(smiles_or_mol, str): if len(smiles_or_mol) == 0: return None mol = Chem.MolFromSmiles(smiles_or_mol) if mol is None: return None try: Chem.SanitizeMol(mol) except ValueError: return None return mol return smiles_or_mol def mapper(n_jobs): """ Returns a mapping function suitable for parallel or sequential execution based on the value of n_jobs. Parameters: - n_jobs (int or Pool): Number of jobs for parallel execution or a multiprocessing Pool object. Returns: - Function: A mapping function that can be used for applying a function over a sequence. """ if n_jobs == 1: def _mapper(*args, **kwargs): return list(map(*args, **kwargs)) return _mapper if isinstance(n_jobs, int): pool = Pool(n_jobs) def _mapper(*args, **kwargs): try: result = pool.map(*args, **kwargs) finally: pool.terminate() return result return _mapper return n_jobs.map def fraction_valid(gen, n_jobs=1): """ Calculates the fraction of valid molecules in a list of SMILES strings. Parameters: - gen (list of str): List of SMILES strings. - n_jobs (int): Number of parallel jobs to use for computation. Returns: - float: Fraction of valid molecules. """ gen = mapper(n_jobs)(get_mol, gen) return 1 - gen.count(None) / len(gen) def canonic_smiles(smiles_or_mol): """ Converts a molecule into its canonical SMILES representation. Parameters: - smiles_or_mol (str or Mol): SMILES string or RDKit molecule object. Returns: - str or None: Canonical SMILES string, or None if conversion fails. """ mol = get_mol(smiles_or_mol) if mol is None: return None return Chem.MolToSmiles(mol) def fraction_unique(gen, k=None, n_jobs=1, check_validity=False): """ Calculates the fraction of unique molecules in a list of SMILES strings. Parameters: - gen (list of str): List of SMILES strings. - k (int, optional): Number of top molecules to consider for uniqueness. If None, considers all. - n_jobs (int): Number of parallel jobs to use for computation. - check_validity (bool): If True, checks for the validity of molecules. Returns: - float: Fraction of unique molecules. """ if k is not None: if len(gen) < k: warnings.warn( "Can't compute unique@{}.".format(k) + "gen contains only {} molecules".format(len(gen)) ) gen = gen[:k] canonic = set(mapper(n_jobs)(canonic_smiles, gen)) if None in canonic and check_validity: raise ValueError("Invalid molecule passed to unique@k") return len(canonic) / len(gen) def novelty(gen, train, n_jobs=1): """ Computes the novelty of generated molecules compared to a training set. Parameters: - gen (List[str]): List of generated SMILES strings. - train (List[str]): List of SMILES strings from the training set. - n_jobs (int): Number of parallel jobs to use for computation. Returns: - float: Novelty score. """ gen_smiles = mapper(n_jobs)(canonic_smiles, gen) gen_smiles_set = set(gen_smiles) - {None} train_set = set(train) return len(gen_smiles_set - train_set) / len(gen_smiles_set) def synthetic_complexity_score(gen): """ Calculate the average Synthetic Complexity Score (SCScore) for a list of molecules represented by their SMILES strings. The SCScore model rates the synthetic complexity of molecules on a scale from 1 to 5. Based on the premise that on average, the products of published chemical reactions should be more synthetically complex than their corresponding reactants Parameters: - gen (list of str): A list containing the SMILES representations of the molecules. Returns: - float: The average Synthetic Accessibility Score for the valid molecules in the list. Returns None if no valid molecules are found. """ model = SCScorer() model.restore() average_score = model.get_avg_score(gen) return average_score def calculate_sa_score(smiles): """ Calculates the SA score for a single SMILES string. Evaluates the ease of synthesizing drug-like molecules in virtual screening. Ranges from 1 (easy to synthesize) to 10 (hard to synthesize) This score reflects the presence of common fragments in a molecule and structural complexities. Parameters: - smiles (str): SMILES string of the molecule. Returns: - float: SA score of the molecule, or None if the molecule couldn't be created. """ mol = get_mol(smiles) if mol: return sascorer.calculateScore(mol) else: return None def average_sascore(gen, n_jobs=1): """ Computes the average synthetic accessibility score for a list of molecules. Parameters: - gen (List[str]): List of generated SMILES strings. - n_jobs (int): Number of parallel jobs to use for computation. Returns: - float: Average SA score, or None if no scores could be computed. """ scores = mapper(n_jobs)(calculate_sa_score, gen) # Filter out None values which indicate failed molecule creation valid_scores = [score for score in scores if score is not None] if valid_scores: return sum(valid_scores) / len(valid_scores) else: return None def average_agg_tanimoto(stock_vecs, gen_vecs, batch_size=5000, agg='max', device=None, p=1): """ Calculates the average aggregate Tanimoto similarity between two sets of molecule fingerprints. Parameters: - stock_vecs (numpy array): Fingerprint vectors for the reference molecule set. - gen_vecs (numpy array): Fingerprint vectors for the generated molecule set. - batch_size (int): The size of batches to process similarities (reduces memory usage). - agg (str): Aggregation method, either 'max' or 'mean'. - device (str or None): The computation device ('cpu' or 'cuda:0', etc.). If None, automatically detect. - p (float): The power for averaging, used in generalized mean calculation. Returns: - float: Average aggregate Tanimoto similarity score. """ assert agg in ['max', 'mean'], "Can aggregate only max or mean" # Automatically detect device if not provided if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' agg_tanimoto = np.zeros(len(gen_vecs)) total = np.zeros(len(gen_vecs)) for j in range(0, stock_vecs.shape[0], batch_size): x_stock = torch.tensor(stock_vecs[j:j + batch_size]).to(device).float() for i in range(0, gen_vecs.shape[0], batch_size): y_gen = torch.tensor(gen_vecs[i:i + batch_size]).to(device).float() y_gen = y_gen.transpose(0, 1) tp = torch.mm(x_stock, y_gen) jac = (tp / (x_stock.sum(1, keepdim=True) + y_gen.sum(0, keepdim=True) - tp)).cpu().numpy() jac[np.isnan(jac)] = 1 if p != 1: jac = jac**p if agg == 'max': agg_tanimoto[i:i + y_gen.shape[1]] = np.maximum( agg_tanimoto[i:i + y_gen.shape[1]], jac.max(0)) elif agg == 'mean': agg_tanimoto[i:i + y_gen.shape[1]] += jac.sum(0) total[i:i + y_gen.shape[1]] += jac.shape[0] if agg == 'mean': agg_tanimoto /= total if p != 1: agg_tanimoto = (agg_tanimoto)**(1/p) return np.mean(agg_tanimoto) def fingerprint(smiles_or_mol, fp_type='maccs', dtype=None, morgan__r=2, morgan__n=1024, *args, **kwargs): """ Generates fingerprint for SMILES If smiles is invalid, returns None Returns numpy array of fingerprint bits Parameters: smiles: SMILES string type: type of fingerprint: [MACCS|morgan] dtype: if not None, specifies the dtype of returned array """ fp_type = fp_type.lower() molecule = get_mol(smiles_or_mol, *args, **kwargs) if molecule is None: return None if fp_type == 'maccs': keys = MACCSkeys.GenMACCSKeys(molecule) keys = np.array(keys.GetOnBits()) fingerprint = np.zeros(166, dtype='uint8') if len(keys) != 0: fingerprint[keys - 1] = 1 # We drop 0-th key that is always zero elif fp_type == 'morgan': fingerprint = np.asarray(Morgan(molecule, morgan__r, nBits=morgan__n), dtype='uint8') else: raise ValueError("Unknown fingerprint type {}".format(fp_type)) if dtype is not None: fingerprint = fingerprint.astype(dtype) return fingerprint def fingerprints(smiles_mols_array, n_jobs=1, already_unique=False, *args, **kwargs): ''' Computes fingerprints of smiles np.array/list/pd.Series with n_jobs workers e.g.fingerprints(smiles_mols_array, type='morgan', n_jobs=10) Inserts np.NaN to rows corresponding to incorrect smiles. IMPORTANT: if there is at least one np.NaN, the dtype would be float Parameters: smiles_mols_array: list/array/pd.Series of smiles or already computed RDKit molecules n_jobs: number of parralel workers to execute already_unique: flag for performance reasons, if smiles array is big and already unique. Its value is set to True if smiles_mols_array contain RDKit molecules already. ''' if isinstance(smiles_mols_array, pd.Series): smiles_mols_array = smiles_mols_array.values else: smiles_mols_array = np.asarray(smiles_mols_array) if not isinstance(smiles_mols_array[0], str): already_unique = True if not already_unique: smiles_mols_array, inv_index = np.unique(smiles_mols_array, return_inverse=True) fps = mapper(n_jobs)( partial(fingerprint, *args, **kwargs), smiles_mols_array ) length = 1 for fp in fps: if fp is not None: length = fp.shape[-1] first_fp = fp break fps = [fp if fp is not None else np.array([np.NaN]).repeat(length)[None, :] for fp in fps] if scipy.sparse.issparse(first_fp): fps = scipy.sparse.vstack(fps).tocsr() else: fps = np.vstack(fps) if not already_unique: return fps[inv_index] return fps def internal_diversity(gen, n_jobs=1, device='cpu', fp_type='morgan', gen_fps=None, p=1): """ Computes internal diversity as: 1/|A|^2 sum_{x, y in AxA} (1-tanimoto(x, y)) Parameters: - gen (List[str]): List of generated SMILES strings. - n_jobs (int): Number of parallel jobs for fingerprint computation. - device (str): Computation device ('cpu' or 'cuda:0', etc.). - fp_type (str): Type of fingerprint to use ('morgan', etc.). - gen_fps (Optional[np.ndarray]): Precomputed fingerprints of generated molecules. If None, will be computed. Returns: - float: Internal diversity score. """ if gen_fps is None: gen_fps = fingerprints(gen, fp_type=fp_type, n_jobs=n_jobs) return 1 - (average_agg_tanimoto(gen_fps, gen_fps, agg='mean', device=device, p=p)).mean() def fcd_metric(gen, train, n_jobs = 1, device = None): """ Computes the Fréchet ChemNet Distance (FCD) between two sets of molecules. FCD is calculated using the Fréchet Distance between feature vectors of generated and real molecules obtained from ChemNet. A lower FCD score indicates higher chemical realism and diversity in the molecules generated by a model. Parameters: - gen (List[str]): List of generated SMILES strings. - train (List[str]): List of training set SMILES strings. - n_jobs (int): Number of parallel jobs for computation. - device (str): Computation device for the FCD calculation. Returns: - float: FCD score. """ # Determine the device dynamically based on CUDA availability if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' else: device = torch.device(device if torch.cuda.is_available() and 'cuda' in device else 'cpu') fcd = FCD(device=device, n_jobs= n_jobs) return fcd(gen, train) def SYBAscore(gen): """ Compute the average SYBA (SYnthetic Bayesian Accessibility) score for a list of SMILES strings. It is a fragment-based method for the rapid classification of organic compounds as easy- (ES) or hard-to-synthesize (HS). Based on a Bernoulli naïve Bayes classifier that is used to assign SYBA score contributions to individual fragments based on their frequencies in the database of ES and HS molecules. Trained on ES molecules available in the ZINC15 database and on HS molecules generated by the Nonpher methodology Parameters: - gen (List[str]): List of generated SMILES strings. Returns: - float: The average SYBA score for the list of molecules. """ syba = SybaClassifier() syba.fitDefaultScore() scores = [] for smiles in gen: try: score = syba.predict(smi=smiles) scores.append(score) except Exception as e: print(f"Error processing SMILES '{smiles}': {e}") continue if scores: return sum(scores) / len(scores) else: return None # Or handle empty list or all failed predictions as needed def qed_metric(gen): """ Computes RDKit's QED score. A [0,1] value estimating how likely a molecule is a viable candidate for a drug. QED is meant to capture certain desirable traits that successful drug molecules tend to possess Parameters: - gen (List[str]): List of generated SMILES strings. Returns: - float: The average QED score for the list of molecules. """ if not gen: return 0.0 # Return 0 or suitable value for empty list # Convert SMILES strings to RDKit molecule objects and calculate QED scores qed_scores = [] for smiles in gen: try: mol = get_mol(smiles) if mol: # Ensure molecule is valid qed_scores.append(qed(mol)) except Exception as e: print(f"Error processing molecule {smiles}: {str(e)}") # Calculate the average QED score if qed_scores: return sum(qed_scores) / len(qed_scores) else: return 0.0 # Return 0 or suitable value if no valid molecules are processed def logP_metric(gen): """ Computes the average RDKit's logP value for a list of SMILES strings. LogP is the log of the partition coefficient of a solute between octanol and water, at near infinite dilution. It is stated that LogP should be between 0 and 5 for a small molecule drug to be a candidate for oral administration. Computed with RDKit's Crippen (Wildman and Crippen, 1999) estimation. Parameters: - gen (List[str]): List of generated SMILES strings. Returns: - float: Average logP value for the list of molecules. """ # Check if the input list is empty if not gen: return 0.0 # Return 0 or suitable value for empty list # Convert SMILES strings to RDKit molecule objects and calculate logP values logP_values = [] for smiles in gen: try: mol = get_mol(smiles) if mol: # Ensure molecule is valid logP_values.append(MolLogP(mol)) except Exception as e: print(f"Error processing molecule {smiles}: {str(e)}") # Calculate the average logP value if logP_values: return sum(logP_values) / len(logP_values) else: return 0.0 # Return 0 or suitable value if no valid molecules are processed def penalized_logp(gen): """ Computes the average PyTDC's penalized logP value for a list of SMILES strings. Captures LogP, SA and penalty for number of rings. Parameters: - gen (List[str]): List of generated SMILES strings. Returns: - float: Average penalized logP value for the list of molecules. """ oracle = Oracle('LogP') score = oracle(gen) if isinstance(score, list): score = sum(score) / len(score) return score _DESCRIPTION = """ Comprehensive suite of metrics designed to assess the performance of molecular generation models, for understanding how well a model can produce novel, chemically valid molecules that are relevant to specific research objectives. """ _KWARGS_DESCRIPTION = """ Args: generated_smiles (`list` of `string`): A collection of SMILES (Simplified Molecular Input Line Entry System) strings generated by the model, ideally encompassing more than 30,000 samples. train_smiles (`list` of `string`): The dataset of SMILES strings used to train the model, serving as a reference to evaluate the novelty and diversity of the generated molecules. Returns: Dectionary item containing various metrics to evaluate model performance """ _CITATION = """ """ @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION) class molgenevalmetric(evaluate.Metric): def _info(self): return evaluate.MetricInfo( description=_DESCRIPTION, citation=_CITATION, inputs_description=_KWARGS_DESCRIPTION, features=datasets.Features( { "gensmi": datasets.Sequence(datasets.Value("string")), "trainsmi": datasets.Sequence(datasets.Value("string")), } if self.config_name == "multilabel" else { "gensmi": datasets.Value("string"), "trainsmi": datasets.Value("string"), } ), reference_urls=["https://github.com/molecularsets/moses", "https://tdcommons.ai/functions/oracles/", "https://github.com/lich-uct/syba", "https://github.com/connorcoley/scscore"], ) def _compute(self, gensmi, trainsmi): metrics = {} metric_functions = { 'Novelty': lambda: novelty(gen=gensmi, train=trainsmi), 'Valid': lambda: fraction_valid(gen=gensmi), 'Unique': lambda: fraction_unique(gen=gensmi), 'IntDiv': lambda: internal_diversity(gen=gensmi), 'FCD': lambda: fcd_metric(gen=gensmi, train=trainsmi), 'QED': lambda: qed_metric(gen=gensmi), 'LogP': lambda: logP_metric(gen=gensmi), 'Penalized LogP': lambda: penalized_logp(gen=gensmi), 'SA': lambda: average_sascore(gen=gensmi), 'SCScore': lambda: synthetic_complexity_score(gen=gensmi), 'SYBA': lambda: SYBAscore(gen=gensmi), # 'Oracles': lambda: oracles(gen=gensmi, train=trainsmi) } for metric_name, compute_func in metric_functions.items(): print(f"Computing {metric_name}...") try: metrics[metric_name] = compute_func() except Exception as e: print(f"Error computing {metric_name}: {e}") metrics[metric_name] = 0 return metrics # def get_n_rings(mol): # """ # Computes the number of rings in a molecule # """ # return mol.GetRingInfo().NumRings() # def fragmenter(mol): # """ # fragment mol using BRICS and return smiles list # """ # fgs = AllChem.FragmentOnBRICSBonds(get_mol(mol)) # fgs_smi = Chem.MolToSmiles(fgs).split(".") # return fgs_smi # def compute_fragments(mol_list, n_jobs=1): # """ # fragment list of mols using BRICS and return smiles list # """ # fragments = Counter() # for mol_frag in mapper(n_jobs)(fragmenter, mol_list): # fragments.update(mol_frag) # return fragments # def compute_scaffolds(mol_list, n_jobs=1, min_rings=2): # """ # Extracts a scafold from a molecule in a form of a canonic SMILES # """ # scaffolds = Counter() # map_ = mapper(n_jobs) # scaffolds = Counter( # map_(partial(compute_scaffold, min_rings=min_rings), mol_list)) # if None in scaffolds: # scaffolds.pop(None) # return scaffolds # def compute_scaffold(mol, min_rings=2): # mol = get_mol(mol) # try: # scaffold = MurckoScaffold.GetScaffoldForMol(mol) # except (ValueError, RuntimeError): # return None # n_rings = get_n_rings(scaffold) # scaffold_smiles = Chem.MolToSmiles(scaffold) # if scaffold_smiles == '' or n_rings < min_rings: # return None # return scaffold_smiles # class Metric: # def __init__(self, n_jobs=1, device='cpu', batch_size=512, **kwargs): # self.n_jobs = n_jobs # self.device = device # self.batch_size = batch_size # for k, v in kwargs.values(): # setattr(self, k, v) # def __call__(self, ref=None, gen=None, pref=None, pgen=None): # assert (ref is None) != (pref is None), "specify ref xor pref" # assert (gen is None) != (pgen is None), "specify gen xor pgen" # if pref is None: # pref = self.precalc(ref) # if pgen is None: # pgen = self.precalc(gen) # return self.metric(pref, pgen) # def precalc(self, moleclues): # raise NotImplementedError # def metric(self, pref, pgen): # raise NotImplementedError # class SNNMetric(Metric): # """ # Computes average max similarities of gen SMILES to ref SMILES # """ # def __init__(self, fp_type='morgan', **kwargs): # self.fp_type = fp_type # super().__init__(**kwargs) # def precalc(self, mols): # return {'fps': fingerprints(mols, n_jobs=self.n_jobs, # fp_type=self.fp_type)} # def metric(self, pref, pgen): # return average_agg_tanimoto(pref['fps'], pgen['fps'], # device=self.device) # def cos_similarity(ref_counts, gen_counts): # """ # Computes cosine similarity between # dictionaries of form {name: count}. Non-present # elements are considered zero: # sim = / ||r|| / ||g|| # """ # if len(ref_counts) == 0 or len(gen_counts) == 0: # return np.nan # keys = np.unique(list(ref_counts.keys()) + list(gen_counts.keys())) # ref_vec = np.array([ref_counts.get(k, 0) for k in keys]) # gen_vec = np.array([gen_counts.get(k, 0) for k in keys]) # return 1 - cos_distance(ref_vec, gen_vec) # class FragMetric(Metric): # def precalc(self, mols): # return {'frag': compute_fragments(mols, n_jobs=self.n_jobs)} # def metric(self, pref, pgen): # return cos_similarity(pref['frag'], pgen['frag']) # class ScafMetric(Metric): # def precalc(self, mols): # return {'scaf': compute_scaffolds(mols, n_jobs=self.n_jobs)} # def metric(self, pref, pgen): # return cos_similarity(pref['scaf'], pgen['scaf']) # class WassersteinMetric(Metric): # def __init__(self, func=None, **kwargs): # self.func = func # super().__init__(**kwargs) # def precalc(self, mols): # if self.func is not None: # values = mapper(self.n_jobs)(self.func, mols) # else: # values = mols # return {'values': values} # def metric(self, pref, pgen): # return wasserstein_distance( # pref['values'], pgen['values'] # ) # def get_frag(gen): # mols = mapper(pool)(get_mol, gen) # kwargs = {'n_jobs': pool, 'device': device, 'batch_size': batch_size}