Source code for gromacs.collections

# GromacsWrapper: collections.py
# Copyright (c) 2009 Oliver Beckstein <orbeckst@gmail.com>
# Released under GPL3+

"""
:mod:`analysis.collections` -- Handling of groups of simulation instances
=========================================================================

This module contains classes and functions that combine multiple
:class:`gromacs.analysis.core.Simulation` objects. In this way the
same kind of analysis or plotting task can be carried out
simultaneously for all simulations in the collection.

.. autoclass:: Collection
"""

import os.path
import pickle
from numpy import all, any


[docs] class Collection(list): """Multiple objects (organized as a list). Methods are applied to all objects in the Collection and returned as new Collection: >>> from gromacs.analysis.collections import Collection >>> animals = Collection(['ant', 'boar', 'ape', 'gnu']) >>> animals.startswith('a') Collection([True, False, True, False]) Similarly, attributes are returned as a Collection. Using :meth:`Collection.save` one can save the whole collection to disk and restore it later with the :meth:`Collection.load` method >>> animals.save('zoo') >>> arc = Collection() >>> arc.load('zoo') >>> arc.load('zoo', append=True) >>> arc ['ant', 'boar', 'ape', 'gnu', 'ant', 'boar', 'ape', 'gnu'] """ # note: do not use with multiple inheritance -- why, I think it could work now... def save(self, filename): """Pickle the whole collection to *filename*. If no extension is provided, ".collection" is appended. """ pickle.dump( self, open(self._canonicalize(filename), "wb"), protocol=pickle.HIGHEST_PROTOCOL, ) def load(self, filename, append=False): """Load collection from pickled file *filename*. *append* determines if the saved collection is added to the current one or if it replaces the current content. If no extension is provided, ".collection" is appended. """ tmp = pickle.load(open(self._canonicalize(filename), "rb")) if append: self.extend(tmp) else: self[:] = tmp[:] del tmp def tolist(self): """Return contents as a simple list.""" return self[:] def _canonicalize(self, filename): """Use .collection as extension unless provided""" path, ext = os.path.splitext(filename) if not ext: ext = ".collection" return path + ext def __getnewargs__(self, *args, **kwargs): """Provide proper initialization to make pickling with protocol 2 work""" return (self.tolist(),) def __getattribute__(self, attr): try: return super(Collection, self).__getattribute__(attr) except AttributeError: pass for o in self: failures = [] if not hasattr(o, attr): failures.append(o) if len(failures) > 0: raise AttributeError( "The following members of the collection do not " "implement the attribute %(attr)r:\n%(failures)r\n" % vars() ) # analyze attribute: functions (the ones with __call__) get delayed, simple # attributes are looked up immediately iscallable = [hasattr(o.__getattribute__(attr), "__call__") for o in self] if all(iscallable): def runall(*args, **kwargs): """Apply function to all members and return a new Collection""" return Collection( [o.__getattribute__(attr)(*args, **kwargs) for o in self] ) runall.__name__ = attr return runall elif any(iscallable): raise TypeError( "Attribute {0!r} is callable only for some objects".format(attr) ) return Collection([o.__getattribute__(attr) for o in self]) def __add__(self, x): return Collection(super(Collection, self).__add__(x)) def __repr__(self): return self.__class__.__name__ + "({0!s})".format( super(Collection, self).__repr__() )