Source Code for Module Biskit.Pymoler

  1  ## 
  2  ## Biskit, a toolkit for the manipulation of macromolecular structures 
  3  ## Copyright (C) 2004-2005 Raik Gruenberg & Johan Leckner 
  4  ## 
  5  ## This program is free software; you can redistribute it and/or 
  6  ## modify it under the terms of the GNU General Public License as 
  7  ## published by the Free Software Foundation; either version 2 of the 
  8  ## License, or any later version. 
  9  ## 
 10  ## This program is distributed in the hope that it will be useful, 
 11  ## but WITHOUT ANY WARRANTY; without even the implied warranty of 
 12  ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
 13  ## General Public License for more details. 
 14  ## 
 15  ## You find a copy of the GNU General Public License in the file 
 16  ## license.txt along with this program; if not, write to the Free 
 17  ## Software Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA. 
 18  ## 
 19  ## 
 20  ## last $Date: 2006/12/21 09:33:09 $ 
 21  ## $Revision: 2.7 $ 
 22  ## last $Author: leckner  
 23  """ 
 24  Display Structures with Pymol 
 25  """ 
 26   
 27   
 28  from PDBModel import PDBModel 
 29  import tools as T 
 30  import os 
 31  import tempfile 
 32  import settings 
 33  from Biskit import Executor, TemplateError 
 34   
 35   
 36  ## ================= single Pymol model =========================== 
 37   
 38 -class PymolModel: 
 39   
 40 -    def __init__( self, model, modName ): 
 41          """ 
 42          @param model: model to view 
 43          @type  model: PDBModel 
 44          @param modName: model name, will show up in PyMol 
 45          @type  modName: str 
 46          """ 
 47          self.fname = '' 
 48          self.temporary = 0 
 49          self.struct = None 
 50   
 51          if type( model ) is str: 
 52              self.fname = model 
 53          else: 
 54              self.struct = model.clone( deepcopy=1 ) ## clone atom dicts 
 55              self.temporary = 1 
 56   
 57          self.modName = modName 
 58   
 59          if self.fname == '': 
 60              self.fname = tempfile.mktemp( 'model') 
 61   
 62   
 63 -    def addProperty( self, values, key='temperature_factor' ): 
 64          """ 
 65          Add extra value to each atom in Structure. The values 
 66          will be written to either the B- (temperature_factor) 
 67          or Q-factor 'occupancy' column in the temporary pdb-file. 
 68          These values can then be used to display properties in PyMol 
 69          via commands like 'color_b' and 'color_q'. See also 
 70          L{addResProperty}. 
 71   
 72          @param values: list of numbers, len( values ) == number of atoms 
 73          @type  values: [float]       
 74          @param key: key for Atom.properties dictionary 
 75                      ('occupancy' OR 'temperature_factor') 
 76          @type  key: occupancy|temperature_factor 
 77          """ 
 78          if self.struct == None: 
 79              self.struct = PDBModel( self.fname ) 
 80              self.temporary = 1 
 81   
 82          i = 0 
 83          for atom in self.struct.atoms: 
 84              atom[ key ] = values[i] 
 85              i += 1 
 86   
 87   
 88 -    def addResProperty( self, values, key='temperature_factor'): 
 89          """ 
 90          Does the same thing as L{addProperty} but on the residue level, 
 91          i.e adds extra value to each residue in Structure. 
 92          (The same value is added to all atoms of a residue.) 
 93          These values can then be used to display properties in PyMol 
 94          via commands like 'color_b' and 'color_q'. 
 95           
 96          @param values: list of numbers, len( values ) == number of residues 
 97          @type  values: [float]       
 98          @param key: key for Atom.properties dictionary 
 99                      ('occupancy' OR 'temperature_factor') 
100          @type  key: occupancy|temperature_factor 
101          """ 
102          try: 
103              if self.struct == None: 
104                  self.struct = PDBModel( self.fname ) 
105                  self.temporary = 1 
106   
107              i = 0 
108              resMap = self.struct.resMap() 
109   
110              for atom in self.struct.atoms: 
111                  atom[ key ] = values[ resMap[i] ] 
112                  i += 1 
113          except: 
114              print T.lastError() 
115   
116   
117 -    def writeIfNeeded( self ): 
118          """ 
119          Create pdb on disc if it's not already there or if it 
120          has to be changed. 
121           
122          @return: filename of new or existing pdb 
123          @rtype: str 
124          """ 
125          if self.temporary: 
126              self.struct.writePdb( self.fname, 2 ) 
127          return self.fname 
128   
129   
130  ## ========================= Pymoler ============================ 
131  ## create Pymol input script and display in Pymol ## 
132   
133 -class Pymoler( Executor ): 
134      """ 
135      Run Pymol 
136      ========= 
137      Create a pymol script file (.pml) and execule it using Pymol. 
138   
139      Example usage 
140      ------------- 
141          >>> pm = Pymoler() 
142          >>> pm.addPdb( model ) 
143          >>> pm.run() 
144   
145      References 
146      ---------- 
147         - U{http://pymol.sourceforge.net/} 
148       
149      @note: the file is only flushed to disc when the object 
150             is destructed, or the flush() method is called! 
151      @note: since version 2.6 Pymoler is using Biskit.Executor 
152      """ 
153   
154 -    def __init__(self, full=0, mode='w', verbose=1, **kw ): 
155          """ 
156          @param mode: open file with this mode, w=override, a=append 
157          @type  mode: str 
158          @param full: dispaly pymol structures in fill screen mode:: 
159                         0 - normal mode 
160                         1 - full screen mode 
161                         2 - full screen and no menues 
162          @type  full: 0|1|2         
163          """ 
164          self.verbose = verbose 
165           
166          # name of .pml file 
167          self.foutName = tempfile.mktemp() + '.pml' 
168   
169          # open for <appending|writing|reading> 
170          self.fgenerate = open(self.foutName, mode)  
171   
172          # will contain PymolModels or lists of PymolModels 
173          self.dic = {} 
174   
175          ## add startup commands 
176          self.initPymol() 
177   
178          ## set arguments for display options (normal, full, all) 
179          arg = args='-q %s'%self.foutName 
180          if full == 1: 
181              arg = args='-qe %s'%self.foutName 
182          if full == 2: 
183              arg = args='-qei %s'%self.foutName 
184               
185          Executor.__init__( self, 'pymol', args=arg, 
186                             catch_err=1, catch_out=1, **kw ) 
187           
188 -    def __del__(self): 
189          self.fgenerate.close() 
190   
191   
192 -    def flush(self): 
193          """ 
194          Flush output file (but keep it open). 
195          """ 
196          self.fgenerate.flush() 
197   
198   
199 -    def add(self, str): 
200          """ 
201          Add String str and line break to file. 
202   
203          @param str: string to add to pml file 
204          @type  str: str         
205          """ 
206          try: 
207              self.fgenerate.write(str + '\n') 
208          except (IOError): 
209              T.errWriteln( 
210                  "PymolInput.add(): Error adding string to pymol script file.") 
211              T.errWriteln( T.lastError() ) 
212   
213   
214 -    def addMovie( self, pdb, modName=None ): 
215          """ 
216          Add one or several existing pdb files or Structure objects 
217          to one model. Several files will hence end up as single movie 
218          (i.e. as frames of a model in PyMol). 
219    
220          @param pdb: file name or a list of file names OR 
221                      PDBModel or list of PDBModels 
222          @type  pdb: str or [str] OR PDBModel or [PDBModel] 
223          @param modName: model name, will show up in PyMol. If 'None' a 
224                          model name will be created from the source file 
225                          name and a serial number. 
226          @type  modName: str OR None 
227   
228          @return: the modName of the added model 
229          @rtype: str         
230          """ 
231          if type( pdb ) is not list: 
232              pdb = [pdb] 
233   
234          ## dream up a nice model name 
235          if modName == None: 
236   
237              if type( pdb[0]) is str: 
238                  modName = T.stripFilename( pdb[0] ) 
239                  modName = self._getFreeModName( modName, 0 ) 
240   
241              else: 
242                  modName = self._getFreeModName( 'models', 0 ) 
243   
244          ## create new empty list 
245          if not self.dic.has_key( modName ): 
246              self.dic[ modName ] = [] 
247   
248          ## create model object for each file and append it to dic 
249          for f in pdb: 
250   
251              ## create model from Structure or file name 
252              model = PymolModel( f, modName ) 
253              self.dic[ modName ].append( model ) 
254   
255              ## add load statement to Pymol script 
256              self.add( 'load '+ model.fname + ',' + modName ) 
257   
258          return modName 
259   
260   
261 -    def addFrame( self, frame, modName ): 
262          """ 
263          Add file(s) or Structure(s) to an EXISTING model as movie. 
264           
265          @param frame: the structure to add to an existing model 
266          @type  frame: str(s) or PDBModel(s)  
267          @param modName: model name, must be existing 
268                          (i.e. an already added model) 
269          @type  modName: str 
270          """ 
271          self.addMovie( frame, modName ) 
272   
273   
274 -    def _getFreeModName( self, base, index ): 
275          """ 
276          Return next free model name in dictionary, made up 
277          from base + index 
278           
279          @param base: name base 
280          @type  base: str 
281          @param index: name suffix  
282          @type  index: int 
283           
284          @return: name composed of base+suffix 
285          @rtype: str 
286          """ 
287          if self.dic.has_key( base ): 
288   
289              if self.dic.has_key( base + str(index) ): 
290   
291                  return self._getFreeModName( base, index + 1 ) 
292   
293              else: 
294                  return base + str( index ) 
295   
296          else: 
297              return base 
298   
299   
300 -    def addPdb( self, pdb, modName=None ): 
301          """ 
302          Add one or several existing pdbs. Make sure all go into 
303          different models (have different model names). 
304           
305          @param pdb: file name or a list of file names OR 
306                      PDBModel or list of PDBModels 
307          @type  pdb: str or [str] OR PDBModel or [PDBModel] 
308   
309          @param modName: force model name, will change for list 
310                          of file names. If 'None' a model name will 
311                          be created from the source file name. 
312          @type  modName: str OR None 
313           
314          @return: model name of first file 
315          @rtype: str 
316          """ 
317   
318          if type( pdb ) is not list: 
319              pdb = [pdb] 
320   
321          ## create one model for each file / Structure 
322          for f in pdb: 
323              result = self.addMovie( f, modName ) 
324   
325          return result 
326   
327   
328 -    def makeSel( self, selDic ): 
329          """ 
330          Make a selection. 
331           
332          @param selDic: a selection dictionary,  that can be of three types: 
333                           1. dictionary with one or more of the keys: 
334                             'model' 'segment', 'chain', 'residue', 'atom' 
335                           2. dictionary with key:   'element' 
336                           3. dictionaty with key:   'expression' 
337          @type  selDic: dict 
338          """ 
339          ## the selection must comply with one of the three selection 
340          ## schemes below 
341          keyList = ['model', 'segment', 'chain', 'residue', 'atom', \ 
342                     'element', 'expression'] 
343   
344          # check for invalid selections 
345          for sel in selDic.keys(): 
346              if sel not in keyList: 
347                  print 'Invalid selection in ' + str(selDic) 
348   
349          # element 
350          if selDic.has_key(keyList[5]): 
351              selection = '( elem ' + str(selDic[keyList[5]]) + ' )' 
352   
353          # expression 
354          elif selDic.has_key(keyList[6]): 
355              selection = '( ' + str(selDic[keyList[6]]) + ' )' 
356   
357          # create selection from selDic 
358          else: 
359              for key in keyList: 
360                  if not selDic.has_key(key): 
361                      selDic[key]='' 
362              selection = '( /' + str(selDic['model']) + \ 
363                          '/' + str(selDic['segment']) + \ 
364                          '/' + str(selDic['chain']) + \ 
365                          '/' + str(selDic['residue']) + \ 
366                          '/' + str(selDic['atom']) + ' )' 
367   
368          return selection 
369   
370   
371 -    def writeStructures( self ): 
372          """ 
373          Write all needed PDB files to disc. 
374          """ 
375          for key in self.dic.keys(): 
376   
377              for model in self.dic[ key ]: 
378   
379                  n = model.writeIfNeeded() 
380                   
381                  if self.verbose: print n 
382   
383   
384 -    def addDeleteScript(self): 
385          """ 
386          Deletes the pymol script file from disc 
387          """ 
388          self.add( "/ import os" ) 
389          self.add("/ os.system('rm " + self.foutName+ "')") 
390   
391   
392 -    def addDeletePdbs(self): 
393          """ 
394          Deletes the pdb-files in the list from disc 
395          """ 
396          self.add( "/ import os" ) 
397          for key in self.dic.keys(): 
398   
399              for model in self.dic[ key ]: 
400                  ## only remove files created by PymolInput! 
401                  if model.temporary: 
402                          self.add( "/ os.system('rm " + model.fname + "')" ) 
403   
404   
405 -    def setAtomValues( self, model, values, key='temperature_factor', 
406                         lastOnly=0 ): 
407          """ 
408          Add numeric value to all atoms of all Structures or the last 
409          Structure in a model.. The values will be written to either 
410          the B- (temperature_factor) or Q-factor 'occupancy' column in 
411          the temporary pdb-file. 
412          These values can then be used to display properties in PyMol 
413          via commands like 'color_b' and 'color_q'. See also 
414          L{setResValues}. 
415   
416          @param model: model name 
417          @type  model: str 
418          @param values: list of numbers, len( values ) == number of atoms 
419          @type  values: [float]       
420          @param key: key for Atom.properties dictionary 
421                      (default: temperature_factor) 
422          @type  key: occupancy|temperature_factor 
423          @param lastOnly: 0 - add to all in model OR 
424                           1 - add only to last Structure (default: 0) 
425          @type  lastOnly: 1|0 
426          """ 
427          if lastOnly: 
428              self.dic[ model ][-1].addProperty( values, key ) 
429   
430          else: 
431              for m in self.dic[ model ]: 
432                  try: 
433                      m.addProperty( values, key ) 
434                  except: 
435                      T.errWriteln( "Warning: error while adding properties.") 
436                      T.errWriteln( "Key: "+str( key )+" values: "+str( values ) ) 
437                      T.errWriteln( T.lastError() ) 
438   
439   
440 -    def setResValues( self, model, values, key='temperature_factor', 
441                        lastOnly=0 ): 
442          """ 
443          Add numeric value per residue to all atoms of all Structures 
444          or the last Structure in a model. The values will be written to 
445          either the B- (temperature_factor) or Q-factor 'occupancy' column 
446          in the temporary pdb-file. 
447          These values can then be used to display properties in PyMol 
448          via commands like 'color_b' and 'color_q'. See also 
449          L{setAtomValues}. 
450           
451          @param model: model name 
452          @type  model: str 
453          @param values: list of numbers, len( values ) == number of residues 
454          @type  values: [float]       
455          @param key: key for Atom.properties dictionary 
456                      (default: temperature_factor) 
457          @type  key: occupancy|temperature_factor 
458          @param lastOnly: 0 - add to all in model OR 
459                           1 - add only to last Structure (default: 0) 
460          @type  lastOnly: 1|0 
461          """ 
462          if lastOnly: 
463              self.dic[ model ][-1].addResProperty( values, key ) 
464   
465          else: 
466              for m in self.dic[ model ]: 
467                  try: 
468                      m.addResProperty( values, key ) 
469                  except: 
470                      T.errWriteln( "Warning: error while adding properties.") 
471                      T.errWriteln( "Key: "+str( key )+" values: "+str( values ) ) 
472                      T.errWriteln( T.lastError() ) 
473   
474   
475 -    def colorAtoms( self, model, values, lastOnly=0 ): 
476          """ 
477          Color atoms of this model by list of values. 
478           
479          @param model: model name 
480          @type  model: str 
481          @param values: len == number of atoms 
482          @type  values: [float] 
483          @param lastOnly: 0 - add to all in model OR 
484                           1 - add only to last Structure (default: 0) 
485          @type  lastOnly: 1|0          
486          """ 
487          self.setAtomValues( model, values, lastOnly=lastOnly ) 
488          self.add( "color_b('%s')" % model ) 
489   
490   
491 -    def colorRes( self, model, values, lastOnly=0 ): 
492          """ 
493          Color residues of this model by list of values. 
494           
495          @param model: model name 
496          @type  model: str 
497          @param values: len == number of residues 
498          @type  values: list of numbers 
499          @param lastOnly: 0 .. add to all in model 
500                           1 .. add only to last Structure 
501          @type  lastOnly:  
502          """ 
503          self.setResValues( model, values, lastOnly=lastOnly ) 
504          self.add( "color_b('%s')" % model ) 
505   
506   
507   
508 -    def addColors( self, nColors, firstColor=[1.0, 0.0, 0.0], 
509                     lastColor=[0.0, 1.0, 0.0] ): 
510          """ 
511          Define a range of colors that can be called in PyMol by name. 
512           
513          @param nColors: numbers of colors generated 
514          @type  nColors: int 
515          @param firstColor: first rgb color (default: [1.0, 0.0, 0.0]) 
516          @type  firstColor: [float] 
517          @param lastColor: last rgb color (default: [0.0, 1.0, 0.0]) 
518          @type  lastColor: [float] 
519          @return: a list of the color names 
520          @rtype: [str] 
521          """    
522          spectrum = T.hexColors( nColors, T.rgb2hex( firstColor ), 
523                                T.rgb2hex( lastColor ) ) 
524          rgb = [] 
525          colorNames = [] 
526   
527          for c in range(0, nColors): 
528              rgb = [ float( spectrum[c][:4] ) / 255, \ 
529                    float( '0x' + spectrum[c][4:6] ) / 255, \ 
530                    float( '0x' + spectrum[c][6:] ) / 255 ] 
531   
532              cName = 'c' + str(c) 
533              colorNames += cName 
534   
535              self.add( 'set_color '+ cName + ', ' + str(rgb) ) 
536   
537          return colorNames 
538   
539   
540 -    def initPymol( self ): 
541          """ 
542          Do some stuff always first... 
543          """ 
544          ## import dssp command and other commands 
545          if os.path.isdir( settings.pymol_scripts ): 
546              for script in os.listdir( settings.pymol_scripts ): 
547                  if not os.path.isdir( script ): 
548                      self.add( 'run ' + settings.pymol_scripts + script ) 
549                      if self.verbose: 
550                          print 'Adding %s script as a PyMol command'%script 
551          else: 
552              print '\n\nWARNING: No external PyMol scripts added\n\n' 
553   
554   
555  ##      # MAKE default SELECTIONS 
556  ##      self.add('select bb, ' + self.makeSel({'expression':'*/c,ca,o,n'} )) 
557  ##      self.add('select sc, ' +\ 
558  ##               self.makeSel({'expression':'!(*/c,o,ca,n) and !element h'})) 
559  ##         self.add('select none') 
560   
561   
562 -    def prepare( self, cleanUp=1 ): 
563          """ 
564          Overrides Executor method. 
565          """ 
566          # clean up files from disc 
567          if cleanUp: 
568              self.addDeletePdbs() 
569              self.addDeleteScript() 
570   
571          self.flush() 
572   
573          ## Write PDB's to disc if needed 
574          self.writeStructures() 
575   
576   
577 -    def show( self ): 
578          """ 
579          Backward compatability with old scripts. 
580          """ 
581          self.run() 
582   
583   
584  ############# 
585  ##  TESTING         
586  ############# 
587           
588 -class Test: 
589      """ 
590      Test class 
591      """ 
592       
593 -    def run( self, local=0 ): 
594          """ 
595          run function test 
596           
597          @param local: transfer local variables to global and perform 
598                        other tasks only when run locally 
599          @type  local: 1|0 
600           
601          @return: 1 
602          @rtype: int 
603          """ 
604          traj = T.Load( T.testRoot() + '/lig_pcr_00/traj.dat' ) 
605   
606          pm = Pymoler( full=0, verbose=local ) 
607           
608          mname = pm.addMovie( [ traj[i] for i in range(0,100,20) ] ) 
609   
610          sel = pm.makeSel({'residue':29}) 
611          pm.add('show stick, %s'%sel) 
612   
613          pm.add('mplay') 
614           
615          if local: 
616              globals().update( locals() ) 
617          else: 
618              pm.add('quit') 
619               
620          pm.run() ## old style call "pm.show()" also works 
621   
622          return 1 
623   
624   
625 -    def expected_result( self ): 
626          """ 
627          Precalculated result to check for consistent performance. 
628   
629          @return: 1 
630          @rtype:  int 
631          """ 
632          return 1 
633       
634   
635  if __name__ == '__main__': 
636   
637      test = Test() 
638   
639      assert test.run( local=1 ) == test.expected_result() 
640