molgenevalmetric.py

import evaluate
import datasets
# import moses
# from moses import metrics
import pandas as pd
# from tdc import Evaluator
# from tdc import Oracle
# from metrics import novelty, fraction_valid, fraction_unique, SAscore, internal_diversity,fcd_metric, SYBAscore, oracles

import os
from collections import Counter
from functools import partial
import numpy as np
import pandas as pd
import scipy.sparse
import torch
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import MACCSkeys
from rdkit.Chem.AllChem import GetMorganFingerprintAsBitVect as Morgan
from rdkit.Chem.QED import qed
from rdkit.Chem.Scaffolds import MurckoScaffold
from rdkit.Chem import Descriptors
from multiprocessing import Pool
from collections import UserList, defaultdict
import numpy as np
import pandas as pd
from rdkit import rdBase
import sys

from rdkit.Chem import RDConfig
import os
# sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
# import sascorer
import pandas as pd
from fcd_torch import FCD
# from syba.syba import SybaClassifier

# from SCScore import SCScorer

'''
This is a standalone, importable SCScorer model. It does not have tensorflow as a
dependency and is a more attractive option for deployment. The calculations are
fast enough that there is no real reason to use GPUs (via tf) instead of CPUs (via np)
'''
import numpy as np
import time
import rdkit.Chem as Chem
import rdkit.Chem.AllChem as AllChem
import json
import gzip
import six

from types import MethodType


import os
project_root = os.path.dirname(os.path.dirname(__file__))

score_scale = 5.0
min_separation = 0.25

FP_len = 1024
FP_rad = 2

def sigmoid(x):
  return 1 / (1 + np.exp(-x))

class SCScorer():
    self.mol_to_fp = MethodType(mol_to_fp, self)

    def __init__(self, score_scale=score_scale):
        self.vars = []
        self.score_scale = score_scale
        self._restored = False

    def restore(self, weight_path='model.ckpt-10654.as_numpy.json.gz', FP_rad, FP_len):
        self.FP_len = FP_len
        self.FP_rad = FP_rad
        self._load_vars(weight_path)

        if 'uint8' in weight_path or 'counts' in weight_path:
            def mol_to_fp(mol):
                if mol is None:
                    return np.array((self.FP_len,), dtype=np.uint8)
                fp = AllChem.GetMorganFingerprint(mol, self.FP_rad, useChirality=True)  # uint sparse vector
                fp_folded = np.zeros((self.FP_len,), dtype=np.uint8)
                for k, v in six.iteritems(fp.GetNonzeroElements()):
                    fp_folded[k % self.FP_len] += v
                return np.array(fp_folded)
        else:
            def mol_to_fp(mol):
                if mol is None:
                    return np.zeros((self.FP_len,), dtype=np.float32)
                return np.array(AllChem.GetMorganFingerprintAsBitVect(mol, self.FP_rad, nBits=self.FP_len, useChirality=True), dtype=np.bool_)

        # Set the mol_to_fp method for the instance
        self.mol_to_fp = MethodType(mol_to_fp, self)

        self._restored = True
        return self

    def smi_to_fp(self, smi):
        if not smi:
            return np.zeros((self.FP_len,), dtype=np.float32)
        return self.mol_to_fp(self, Chem.MolFromSmiles(smi))

    def apply(self, x):
        if not self._restored:
            raise ValueError('Must restore model weights!')
        # Each pair of vars is a weight and bias term
        for i in range(0, len(self.vars), 2):
            last_layer = (i == len(self.vars)-2)
            W = self.vars[i]
            b = self.vars[i+1]
            x = np.matmul(x, W) + b
            if not last_layer:
                x = x * (x > 0) # ReLU
        x = 1 + (score_scale - 1) * sigmoid(x)
        return x

    def get_score_from_smi(self, smi='', v=False):
        if not smi:
            return ('', 0.)
        fp = np.array((self.smi_to_fp(smi)), dtype=np.float32)
        if sum(fp) == 0:
            if v: print('Could not get fingerprint?')
            cur_score = 0.
        else:
            # Run
            cur_score = self.apply(fp)
            if v: print('Score: {}'.format(cur_score))
        mol = Chem.MolFromSmiles(smi)
        if mol:
            smi = Chem.MolToSmiles(mol, isomericSmiles=True, kekuleSmiles=True)
        else:
            smi = ''
        return (smi, cur_score)
    
    def get_avg_score(self, smis):
        """
        Compute the average score for a list of SMILES strings.

        Args:
            smis (list of str): A list of SMILES strings.

        Returns:
            float: The average score of the given SMILES strings.
        """
        if not smis:  # Check if the list is empty
            return 0.0
        
        total_score = 0.0
        valid_smiles_count = 0
        
        for smi in smis:
            _, score = self.get_score_from_smi(smi)
            if score > 0:  # Assuming only positive scores are valid
                total_score += score
                valid_smiles_count += 1

        # Avoid division by zero
        if valid_smiles_count == 0:
            return 0.0
        else:
            return total_score / valid_smiles_count

    def _load_vars(self, weight_path):
        if weight_path.endswith('pickle'):
            import pickle
            with open(weight_path, 'rb') as fid:
                self.vars = pickle.load(fid)
                self.vars = [x.tolist() for x in self.vars]
        elif weight_path.endswith('json.gz'):
            with gzip.GzipFile(weight_path, 'r') as fin:    # 4. gzip
                json_bytes = fin.read()                      # 3. bytes (i.e. UTF-8)
                json_str = json_bytes.decode('utf-8')            # 2. string (i.e. JSON)
                self.vars = json.loads(json_str)
                self.vars = [np.array(x) for x in self.vars]



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=True):
    """
    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 SAscore(gen):
#     """
#     Calculate the average Synthetic Accessibility Score (SAscore) for a list of molecules represented by their SMILES strings.

