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