def compute_intermediate_statistics(smiles, n_jobs=1, device='cpu',
batch_size=512, pool=None):
"""
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.
"""
close_pool = False
if pool is None:
if n_jobs != 1:
pool = Pool(n_jobs)
close_pool = True
else:
pool = 1
statistics = {}
mols = mapper(pool)(get_mol, smiles)
kwargs = {'n_jobs': pool, 'device': device, 'batch_size': batch_size}
kwargs_fcd = {'n_jobs': n_jobs, 'device': device, 'batch_size': batch_size}
statistics['FCD'] = FCDMetric(**kwargs_fcd).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 close_pool:
pool.terminate()
return statistics
def _get_metrics(self, ref, ref_mols, rollout):
rollout_mols = mapper(self.n_jobs)(get_mol, rollout)
result = [[0 if m is None else 1] for m in rollout_mols]
if sum([r[0] for r in result], 0) == 0:
return result
rollout = remove_invalid(rollout, canonize=True, n_jobs=self.n_jobs)
rollout_mols = mapper(self.n_jobs)(get_mol, rollout)
if len(rollout) < 2:
return result
if len(self.metrics):
for metric_name in self.metrics:
if metric_name == 'fcd':
m = FCDMetric(n_jobs=self.n_jobs)(ref, rollout)
elif metric_name == 'morgan':
m = SNNMetric(n_jobs=self.n_jobs)(ref_mols, rollout_mols)
elif metric_name == 'fragments':
m = FragMetric(n_jobs=self.n_jobs)(ref_mols, rollout_mols)
elif metric_name == 'scaffolds':
m = ScafMetric(n_jobs=self.n_jobs)(ref_mols, rollout_mols)
elif metric_name == 'internal_diversity':
m = internal_diversity(rollout_mols, n_jobs=self.n_jobs)
elif metric_name == 'filters':
m = fraction_passes_filters(
rollout_mols, n_jobs=self.n_jobs
)
elif metric_name == 'logp':
m = -FrechetMetric(func=logP, n_jobs=self.n_jobs)(
ref_mols, rollout_mols
)
elif metric_name == 'sa':
m = -FrechetMetric(func=SA, n_jobs=self.n_jobs)(
ref_mols, rollout_mols
)
elif metric_name == 'qed':
m = -FrechetMetric(func=QED, n_jobs=self.n_jobs)(
ref_mols, rollout_mols
)
elif metric_name == 'np':
m = -FrechetMetric(func=NP, n_jobs=self.n_jobs)(
ref_mols, rollout_mols
)
elif metric_name == 'weight':
m = -FrechetMetric(func=weight, n_jobs=self.n_jobs)(
ref_mols, rollout_mols
)
m = MetricsReward._nan2zero(m)
for i in range(len(rollout)):
result[i].append(m)
return result
def get_all_metrics(test, gen, k=None, n_jobs=1, device='cpu',
batch_size=512, test_scaffolds=None,
ptest=None, ptest_scaffolds=None,
pool=None, gpu=None, train=None):
"""
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: int or list with values for unique@k. Will calculate number of
unique molecules in the first k molecules. Default [1000, 10000]
n_jobs: number of workers for parallel processing
device: 'cpu' or 'cuda:n', where n is GPU device number
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
pool: optional multiprocessing pool to use for parallelization
gpu: deprecated, use `device`
train: list of train SMILES
Available metrics:
* %valid
* %unique@k
* Frechet ChemNet Distance (FCD)
* Fragment similarity (Frag)
* Scaffold similarity (Scaf)
* Similarity to nearest neighbour (SNN)
* Internal diversity (IntDiv)
* Internal diversity 2: using square root of mean squared
Tanimoto similarity (IntDiv2)
* %passes filters (Filters)
* Distribution difference for logP, SA, QED, NP, weight
* Novelty (molecules not present in train)
"""
if k is None:
k = [1000, 10000]
disable_rdkit_log()
metrics = {}
if gpu is not None:
warnings.warn(
"parameter `gpu` is deprecated. Use `device`",
DeprecationWarning
)
if gpu == -1:
device = 'cpu'
else:
device = 'cuda:{}'.format(gpu)
close_pool = False
if pool is None:
if n_jobs != 1:
pool = Pool(n_jobs)
close_pool = True
else:
pool = 1
metrics['valid'] = fraction_valid(gen, n_jobs=pool)
gen = remove_invalid(gen, canonize=True)
if not isinstance(k, (list, tuple)):
k = [k]
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,
device=device,
batch_size=batch_size,
pool=pool)
if test_scaffolds is not None and ptest_scaffolds is None:
ptest_scaffolds = compute_intermediate_statistics(
test_scaffolds, n_jobs=n_jobs,
device=device, batch_size=batch_size,
pool=pool
)
mols = mapper(pool)(get_mol, gen)
kwargs = {'n_jobs': pool, 'device': device, 'batch_size': batch_size}
kwargs_fcd = {'n_jobs': n_jobs, 'device': device, 'batch_size': batch_size}
metrics['FCD/Test'] = FCDMetric(**kwargs_fcd)(gen=gen, pref=ptest['FCD'])
metrics['SNN/Test'] = SNNMetric(**kwargs)(gen=mols, pref=ptest['SNN'])
metrics['Frag/Test'] = FragMetric(**kwargs)(gen=mols, pref=ptest['Frag'])
metrics['Scaf/Test'] = ScafMetric(**kwargs)(gen=mols, pref=ptest['Scaf'])
if ptest_scaffolds is not None:
metrics['FCD/TestSF'] = FCDMetric(**kwargs_fcd)(
gen=gen, pref=ptest_scaffolds['FCD']
)
metrics['SNN/TestSF'] = SNNMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['SNN']
)
metrics['Frag/TestSF'] = FragMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['Frag']
)
metrics['Scaf/TestSF'] = ScafMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['Scaf']
)
metrics['IntDiv'] = internal_diversity(mols, pool, device=device)
metrics['IntDiv2'] = internal_diversity(mols, pool, device=device, 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,
pref=ptest[name])
if train is not None:
metrics['Novelty'] = novelty(mols, train, pool)
enable_rdkit_log()
if close_pool:
pool.close()
pool.join()
return metrics
def get_all_metrics(gen,
k=None,
n_jobs=1,
device='cpu',
batch_size=512,
pool=None,
test=None,
test_scaffolds=None,
ptest=None,
ptest_scaffolds=None,
train=None):
"""
Computes all available metrics between test (scaffold test)
and generated sets of SMILES.
Parameters:
gen: list of generated SMILES
k: int or list with values for unique@k. Will calculate number of
unique molecules in the first k molecules. Default [1000, 10000]
n_jobs: number of workers for parallel processing
device: 'cpu' or 'cuda:n', where n is GPU device number
batch_size: batch size for FCD metric
pool: optional multiprocessing pool to use for parallelization
test (None or list): test SMILES. If None, will load
a default test set
test_scaffolds (None or list): scaffold test SMILES. If None, will
load a default scaffold test set
ptest (None or dict): precalculated statistics of the test set. If
None, will load default test statistics. If you specified a custom
test set, default test statistics will be ignored
ptest_scaffolds (None or dict): precalculated statistics of the
scaffold test set If None, will load default scaffold test
statistics. If you specified a custom test set, default test
statistics will be ignored
train (None or list): train SMILES. If None, will load a default
train set
Available metrics:
* %valid
* %unique@k
* Frechet ChemNet Distance (FCD)
* Fragment similarity (Frag)
* Scaffold similarity (Scaf)
* Similarity to nearest neighbour (SNN)
* Internal diversity (IntDiv)
* Internal diversity 2: using square root of mean squared
Tanimoto similarity (IntDiv2)
* %passes filters (Filters)
* Distribution difference for logP, SA, QED, weight
* Novelty (molecules not present in train)
"""
if test is None:
if ptest is not None:
raise ValueError("You cannot specify custom test "
"statistics for default test set")
test = get_dataset('test')
ptest = get_statistics('test')
if test_scaffolds is None:
if ptest_scaffolds is not None:
raise ValueError("You cannot specify custom scaffold test "
"statistics for default scaffold test set")
test_scaffolds = get_dataset('test_scaffolds')
ptest_scaffolds = get_statistics('test_scaffolds')
train = train or get_dataset('train')
if k is None:
k = [1000, 10000]
disable_rdkit_log()
metrics = {}
close_pool = False
if pool is None:
if n_jobs != 1:
pool = Pool(n_jobs)
close_pool = True
else:
pool = 1
metrics['valid'] = fraction_valid(gen, n_jobs=pool)
gen = remove_invalid(gen, canonize=True)
if not isinstance(k, (list, tuple)):
k = [k]
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,
device=device,
batch_size=batch_size,
pool=pool)
if test_scaffolds is not None and ptest_scaffolds is None:
ptest_scaffolds = compute_intermediate_statistics(
test_scaffolds,
n_jobs=n_jobs,
device=device,
batch_size=batch_size,
pool=pool)
mols = mapper(pool)(get_mol, gen)
kwargs = {'n_jobs': pool, 'device': device, 'batch_size': batch_size}
kwargs_fcd = {'n_jobs': n_jobs, 'device': device, 'batch_size': batch_size}
metrics['FCD/Test'] = FCDMetric(**kwargs_fcd)(gen=gen, pref=ptest['FCD'])
metrics['SNN/Test'] = SNNMetric(**kwargs)(gen=mols, pref=ptest['SNN'])
metrics['Frag/Test'] = FragMetric(**kwargs)(gen=mols, pref=ptest['Frag'])
metrics['Scaf/Test'] = ScafMetric(**kwargs)(gen=mols, pref=ptest['Scaf'])
if ptest_scaffolds is not None:
metrics['FCD/TestSF'] = FCDMetric(**kwargs_fcd)(
gen=gen, pref=ptest_scaffolds['FCD'])
metrics['SNN/TestSF'] = SNNMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['SNN'])
metrics['Frag/TestSF'] = FragMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['Frag'])
metrics['Scaf/TestSF'] = ScafMetric(**kwargs)(
gen=mols, pref=ptest_scaffolds['Scaf'])
metrics['IntDiv'] = internal_diversity(mols, pool, device=device)
metrics['IntDiv2'] = internal_diversity(mols, pool, device=device, p=2)
metrics['Filters'] = fraction_passes_filters(mols, pool)
# Properties
for name, func in [('logP', logP), ('SA', SA), ('QED', QED),
('weight', weight)]:
metrics[name] = WassersteinMetric(func, **kwargs)(gen=mols,
pref=ptest[name])
if train is not None:
metrics['Novelty'] = novelty(mols, train, pool)
enable_rdkit_log()
if close_pool:
pool.close()
pool.join()
return metrics
def calculate(self,
gen,
calc_valid=False,
calc_unique=False,
unique_k=None,
se_k=None):
metrics = {}
metrics['#'] = len(gen)
# Calculate validity
if calc_valid:
metrics['Validity'] = fraction_valid(gen, self.pool)
gen = remove_invalid(gen, canonize=True)
mols = mapper(self.pool)(get_mol, gen)
metrics['# valid'] = len(gen)
# Calculate Uniqueness
if calc_unique:
metrics['Uniqueness'] = fraction_unique(gen=gen,
k=None,
n_jobs=self.pool)
if unique_k is not None:
metrics[f'Unique@{unique_k/1000:.0f}k'] = fraction_unique(
gen=gen, k=unique_k, n_jobs=self.pool)
# Now subset only unique molecules
gen = list(set(gen))
mols = mapper(self.pool)(get_mol, gen)
# Precalculate some things
mol_fps = fingerprints(mols,
self.pool,
already_unique=True,
fp_type='morgan')
scaffs = compute_scaffolds(mols, n_jobs=self.n_jobs)
scaff_gen = list(scaffs.keys())
scaff_mols = mapper(self.pool)(get_mol, scaff_gen)
metrics['# valid & unique'] = len(gen)
# Calculate diversity related metrics
if self.train is not None:
metrics['Novelty'] = novelty(gen, self.train, self.pool)
metrics['IntDiv1'] = internal_diversity(gen=mol_fps,
n_jobs=self.pool,
device=self.device)
metrics['IntDiv2'] = internal_diversity(gen=mol_fps,
n_jobs=self.pool,
device=self.device,
p=2)
metrics['SEDiv'] = se_diversity(gen=mols, n_jobs=self.pool)
if se_k is not None:
metrics[f'SEDiv@{se_k/1000:.0f}k'] = se_diversity(gen=mols,
k=se_k,
n_jobs=self.pool,
normalize=True)
metrics['ScaffDiv'] = internal_diversity(gen=scaff_mols,
n_jobs=self.pool,
device=self.device,
fp_type='morgan')
metrics['Scaff uniqueness'] = len(scaff_gen) / len(gen)
# Calculate % pass filters
metrics['Filters'] = fraction_passes_filters(mols, self.pool)
# Calculate FCD
pgen = FCDMetric(**self.kwargs_fcd).precalc(gen)
if self.ptrain:
metrics['FCD_train'] = FCDMetric(**self.kwargs_fcd)(
pgen=pgen, pref=self.ptrain)
if self.ptest:
metrics['FCD_test'] = FCDMetric(**self.kwargs_fcd)(pgen=pgen,
pref=self.ptest)
if self.ptest_scaffolds:
metrics['FCD_testSF'] = FCDMetric(**self.kwargs_fcd)(
pgen=pgen, pref=self.ptest_scaffolds)
if self.ptarget:
metrics['FCD_target'] = FCDMetric(**self.kwargs_fcd)(
pgen=pgen, pref=self.ptarget)
# Test metrics
if self.test_int is not None:
metrics['SNN_test'] = SNNMetric(**self.kwargs)(
pgen={
'fps': mol_fps
}, pref=self.test_int['SNN'])
metrics['Frag_test'] = FragMetric(**self.kwargs)(
gen=mols, pref=self.test_int['Frag'])
metrics['Scaf_test'] = ScafMetric(**self.kwargs)(
pgen={
'scaf': scaffs
}, pref=self.test_int['Scaf'])
# Test scaff metrics
if self.test_scaffolds_int is not None:
metrics['SNN_testSF'] = SNNMetric(**self.kwargs)(
pgen={
'fps': mol_fps
}, pref=self.test_scaffolds_int['SNN'])
metrics['Frag_testSF'] = FragMetric(**self.kwargs)(
gen=mols, pref=self.test_scaffolds_int['Frag'])
metrics['Scaf_testSF'] = ScafMetric(**self.kwargs)(
pgen={
'scaf': scaffs
}, pref=self.test_scaffolds_int['Scaf'])
# Target metrics
if self.target_int is not None:
metrics['SNN_target'] = SNNMetric(**self.kwargs)(
pgen={
'fps': mol_fps
}, pref=self.target_int['SNN'])
metrics['Frag_target'] = FragMetric(**self.kwargs)(
gen=mols, pref=self.target_int['Frag'])
metrics['Scaf_target'] = ScafMetric(**self.kwargs)(
pgen={
'scaf': scaffs
}, pref=self.target_int['Scaf'])
return metrics