#     Parameters:
#     - smiles_list (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.
#     """
#     scores = []
#     for smiles in gen:
#         mol = Chem.MolFromSmiles(smiles)
#         if mol:  # Ensures the molecule could be parsed from the SMILES string
#             score = sascorer.calculateScore(mol)
#             scores.append(score)
    
#     if scores:  # Checks if there are any scores calculated
#         return np.mean(scores)
#     else:
#         return None


def synthetic_complexity_score(gen):
    """
    Calculate the average Synthetic Complexity Score (SCScore) for a list of molecules represented by their SMILES strings.

    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()
    average_score = model.get_avg_score(gen)
    return average_score


def average_agg_tanimoto(stock_vecs, gen_vecs,
                         batch_size=5000, agg='max',
                         device='cpu', 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): The computation device ('cpu' or 'cuda:0', etc.).
    - 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"
    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 = 8, device = 'cuda:0'):
    """
    Computes the Fréchet ChemNet Distance (FCD) between two sets of molecules.

    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.
    """

    fcd = FCD(device=device, n_jobs= n_jobs)
    return fcd(gen, train)

# def SYBAscore(gen):
#     """
#     Compute the average SYBA score for a list of SMILES strings.

#     Parameters:
#     - smiles_list (list of str): A list of SMILES strings representing molecules.

#     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 oracles(gen, train):

    """
    Computes scores from various oracles for a list of generated molecules.

    Parameters:
    - gen (List[str]): List of generated SMILES strings.
    - train (List[str]): List of training set SMILES strings.

    Returns:
    - Dict[str, Any]: A dictionary with oracle names as keys and their corresponding scores as values.
    """

    Result = {}
    evaluator = Evaluator(name = 'KL_Divergence')
    KL_Divergence = evaluator(gen, train)
            
    Result["KL_Divergence"] = KL_Divergence


    oracle_list = [
    'QED', 'SA', 'MPO', 'GSK3B', 'JNK3',
    'DRD2', 'LogP', 'Rediscovery', 'Similarity',
    'Median', 'Isomers', 'Valsartan_SMARTS', 'Hop'
    ]

    for oracle_name in oracle_list:
        oracle = Oracle(name=oracle_name)
        if oracle_name in ['Rediscovery', 'MPO', 'Similarity', 'Median', 'Isomers', 'Hop']:
            score = oracle(gen)
            if isinstance(score, dict):
                score = {key: sum(values)/len(values) for key, values in score.items()}
        else:
            score = oracle(gen)
            if isinstance(score, list):
                score = sum(score) / len(score)
        
        Result[f"{oracle_name}"] = score
    
    return Result



