【原】RDKit | 基于RDKit的指定原子或键高亮

指定原子或键高亮在某些场景下有需求，此处指定原子或键高亮方法供参考。

脚本文件：

HilightChemAtom.py

from rdkit import Chemfrom rdkit.Chem.Draw import rdMolDraw2Dfrom IPython.display import SVGfrom io import BytesIOfrom PIL import Imagefrom cairosvg import svg2pngimport argparse  def generate_image(mol, highlight_atoms, highlight_bonds, atomColors, bondColors, radii, size, output, isNumber=False):     print(highlight_atoms)    print(highlight_bonds)    print(atomColors)    print(bondColors)     image_data = BytesIO()    view = rdMolDraw2D.MolDraw2DSVG(size[0], size[1])    tm = rdMolDraw2D.PrepareMolForDrawing(mol)     option = view.drawOptions()    if isNumber:        for atom in mol.GetAtoms():            option.atomLabels[atom.GetIdx()] = atom.GetSymbol() + str(atom.GetIdx() + 1)     view.DrawMolecule(tm, highlightAtoms=highlight_atoms, highlightBonds=highlight_bonds,                      highlightAtomColors=atomColors, highlightBondColors=bondColors, highlightAtomRadii=radii)    view.FinishDrawing()    svg = view.GetDrawingText()    SVG(svg.replace('svg:', ''))    svg2png(bytestring=svg, write_to=output)    img = Image.open(output)    img.save(image_data, format='PNG')     return image_data  def main():     parser = argparse.ArgumentParser()    parser.add_argument("-i", type=str, required=True)    parser.add_argument("-o", type=str, required=True)    parser.add_argument("-a", type=str, action="append")    parser.add_argument("-b", type=str, action="append")    args = parser.parse_args()     atom_hilights = {}    bond_hilights = {}    radii = {}    with open(args.i, "r") as f:        string_mol = f.read()        mol = Chem.MolFromMolBlock(string_mol)         if args.a:            for atom in args.a:                print(atom)                datas = atom.split(",")                atom_hilights[int(datas[0])] = (int(datas[1]), int(datas[2]), int(datas[3]))                radii[int(datas[0])] = 0.20        if args.b:            for bond in args.b:                datas = bond.split(",")                bond_hilights[int(datas[0])] = (int(datas[1]), int(datas[2]), int(datas[3]))         generate_image(mol,list(atom_hilights.keys()), list(bond_hilights.keys()), atom_hilights, bond_hilights, radii, (400, 200), args.o, False) if __name__ == "__main__":    main()

使用方法

python HilightChemAtom.py  -i test.mol -o test.png [-a A] [-b B]usage: HilightChemExample.py [-h] -i I -o O [-a A] [-b B]

例

python HilightChemAtom.py -i test.mol -o test.png -a 0,1,0,1

python HilightChemAtom.py -i test.mol -o test1.png -a 0,1,0,1 -a 8,1,0,0 -b 9,0,0,1

python HilightChemAtom.py -i test.mol -o test2.png -a 0,1,0,1 -a 8,1,0,0 -b 0,0,0,1

DrugAI