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