Chem.Draw.MolDrawing

Package Chem :: Package Draw :: Module MolDrawing |
[hide private]
[frames] | no frames]

# Source Code for Module Chem.Draw.MolDrawing

      1  # $Id: MolDrawing.py 680 2008-05-19 04:35:52Z glandrum $
      2  #
      3  #  Copyright (C) 2008 Greg Landrum
      4  #
      5  #   @@ All Rights Reserved  @@
      6  #
      7  import Chem
      8  import RDConfig
      9  import math
     10
     11  elemDict={
     12    7:(0,0,1),
     13    8:(1,0,0),
     14    9:(.2,.8,.8),
     15    15:(1,.5,0),
     16    16:(.8,.8,0),
     17    17:(0,.8,0),
     18    35:(.5,.3,.1),
     19    0:(.5,.5,.5),
     20    }
     21
     22 -class Font(object):
     23    face='sans'
     24    size='12'
     25    weight='normal'
     26    name=None
     27 -  def __init__(self,face=None,size=None,name=None,weight=None):
     28      if face: self.face=face
     29      if size: self.size=size
     30      if name: self.name=name
     31      if weight: self.weight=weight
     32
     33 -class MolDrawing(object):
     34    dotsPerAngstrom = 30
     35
     36    atomLabelFontFace = "sans"
     37    atomLabelFontSize = 12
     38    atomLabelMinFontSize = 10
     39
     40    bondLineWidth = 1.2
     41    dblBondOffset = .2
     42    dblBondLengthFrac = .8
     43
     44    defaultColor = (1,0,0)
     45    selectColor = (1,0,0)
     46
     47    colorBonds=True
     48    noCarbonSymbols=True
     49    includeAtomNumbers=False
     50    atomNumberOffset=0
     51
     52    dash = (1,1)
     53    atomPs = None
     54    canvas = None
     55    canvasSize=None
     56
     57    wedgeDashedBonds=False
     58
     59 -  def __init__(self,canvas=None):
     60      self.canvas = canvas
     61      if canvas:
     62        self.canvasSize=canvas.size
     63      self.atomPs = {}
     64
     65 -  def transformPoint(self,pos):
     66      res = [0,0]
     67      res[0] = (pos[0] + self.molTrans[0])*self.dotsPerAngstrom + self.drawingTrans[0]
     68      res[1] = self.canvasSize[1]-((pos[1] + self.molTrans[1])*self.dotsPerAngstrom + \
     69                                   self.drawingTrans[1])
     70      return res
     71
     72
     73 -  def _getBondOffset(self,p1,p2):
     74      # get the vector between the points:
     75      dx = p2[0]-p1[0]
     76      dy = p2[1]-p1[1]
     77
     78      # figure out the angle and the perpendicular:
     79      ang = math.atan2(dy,dx)
     80      perp = ang + math.pi/2.
     81
     82      # here's the offset for the parallel bond:
     83      offsetX = math.cos(perp)*self.dblBondOffset*self.dotsPerAngstrom
     84      offsetY = math.sin(perp)*self.dblBondOffset*self.dotsPerAngstrom
     85
     86      return perp,offsetX,offsetY
     87
     88 -  def _getOffsetBondPts(self,p1,p2,
     89                          offsetX,offsetY,
     90                          lenFrac=None):
     91      if not lenFrac:
     92        lenFrac = self.dblBondLengthFrac
     93
     94      dx = p2[0]-p1[0]
     95      dy = p2[1]-p1[1]
     96      # ----
     97      # now figure out where to start and end it:
     98
     99      # offset the start point:
    100      fracP1 = p1[0]+offsetX,p1[1]+offsetY
    101
    102      # now move a portion of the way along the line to the neighbor:
    103      frac = (1.-lenFrac)/2
    104      fracP1 = fracP1[0]+dx*frac,\
    105               fracP1[1]+dy*frac
    106
    107      fracP2 = fracP1[0]+dx*lenFrac,\
    108               fracP1[1]+dy*lenFrac
    109      return fracP1,fracP2
    110
    111 -  def _offsetDblBond(self,p1,p2,bond,a1,a2,conf,dir=1,
    112                       lenFrac=None):
    113      perp,offsetX,offsetY = self._getBondOffset(p1,p2)
    114      offsetX = offsetX*dir
    115      offsetY = offsetY*dir
    116
    117      # if we're a ring bond, we may need to flip over to the other side:
    118      if bond.IsInRing():
    119        bondIdx = bond.GetIdx()
    120        a1Idx = a1.GetIdx()
    121        a2Idx = a2.GetIdx()
    122        # find a ring bond from a1 to an atom other than a2:
    123        for otherBond in a1.GetBonds():
    124          if otherBond.GetIdx()!=bondIdx and \
    125             otherBond.IsInRing():
    126            sharedRing=False
    127            for ring in self.bondRings:
    128              if bondIdx in ring and otherBond.GetIdx() in ring:
    129                sharedRing=True
    130                break
    131            if not sharedRing:
    132              continue
    133            a3 = otherBond.GetOtherAtom(a1)
    134            if a3.GetIdx() != a2Idx:
    135              p3 = self.transformPoint(conf.GetAtomPosition(a3.GetIdx()))
    136              dx2 = p3[0] - p1[0]
    137              dy2 = p3[1] - p1[1]
    138              dotP = dx2*offsetX + dy2*offsetY
    139              if dotP < 0:
    140                perp += math.pi
    141                offsetX = math.cos(perp)*self.dblBondOffset*self.dotsPerAngstrom
    142                offsetY = math.sin(perp)*self.dblBondOffset*self.dotsPerAngstrom
    143
    144      fracP1,fracP2 = self._getOffsetBondPts(p1,p2,
    145                                             offsetX,offsetY,
    146                                             lenFrac=lenFrac)
    147      return fracP1,fracP2
    148
    149 -  def _drawWedgedBond(self,canvas,bond,pos,nbrPos,
    150                        width=bondLineWidth,color=defaultColor,
    151                        dash=None):
    152      perp,offsetX,offsetY = self._getBondOffset(pos,nbrPos)
    153      offsetX *=.75
    154      offsetY *=.75
    155      poly = ((pos[0],pos[1]),
    156              (nbrPos[0]+offsetX,nbrPos[1]+offsetY),
    157              (nbrPos[0]-offsetX,nbrPos[1]-offsetY))
    158      #canvas.drawPolygon(poly,edgeColor=color,edgeWidth=1,fillColor=color,closed=1)
    159      if not dash:
    160        addCanvasPolygon(canvas,poly,color=color)
    161      elif self.wedgeDashedBonds and addCanvasDashedWedge:
    162        addCanvasDashedWedge(canvas,poly[0],poly[1],poly[2],color=color)
    163      else:
    164        addCanvasLine(canvas,pos,nbrPos,linewidth=width*2,color=color,
    165                      dashes=dash)
    166
    167 -  def _drawBond(self,canvas,bond,atom,nbr,pos,nbrPos,conf,
    168                  width=bondLineWidth,color=defaultColor,color2=None):
    169      if bond.GetBondType() == Chem.BondType.SINGLE:
    170        bDir = bond.GetBondDir()
    171        if bDir in (Chem.BondDir.BEGINWEDGE,Chem.BondDir.BEGINDASH):
    172          # if the bond is "backwards", change the drawing direction:
    173          if bond.GetBeginAtom().GetChiralTag() in (Chem.ChiralType.CHI_TETRAHEDRAL_CW,
    174                                                    Chem.ChiralType.CHI_TETRAHEDRAL_CCW):
    175            p1,p2 = pos,nbrPos
    176          else:
    177            p2,p1 = pos,nbrPos
    178          if bDir==Chem.BondDir.BEGINWEDGE:
    179            self._drawWedgedBond(canvas,bond,p1,p2,color=(0,0,0),width=width)
    180          elif bDir==Chem.BondDir.BEGINDASH:
    181            self._drawWedgedBond(canvas,bond,p1,p2,color=(0,0,0),width=width,
    182                                 dash=self.dash)
    183        else:
    184          addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    185      elif bond.GetBondType() == Chem.BondType.DOUBLE:
    186        if bond.IsInRing() or (atom.GetDegree()!=1 and bond.GetOtherAtom(atom).GetDegree()!=1):
    187          addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    188          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf)
    189          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    190        else:
    191          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf,dir=.5,
    192                                        lenFrac=1.0)
    193          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    194          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf,dir=-.5,
    195                                        lenFrac=1.0)
    196          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    197      elif bond.GetBondType() == Chem.BondType.AROMATIC:
    198        addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    199        fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf)
    200        addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2,
    201                      dash=self.dash)
    202        #DASH
    203      elif bond.GetBondType() == Chem.BondType.TRIPLE:
    204        addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    205        fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf)
    206        addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    207        fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf,dir=-1)
    208        addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    209
    210 -  def _scaleAndCenter(self,mol,conf,coordCenter=False):
    211      canvasSize = self.canvasSize
    212      xAccum = 0
    213      yAccum = 0
    214      nAts = 0
    215      minX = 1e8
    216      minY = 1e8
    217      maxX = -1e8
    218      maxY = -1e8
    219
    220      for atom in mol.GetAtoms():
    221        pos = conf.GetAtomPosition(atom.GetIdx())
    222        xAccum += pos[0]
    223        yAccum += pos[1]
    224        minX = min(minX,pos[0])
    225        minY = min(minY,pos[1])
    226        maxX = max(maxX,pos[0])
    227        maxY = max(maxY,pos[1])
    228        nAts += 1
    229
    230      dx = abs(maxX-minX)
    231      dy = abs(maxY-minY)
    232      xSize = dx*self.dotsPerAngstrom
    233      ySize = dy*self.dotsPerAngstrom
    234
    235      if coordCenter:
    236        molTrans = -xAccum/nAts,-yAccum/nAts
    237      else:
    238        molTrans = -(minX+(maxX-minX)/2),-(minY+(maxY-minY)/2)
    239      self.dotsPerAngstrom=30.0
    240      self.molTrans = molTrans
    241      self.drawingTrans=(0,0)
    242      tMin = self.transformPoint((minX,minY))
    243      tMax = self.transformPoint((maxX,maxY))
    244
    245      if xSize>=.95*canvasSize[0]:
    246        scale = .9*canvasSize[0]/xSize
    247        xSize*=scale
    248        ySize*=scale
    249        self.dotsPerAngstrom*=scale
    250        self.atomLabelFontSize = max(self.atomLabelFontSize*scale,
    251                                     self.atomLabelMinFontSize)
    252      if ySize>=.95*canvasSize[1]:
    253        scale = .9*canvasSize[1]/ySize
    254        xSize*=scale
    255        ySize*=scale
    256        self.dotsPerAngstrom*=scale
    257        self.atomLabelFontSize = max(self.atomLabelFontSize*scale,
    258                                     self.atomLabelMinFontSize)
    259      drawingTrans = canvasSize[0]/2,canvasSize[1]/2
    260      self.drawingTrans = drawingTrans
    261      tMax = self.transformPoint((maxX,maxY))
    262      tMin = self.transformPoint((minX,minY))
    263
    264 -  def _drawLabel(self,canvas,label,pos,font,color=None,
    265                   highlightIt=False):
    266      if highlightIt:
    267        color = self.selectColor
    268      elif not color:
    269        color = self.defaultColor
    270      x1 = pos[0]
    271      y1 = pos[1]
    272      labelP = x1,y1
    273      addCanvasText(canvas,label,(x1,y1),font,color)
    274
    275 -  def AddMol(self,mol,canvas=None,centerIt=True,molTrans=(0,0),drawingTrans=(0,0),
    276               highlightAtoms=[],confId=-1):
    277      """
    278
    279      Notes:
    280        - specifying centerIt will cause molTrans and drawingTrans to be ignored
    281
    282      """
    283      try:
    284        dl = addCanvasLine
    285      except NameError:
    286        registerCanvas('sping')
    287      if canvas is None:
    288        canvas = self.canvas
    289      else:
    290        self.canvas = canvas
    291      self.canvasSize=canvas.size
    292
    293      conf = mol.GetConformer(confId)
    294
    295      if centerIt:
    296        self._scaleAndCenter(mol,conf)
    297      else:
    298        self.molTrans = molTrans
    299        self.drawingTrans = drawingTrans
    300      font = Font(face=self.atomLabelFontFace,size=self.atomLabelFontSize)
    301
    302      if not mol.HasProp('_drawingBondsWedged'):
    303        Chem.WedgeMolBonds(mol,conf)
    304        mol.SetProp('_drawingBondsWedged','set',True)
    305
    306      self.atomPs[mol] = {}
    307      self.activeMol = mol
    308      self.bondRings = mol.GetRingInfo().BondRings()
    309      for atom in mol.GetAtoms():
    310        idx = atom.GetIdx()
    311        pos = self.atomPs[mol].get(idx,None)
    312        if pos is None:
    313          pos = self.transformPoint(conf.GetAtomPosition(idx))
    314          self.atomPs[mol][idx] = pos
    315        nbrSum = [0,0]
    316        for bond in atom.GetBonds():
    317          nbr = bond.GetOtherAtom(atom)
    318          nbrIdx = nbr.GetIdx()
    319          if nbrIdx > idx:
    320            nbrPos = self.atomPs[mol].get(nbrIdx,None)
    321            if nbrPos is None:
    322              nbrPos = self.transformPoint(conf.GetAtomPosition(nbrIdx))
    323              self.atomPs[mol][nbrIdx] = nbrPos
    324
    325            if highlightAtoms and idx in highlightAtoms and nbrIdx in highlightAtoms:
    326              width=2.0*self.bondLineWidth
    327              color = self.selectColor
    328              color2 = self.selectColor
    329            else:
    330              width=self.bondLineWidth
    331              if self.colorBonds:
    332                color = elemDict.get(atom.GetAtomicNum(),(0,0,0))
    333                color2 = elemDict.get(nbr.GetAtomicNum(),(0,0,0))
    334              else:
    335                color = self.defaultColor
    336                color2= color
    337
    338            # make sure we draw from the beginning to the end
    339            # (this was Issue400)
    340            if atom.GetIdx()==bond.GetBeginAtomIdx():
    341              self._drawBond(canvas,bond,atom,nbr,pos,nbrPos,conf,
    342                             color=color,width=width,color2=color2)
    343            else:
    344              self._drawBond(canvas,bond,nbr,atom,nbrPos,pos,conf,
    345                             color=color2,width=width,color2=color)
    346          else:
    347            nbrPos = self.atomPs[mol][nbrIdx]
    348          nbrSum[0] += nbrPos[0]-pos[0]
    349          nbrSum[1] += nbrPos[1]-pos[1]
    350
    351
    352        base = atom.GetSymbol()
    353        nHs = atom.GetTotalNumHs()
    354        if nHs>0:
    355          if nHs>1:
    356            hs='H%d'%nHs
    357          else:
    358            hs ='H'
    359        else:
    360          hs = ''
    361        chg = atom.GetFormalCharge()
    362        if chg!=0:
    363          if chg==1:
    364            chg = '+'
    365          elif chg==-1:
    366            chg = '-'
    367          elif chg>1:
    368            chg = '+%d'%chg
    369          elif chg<-1:
    370            chg = '-%d'%chg
    371        else:
    372          chg = ''
    373        if nbrSum[0]<=0:
    374          symbol = '%s%s%s'%(base,hs,chg)
    375        else:
    376          symbol = '%s%s%s'%(chg,hs,base)
    377
    378        #symbol = str(atom.GetIdx()+1)
    379        labelIt= not self.noCarbonSymbols or \
    380                 atom.GetAtomicNum()!=6 or \
    381                 atom.GetFormalCharge()!=0 or \
    382                 self.includeAtomNumbers
    383        if self.includeAtomNumbers:
    384          symbol = str(atom.GetIdx())
    385        if labelIt:
    386          color = elemDict.get(atom.GetAtomicNum(),(0,0,0))
    387          self._drawLabel(canvas,symbol,pos,font,color=color,
    388                          highlightIt=(highlightAtoms and idx in highlightAtoms))
    389
    390 -def registerCanvas(canvasNm):
    391    g= globals()
    392    if canvasNm in ('sping','SPING'):
    393      from spingCanvas import addCanvasLine,addCanvasText,addCanvasPolygon,addCanvasDashedWedge
    394    elif canvasNm in ('agg','AGG'):
    395      from aggCanvas import addCanvasLine,addCanvasText,addCanvasPolygon,addCanvasDashedWedge
    396    elif canvasNm in ('mpl','MPL'):
    397      from mplCanvas import addCanvasLine,addCanvasText,addCanvasPolygon
    398      addCanvasDashedWedge=None
    399    else:
    400      raise ValueError,'unrecognized canvas type'
    401    g['addCanvasLine']=addCanvasLine
    402    g['addCanvasText']=addCanvasText
    403    g['addCanvasPolygon']=addCanvasPolygon
    404    g['addCanvasDashedWedge']=addCanvasDashedWedge
    405
    406  if __name__=='__main__':
    407    import sys
    408    if len(sys.argv)<2:
    409      mol = Chem.MolFromSmiles('O=C1C([C@@H](F)C=CN[C@H](Cl)Br)C(c2c(O)c(NN)ccc2)=C1C#N')
    410    else:
    411      mol = Chem.MolFromSmiles(sys.argv[1])
    412
    413    #mol = Chem.MolFromSmiles('OC1C(O)CC1')
    414    Chem.Kekulize(mol)
    415    from Chem import rdDepictor
    416    rdDepictor.Compute2DCoords(mol)
    417
    418    if 1:
    419      from aggdraw import Draw
    420      registerCanvas('agg')
    421      from PIL import Image
    422      img = Image.new("RGBA",(300,300),"white")
    423      canvas=Draw(img)
    424      canvas.setantialias(True)
    425      drawer = MolDrawing(canvas)
    426      drawer.AddMol(mol)
    427      canvas.flush()
    428      img.save("foo.png")
    429    elif 0:
    430      from matplotlib.figure import Figure
    431      from matplotlib.backends.backend_agg import FigureCanvasAgg
    432      registerCanvas('mpl')
    433      fig = Figure(figsize=(3,3))
    434      ax = fig.add_axes((0,0,1,1),xticks=[],yticks=[],frame_on=False)
    435      xd = ax.get_xlim()
    436      xd = xd[1]-xd[0]
    437      yd = ax.get_ylim()
    438      yd = yd[1]-yd[0]
    439      ax.size=(xd,yd)
    440      drawer = MolDrawing(ax)
    441      drawer.AddMol(mol)
    442      canv = FigureCanvasAgg(fig)
    443      canv.print_figure("foo.png",dpi=80)
    444    else:
    445      from sping.PDF.pidPDF import PDFCanvas as Canvas
    446      canvas = Canvas(size=(300,300),name='test.pdf')
    447      registerCanvas('sping')
    448      drawer = MolDrawing(canvas)
    449      drawer.AddMol(mol)
    450      canvas.save()
    451