Package rdkit :: Package Chem :: Package Draw
[hide private]
[frames] | no frames]

Source Code for Package rdkit.Chem.Draw

  1  # 
  2  # Copyright (C) 2006-2016 Greg Landrum 
  3  #  All Rights Reserved 
  4  # 
  5  import os 
  6  import re 
  7   
  8  from rdkit.Chem.Draw import rdMolDraw2D 
  9  from rdkit.Chem.Draw.MolDrawing import MolDrawing, DrawingOptions 
 10  from rdkit.Chem.Draw.rdMolDraw2D import * 
 11  from rdkit.six import iteritems 
 12   
 13   
14 -def _getCanvas():
15 useAGG = False 16 useCairo = False 17 useSping = False 18 Canvas = None 19 if not os.environ.get('RDKIT_CANVAS', ''): 20 try: 21 from rdkit.Chem.Draw.cairoCanvas import Canvas 22 useCairo = True 23 except ImportError: 24 try: 25 from rdkit.Chem.Draw.aggCanvas import Canvas 26 useAGG = True 27 except ImportError: 28 from rdkit.Chem.Draw.spingCanvas import Canvas 29 useSping = True 30 else: 31 canv = os.environ['RDKIT_CANVAS'].lower() 32 if canv == 'cairo': 33 from rdkit.Chem.Draw.cairoCanvas import Canvas 34 useCairo = True 35 elif canv == 'agg': 36 from rdkit.Chem.Draw.aggCanvas import Canvas 37 useAGG = True 38 else: 39 from rdkit.Chem.Draw.spingCanvas import Canvas 40 useSping = True 41 if useSping: 42 # <- the sping canvas doesn't support unicode well 43 DrawingOptions.radicalSymbol = '.' 44 return useAGG, useCairo, Canvas
45 46
47 -def _createCanvas(size):
48 useAGG, useCairo, Canvas = _getCanvas() 49 if useAGG or useCairo: 50 try: 51 import Image 52 except ImportError: 53 from PIL import Image 54 img = Image.new("RGBA", size, (0, 0, 0, 0)) 55 canvas = Canvas(img) 56 else: 57 from rdkit.Chem.Draw.spingCanvas import Canvas 58 canvas = Canvas(size=size, name='MolToImageFile') 59 img = canvas._image 60 return img, canvas
61 62
63 -def MolToImage(mol, size=(300, 300), kekulize=True, wedgeBonds=True, fitImage=False, options=None, 64 canvas=None, **kwargs):
65 """Returns a PIL image containing a drawing of the molecule 66 67 ARGUMENTS: 68 69 - kekulize: run kekulization routine on input `mol` (default True) 70 71 - size: final image size, in pixel (default (300,300)) 72 73 - wedgeBonds: draw wedge (stereo) bonds (default True) 74 75 - highlightAtoms: list of atoms to highlight (default []) 76 77 - highlightMap: dictionary of (atom, color) pairs (default None) 78 79 - highlightBonds: list of bonds to highlight (default []) 80 81 - highlightColor: RGB color as tuple (default [1, 0, 0]) 82 83 NOTE: 84 85 use 'matplotlib.colors.to_rgb()' to convert string and 86 HTML color codes into the RGB tuple representation, eg. 87 88 from matplotlib.colors import ColorConverter 89 img = Draw.MolToImage(m, highlightAtoms=[1,2], highlightColor=ColorConverter().to_rgb('aqua')) 90 img.save("molecule.png") 91 92 RETURNS: 93 94 a PIL Image object 95 """ 96 97 if not mol: 98 raise ValueError('Null molecule provided') 99 if canvas is None: 100 img, canvas = _createCanvas(size) 101 else: 102 img = None 103 104 options = options or DrawingOptions() 105 if fitImage: 106 options.dotsPerAngstrom = int(min(size) / 10) 107 options.wedgeDashedBonds = wedgeBonds 108 if 'highlightColor' in kwargs: 109 color = kwargs.pop('highlightColor', (1, 0, 0)) 110 options.selectColor = color 111 112 drawer = MolDrawing(canvas=canvas, drawingOptions=options) 113 114 if kekulize: 115 from rdkit import Chem 116 mol = Chem.Mol(mol.ToBinary()) 117 Chem.Kekulize(mol) 118 119 if not mol.GetNumConformers(): 120 from rdkit.Chem import AllChem 121 AllChem.Compute2DCoords(mol) 122 123 if 'legend' in kwargs: 124 legend = kwargs['legend'] 125 del kwargs['legend'] 126 else: 127 legend = '' 128 129 drawer.AddMol(mol, **kwargs) 130 131 if legend: 132 from rdkit.Chem.Draw.MolDrawing import Font 133 bbox = drawer.boundingBoxes[mol] 134 pos = size[0] / 2, int(.94 * size[1]), 0 # the 0.94 is extremely empirical 135 # canvas.addCanvasPolygon(((bbox[0],bbox[1]),(bbox[2],bbox[1]),(bbox[2],bbox[3]),(bbox[0],bbox[3])), 136 # color=(1,0,0),fill=False,stroke=True) 137 # canvas.addCanvasPolygon(((0,0),(0,size[1]),(size[0],size[1]),(size[0],0) ), 138 # color=(0,0,1),fill=False,stroke=True) 139 font = Font(face='sans', size=12) 140 canvas.addCanvasText(legend, pos, font) 141 142 if kwargs.get('returnCanvas', False): 143 return img, canvas, drawer 144 else: 145 canvas.flush() 146 return img
147 148
149 -def MolToFile(mol, fileName, size=(300, 300), kekulize=True, wedgeBonds=True, imageType=None, 150 fitImage=False, options=None, **kwargs):
151 """ Generates a drawing of a molecule and writes it to a file 152 """ 153 # original contribution from Uwe Hoffmann 154 if not fileName: 155 raise ValueError('no fileName provided') 156 if not mol: 157 raise ValueError('Null molecule provided') 158 159 if imageType is None: 160 imageType = os.path.splitext(fileName)[1][1:] 161 162 if options is None: 163 options = DrawingOptions() 164 useAGG, useCairo, Canvas = _getCanvas() 165 if fitImage: 166 options.dotsPerAngstrom = int(min(size) / 10) 167 options.wedgeDashedBonds = wedgeBonds 168 if useCairo or useAGG: 169 canvas = Canvas(size=size, imageType=imageType, fileName=fileName) 170 else: 171 options.radicalSymbol = '.' # <- the sping canvas doesn't support unicode well 172 canvas = Canvas(size=size, name=fileName, imageType=imageType) 173 drawer = MolDrawing(canvas=canvas, drawingOptions=options) 174 if kekulize: 175 from rdkit import Chem 176 mol = Chem.Mol(mol.ToBinary()) 177 Chem.Kekulize(mol) 178 179 if not mol.GetNumConformers(): 180 from rdkit.Chem import AllChem 181 AllChem.Compute2DCoords(mol) 182 183 drawer.AddMol(mol, **kwargs) 184 if useCairo or useAGG: 185 canvas.flush() 186 else: 187 canvas.save()
188 189
190 -def MolToImageFile(mol, filename, size=(300, 300), kekulize=True, wedgeBonds=True, **kwargs):
191 """ DEPRECATED: please use MolToFile instead 192 193 """ 194 img = MolToImage(mol, size=size, kekulize=kekulize, wedgeBonds=wedgeBonds, **kwargs) 195 img.save(filename)
196 197 198 tkRoot = None 199 tkLabel = None 200 tkPI = None 201 202
203 -def ShowMol(mol, size=(300, 300), kekulize=True, wedgeBonds=True, title='RDKit Molecule', **kwargs):
204 """ Generates a picture of a molecule and displays it in a Tkinter window 205 """ 206 global tkRoot, tkLabel, tkPI 207 try: 208 import Tkinter 209 except ImportError: 210 import tkinter as Tkinter 211 try: 212 import ImageTk 213 except ImportError: 214 from PIL import ImageTk 215 216 img = MolToImage(mol, size, kekulize, wedgeBonds, **kwargs) 217 218 if not tkRoot: 219 tkRoot = Tkinter.Tk() 220 tkRoot.title(title) 221 tkPI = ImageTk.PhotoImage(img) 222 tkLabel = Tkinter.Label(tkRoot, image=tkPI) 223 tkLabel.place(x=0, y=0, width=img.size[0], height=img.size[1]) 224 else: 225 tkPI.paste(img) 226 tkRoot.geometry('%dx%d' % (img.size))
227 228
229 -def MolToMPL(mol, size=(300, 300), kekulize=True, wedgeBonds=True, imageType=None, fitImage=False, 230 options=None, **kwargs):
231 """ Generates a drawing of a molecule on a matplotlib canvas 232 """ 233 if not mol: 234 raise ValueError('Null molecule provided') 235 from rdkit.Chem.Draw.mplCanvas import Canvas 236 canvas = Canvas(size) 237 if options is None: 238 options = DrawingOptions() 239 options.bgColor = None 240 if fitImage: 241 options.dotsPerAngstrom = int(min(size) / 10) 242 options.wedgeDashedBonds = wedgeBonds 243 drawer = MolDrawing(canvas=canvas, drawingOptions=options) 244 omol = mol 245 if kekulize: 246 from rdkit import Chem 247 mol = Chem.Mol(mol.ToBinary()) 248 Chem.Kekulize(mol) 249 250 if not mol.GetNumConformers(): 251 from rdkit.Chem import AllChem 252 AllChem.Compute2DCoords(mol) 253 254 drawer.AddMol(mol, **kwargs) 255 omol._atomPs = drawer.atomPs[mol] 256 for k, v in iteritems(omol._atomPs): 257 omol._atomPs[k] = canvas.rescalePt(v) 258 canvas._figure.set_size_inches(float(size[0]) / 100, float(size[1]) / 100) 259 return canvas._figure
260 261
262 -def calcAtomGaussians(mol, a=0.03, step=0.02, weights=None):
263 """ 264 useful things to do with these: 265 fig.axes[0].imshow(z,cmap=cm.gray,interpolation='bilinear',origin='lower',extent=(0,1,0,1)) 266 fig.axes[0].contour(x,y,z,20,colors='k') 267 268 fig=Draw.MolToMPL(m); 269 contribs=Crippen.rdMolDescriptors._CalcCrippenContribs(m) 270 logps,mrs=zip(*contribs) 271 x,y,z=Draw.calcAtomGaussians(m,0.03,step=0.01,weights=logps) 272 fig.axes[0].imshow(z,cmap=cm.jet,interpolation='bilinear',origin='lower',extent=(0,1,0,1)) 273 fig.axes[0].contour(x,y,z,20,colors='k',alpha=0.5) 274 fig.savefig('coumlogps.colored.png',bbox_inches='tight') 275 276 277 """ 278 import numpy 279 from matplotlib import mlab 280 x = numpy.arange(0, 1, step) 281 y = numpy.arange(0, 1, step) 282 X, Y = numpy.meshgrid(x, y) 283 if weights is None: 284 weights = [1.] * mol.GetNumAtoms() 285 Z = mlab.bivariate_normal(X, Y, a, a, mol._atomPs[0][0], mol._atomPs[0][1]) * weights[0] 286 for i in range(1, mol.GetNumAtoms()): 287 Zp = mlab.bivariate_normal(X, Y, a, a, mol._atomPs[i][0], mol._atomPs[i][1]) 288 Z += Zp * weights[i] 289 return X, Y, Z
290 291
292 -def MolsToImage(mols, subImgSize=(200, 200), legends=None, **kwargs):
293 """ 294 """ 295 try: 296 import Image 297 except ImportError: 298 from PIL import Image 299 if legends is None: 300 legends = [None] * len(mols) 301 res = Image.new("RGBA", (subImgSize[0] * len(mols), subImgSize[1])) 302 for i, mol in enumerate(mols): 303 res.paste(MolToImage(mol, subImgSize, legend=legends[i], **kwargs), (i * subImgSize[0], 0)) 304 return res
305 306 from io import BytesIO
307 -def _drawerToImage(d2d):
308 try: 309 import Image 310 except ImportError: 311 from PIL import Image 312 sio = BytesIO(d2d.GetDrawingText()) 313 return Image.open(sio)
314
315 -def _okToKekulizeMol(mol,kekulize):
316 if kekulize: 317 for bond in mol.GetBonds(): 318 if bond.GetIsAromatic() and bond.HasQuery(): 319 return False 320 return True 321 return kekulize
322
323 -def _moltoimg(mol, sz, highlights, legend, returnPNG=False, **kwargs):
324 try: 325 mol.GetAtomWithIdx(0).GetExplicitValence() 326 except RuntimeError: 327 mol.UpdatePropertyCache(False) 328 329 kekulize=_okToKekulizeMol(mol,kwargs.get('kekulize', True)) 330 331 try: 332 mc = rdMolDraw2D.PrepareMolForDrawing(mol,kekulize=kekulize) 333 except ValueError: # <- can happen on a kekulization failure 334 mc = rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=False) 335 if not hasattr(rdMolDraw2D, 'MolDraw2DCairo'): 336 img = MolToImage(mc, sz, legend=legend, highlightAtoms=highlights, **kwargs) 337 if returnPNG: 338 bio = BytesIO() 339 img.save(bio, format='PNG') 340 img = bio.getvalue() 341 else: 342 d2d = rdMolDraw2D.MolDraw2DCairo(sz[0], sz[1]) 343 d2d.DrawMolecule(mc, legend=legend, highlightAtoms=highlights) 344 d2d.FinishDrawing() 345 if returnPNG: 346 img = d2d.GetDrawingText() 347 else: 348 img = _drawerToImage(d2d) 349 return img
350
351 -def _moltoSVG(mol, sz, highlights, legend, kekulize, **kwargs):
352 try: 353 mol.GetAtomWithIdx(0).GetExplicitValence() 354 except RuntimeError: 355 mol.UpdatePropertyCache(False) 356 357 kekulize=_okToKekulizeMol(mol,kekulize) 358 359 try: 360 mc = rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=kekulize) 361 except ValueError: # <- can happen on a kekulization failure 362 mc = rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=False) 363 d2d = rdMolDraw2D.MolDraw2DSVG(sz[0], sz[1]) 364 d2d.DrawMolecule(mc, legend=legend, highlightAtoms=highlights) 365 d2d.FinishDrawing() 366 svg = d2d.GetDrawingText() 367 return svg.replace("svg:", "")
368 369 370
371 -def _MolsToGridImage(mols, molsPerRow=3, subImgSize=(200, 200), legends=None, 372 highlightAtomLists=None, highlightBondLists=None, **kwargs):
373 """ returns a PIL Image of the grid 374 """ 375 if legends is None: 376 legends = [''] * len(mols) 377 378 nRows = len(mols) // molsPerRow 379 if len(mols) % molsPerRow: 380 nRows += 1 381 382 if not hasattr(rdMolDraw2D, 'MolDraw2DCairo'): 383 try: 384 import Image 385 except ImportError: 386 from PIL import Image 387 res = Image.new("RGBA", (molsPerRow * subImgSize[0], nRows * subImgSize[1]), (255, 255, 255, 0)) 388 for i, mol in enumerate(mols): 389 row = i // molsPerRow 390 col = i % molsPerRow 391 highlights = None 392 if highlightAtomLists and highlightAtomLists[i]: 393 highlights = highlightAtomLists[i] 394 if mol is not None: 395 img = _moltoimg(mol, subImgSize, highlights, legends[i], **kwargs) 396 res.paste(img, (col * subImgSize[0], row * subImgSize[1])) 397 else: 398 fullSize = (molsPerRow * subImgSize[0], nRows * subImgSize[1]) 399 d2d = rdMolDraw2D.MolDraw2DCairo(fullSize[0],fullSize[1],subImgSize[0], subImgSize[1]) 400 d2d.DrawMolecules(list(mols),legends=legends,highlightAtoms=highlightAtomLists, 401 highlightBonds=highlightBondLists,**kwargs) 402 d2d.FinishDrawing() 403 res = _drawerToImage(d2d) 404 405 return res
406 407
408 -def _MolsToGridSVG(mols, molsPerRow=3, subImgSize=(200, 200), legends=None, 409 highlightAtomLists=None, highlightBondLists=None, 410 stripSVGNamespace=True, **kwargs):
411 """ returns an SVG of the grid 412 """ 413 if legends is None: 414 legends = [''] * len(mols) 415 416 nRows = len(mols) // molsPerRow 417 if len(mols) % molsPerRow: 418 nRows += 1 419 420 blocks = [''] * (nRows * molsPerRow) 421 422 fullSize = (molsPerRow * subImgSize[0], nRows * subImgSize[1]) 423 424 d2d = rdMolDraw2D.MolDraw2DSVG(fullSize[0],fullSize[1],subImgSize[0], subImgSize[1]) 425 d2d.DrawMolecules(mols,legends=legends,highlightAtoms=highlightAtomLists, 426 highlightBonds=highlightBondLists,**kwargs) 427 d2d.FinishDrawing() 428 res = d2d.GetDrawingText() 429 if stripSVGNamespace: 430 res = res.replace('svg:', '') 431 return res
432 433
434 -def MolsToGridImage(mols, molsPerRow=3, subImgSize=(200, 200), legends=None, 435 highlightAtomLists=None, highlightBondLists=None, 436 useSVG=False, **kwargs):
437 if legends and len(legends)>len(mols): 438 legends = legends[:len(mols)] 439 if highlightAtomLists and len(highlightAtomLists)>len(mols): 440 highlightAtomLists = highlightAtomLists[:len(mols)] 441 if highlightBondLists and len(highlightBondLists)>len(mols): 442 highlightBondLists = highlightBondLists[:len(mols)] 443 444 if useSVG: 445 return _MolsToGridSVG(mols, molsPerRow=molsPerRow, subImgSize=subImgSize, legends=legends, 446 highlightAtomLists=highlightAtomLists, 447 highlightBondLists=highlightBondLists, 448 **kwargs) 449 else: 450 return _MolsToGridImage(mols, molsPerRow=molsPerRow, subImgSize=subImgSize, legends=legends, 451 highlightAtomLists=highlightAtomLists, 452 highlightBondLists=highlightBondLists, 453 **kwargs)
454 455
456 -def _legacyReactionToImage(rxn, subImgSize=(200, 200), **kwargs):
457 try: 458 import Image 459 except ImportError: 460 from PIL import Image 461 462 mols = [] 463 for i in range(rxn.GetNumReactantTemplates()): 464 tmpl = rxn.GetReactantTemplate(i) 465 tmpl.UpdatePropertyCache(False) 466 mols.append(tmpl) 467 mols.append(None) 468 for i in range(rxn.GetNumProductTemplates()): 469 tmpl = rxn.GetProductTemplate(i) 470 tmpl.UpdatePropertyCache(False) 471 mols.append(tmpl) 472 473 res = Image.new("RGBA", (subImgSize[0] * len(mols), subImgSize[1]), (255, 255, 255, 0)) 474 for i, mol in enumerate(mols): 475 if mol is not None: 476 nimg = MolToImage(mol, subImgSize, kekulize=False, **kwargs) 477 else: 478 nimg, canvas = _createCanvas(subImgSize) 479 p0 = (10, subImgSize[1] // 2) 480 p1 = (subImgSize[0] - 10, subImgSize[1] // 2) 481 p3 = (subImgSize[0] - 20, subImgSize[1] // 2 - 10) 482 p4 = (subImgSize[0] - 20, subImgSize[1] // 2 + 10) 483 canvas.addCanvasLine(p0, p1, lineWidth=2, color=(0, 0, 0)) 484 canvas.addCanvasLine(p3, p1, lineWidth=2, color=(0, 0, 0)) 485 canvas.addCanvasLine(p4, p1, lineWidth=2, color=(0, 0, 0)) 486 if hasattr(canvas, 'flush'): 487 canvas.flush() 488 else: 489 canvas.save() 490 res.paste(nimg, (i * subImgSize[0], 0)) 491 return res
492 493
494 -def ReactionToImage(rxn, subImgSize=(200, 200), useSVG=False, **kwargs):
495 if not useSVG and not hasattr(rdMolDraw2D, 'MolDraw2DCairo'): 496 return _legacyReactionToImage(rxn,subImgSize=subImgSize,**kwargs) 497 else: 498 width = subImgSize[0] * (rxn.GetNumReactantTemplates() + rxn.GetNumProductTemplates() + 1) 499 if useSVG: 500 d = rdMolDraw2D.MolDraw2DSVG(width, subImgSize[1]) 501 else: 502 d = rdMolDraw2D.MolDraw2DCairo(width, subImgSize[1]) 503 d.DrawReaction(rxn,**kwargs) 504 d.FinishDrawing() 505 if useSVG: 506 return d.GetDrawingText().replace('svg:', '') 507 else: 508 return _drawerToImage(d)
509
510 -def MolToQPixmap(mol, size=(300, 300), kekulize=True, wedgeBonds=True, fitImage=False, options=None, 511 **kwargs):
512 """ Generates a drawing of a molecule on a Qt QPixmap 513 """ 514 if not mol: 515 raise ValueError('Null molecule provided') 516 from rdkit.Chem.Draw.qtCanvas import Canvas 517 canvas = Canvas(size) 518 if options is None: 519 options = DrawingOptions() 520 options.bgColor = None 521 if fitImage: 522 options.dotsPerAngstrom = int(min(size) / 10) 523 options.wedgeDashedBonds = wedgeBonds 524 if kekulize: 525 from rdkit import Chem 526 mol = Chem.Mol(mol.ToBinary()) 527 Chem.Kekulize(mol) 528 if not mol.GetNumConformers(): 529 from rdkit.Chem import AllChem 530 AllChem.Compute2DCoords(mol) 531 drawer = MolDrawing(canvas=canvas, drawingOptions=options) 532 drawer.AddMol(mol, **kwargs) 533 canvas.flush() 534 return canvas.pixmap
535