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