指定原子或键高亮在某些场景下有需求,此处指定原子或键高亮方法供参考。
脚本文件:
HilightChemAtom.py
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()
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