documentation

molgenevalmetric.py

@@ -43,9 +43,18 @@ from fcd_torch import FCD
 
 
 def get_mol(smiles_or_mol):
-    '''
-    Loads SMILES/molecule into RDKit's object
-    '''
+    """
+    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
@@ -60,12 +69,17 @@ def get_mol(smiles_or_mol):
     return smiles_or_mol
 
 def mapper(n_jobs):
-    '''
-    Returns function for map call.
-    If n_jobs == 1, will use standard map
-    If n_jobs > 1, will use multiprocessing pool
-    If n_jobs is a pool object, will return its map function
-    '''
+    """
+    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))
@@ -86,16 +100,30 @@ def mapper(n_jobs):
 
 def fraction_valid(gen, n_jobs=1):
     """
-    Computes a number of valid molecules
+    Calculates the fraction of valid molecules in a list of SMILES strings.
+
     Parameters:
-        gen: list of SMILES
-        n_jobs: number of threads for calculation
+    - 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
@@ -103,12 +131,16 @@ def canonic_smiles(smiles_or_mol):
 
 def fraction_unique(gen, k=None, n_jobs=1, check_validity=True):
     """
-    Computes a number of unique molecules
+    Calculates the fraction of unique molecules in a list of SMILES strings.
+
     Parameters:
-        gen: list of SMILES
-        k: compute unique@k
-        n_jobs: number of threads for calculation
-        check_validity: raises ValueError if invalid molecules are present
+    - 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:
@@ -124,6 +156,18 @@ def fraction_unique(gen, k=None, n_jobs=1, check_validity=True):
     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)
@@ -158,15 +202,20 @@ def average_agg_tanimoto(stock_vecs, gen_vecs,
                          batch_size=5000, agg='max',
                          device='cpu', p=1):
     """
-    For each molecule in gen_vecs finds closest molecule in stock_vecs.
-    Returns average tanimoto score for between these molecules
+    Calculates the average aggregate Tanimoto similarity between two sets of molecule fingerprints.
 
     Parameters:
-        stock_vecs: numpy array <n_vectors x dim>
-        gen_vecs: numpy array <n_vectors' x dim>
-        agg: max or mean
-        p: power for averaging: (mean x^p)^(1/p)
+    - 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))
@@ -276,6 +325,17 @@ def internal_diversity(gen, n_jobs=1, device='cpu', fp_type='morgan',
     """
     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)
@@ -284,6 +344,19 @@ def internal_diversity(gen, n_jobs=1, device='cpu', fp_type='morgan',
 
 
 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)
 
@@ -315,33 +388,45 @@ def fcd_metric(gen, train, n_jobs = 8, device = 'cuda:0'):
 #         return None  # Or handle empty list or all failed predictions as needed
 
 def oracles(gen, train):
-        Result = {}
-        # evaluator = Evaluator(name = 'KL_Divergence')
-        # KL_Divergence = evaluator(gen, train)
-                
-        # Result["KL_Divergence"]: KL_Divergence
 
+    """
+    Computes scores from various oracles for a list of generated molecules.
 
-        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)
+    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[f"{oracle_name}"] = score
+    # 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)
         
-        return Result
+        Result[f"{oracle_name}"] = score
+    
+    return Result
 
 
 
@@ -424,7 +509,7 @@ class molgenevalmetric(evaluate.Metric):
         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'] = SAscore(gen=trainsmi)