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