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 795 2008-08-19 03:11:01Z 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=True
     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      bType=bond.GetBondType()
    170      if bType == Chem.BondType.SINGLE:
    171        bDir = bond.GetBondDir()
    172        if bDir in (Chem.BondDir.BEGINWEDGE,Chem.BondDir.BEGINDASH):
    173          # if the bond is "backwards", change the drawing direction:
    174          if bond.GetBeginAtom().GetChiralTag() in (Chem.ChiralType.CHI_TETRAHEDRAL_CW,
    175                                                    Chem.ChiralType.CHI_TETRAHEDRAL_CCW):
    176            p1,p2 = pos,nbrPos
    177          else:
    178            p2,p1 = pos,nbrPos
    179          if bDir==Chem.BondDir.BEGINWEDGE:
    180            self._drawWedgedBond(canvas,bond,p1,p2,color=(0,0,0),width=width)
    181          elif bDir==Chem.BondDir.BEGINDASH:
    182            self._drawWedgedBond(canvas,bond,p1,p2,color=(0,0,0),width=width,
    183                                 dash=self.dash)
    184        else:
    185          addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    186      elif bType == Chem.BondType.DOUBLE:
    187        if bond.IsInRing() or (atom.GetDegree()!=1 and bond.GetOtherAtom(atom).GetDegree()!=1):
    188          addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    189          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf)
    190          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    191        else:
    192          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf,dir=.5,
    193                                        lenFrac=1.0)
    194          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    195          fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf,dir=-.5,
    196                                        lenFrac=1.0)
    197          addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2)
    198      elif bType == Chem.BondType.AROMATIC:
    199        addCanvasLine(canvas,pos,nbrPos,linewidth=width,color=color,color2=color2)
    200        fp1,fp2 = self._offsetDblBond(pos,nbrPos,bond,atom,nbr,conf)
    201        addCanvasLine(canvas,fp1,fp2,linewidth=width,color=color,color2=color2,
    202                      dash=self.dash)
    203      elif bType == 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      minX = 1e8
    215      minY = 1e8
    216      maxX = -1e8
    217      maxY = -1e8
    218
    219      nAts = mol.GetNumAtoms()
    220      for i in range(nAts):
    221        pos = conf.GetAtomPosition(i)
    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
    229      dx = abs(maxX-minX)
    230      dy = abs(maxY-minY)
    231      xSize = dx*self.dotsPerAngstrom
    232      ySize = dy*self.dotsPerAngstrom
    233
    234      if coordCenter:
    235        molTrans = -xAccum/nAts,-yAccum/nAts
    236      else:
    237        molTrans = -(minX+(maxX-minX)/2),-(minY+(maxY-minY)/2)
    238      self.dotsPerAngstrom=30.0
    239      self.molTrans = molTrans
    240
    241      if xSize>=.95*canvasSize[0]:
    242        scale = .9*canvasSize[0]/xSize
    243        xSize*=scale
    244        ySize*=scale
    245        self.dotsPerAngstrom*=scale
    246        self.atomLabelFontSize = max(self.atomLabelFontSize*scale,
    247                                     self.atomLabelMinFontSize)
    248      if ySize>=.95*canvasSize[1]:
    249        scale = .9*canvasSize[1]/ySize
    250        xSize*=scale
    251        ySize*=scale
    252        self.dotsPerAngstrom*=scale
    253        self.atomLabelFontSize = max(self.atomLabelFontSize*scale,
    254                                     self.atomLabelMinFontSize)
    255      drawingTrans = canvasSize[0]/2,canvasSize[1]/2
    256      self.drawingTrans = drawingTrans
    257      tMax = self.transformPoint((maxX,maxY))
    258      tMin = self.transformPoint((minX,minY))
    259
    260 -  def _drawLabel(self,canvas,label,pos,font,color=None,
    261                   highlightIt=False):
    262      if highlightIt:
    263        color = self.selectColor
    264      elif not color:
    265        color = self.defaultColor
    266      x1 = pos[0]
    267      y1 = pos[1]
    268      labelP = x1,y1
    269      addCanvasText(canvas,label,(x1,y1),font,color)
    270
    271 -  def AddMol(self,mol,canvas=None,centerIt=True,molTrans=(0,0),drawingTrans=(0,0),
    272               highlightAtoms=[],confId=-1):
    273      """
    274
    275      Notes:
    276        - specifying centerIt will cause molTrans and drawingTrans to be ignored
    277
    278      """
    279      try:
    280        dl = addCanvasLine
    281      except NameError:
    282        registerCanvas('sping')
    283      if canvas is None:
    284        canvas = self.canvas
    285      else:
    286        self.canvas = canvas
    287      self.canvasSize=canvas.size
    288
    289      conf = mol.GetConformer(confId)
    290
    291      if centerIt:
    292        self._scaleAndCenter(mol,conf)
    293      else:
    294        self.molTrans = molTrans
    295        self.drawingTrans = drawingTrans
    296      font = Font(face=self.atomLabelFontFace,size=self.atomLabelFontSize)
    297
    298      if not mol.HasProp('_drawingBondsWedged'):
    299        Chem.WedgeMolBonds(mol,conf)
    300
    301      self.atomPs[mol] = {}
    302      self.activeMol = mol
    303      self.bondRings = mol.GetRingInfo().BondRings()
    304      for atom in mol.GetAtoms():
    305        idx = atom.GetIdx()
    306        pos = self.atomPs[mol].get(idx,None)
    307        if pos is None:
    308          pos = self.transformPoint(conf.GetAtomPosition(idx))
    309          self.atomPs[mol][idx] = pos
    310        nbrSum = [0,0]
    311        for bond in atom.GetBonds():
    312          nbr = bond.GetOtherAtom(atom)
    313          nbrIdx = nbr.GetIdx()
    314          if nbrIdx > idx:
    315            nbrPos = self.atomPs[mol].get(nbrIdx,None)
    316            if nbrPos is None:
    317              nbrPos = self.transformPoint(conf.GetAtomPosition(nbrIdx))
    318              self.atomPs[mol][nbrIdx] = nbrPos
    319
    320            if highlightAtoms and idx in highlightAtoms and nbrIdx in highlightAtoms:
    321              width=2.0*self.bondLineWidth
    322              color = self.selectColor
    323              color2 = self.selectColor
    324            else:
    325              width=self.bondLineWidth
    326              if self.colorBonds:
    327                color = elemDict.get(atom.GetAtomicNum(),(0,0,0))
    328                color2 = elemDict.get(nbr.GetAtomicNum(),(0,0,0))
    329              else:
    330                color = self.defaultColor
    331                color2= color
    332
    333            # make sure we draw from the beginning to the end
    334            # (this was Issue400)
    335            if idx==bond.GetBeginAtomIdx():
    336              self._drawBond(canvas,bond,atom,nbr,pos,nbrPos,conf,
    337                             color=color,width=width,color2=color2)
    338            else:
    339              self._drawBond(canvas,bond,nbr,atom,nbrPos,pos,conf,
    340                             color=color2,width=width,color2=color)
    341          else:
    342            nbrPos = self.atomPs[mol][nbrIdx]
    343          nbrSum[0] += nbrPos[0]-pos[0]
    344          nbrSum[1] += nbrPos[1]-pos[1]
    345
    346
    347        labelIt= not self.noCarbonSymbols or \
    348                 atom.GetAtomicNum()!=6 or \
    349                 atom.GetFormalCharge()!=0 or \
    350                 self.includeAtomNumbers
    351        if labelIt:
    352          if self.includeAtomNumbers:
    353            symbol = str(atom.GetIdx())
    354          else:
    355            base = atom.GetSymbol()
    356            nHs = atom.GetTotalNumHs()
    357            if nHs>0:
    358              if nHs>1:
    359                hs='H%d'%nHs
    360              else:
    361                hs ='H'
    362            else:
    363              hs = ''
    364            chg = atom.GetFormalCharge()
    365            if chg!=0:
    366              if chg==1:
    367                chg = '+'
    368              elif chg==-1:
    369                chg = '-'
    370              elif chg>1:
    371                chg = '+%d'%chg
    372              elif chg<-1:
    373                chg = '-%d'%chg
    374            else:
    375              chg = ''
    376            if nbrSum[0]<=0:
    377              symbol = '%s%s%s'%(base,hs,chg)
    378            else:
    379              symbol = '%s%s%s'%(chg,hs,base)
    380
    381          color = elemDict.get(atom.GetAtomicNum(),(0,0,0))
    382          self._drawLabel(canvas,symbol,pos,font,color=color,
    383                          highlightIt=(highlightAtoms and idx in highlightAtoms))
    384
    385 -def registerCanvas(canvasNm):
    386    g= globals()
    387    if canvasNm in ('sping','SPING'):
    388      from spingCanvas import addCanvasLine,addCanvasText,addCanvasPolygon,addCanvasDashedWedge
    389    elif canvasNm in ('agg','AGG'):
    390      from aggCanvas import addCanvasLine,addCanvasText,addCanvasPolygon,addCanvasDashedWedge
    391    elif canvasNm in ('mpl','MPL'):
    392      from mplCanvas import addCanvasLine,addCanvasText,addCanvasPolygon
    393      addCanvasDashedWedge=None
    394    else:
    395      raise ValueError,'unrecognized canvas type'
    396    g['addCanvasLine']=addCanvasLine
    397    g['addCanvasText']=addCanvasText
    398    g['addCanvasPolygon']=addCanvasPolygon
    399    g['addCanvasDashedWedge']=addCanvasDashedWedge
    400
    401  if __name__=='__main__':
    402    import sys
    403    if len(sys.argv)<2:
    404      mol = Chem.MolFromSmiles('O=C1C([C@@H](F)C=CN[C@H](Cl)Br)C(c2c(O)c(NN)ccc2)=C1C#N')
    405    else:
    406      mol = Chem.MolFromSmiles(sys.argv[1])
    407
    408    #mol = Chem.MolFromSmiles('OC1C(O)CC1')
    409    Chem.Kekulize(mol)
    410    from Chem import rdDepictor
    411    rdDepictor.Compute2DCoords(mol)
    412
    413    if 1:
    414      from aggdraw import Draw
    415      registerCanvas('agg')
    416      from PIL import Image
    417      img = Image.new("RGBA",(300,300),"white")
    418      canvas=Draw(img)
    419      canvas.setantialias(True)
    420      drawer = MolDrawing(canvas)
    421      drawer.AddMol(mol)
    422      canvas.flush()
    423      img.save("foo.png")
    424    elif 0:
    425      from matplotlib.figure import Figure
    426      from matplotlib.backends.backend_agg import FigureCanvasAgg
    427      registerCanvas('mpl')
    428      fig = Figure(figsize=(3,3))
    429      ax = fig.add_axes((0,0,1,1),xticks=[],yticks=[],frame_on=False)
    430      xd = ax.get_xlim()
    431      xd = xd[1]-xd[0]
    432      yd = ax.get_ylim()
    433      yd = yd[1]-yd[0]
    434      ax.size=(xd,yd)
    435      drawer = MolDrawing(ax)
    436      drawer.AddMol(mol)
    437      canv = FigureCanvasAgg(fig)
    438      canv.print_figure("foo.png",dpi=80)
    439    else:
    440      from sping.PDF.pidPDF import PDFCanvas as Canvas
    441      canvas = Canvas(size=(300,300),name='test.pdf')
    442      registerCanvas('sping')
    443      drawer = MolDrawing(canvas)
    444      drawer.AddMol(mol)
    445      canvas.save()
    446