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