Precompute metrics on test splits (#14)

* metrics precalculation

* run.py supports precalculated statistics

* GPU from run.py

* imports

README.md

@@ -247,9 +247,9 @@ python scripts/run.py
 ```
 This will **download** the dataset, **train** the models, **generate** new molecules, and **calculate** the metrics. Evaluation results will be saved in `metrics.csv`.
 
-You can specify the device and/or model by running:
+You can specify the GPU index (-1 for CPU) and/or model by running:
 ```
-python scripts/run.py --device cuda:5 --model aae
+python scripts/run.py --gpu 1 --model aae
 ```
 
 For more details run `python scripts/run.py --help`.

moses/metrics/metrics.py

@@ -10,51 +10,112 @@
 from moses.utils import mapper
 from .utils_fcd import get_predictions, calculate_frechet_distance
 from multiprocessing import Pool
+from moses.utils import disable_rdkit_log, enable_rdkit_log
 
 
-def get_all_metrics(ref, gen, k=[1000, 10000], n_jobs=1, gpu=-1, batch_size=512):
+def get_all_metrics(test, gen, k=[1000, 10000], n_jobs=1, gpu=-1,
+                    batch_size=512, test_scaffolds=None,
+                    ptest=None, ptest_scaffolds=None):
     '''
-    Computes all available metrics between two lists of SMILES:
-    * %valid
-    ----- Next metrics are only computed for valid molecules -----
-    * %unique@k
-    * Frechet ChemNet Distance (FCD)
-    * fragment similarity
-    * scaffold similarity
-    * morgan similarity
+    Computes all available metrics between test (scaffold test) and generated sets of SMILES.
+    Parameters:
+        test: list of test SMILES
+        gen: list of generated SMILES
+        k: list with values for unique@k.
+            Will calculate number of unique molecules in the first k molecules.
+        n_jobs: number of workers for parallel processing
+        gpu: index of GPU for FCD metric
+        batch_size: batch size for FCD metric
+        test_scaffolds: list of scaffold test SMILES
+            Will compute only on the general test set if not specified
+        ptest: dict with precalculated statistics of the test set
+        ptest_scaffolds: dict with precalculated statistics of the scaffold test set
+        
+    
+    Available metrics:
+        * %valid
+        * %unique@k
+        * Frechet ChemNet Distance (FCD)
+        * Fragment similarity (Frag)
+        * Scaffold similarity (Scaf)
+        * Similarity to nearest neighbour (SNN)
+        * Internal diversity (IntDiv)
+        * %passes filters (Filters)
+        * Distribution difference for logP, SA, QED, NP, weight
     '''
+    disable_rdkit_log()
     metrics = {}
     if n_jobs != 1:
         pool = Pool(n_jobs)
     else:
         pool = 1
     metrics['valid'] = fraction_valid(gen, n_jobs=n_jobs)
     gen = remove_invalid(gen, canonize=True)
-    ref = remove_invalid(ref, canonize=True)
-    gen_mols = mapper(pool)(get_mol, gen)
-    ref_mols = mapper(pool)(get_mol, ref)
-
     if not isinstance(k, (list, tuple)):
         k = [k]
-    for k_ in k:
-        metrics['unique@{}'.format(k_)] = fraction_unique(gen, k_, pool)
-
-    metrics['FCD'] = frechet_chemnet_distance(ref, gen, gpu=gpu, batch_size=batch_size)
-    metrics['SNN'] = morgan_similarity(ref_mols, gen_mols, pool, gpu=gpu)
-    metrics['Frag'] = fragment_similarity(ref_mols, gen_mols, pool)
-    metrics['Scaf'] = scaffold_similarity(ref_mols, gen_mols, pool)
-    metrics['IntDiv'] = internal_diversity(gen_mols, pool)
-    metrics['Filters'] = fraction_passes_filters(gen_mols, pool)
-    metrics['logP'] = frechet_distance(ref_mols, gen_mols, logP, pool)
-    metrics['SA'] = frechet_distance(ref_mols, gen_mols, SA, pool)
-    metrics['QED'] = frechet_distance(ref_mols, gen_mols, QED, pool)
-    metrics['NP'] = frechet_distance(ref_mols, gen_mols, NP, pool)
-    metrics['weight'] = frechet_distance(ref_mols, gen_mols, weight, pool)
+    for _k in k:
+        metrics['unique@{}'.format(_k)] = fraction_unique(gen, _k, pool)
+
+    if ptest is None:
+        ptest = compute_intermediate_statistics(test, n_jobs=n_jobs, gpu=gpu, batch_size=batch_size)
+    if test_scaffolds is not None and ptest_scaffolds is None:
+        ptest_scaffolds = compute_intermediate_statistics(test_scaffolds, n_jobs=n_jobs,
+                                                                 gpu=gpu, batch_size=batch_size)
+    mols = mapper(pool)(get_mol, gen)
+    kwargs = {'n_jobs': pool, 'gpu': gpu, 'batch_size': batch_size}
+    metrics['FCD/Test'] = FCDMetric(**kwargs)(gen=gen, ptest=ptest['FCD'])
+    metrics['SNN/Test'] = SNNMetric(**kwargs)(gen=mols, ptest=ptest['SNN'])
+    metrics['Frag/Test'] = FragMetric(**kwargs)(gen=mols, ptest=ptest['Frag'])
+    metrics['Scaf/Test'] = ScafMetric(**kwargs)(gen=mols, ptest=ptest['Scaf'])
+    if ptest_scaffolds is not None:
+        metrics['FCD/TestSF'] = FCDMetric(**kwargs)(gen=gen, ptest=ptest_scaffolds['FCD'])
+        metrics['SNN/TestSF'] = SNNMetric(**kwargs)(gen=mols, ptest=ptest_scaffolds['SNN'])
+        metrics['Frag/TestSF'] = FragMetric(**kwargs)(gen=mols, ptest=ptest_scaffolds['Frag'])
+        metrics['Scaf/TestSF'] = ScafMetric(**kwargs)(gen=mols, ptest=ptest_scaffolds['Scaf'])
+
+    metrics['IntDiv'] = internal_diversity(mols, pool)
+    metrics['IntDiv2'] = internal_diversity(mols, pool, p=2)
+    metrics['Filters'] = fraction_passes_filters(mols, pool)
+
+    # Properties
+    for name, func in [('logP', logP), ('SA', SA),
+                       ('QED', QED), ('NP', NP),
+                       ('weight', weight)]:
+        metrics[name] = FrechetMetric(func, **kwargs)(gen=mols, ptest=ptest[name])
+    enable_rdkit_log()
     if n_jobs != 1:
-        pool.close()
+        pool.terminate()
     return metrics
 
 
+
+def compute_intermediate_statistics(smiles, n_jobs=1, gpu=-1, batch_size=512):
+    '''
+    The function precomputes statistics such as mean and variance for FCD, etc.
+    It is useful to compute the statistics for test and scaffold test sets to
+        speedup metrics calculation.
+    '''
+
+    if n_jobs != 1:
+        pool = Pool(n_jobs)
+    else:
+        pool = 1
+    statistics = {}
+    mols = mapper(pool)(get_mol, smiles)
+    kwargs = {'n_jobs': n_jobs, 'gpu': gpu, 'batch_size': batch_size}
+    statistics['FCD'] = FCDMetric(**kwargs).precalc(smiles)
+    statistics['SNN'] = SNNMetric(**kwargs).precalc(mols)
+    statistics['Frag'] = FragMetric(**kwargs).precalc(mols)
+    statistics['Scaf'] = ScafMetric(**kwargs).precalc(mols)
+    for name, func in [('logP', logP), ('SA', SA),
+                       ('QED', QED), ('NP', NP),
+                       ('weight', weight)]:
+        statistics[name] = FrechetMetric(func, **kwargs).precalc(mols)
+    if n_jobs != 1:
+        pool.terminate()
+    return statistics
+
+
 def fraction_passes_filters(gen, n_jobs=1):
     '''
     Computes the fraction of molecules that pass filters:
@@ -67,14 +128,15 @@ def fraction_passes_filters(gen, n_jobs=1):
     return np.mean(passes)
 
 
-def internal_diversity(gen, n_jobs=1, gpu=-1, fp_type='morgan'):
+def internal_diversity(gen, n_jobs=1, gpu=-1, fp_type='morgan', gen_fps=None, p=1):
     '''
     Computes internal diversity as:
     1/|A|^2 sum_{x, y in AxA} (1-tanimoto(x, y))
     '''
-    gen_fps = fingerprints(gen, fp_type=fp_type, n_jobs=n_jobs)
+    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', gpu=gpu)).mean()
+                                     agg='mean', gpu=gpu, p=p)).mean()
 
 
 def fraction_unique(gen, k=None, n_jobs=1, check_validity=True):
@@ -117,75 +179,108 @@ def remove_invalid(gen, canonize=True, n_jobs=1):
                 x is not None]
 
 
-def morgan_similarity(ref, gen, n_jobs=1, gpu=-1):
-    return fingerprint_similarity(ref, gen, 'morgan', n_jobs=n_jobs, gpu=gpu)
+class Metric:
+    def __init__(self, n_jobs=1, gpu=-1, batch_size=512, **kwargs):
+        self.n_jobs = n_jobs
+        self.gpu = gpu
+        self.batch_size = batch_size
+        for k, v in kwargs.values():
+            setattr(self, k, v)
 
+    def __call__(self, test=None, gen=None, ptest=None, pgen=None):
+        assert (test is None) != (ptest is None), "specify test xor ptest"
+        assert (gen is None) != (pgen is None), "specify gen xor pgen"
+        if ptest is None:
+            ptest = self.precalc(test)
+        if pgen is None:
+            pgen = self.precalc(gen)
+        return self.metric(ptest, pgen)
 
-def frechet_chemnet_distance(ref, gen, gpu=-1, batch_size=512):
-    '''
-    Computes Frechet ChemNet Distance between two lists of SMILES
+    def precalc(self, moleclues):
+        raise NotImplementedError
+
+    def metric(self, ptest, pgen):
+        raise NotImplementedError
+
+
+class FCDMetric(Metric):
     '''
-    if len(ref) < 2 or len(gen) < 2:
-        warnings.warn("Can't compute FCD for less than 2 molecules")
-        return np.nan
-    gen_activations, ref_activations = get_predictions(gen, ref, gpu=gpu,
-                                                       batch_size=batch_size)
-    mu1 = gen_activations.mean(0)
-    mu2 = ref_activations.mean(0)
-    sigma1 = np.cov(gen_activations.T)
-    sigma2 = np.cov(ref_activations.T)
-    fcd = calculate_frechet_distance(mu1, sigma1, mu2, sigma2)
-    return fcd
+    Computes Frechet ChemNet Distance
+    '''        
+    def precalc(self, smiles):
+        if len(smiles) < 2:
+            warnings.warn("Can't compute FCD for less than 2 molecules")
+            return np.nan
 
+        chemnet_activations = get_predictions(smiles, gpu=self.gpu,
+                                              batch_size=self.batch_size)
+        mu = chemnet_activations.mean(0)
+        sigma = np.cov(chemnet_activations.T)
+        return {'mu': mu, 'sigma': sigma}
+
+    def metric(self, ptest, pgen):
+        return calculate_frechet_distance(ptest['mu'], ptest['sigma'],
+                                          pgen['mu'], pgen['sigma'])
 
-def fingerprint_similarity(ref, gen, fp_type='morgan', n_jobs=1, gpu=-1):
+
+
+class SNNMetric(Metric):
     '''
-    Computes average max similarities of gen SMILES to ref SMILES
+    Computes average max similarities of gen SMILES to test SMILES
     '''
-    ref_fp = fingerprints(ref, n_jobs=n_jobs, fp_type=fp_type, morgan__r=2,
-                          morgan__n=1024)
-    gen_fp = fingerprints(gen, n_jobs=n_jobs, fp_type=fp_type, morgan__r=2,
-                          morgan__n=1024)
-    similarity = average_agg_tanimoto(ref_fp, gen_fp, gpu=gpu)
-    return similarity
+    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, ptest, pgen):
+        return average_agg_tanimoto(ptest['fps'], pgen['fps'], gpu=self.gpu)
 
 
-def count_distance(ref_counts, gen_counts):
+def cos_distance(test_counts, gen_counts):
     '''
     Computes 1 - cosine similarity between
-     dictionaries of form {type: count}. Non-present
+     dictionaries of form {name: count}. Non-present
      elements are considered zero
     '''
-    if len(ref_counts) == 0 or len(gen_counts) == 0:
+    if len(test_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])
+    keys = np.unique(list(test_counts.keys()) + list(gen_counts.keys()))
+    test_vec = np.array([test_counts.get(k, 0) for k in keys])
     gen_vec = np.array([gen_counts.get(k, 0) for k in keys])
-    return 1 - cosine(ref_vec, gen_vec)
+    return 1 - cosine(test_vec, gen_vec)
 
 
-def fragment_similarity(ref, gen, n_jobs=1):
-    ref_fragments = compute_fragments(ref, n_jobs=n_jobs)
-    gen_fragments = compute_fragments(gen, n_jobs=n_jobs)
-    return count_distance(ref_fragments, gen_fragments)
+class FragMetric(Metric):
+    def precalc(self, mols):
+        return {'frag': compute_fragments(mols, n_jobs=self.n_jobs)}
 
+    def metric(self, ptest, pgen):
+        return cos_distance(ptest['frag'], pgen['frag'])
 
-def scaffold_similarity(ref, gen, n_jobs=1):
-    ref_scaffolds = compute_scaffolds(ref, n_jobs=n_jobs)
-    gen_scaffolds = compute_scaffolds(gen, n_jobs=n_jobs)
-    return count_distance(ref_scaffolds, gen_scaffolds)
 
+class ScafMetric(Metric):
+    def precalc(self, mols):
+        return {'scaf': compute_scaffolds(mols, n_jobs=self.n_jobs)}
 
-def frechet_distance(ref, gen, func=None, n_jobs=1):
-    if func is not None:
-        ref_values = mapper(n_jobs)(func, ref)
-        gen_values = mapper(n_jobs)(func, gen)
-    else:
-        ref_values = ref
-        gen_values = gen
-    ref_mean = np.mean(ref_values)
-    ref_var = np.var(ref_values)
-    gen_mean = np.mean(gen_values)
-    gen_var = np.var(gen_values)
-    return calculate_frechet_distance(ref_mean, ref_var,
-                                      gen_mean, gen_var)
+    def metric(self, ptest, pgen):
+        return cos_distance(ptest['scaf'], pgen['scaf'])
+
+
+class FrechetMetric(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 {'mu': np.mean(values), 'var': np.var(values)}
+
+    def metric(self, ptest, pgen):
+        return calculate_frechet_distance(ptest['mu'], ptest['var'],
+                                          pgen['mu'], pgen['var'])

moses/metrics/utils.py

@@ -138,13 +138,14 @@ def compute_scaffold(mol, min_rings=2):
 
 def average_agg_tanimoto(stock_vecs, gen_vecs,
                          batch_size=5000, agg='max',
-                         gpu=-1):
+                         gpu=-1, p=1):
     '''
     For each molecule in gen_vecs finds closest molecule in stock_vecs.
     Returns average tanimoto score for between these molecules
     :param stock_vecs: numpy array <n_vectors x dim>
     :param gen_vecs: numpy array <n_vectors' x dim>
     :param agg: max or mean
+    :param p: power for averaging: (mean x^p)^(1/p)
     '''
     assert agg in ['max', 'mean'], "Can aggregate only max or mean"
     if gpu != -1:
@@ -162,6 +163,8 @@ def average_agg_tanimoto(stock_vecs, gen_vecs,
             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))
@@ -170,6 +173,8 @@ def average_agg_tanimoto(stock_vecs, gen_vecs,
                 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)