_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 = """
@article{DBLP:journals/corr/abs-1811-12823,
  author       = {Daniil Polykovskiy and
                  Alexander Zhebrak and
                  Benjam{\'{\i}}n S{\'{a}}nchez{-}Lengeling and
                  Sergey Golovanov and
                  Oktai Tatanov and
                  Stanislav Belyaev and
                  Rauf Kurbanov and
                  Aleksey Artamonov and
                  Vladimir Aladinskiy and
                  Mark Veselov and
                  Artur Kadurin and
                  Sergey I. Nikolenko and
                  Al{\'{a}}n Aspuru{-}Guzik and
                  Alex Zhavoronkov},
  title        = {Molecular Sets {(MOSES):} {A} Benchmarking Platform for Molecular
                  Generation Models},
  journal      = {CoRR},
  volume       = {abs/1811.12823},
  year         = {2018},
  url          = {http://arxiv.org/abs/1811.12823},
  eprinttype    = {arXiv},
  eprint       = {1811.12823},
  timestamp    = {Fri, 26 Nov 2021 15:34:30 +0100},
  biburl       = {https://dblp.org/rec/journals/corr/abs-1811-12823.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}
"""


@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/"],
        )

    def _compute(self, gensmi, trainsmi):

        metrics = {}
        metrics['novelty'] = novelty(gen = gensmi, train = trainsmi)
        metrics['valid'] = fraction_valid(gen=gensmi)
        metrics['unique'] = fraction_unique(gen=gensmi)
        metrics['IntDiv'] = internal_diversity(gen=gensmi)
        metrics['FCD'] = fcd_metric(gen = gensmi, train = trainsmi)
        # metrics['Oracles'] = oracles(gen = gensmi, train = trainsmi)

        # metrics['SA'] = SAscore(gen=gensmi)
        metrics['SCS'] = synthetic_complexity_score(gen=gensmi)

        return metrics

        
        # generated_smiles = [s for s in generated_smiles if s != '']

        # evaluator = Evaluator(name = 'KL_Divergence')
        # KL_Divergence = evaluator(generated_smiles, train_smiles)
                
        # Results.update({
        #     "KL_Divergence": KL_Divergence,
        # })


        # oracle_list = [
        # 'QED', 'SA', 'MPO', 'GSK3B', 'JNK3',
        # 'DRD2', 'LogP', 'Rediscovery', 'Similarity',
        # 'Median', 'Isomers', 'Valsartan_SMARTS', 'Hop'
        # ]
    
        # for oracle_name in oracle_list:
        #     oracle = Oracle(name=oracle_name)
        #     if oracle_name in ['Rediscovery', 'MPO', 'Similarity', 'Median', 'Isomers', 'Hop']:
        #         score = oracle(generated_smiles)
        #         if isinstance(score, dict):
        #             score = {key: sum(values)/len(values) for key, values in score.items()}
        #     else:
        #         score = oracle(generated_smiles)
        #         if isinstance(score, list):
        #             score = sum(score) / len(score)
            
        #     Results.update({f"{oracle_name}": score})

        # # keys_to_remove = ["FCD/TestSF", "SNN/TestSF", "Frag/TestSF", "Scaf/TestSF"]
        # # for key in keys_to_remove:
        # #     Results.pop(key, None)    

        # return {"results": Results}