Source Code for Module gromacs.analysis.collections

  1  # GromacsWrapper: collections.py 
  2  # Copyright (c) 2009 Oliver Beckstein <orbeckst@gmail.com> 
  3  # Released under GPL3+ 
  4   
  5  """ 
  6  :mod:`analysis.collections` -- Handling of groups of simulation instances 
  7  ========================================================================= 
  8   
  9  This module contains classes and functions that combine multiple 
 10  :class:`gromacs.analysis.core.Simulation` objects. In this way the 
 11  same kind of analysis or plotting task can be carried out 
 12  simultaneously for all simulations in the collection. 
 13   
 14  .. autoclass:: Collection 
 15  """ 
 16   
 17  import os.path 
 18  import cPickle 
 19  from numpy import all, any 
 20   
 21 -class Collection(list): 
 22      """Multiple objects (organized as a list). 
 23   
 24      Methods are applied to all objects in the Collection and returned 
 25      as new Collection: 
 26   
 27        >>> from gromacs.analysis.collections import Collection 
 28        >>> animals = Collection(['ant', 'boar', 'ape', 'gnu']) 
 29        >>> animals.startswith('a') 
 30        Collection([True, False, True, False]) 
 31   
 32      Similarly, attributes are returned as a Collection. 
 33   
 34      Using :meth:`Collection.save` one can save the whole collection to 
 35      disk and restore it later with the :meth:`Collection.load` method 
 36   
 37        >>> animals.save('zoo') 
 38        >>> arc = Collection() 
 39        >>> arc.load('zoo') 
 40        >>> arc.load('zoo', append=True) 
 41        >>> arc 
 42        ['ant', 'boar', 'ape', 'gnu', 'ant', 'boar', 'ape', 'gnu'] 
 43      """ 
 44      # note: do not use with multiple inheritance -- why, I think it could work now... 
 45   
 46 -    def save(self, filename): 
 47          """Pickle the whole collection to *filename*. 
 48           
 49          If no extension is provided, ".collection" is appended. 
 50          """ 
 51          cPickle.dump(self, open(self._canonicalize(filename), 'wb'),  
 52                       protocol=cPickle.HIGHEST_PROTOCOL) 
 53   
 54 -    def load(self, filename, append=False): 
 55          """Load collection from pickled file *filename*. 
 56   
 57          *append* determines if the saved collection is added to the current one 
 58          or if it replaces the current content. 
 59   
 60          If no extension is provided, ".collection" is appended. 
 61          """ 
 62          tmp = cPickle.load(open(self._canonicalize(filename), 'rb')) 
 63          if append: 
 64              self.extend(tmp) 
 65          else: 
 66              self[:] = tmp[:] 
 67          del tmp 
 68   
 69 -    def tolist(self): 
 70          """Return contents as a simple list.""" 
 71          return self[:] 
 72   
 73 -    def _canonicalize(self, filename): 
 74          """Use .collection as extension unless provided""" 
 75          path, ext = os.path.splitext(filename) 
 76          if not ext: 
 77              ext = ".collection" 
 78          return path + ext 
 79   
 80 -    def __getnewargs__(self, *args, **kwargs): 
 81          """Provide proper initialization to make pickling with protocol 2 work""" 
 82          return (self.tolist(),) 
 83   
 84 -    def __getattribute__(self, attr): 
 85          try: 
 86              return super(Collection, self).__getattribute__(attr) 
 87          except AttributeError: 
 88              pass 
 89           
 90          for o in self: 
 91              failures = [] 
 92              if not hasattr(o, attr): 
 93                  failures.append(o) 
 94          if len(failures) > 0: 
 95              raise AttributeError("The following members of the collection do not " 
 96                                   "implement the attribute %(attr)r:\n%(failures)r\n" 
 97                                   % vars()) 
 98   
 99          # analyze attribute: functions (the ones with __call__) get delayed, simple 
100          # attributes are looked up immediately 
101          iscallable = [hasattr(o.__getattribute__(attr), '__call__') for o in self] 
102          if all(iscallable): 
103              def runall(*args, **kwargs): 
104                  """Apply function to all members and return a new Collection""" 
105                  return Collection([o.__getattribute__(attr)(*args, **kwargs) for o in self]) 
106              runall.__name__ = attr 
107              return runall 
108          elif any(iscallable): 
109              raise TypeError("Attribute %r is callable only for some objects" % attr) 
110   
111          return Collection([o.__getattribute__(attr) for o in self]) 
112   
113 -    def __add__(self, x): 
114          return Collection(super(Collection, self).__add__(x)) 
115   
116 -    def __repr__(self): 
117          return self.__class__.__name__+"(%s)" % super(Collection, self).__repr__() 
118