rdkit.Chem.PyMol

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, sys
     16
     17  # Python3 compatibility
     18  try:
     19    from xmlrpclib import Server
     20  except ImportError:
     21    from xmlrpc.client import Server
     22
     23  _server = None
     24
     25
     26 -class MolViewer(object):
     27
     28 -  def __init__(self, host=None, port=9123, force=0, **kwargs):
     29      global _server
     30      if not force and _server is not None:
     31        self.server = _server
     32      else:
     33        if not host:
     34          host = os.environ.get('PYMOL_RPCHOST', 'localhost')
     35        _server = None
     36        serv = Server('http://%s:%d' % (host, port))
     37        serv.ping()
     38        _server = serv
     39        self.server = serv
     40      self.InitializePyMol()
     41
     42 -  def InitializePyMol(self):
     43      """ does some initializations to set up PyMol according to our
     44      tastes
     45
     46      """
     47      self.server.do('set valence,1')
     48      self.server.do('set stick_rad,0.15')
     49      self.server.do('set mouse_selection_mode,0')
     50      self.server.do('set line_width,2')
     51      self.server.do('set selection_width,10')
     52      self.server.do('set auto_zoom,0')
     53
     54 -  def DeleteAll(self):
     55      " blows out everything in the viewer "
     56      self.server.deleteAll()
     57
     58 -  def DeleteAllExcept(self, excludes):
     59      " deletes everything except the items in the provided list of arguments "
     60      allNames = self.server.getNames('*', False)
     61      for nm in allNames:
     62        if nm not in excludes:
     63          self.server.deleteObject(nm)
     64
     65 -  def LoadFile(self, filename, name, showOnly=False):
     66      """ calls pymol's "load" command on the given filename; the loaded object
     67      is assigned the name "name"
     68      """
     69      if showOnly:
     70        self.DeleteAll()
     71      id = self.server.loadFile(filename, name)
     72      return id
     73
     74 -  def ShowMol(self, mol, name='molecule', showOnly=True, highlightFeatures=[], molB="", confId=-1,
     75                zoom=True, forcePDB=False, showSticks=False):
     76      """ special case for displaying a molecule or mol block """
     77
     78      server = self.server
     79      if not zoom:
     80        self.server.do('view rdinterface,store')
     81      if showOnly:
     82        self.DeleteAll()
     83
     84      if not forcePDB and mol.GetNumAtoms() < 999:
     85        if not molB:
     86          molB = Chem.MolToMolBlock(mol, confId=confId)
     87        mid = server.loadMolBlock(molB, name)
     88      else:
     89        if not molB:
     90          molB = Chem.MolToPDBBlock(mol, confId=confId)
     91        mid = server.loadPDB(molB, name)
     92
     93      if highlightFeatures:
     94        nm = name + '-features'
     95        conf = mol.GetConformer(confId)
     96        for feat in highlightFeatures:
     97          pt = [0.0, 0.0, 0.0]
     98          for idx in feat:
     99            loc = conf.GetAtomPosition(idx)
    100            pt[0] += loc[0] / len(feat)
    101            pt[1] += loc[1] / len(feat)
    102            pt[2] += loc[2] / len(feat)
    103          server.sphere(pt, 0.2, (1, 1, 1), nm)
    104      if zoom:
    105        server.zoom('visible')
    106      else:
    107        self.server.do('view rdinterface,recall')
    108      if showSticks:  # show molecule in stick view
    109        self.server.do('show sticks, {}'.format(name))
    110      return mid
    111
    112 -  def GetSelectedAtoms(self, whichSelection=None):
    113      " returns the selected atoms "
    114      if not whichSelection:
    115        sels = self.server.getNames('selections')
    116        if sels:
    117          whichSelection = sels[-1]
    118        else:
    119          whichSelection = None
    120      if whichSelection:
    121        items = self.server.index(whichSelection)
    122      else:
    123        items = []
    124      return items
    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, name='neighborhood',
    150                                  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' %
    154                     locals())
    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 -  def SetDisplayUpdate(self, val):
    170      if not val:
    171        self.server.do('set defer_update,1')
    172      else:
    173        self.server.do('set defer_update,0')
    174
    175 -  def GetAtomCoords(self, sels):
    176      " returns the coordinates of the selected atoms "
    177      res = {}
    178      for label, idx in sels:
    179        coords = self.server.getAtomCoords('(%s and id %d)' % (label, idx))
    180        res[(label, idx)] = coords
    181      return res
    182
    183 -  def HideAll(self):
    184      self.server.do('disable all')
    185
    186 -  def HideObject(self, objName):
    187      self.server.do('disable %s' % objName)
    188
    189 -  def DisplayObject(self, objName):
    190      self.server.do('enable %s' % objName)
    191
    192 -  def Redraw(self):
    193      self.server.do('refresh')
    194
    195 -  def Zoom(self, objName):
    196      self.server.zoom(objName)
    197
    198 -  def DisplayHBonds(self, objName, molName, proteinName, molSelText='(%(molName)s)',
    199                      proteinSelText='(%(proteinName)s and not het)'):
    200      " toggles display of h bonds between the protein and a specified molecule "
    201      cmd = "delete %(objName)s;\n"
    202      cmd += "dist %(objName)s," + molSelText + "," + proteinSelText + ",mode=2;\n"
    203      cmd += "enable %(objName)s;"
    204      cmd = cmd % locals()
    205
    206      self.server.do(cmd)
    207
    208 -  def DisplayCollisions(self, objName, molName, proteinName, distCutoff=3.0, color='red',
    209                          molSelText='(%(molName)s)', proteinSelText='(%(proteinName)s and not het)'):
    210      " toggles display of collisions between the protein and a specified molecule "
    211      cmd = "delete %(objName)s;\n"
    212      cmd += "dist %(objName)s," + molSelText + "," + proteinSelText + ",%(distCutoff)f,mode=0;\n"
    213      cmd += """enable %(objName)s
    214      color %(color)s, %(objName)s"""
    215      cmd = cmd % locals()
    216      self.server.do(cmd)
    217
    218 -  def SaveFile(self, filename):
    219      # PyMol will interpret the path to be relative to where it was started
    220      # from. Remedy that.
    221      if not filename:
    222        raise ValueError('empty filename')
    223      filename = os.path.abspath(filename)
    224      self.server.save(filename)
    225
    226 -  def GetPNG(self, h=None, w=None, preDelay=0):
    227      try:
    228        import Image
    229      except ImportError:
    230        from PIL import Image
    231      import time
    232      if preDelay > 0:
    233        time.sleep(preDelay)
    234      fd = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
    235      fd.close()
    236      self.server.do('png %s' % fd.name)
    237      time.sleep(0.2)  # <- wait a short period so that PyMol can finish
    238      for i in range(10):
    239        try:
    240          img = Image.open(fd.name)
    241          break
    242        except IOError:
    243          time.sleep(0.1)
    244      try:
    245        os.unlink(fd.name)
    246      except (OSError, PermissionError):
    247        # happens sometimes on Windows. Not going to worry about this too deeply since
    248        # the files are in a temp dir anyway. This was github #936
    249        pass
    250      fd = None
    251      if h is not None or w is not None:
    252        sz = img.size
    253        if h is None:
    254          h = sz[1]
    255        if w is None:
    256          w = sz[0]
    257        if h < sz[1]:
    258          frac = float(h) / sz[1]
    259          w *= frac
    260          w = int(w)
    261          img = img.resize((w, h), True)
    262        elif w < sz[0]:
    263          frac = float(w) / sz[0]
    264          h *= frac
    265          h = int(h)
    266          img = img.resize((w, h), True)
    267      return img
    268