def __init__(self, ref, n_ref_subsample, n_rollouts, n_jobs, metrics=[]):
assert all([m in MetricsReward.supported_metrics for m in metrics])
self.ref = remove_invalid(ref, canonize=True, n_jobs=n_jobs)
self.ref_mols = mapper(n_jobs)(get_mol, self.ref)
self.n_ref_subsample = n_ref_subsample
self.n_rollouts = n_rollouts
self.n_jobs = n_jobs
self.metrics = metrics
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)
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 = SNN(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_reference_data(self, data):
ref_smiles = remove_invalid(data, canonize=True, n_jobs=self.n_jobs)
ref_mols = mapper(self.n_jobs)(get_mol, ref_smiles)
return ref_smiles, ref_mols