File size: 3,955 Bytes
5e9bd47 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import numpy as np
import multiprocessing
import rdkit
import rdkit.Chem as Chem
rdkit.RDLogger.DisableLog('rdApp.*')
from SmilesPE.pretokenizer import atomwise_tokenizer
def canonicalize_smiles(smiles, ignore_chiral=False, ignore_cistrans=False, replace_rgroup=True):
if type(smiles) is not str or smiles == '':
return '', False
if ignore_cistrans:
smiles = smiles.replace('/', '').replace('\\', '')
if replace_rgroup:
tokens = atomwise_tokenizer(smiles)
for j, token in enumerate(tokens):
if token[0] == '[' and token[-1] == ']':
symbol = token[1:-1]
if symbol[0] == 'R' and symbol[1:].isdigit():
tokens[j] = f'[{symbol[1:]}*]'
elif Chem.AtomFromSmiles(token) is None:
tokens[j] = '*'
smiles = ''.join(tokens)
try:
canon_smiles = Chem.CanonSmiles(smiles, useChiral=(not ignore_chiral))
success = True
except:
canon_smiles = smiles
success = False
return canon_smiles, success
def convert_smiles_to_canonsmiles(
smiles_list, ignore_chiral=False, ignore_cistrans=False, replace_rgroup=True, num_workers=16):
with multiprocessing.Pool(num_workers) as p:
results = p.starmap(canonicalize_smiles,
[(smiles, ignore_chiral, ignore_cistrans, replace_rgroup) for smiles in smiles_list],
chunksize=128)
canon_smiles, success = zip(*results)
return list(canon_smiles), np.mean(success)
class SmilesEvaluator(object):
def __init__(self, gold_smiles, num_workers=16):
self.gold_smiles = gold_smiles
self.gold_canon_smiles, self.gold_valid = convert_smiles_to_canonsmiles(gold_smiles, num_workers=num_workers)
self.gold_smiles_chiral, _ = convert_smiles_to_canonsmiles(gold_smiles,
ignore_chiral=True, num_workers=num_workers)
self.gold_smiles_cistrans, _ = convert_smiles_to_canonsmiles(gold_smiles,
ignore_cistrans=True, num_workers=num_workers)
self.gold_canon_smiles = self._replace_empty(self.gold_canon_smiles)
self.gold_smiles_chiral = self._replace_empty(self.gold_smiles_chiral)
self.gold_smiles_cistrans = self._replace_empty(self.gold_smiles_cistrans)
def _replace_empty(self, smiles_list):
"""Replace empty SMILES in the gold, otherwise it will be considered correct if both pred and gold is empty."""
return [smiles if smiles is not None and type(smiles) is str and smiles != "" else "<empty>"
for smiles in smiles_list]
def evaluate(self, pred_smiles):
results = {}
results['gold_valid'] = self.gold_valid
# Canon SMILES
pred_canon_smiles, pred_valid = convert_smiles_to_canonsmiles(pred_smiles)
results['canon_smiles_em'] = (np.array(self.gold_canon_smiles) == np.array(pred_canon_smiles)).mean()
results['pred_valid'] = pred_valid
# Ignore chirality (Graph exact match)
pred_smiles_chiral, _ = convert_smiles_to_canonsmiles(pred_smiles, ignore_chiral=True)
results['graph'] = (np.array(self.gold_smiles_chiral) == np.array(pred_smiles_chiral)).mean()
# Ignore double bond cis/trans
pred_smiles_cistrans, _ = convert_smiles_to_canonsmiles(pred_smiles, ignore_cistrans=True)
results['canon_smiles'] = (np.array(self.gold_smiles_cistrans) == np.array(pred_smiles_cistrans)).mean()
# Evaluate on molecules with chiral centers
chiral = np.array([[g, p] for g, p in zip(self.gold_smiles_cistrans, pred_smiles_cistrans) if '@' in g])
results['chiral_ratio'] = len(chiral) / len(self.gold_smiles)
results['chiral'] = (chiral[:, 0] == chiral[:, 1]).mean() if len(chiral) > 0 else -1
return results
|