Package rdkit :: Package Chem :: Package Draw

[hide private]
[frames] | no frames]

# Source Code for Package rdkit.Chem.Draw

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