File: PyMol.py

package info (click to toggle)
rdkit 202009.4-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 129,624 kB
  • sloc: cpp: 288,030; python: 75,571; java: 6,999; ansic: 5,481; sql: 1,968; yacc: 1,842; lex: 1,254; makefile: 572; javascript: 461; xml: 229; fortran: 183; sh: 134; cs: 93
file content (267 lines) | stat: -rw-r--r-- 8,405 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
# $Id$
#
# Copyright (C) 2004-2012 Greg Landrum and Rational Discovery LLC
#
#   @@ All Rights Reserved @@
#  This file is part of the RDKit.
#  The contents are covered by the terms of the BSD license
#  which is included in the file license.txt, found at the root
#  of the RDKit source tree.
#
""" uses pymol to interact with molecules

"""
from rdkit import Chem
import os, tempfile, sys

# Python3 compatibility
try:
  from xmlrpclib import Server
except ImportError:
  from xmlrpc.client import Server

_server = None


class MolViewer(object):

  def __init__(self, host=None, port=9123, force=0, **kwargs):
    global _server
    if not force and _server is not None:
      self.server = _server
    else:
      if not host:
        host = os.environ.get('PYMOL_RPCHOST', 'localhost')
      _server = None
      serv = Server('http://%s:%d' % (host, port))
      serv.ping()
      _server = serv
      self.server = serv
    self.InitializePyMol()

  def InitializePyMol(self):
    """ does some initializations to set up PyMol according to our
    tastes

    """
    self.server.do('set valence,1')
    self.server.do('set stick_rad,0.15')
    self.server.do('set mouse_selection_mode,0')
    self.server.do('set line_width,2')
    self.server.do('set selection_width,10')
    self.server.do('set auto_zoom,0')

  def DeleteAll(self):
    " blows out everything in the viewer "
    self.server.deleteAll()

  def DeleteAllExcept(self, excludes):
    " deletes everything except the items in the provided list of arguments "
    allNames = self.server.getNames('*', False)
    for nm in allNames:
      if nm not in excludes:
        self.server.deleteObject(nm)

  def LoadFile(self, filename, name, showOnly=False):
    """ calls pymol's "load" command on the given filename; the loaded object
    is assigned the name "name"
    """
    if showOnly:
      self.DeleteAll()
    id = self.server.loadFile(filename, name)
    return id

  def ShowMol(self, mol, name='molecule', showOnly=True, highlightFeatures=[], molB="", confId=-1,
              zoom=True, forcePDB=False, showSticks=False):
    """ special case for displaying a molecule or mol block """

    server = self.server
    if not zoom:
      self.server.do('view rdinterface,store')
    if showOnly:
      self.DeleteAll()

    if not forcePDB and mol.GetNumAtoms() < 999:
      if not molB:
        molB = Chem.MolToMolBlock(mol, confId=confId)
      mid = server.loadMolBlock(molB, name)
    else:
      if not molB:
        molB = Chem.MolToPDBBlock(mol, confId=confId)
      mid = server.loadPDB(molB, name)

    if highlightFeatures:
      nm = name + '-features'
      conf = mol.GetConformer(confId)
      for feat in highlightFeatures:
        pt = [0.0, 0.0, 0.0]
        for idx in feat:
          loc = conf.GetAtomPosition(idx)
          pt[0] += loc[0] / len(feat)
          pt[1] += loc[1] / len(feat)
          pt[2] += loc[2] / len(feat)
        server.sphere(pt, 0.2, (1, 1, 1), nm)
    if zoom:
      server.zoom('visible')
    else:
      self.server.do('view rdinterface,recall')
    if showSticks:  # show molecule in stick view
      self.server.do('show sticks, {}'.format(name))
    return mid

  def GetSelectedAtoms(self, whichSelection=None):
    " returns the selected atoms "
    if not whichSelection:
      sels = self.server.getNames('selections')
      if sels:
        whichSelection = sels[-1]
      else:
        whichSelection = None
    if whichSelection:
      items = self.server.index(whichSelection)
    else:
      items = []
    return items

  def SelectAtoms(self, itemId, atomIndices, selName='selection'):
    " selects a set of atoms "
    ids = '(id '
    ids += ','.join(['%d' % (x + 1) for x in atomIndices])
    ids += ')'
    cmd = 'select %s,%s and %s' % (selName, ids, itemId)
    self.server.do(cmd)

  def HighlightAtoms(self, indices, where, extraHighlight=False):
    " highlights a set of atoms "
    if extraHighlight:
      idxText = ','.join(['%s and (id %d)' % (where, x) for x in indices])
      self.server.do('edit %s' % idxText)
    else:
      idxText = ' or '.join(['id %d' % x for x in indices])
      self.server.do('select selection, %s and (%s)' % (where, idxText))

  def SetDisplayStyle(self, obj, style=''):
    " change the display style of the specified object "
    self.server.do('hide everything,%s' % (obj, ))
    if style:
      self.server.do('show %s,%s' % (style, obj))

  def SelectProteinNeighborhood(self, aroundObj, inObj, distance=5.0, name='neighborhood',
                                showSurface=False):
    """ selects the area of a protein around a specified object/selection name;
    optionally adds a surface to that """
    self.server.do('select %(name)s,byres (%(aroundObj)s around %(distance)f) and %(inObj)s' %
                   locals())

    if showSurface:
      self.server.do('show surface,%s' % name)
      self.server.do('disable %s' % name)

  def AddPharmacophore(self, locs, colors, label, sphereRad=0.5):
    " adds a set of spheres "
    self.server.do('view rdinterface,store')
    self.server.resetCGO(label)
    for i, loc in enumerate(locs):
      self.server.sphere(loc, sphereRad, colors[i], label, 1)
    self.server.do('enable %s' % label)
    self.server.do('view rdinterface,recall')

  def SetDisplayUpdate(self, val):
    if not val:
      self.server.do('set defer_update,1')
    else:
      self.server.do('set defer_update,0')

  def GetAtomCoords(self, sels):
    " returns the coordinates of the selected atoms "
    res = {}
    for label, idx in sels:
      coords = self.server.getAtomCoords('(%s and id %d)' % (label, idx))
      res[(label, idx)] = coords
    return res

  def HideAll(self):
    self.server.do('disable all')

  def HideObject(self, objName):
    self.server.do('disable %s' % objName)

  def DisplayObject(self, objName):
    self.server.do('enable %s' % objName)

  def Redraw(self):
    self.server.do('refresh')

  def Zoom(self, objName):
    self.server.zoom(objName)

  def DisplayHBonds(self, objName, molName, proteinName, molSelText='(%(molName)s)',
                    proteinSelText='(%(proteinName)s and not het)'):
    " toggles display of h bonds between the protein and a specified molecule "
    cmd = "delete %(objName)s;\n"
    cmd += "dist %(objName)s," + molSelText + "," + proteinSelText + ",mode=2;\n"
    cmd += "enable %(objName)s;"
    cmd = cmd % locals()

    self.server.do(cmd)

  def DisplayCollisions(self, objName, molName, proteinName, distCutoff=3.0, color='red',
                        molSelText='(%(molName)s)', proteinSelText='(%(proteinName)s and not het)'):
    " toggles display of collisions between the protein and a specified molecule "
    cmd = "delete %(objName)s;\n"
    cmd += "dist %(objName)s," + molSelText + "," + proteinSelText + ",%(distCutoff)f,mode=0;\n"
    cmd += """enable %(objName)s
    color %(color)s, %(objName)s"""
    cmd = cmd % locals()
    self.server.do(cmd)

  def SaveFile(self, filename):
    # PyMol will interpret the path to be relative to where it was started
    # from. Remedy that.
    if not filename:
      raise ValueError('empty filename')
    filename = os.path.abspath(filename)
    self.server.save(filename)

  def GetPNG(self, h=None, w=None, preDelay=0):
    try:
      import Image
    except ImportError:
      from PIL import Image
    import time
    if preDelay > 0:
      time.sleep(preDelay)
    fd = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
    fd.close()
    self.server.do('png %s' % fd.name)
    time.sleep(0.2)  # <- wait a short period so that PyMol can finish
    for i in range(10):
      try:
        img = Image.open(fd.name)
        break
      except IOError:
        time.sleep(0.1)
    try:
      os.unlink(fd.name)
    except (OSError, PermissionError):
      # happens sometimes on Windows. Not going to worry about this too deeply since
      # the files are in a temp dir anyway. This was github #936
      pass
    fd = None
    if h is not None or w is not None:
      sz = img.size
      if h is None:
        h = sz[1]
      if w is None:
        w = sz[0]
      if h < sz[1]:
        frac = float(h) / sz[1]
        w *= frac
        w = int(w)
        img = img.resize((w, h), True)
      elif w < sz[0]:
        frac = float(w) / sz[0]
        h *= frac
        h = int(h)
        img = img.resize((w, h), True)
    return img