首页
学习
活动
专区
圈层
工具
发布

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

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

脚本文件:

HilightChemAtom.py

代码语言:javascript
复制
from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D
from IPython.display import SVG
from io import BytesIO
from PIL import Image
from cairosvg import svg2png
import 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()

使用方法

代码语言:javascript
复制
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
下一篇
举报
领券