|
|
from rdkit import Chem
|
|
|
from rdkit.Chem import Draw
|
|
|
from rdkit.Chem import AllChem, MACCSkeys, rdMolDescriptors as rdDesc
|
|
|
from collections import defaultdict
|
|
|
import numpy as np
|
|
|
import os, pickle, hashlib
|
|
|
|
|
|
AllChem.SetPreferCoordGen(True)
|
|
|
|
|
|
FINGERPRINT_DICT = defaultdict(lambda : len(FINGERPRINT_DICT))
|
|
|
|
|
|
ELEMENTS = ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al',
|
|
|
'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn',
|
|
|
'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb',
|
|
|
'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In',
|
|
|
'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm',
|
|
|
'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta',
|
|
|
'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At',
|
|
|
'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk',
|
|
|
'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt',
|
|
|
'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og']
|
|
|
|
|
|
for e in ELEMENTS:
|
|
|
FINGERPRINT_DICT[e]
|
|
|
|
|
|
if os.path.exists('rdkit_fingerprint_list_r1.pkl'):
|
|
|
l = pickle.load(open('rdkit_fingerprint_list_r1.pkl', 'rb'))
|
|
|
|
|
|
for smi in l:
|
|
|
FINGERPRINT_DICT[smi]
|
|
|
|
|
|
print('Len fingerprint_list: %s' %len(FINGERPRINT_DICT)) + len(ELEMENTS)
|
|
|
|
|
|
def mol_with_atom_index(mol):
|
|
|
atoms = mol.GetNumAtoms()
|
|
|
for idx in range(atoms):
|
|
|
mol.GetAtomWithIdx(idx).SetProp('molAtomMapNumber', str(mol.GetAtomWithIdx(idx).GetIdx()))
|
|
|
return mol
|
|
|
|
|
|
def prepare_mol_for_drawing(mol):
|
|
|
try:
|
|
|
mol_draw = Draw.rdMolDraw2D.PrepareMolForDrawing(mol)
|
|
|
except Chem.KekulizeException:
|
|
|
mol_draw = Draw.rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=False)
|
|
|
Chem.SanitizeMol(mol_draw, Chem.SANITIZE_ALL ^ Chem.SANITIZE_KEKULIZE)
|
|
|
return mol_draw
|
|
|
|
|
|
def get_atom_submol_radn(mol, radius, sanitize=True):
|
|
|
atoms = []
|
|
|
submols = []
|
|
|
|
|
|
for atom in mol.GetAtoms():
|
|
|
atoms.append(atom)
|
|
|
r = radius
|
|
|
while r > 0:
|
|
|
try:
|
|
|
env = Chem.FindAtomEnvironmentOfRadiusN(mol, r, atom.GetIdx())
|
|
|
amap={}
|
|
|
submol = Chem.PathToSubmol(mol, env, atomMap=amap)
|
|
|
if sanitize:
|
|
|
Chem.SanitizeMol(submol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_ALL^Chem.SanitizeFlags.SANITIZE_KEKULIZE)
|
|
|
|
|
|
submols.append(submol)
|
|
|
break
|
|
|
except Exception as e:
|
|
|
print(64, e)
|
|
|
r -= 1
|
|
|
|
|
|
return atoms, submols
|
|
|
|
|
|
def gen_fps_from_mol(mol, nbits=256, use_morgan=True, use_macc=False, use_rdkit=False):
|
|
|
|
|
|
fp = []
|
|
|
if use_morgan:
|
|
|
fp_vec = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=nbits)
|
|
|
fp1 = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0')
|
|
|
fp = fp1.tolist()
|
|
|
if use_macc:
|
|
|
|
|
|
fp_vec = MACCSkeys.GenMACCSKeys(mol)
|
|
|
fp1 = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0')
|
|
|
fp.extend(fp1.tolist())
|
|
|
if use_rdkit:
|
|
|
fp_vec = Chem.RDKFingerprint(mol)
|
|
|
fp1 = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0')
|
|
|
fp.extend(fp1.tolist())
|
|
|
|
|
|
return fp
|
|
|
|
|
|
def gen_subgraph_fps_from_str(s, wordsdict={}):
|
|
|
if s in wordsdict:
|
|
|
return [wordsdict[s]]
|
|
|
else:
|
|
|
return [len(wordsdict)]
|
|
|
|
|
|
def gen_subgraph_fps_from_mol(mol, wordsdict={}):
|
|
|
try:
|
|
|
k = Chem.MolToSmiles(mol)
|
|
|
return gen_subgraph_fps_from_str(k, wordsdict)
|
|
|
except Exception as e:
|
|
|
print(e)
|
|
|
return [len(wordsdict)]
|
|
|
|
|
|
def calc_subgraph_fps_from_mol(mol, radius=2, nbits=128, use_macc=True, fptype=1, wordsdict={}):
|
|
|
|
|
|
atoms, submols = get_atom_submol_radn(mol, radius, True)
|
|
|
feats = []
|
|
|
for idx, submol in enumerate(submols):
|
|
|
if fptype == 1:
|
|
|
feat = gen_fps_from_mol(submol, nbits, use_macc)
|
|
|
feats.append(feat)
|
|
|
elif fptype == 2:
|
|
|
feat = gen_subgraph_fps_from_mol(submol, wordsdict)
|
|
|
feats.append(feat)
|
|
|
|
|
|
return np.array(feats)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
smi = 'C=C(S)C(N)(O)C'
|
|
|
smi = 'CC1CCN(CC1N(C)C2=NC=NC3=C2C=CN3)C(=O)CC#N'
|
|
|
|
|
|
mol = Chem.MolFromSmiles(smi, sanitize=False)
|
|
|
|
|
|
print(calc_subgraph_fps_from_mol(mol, 3))
|
|
|
|
|
|
mol = mol_with_atom_index(mol)
|
|
|
submols = get_atom_submol_radn(mol, 3)
|
|
|
submols = [prepare_mol_for_drawing(m) for m in submols]
|
|
|
hl = []
|
|
|
for idx, m in enumerate(submols):
|
|
|
for a in m.GetAtoms():
|
|
|
if int(a.GetProp('molAtomMapNumber')) == idx:
|
|
|
hl.append([a.GetIdx()])
|
|
|
break
|
|
|
|
|
|
draw = Draw.MolsToGridImage([mol] + submols, highlightAtomLists=[[]] + hl, molsPerRow=5)
|
|
|
draw.show()
|
|
|
|
|
|
|
