Package rdkit :: Package Chem :: Module PyMol

[hide private]
[frames] | no frames]

# Source Code for Module rdkit.Chem.PyMol

      1  # $Id$
      2  #
      3  # Copyright (C) 2004-2012 Greg Landrum and Rational Discovery LLC
      4  #
      5  #   @@ All Rights Reserved @@
      6  #  This file is part of the RDKit.
      7  #  The contents are covered by the terms of the BSD license
      8  #  which is included in the file license.txt, found at the root
      9  #  of the RDKit source tree.
     10  #
     11  """ uses pymol to interact with molecules
     12
     13  """
     14  from rdkit import Chem
     15  import os, tempfile
     16
     17  # Python3 compatibility
     18  try:
     19    from xmlrpclib import Server
     20  except ImportError:
     21    from xmlrpc.client import Server
     22
     23
     24  _server=None
     25 -class MolViewer(object):
     26 -  def __init__(self,host=None,port=9123,force=0,**kwargs):
     27      global _server
     28      if not force and _server is not None:
     29        self.server=_server
     30      else:
     31        if not host:
     32          host=os.environ.get('PYMOL_RPCHOST','localhost')
     33        _server=None
     34        serv = Server('http://%s:%d'%(host,port))
     35        serv.ping()
     36        _server = serv
     37        self.server=serv
     38      self.InitializePyMol()
     39
     40 -  def InitializePyMol(self):
     41      """ does some initializations to set up PyMol according to our
     42      tastes
     43
     44      """
     45      self.server.do('set valence,1')
     46      self.server.do('set stick_rad,0.15')
     47      self.server.do('set mouse_selection_mode,0')
     48      self.server.do('set line_width,2')
     49      self.server.do('set selection_width,10')
     50      self.server.do('set auto_zoom,0')
     51
     52
     53 -  def DeleteAll(self):
     54      " blows out everything in the viewer "
     55      self.server.deleteAll()
     56
     57 -  def DeleteAllExcept(self,excludes):
     58      " deletes everything except the items in the provided list of arguments "
     59      allNames = self.server.getNames('*',False)
     60      for nm in allNames:
     61        if nm not in excludes:
     62          self.server.deleteObject(nm)
     63
     64 -  def LoadFile(self,filename,name,showOnly=False):
     65      """ calls pymol's "load" command on the given filename; the loaded object
     66      is assigned the name "name"
     67      """
     68      if showOnly:
     69        self.DeleteAll()
     70      id = self.server.loadFile(filename,name)
     71      return id
     72
     73 -  def ShowMol(self,mol,name='molecule',showOnly=True,highlightFeatures=[],
     74                molB="",confId=-1,zoom=True,forcePDB=False, showSticks=False):
     75      """ special case for displaying a molecule or mol block """
     76
     77      server = self.server
     78      if not zoom:
     79        self.server.do('view rdinterface,store')
     80      if showOnly:
     81        self.DeleteAll()
     82
     83      if not forcePDB and mol.GetNumAtoms()<999 :
     84        if not molB:
     85          molB = Chem.MolToMolBlock(mol,confId=confId)
     86        mid = server.loadMolBlock(molB,name)
     87      else:
     88        if not molB:
     89          molB = Chem.MolToPDBBlock(mol,confId=confId)
     90        mid = server.loadPDB(molB,name)
     91
     92      if highlightFeatures:
     93        nm = name+'-features'
     94        conf = mol.GetConformer(confId)
     95        for feat in highlightFeatures:
     96          pt = [0.0,0.0,0.0]
     97          for idx in feat:
     98            loc = conf.GetAtomPosition(idx)
     99            pt[0] += loc[0]/len(feat)
    100            pt[1] += loc[1]/len(feat)
    101            pt[2] += loc[2]/len(feat)
    102          server.sphere(pt,0.2,(1,1,1),nm)
    103      if zoom:
    104        server.zoom('visible')
    105      else:
    106        self.server.do('view rdinterface,recall')
    107      if showSticks:  # show molecule in stick view
    108        self.server.do('show sticks, {}'.format(name))
    109      return mid
    110
    111 -  def GetSelectedAtoms(self,whichSelection=None):
    112      " returns the selected atoms "
    113      if not whichSelection:
    114        sels = self.server.getNames('selections')
    115        if sels:
    116          whichSelection = sels[-1]
    117        else:
    118          whichSelection=None
    119      if whichSelection:
    120        items = self.server.index(whichSelection)
    121      else:
    122        items = []
    123      return items
    124
    125
    126 -  def SelectAtoms(self,itemId,atomIndices,selName='selection'):
    127      " selects a set of atoms "
    128      ids = '(id '
    129      ids += ','.join(['%d'%(x+1) for x in atomIndices])
    130      ids += ')'
    131      cmd = 'select %s,%s and %s'%(selName,ids,itemId)
    132      self.server.do(cmd)
    133
    134 -  def HighlightAtoms(self,indices,where,extraHighlight=False):
    135      " highlights a set of atoms "
    136      if extraHighlight:
    137        idxText = ','.join(['%s and (id %d)'%(where,x) for x in indices])
    138        self.server.do('edit %s'%idxText)
    139      else:
    140        idxText = ' or '.join(['id %d'%x for x in indices])
    141        self.server.do('select selection, %s and (%s)'%(where,idxText))
    142
    143 -  def SetDisplayStyle(self,obj,style=''):
    144      " change the display style of the specified object "
    145      self.server.do('hide everything,%s'%(obj,))
    146      if style:
    147        self.server.do('show %s,%s'%(style,obj))
    148
    149 -  def SelectProteinNeighborhood(self,aroundObj,inObj,distance=5.0,
    150                                  name='neighborhood',showSurface=False):
    151      """ selects the area of a protein around a specified object/selection name;
    152      optionally adds a surface to that """
    153      self.server.do('select %(name)s,byres (%(aroundObj)s around %(distance)f) and %(inObj)s'%locals())
    154
    155
    156      if showSurface:
    157        self.server.do('show surface,%s'%name)
    158        self.server.do('disable %s'%name)
    159
    160 -  def AddPharmacophore(self,locs,colors,label,sphereRad=0.5):
    161      " adds a set of spheres "
    162      self.server.do('view rdinterface,store')
    163      self.server.resetCGO(label)
    164      for i,loc in enumerate(locs):
    165        self.server.sphere(loc,sphereRad,colors[i],label,1)
    166      self.server.do('enable %s'%label)
    167      self.server.do('view rdinterface,recall')
    168
    169
    170 -  def SetDisplayUpdate(self,val):
    171      if not val:
    172        self.server.do('set defer_update,1')
    173      else:
    174        self.server.do('set defer_update,0')
    175
    176 -  def GetAtomCoords(self,sels):
    177      " returns the coordinates of the selected atoms "
    178      res = {}
    179      for label,idx in sels:
    180        coords = self.server.getAtomCoords('(%s and id %d)'%(label,idx))
    181        res[(label,idx)] = coords
    182      return res
    183
    184 -  def HideAll(self):
    185      self.server.do('disable all')
    186 -  def HideObject(self,objName):
    187      self.server.do('disable %s'%objName)
    188 -  def DisplayObject(self,objName):
    189      self.server.do('enable %s'%objName)
    190
    191 -  def Redraw(self):
    192      self.server.do('refresh')
    193 -  def Zoom(self,objName):
    194      self.server.zoom(objName)
    195
    196 -  def DisplayHBonds(self,objName,molName,proteinName,
    197                      molSelText='(%(molName)s)',
    198                      proteinSelText='(%(proteinName)s and not het)'):
    199      " toggles display of h bonds between the protein and a specified molecule "
    200      cmd = "delete %(objName)s;\n"
    201      cmd += "dist %(objName)s," + molSelText+","+proteinSelText+",mode=2;\n"
    202      cmd += "enable %(objName)s;"
    203      cmd = cmd%locals()
    204
    205      self.server.do(cmd)
    206
    207 -  def DisplayCollisions(self,objName,molName,proteinName,distCutoff=3.0,
    208                          color='red',
    209                          molSelText='(%(molName)s)',
    210                          proteinSelText='(%(proteinName)s and not het)'):
    211      " toggles display of collisions between the protein and a specified molecule "
    212      cmd = "delete %(objName)s;\n"
    213      cmd += "dist %(objName)s," + molSelText+","+proteinSelText+",%(distCutoff)f,mode=0;\n"
    214      cmd += """enable %(objName)s
    215      color %(color)s, %(objName)s"""
    216      cmd = cmd%locals()
    217      self.server.do(cmd)
    218
    219 -  def GetPNG(self,h=None,w=None,preDelay=0):
    220      try:
    221        import Image
    222      except ImportError:
    223        from PIL import Image
    224      import time
    225      if preDelay>0:
    226        time.sleep(preDelay)
    227      fd = tempfile.NamedTemporaryFile(suffix='.png',delete=False)
    228      fd.close()
    229      self.server.do('png %s'%fd.name)
    230      time.sleep(0.2)  # <- wait a short period so that PyMol can finish
    231      for i in range(10):
    232        try:
    233          img = Image.open(fd.name)
    234          break
    235        except IOError:
    236          time.sleep(0.1)
    237      try:
    238          os.unlink(fd.name)
    239      except (OSError,PermissionError):
    240          # happens sometimes on Windows. Not going to worry about this too deeply since
    241          # the files are in a temp dir anyway. This was github #936
    242          pass
    243      fd=None
    244      if h is not None or w is not None:
    245        sz = img.size
    246        if h is None:
    247          h=sz[1]
    248        if w is None:
    249          w=sz[0]
    250        if h<sz[1]:
    251          frac = float(h)/sz[1]
    252          w *= frac
    253          w = int(w)
    254          img=img.resize((w,h),True)
    255        elif w<sz[0]:
    256          frac = float(w)/sz[0]
    257          h *= frac
    258          h = int(h)
    259          img=img.resize((w,h),True)
    260      return img
    261