rdkit.Chem.Draw

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