Source code for MDAnalysis.coordinates.xdrfile.TRR

# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# MDAnalysis --- http://mdanalysis.googlecode.com
# Copyright (c) 2006-2015 Naveen Michaud-Agrawal, Elizabeth J. Denning, Oliver Beckstein
# and contributors (see AUTHORS for the full list)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#


"""
Gromacs TRR trajectory I/O --- :mod:`MDAnalysis.coordinates.xdrfile.TRR`
========================================================================

Classes for reading and writing of `Gromacs TRR trajectories`_
together with supporting code.

.. note:: Users should access classes from :mod:`MDAnalysis.coordinates.TRR`.

.. _Gromacs TRR trajectories: http://www.gromacs.org/Documentation/File_Formats/.trr_File
.. _Gromacs: http://www.gromacs.org


.. SeeAlso:: :mod:`MDAnalysis.coordinates.xdrfile.libxdrfile2` for low-level
   bindings to the Gromacs trajectory file formats

Classes
-------

.. autoclass:: Timestep
   :members:
   :inherited-members:
.. autoclass:: TRRReader
   :members:
   :inherited-members:
.. autoclass:: TRRWriter
   :members:
   :inherited-members:

"""

import numpy

import core
import libxdrfile2
from MDAnalysis import NoDataError


[docs]class Timestep(core.Timestep): """Timestep for a Gromacs_ TRR trajectory. The Timestep can be initialized with *arg* being 1. an integer (the number of atoms) 2. another :class:`Timestep` instance, in which case a copy is made; attention: loss of attributes that do not exist within the TRR :class:`Timestep` may occur; 3. a :class:`numpy.ndarray` of shape ``(numatoms, 3)`` (for positions only) or ``(numatoms, 9)`` (for positions, velocities, and forces): ``positions = arg[:,:3]``, ``velocities = arg[:,3:6]``, and ``forces = arg[:,6:]``. The constructor also takes the named arguments *has_x*, *has_v*, and *has_f*, which are used to set the :class:`Timestep` flags :attr:`~Timestep.has_x`, :attr:`~Timestep.has_v`, and :attr:`~Timestep.has_f`, described below. Depending on the *arg* use-case above, the defaults set for these flags will vary: 1. when *arg* is an integer :attr:`~Timestep.has_x` defaults to ``True`` and :attr:`~Timestep.has_v` and :attr:`~Timestep.has_f` to ``False``. 2. when *arg* is another :class:`Timestep` instance the flags will default to being copied from the passed :class:`Timestep`. If that instance has no 'has_*' flags the behavior is to assign them to ``True`` depending on the existence of :attr:`~Timestep._velocities` and :attr:`~Timestep._forces` (:attr:`~Timestep._pos` is assumed to always be there, so in this case :attr:`~Timestep.has_x` defaults to ``True``). 3. when *arg* is a numpy array, the default flags will reflect what information is passed in the array. TRR :class:`Timestep` objects are now fully aware of the existence or not of coordinate/velocity/force information in frames, reflected in the :attr:`~Timestep.has_x`, :attr:`~Timestep.has_v`, and :attr:`~Timestep.has_f` flags. Accessing either kind of information while the corresponding flag is set to ``False`` wil raise a :exc:`NoDataError`. Internally, however, the arrays are always populated, even when the flags are ``False``; upon creation of a :class:`Timestep` they are zero-filled, but this might not always be the case later on for properties flagged as ``False`` if the same :class:`Timestep` instance is used to read from a TRR frame. When doing low-level writing to :attr:`~Timestep._pos`, :attr:`~Timestep._velocities`, or :attr:`~Timestep._forces:attr:, the corresponding flags must be set beforehand. The TRR :class:`Timestep` constructor allows for the named boolean arguments *has_x*, *has_v*, and *has_f* to be passed for automatic setting of the corresponding flag. An exception to this is assignment to the full property array thus:: ts = MDAnalysis.coordinates.TRR.Timestep(N) # N being the number of atoms ts._velocities = vel_array # Where vel_array is an existing array of shape (N, DIM) # This will also automatically set 'has_v' to True. Attempting to populate the array instead will, however, raise a NoDataError exception:: ts = MDAnalysis.coordinates.TRR.Timestep(N) # N being the number of atoms ts._velocities[:] = vel_array # This will fail if 'has_v' hasn't been set to True. .. versionchanged:: 0.8.0 TRR :class:`Timestep` objects are now fully aware of the existence or not of coordinate/velocity/force information in frames. .. _Gromacs: http://www.gromacs.org """ # The exception error _nodataerr = "You are accessing the %s of a Timestep but there are none in this frame. \ It might be the case that your trajectory doesn't have %s, or not every frame. In the latter case you can \ (1) get rid of %s-less frames before passing the trajectory to MDAnalysis, or \ (2) reflow your code to not access %s when they're not there, by making use of the '%s' flag of Timestep objects." def __init__(self, arg, **kwargs): DIM = libxdrfile2.DIM # compiled-in dimension (most likely 3) # These flags control which attributes were actually initialized in this Timestep. Due to the possibility of TRR # frames with no coords/vels/forces, and the way reading is handled, the _tpos, _tvelocities and _tforces # arrays # may end up containing data from different frames when the relevant attribute is missing in the trajectory. # This is flagged and exceptions can be raised whenever there is an attempt to read flagged data. if numpy.dtype(type(arg)) == numpy.dtype(int): self.has_x = kwargs.pop('has_x', True) self.has_v = kwargs.pop('has_v', False) self.has_f = kwargs.pop('has_f', False) self.frame = 0 self.numatoms = arg # C floats and C-order for arrays (see libxdrfile2.i) self._tpos = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') self._tvelocities = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') self._tforces = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') self._unitcell = numpy.zeros((DIM, DIM), dtype=numpy.float32) # additional data for xtc self.status = libxdrfile2.exdrOK self.step = 0 self.time = 0 self.lmbda = 0 elif isinstance(arg, Timestep): # Copy constructor # This makes a deepcopy of the timestep self.frame = arg.frame self.numatoms = arg.numatoms self._unitcell = numpy.array(arg._unitcell) # The 'has_' flags are set to True if the passed ts has them set to True OR if there is no such flag # and the attributes are there. This is overriden by the has_ named arguments passed to the constructor. # has_x is an exception because we can assume that, if the flag doesn't exist, the _pos attribute will # always be there. self.has_x = kwargs.pop('has_x', arg.__dict__.get("has_x", True)) self.has_v = kwargs.pop('has_v', arg.__dict__.get("has_v", hasattr(arg, "_velocities"))) self.has_f = kwargs.pop('has_f', arg.__dict__.get("has_f", hasattr(arg, "_forces"))) # We now either copy the properties or initialize a zeros array, depending on whether the passed ts # has the appropriate 'has_' flags set or has the attribute. If the attribute is there but the 'has_' # flag is False we initialize to zeros (it's just as slow and it makes no sense to copy data from # undefined behavior). # # Copies are done using the casted numpy.array function, which serves to early identify invalid input data. # COORDINATES: if self.has_x: try: self._tpos = numpy.array(arg._pos, dtype=numpy.float32) except ValueError as err: raise ValueError("Attempted to create a Timestep with invalid coordinate data: " + err.message) else: self._tpos = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') # VELOCITIES: if self.has_v: try: self._tvelocities = numpy.array(arg._velocities, dtype=numpy.float32) except ValueError as err: raise ValueError("Attempted to create a Timestep with invalid velocity data: " + err.message) except AttributeError: raise AttributeError( "Attempted to create a new Timestep from inconsistent Timestep data (the passed Timestep " "object has the 'has_v' flag set, but the '_velocities' attribute is missing.)") else: self._tvelocities = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') # FORCES: if self.has_f: try: self._tforces = numpy.array(arg._forces, dtype=numpy.float32) except ValueError as err: raise ValueError("Attempted to create a Timestep with invalid force data: " + err.message) except AttributeError: raise AttributeError( "Attempted to create a new Timestep from inconsistent Timestep data (the passed Timestep " "object has the 'has_f' flag set, but the '_forces' attribute is missing.)") else: self._tforces = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') ## for attr in ('status', 'step', 'time', 'lmbda', 'prec'): if hasattr(arg, attr): self.__setattr__(attr, arg.__getattribute__(attr)) elif isinstance(arg, numpy.ndarray): # provide packed array shape == (natoms, 3*DIM) # which contains pos = arg[:,0:3], v = arg[:,3:6], f = arg[:, 6:9] # or just positions: pos = arg[:,0:3] == arg # # The 'has_' flags are set from the named arguments passed to the constructor, # but default to what fields are available from the input array. self.has_x = kwargs.pop('has_x', True) self.has_v = kwargs.pop('has_v', arg.shape[1] == 3 * DIM) self.has_f = kwargs.pop('has_f', arg.shape[1] == 3 * DIM) if len(arg.shape) != 2: raise ValueError("packed numpy array (x,v,f) can only have 2 dimensions") self._unitcell = numpy.zeros((DIM, DIM), dtype=numpy.float32) self.frame = 0 if (arg.shape[0] == 3 * DIM and arg.shape[1] != 3 * DIM) or \ (arg.shape[0] == DIM and arg.shape[1] != DIM): # wrong order (but need to exclude case where natoms == DIM or natoms == 3*DIM!) raise ValueError("TRR timestep is to be initialized from an array with dimensions (natoms, " "3*%d) or (natoms, %d). \ If you wish to skip a property set the corresponding 'has_' flag to False on construction." % (DIM, DIM)) self.numatoms = arg.shape[0] if arg.shape[1] != DIM and arg.shape[1] != 3 * DIM: raise ValueError( "TRR timestep doesn't have second dimension %d or 3*%d: shape=%r" % (DIM, DIM, arg.shape)) # COORDINATES self._tpos = arg[:, 0:DIM].copy('C') # C-order # Velocities and forces, if the array seems to have them. # VELOCITIES if arg.shape[1] == DIM: self._tvelocities = arg[:, DIM:2 * DIM].copy('C') # C-order # FORCES self._tforces = arg[:, 2 * DIM:3 * DIM].copy('C') # C-order else: self._tvelocities = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') self._tforces = numpy.zeros((self.numatoms, DIM), dtype=numpy.float32, order='C') # additional data for trr self.status = libxdrfile2.exdrOK self.step = 0 self.time = 0 self.lmbda = 0 else: raise ValueError("Cannot create an empty Timestep") # The entire coordinate/velocity/force information is now hidden behind # existence-checking decorated functions. # COORDINATES @property def _pos(self): if self.has_x: return self._tpos else: raise NoDataError(self._nodataerr % ("coordinates", "coordinates", "coordinate", "coordinates", "has_x")) @_pos.setter def _pos(self, x): if x.shape == (self.numatoms, libxdrfile2.DIM): self._tpos = x self.has_x = True else: raise ValueError("You are attempting to set the positions array of a Timestep with an array \ that doesn't have the same number of atoms or the same number of dimensions. The Timestep \ has number-of-atoms,dimensions %r, and you supplied an array of shape %r." % ((self.numatoms, libxdrfile2.DIM), x.shape)) # @property def _x(self): return self._pos[:, 0] @_x.setter def _x(self, x): self._tpos[:, 0] = x @property def _y(self): return self._pos[:, 1] @_y.setter def _y(self, y): self._tpos[:, 1] = y @property def _z(self): return self._pos[:, 2] @_z.setter def _z(self, z): self._tpos[:, 2] = z # VELOCITIES @property def _velocities(self): if self.has_v: return self._tvelocities else: raise NoDataError(self._nodataerr % ("velocities", "velocities", "velocity", "velocities", "has_v")) @_velocities.setter def _velocities(self, v): if v.shape == (self.numatoms, libxdrfile2.DIM): self._tvelocities = v self.has_v = True else: raise ValueError("You are attempting to set the velocities array of a Timestep with an array \ that doesn't have the same number of atoms or the same number of dimensions. The Timestep \ has number-of-atoms,dimensions %r, and you supplied an array of shape %r." % ((self.numatoms, libxdrfile2.DIM), v.shape)) #FORCES @property def _forces(self): if self.has_f: return self._tforces else: raise NoDataError(self._nodataerr % ("forces", "forces", "force", "forces", "has_f")) @_forces.setter def _forces(self, f): if f.shape == (self.numatoms, libxdrfile2.DIM): self._tforces = f self.has_f = True else: raise ValueError("You are attempting to set the forces array of a Timestep with an array \ that doesn't have the same number of atoms or the same number of dimensions. The Timestep \ has number-of-atoms,dimensions %r, and you supplied an array of shape %r." % ((self.numatoms, libxdrfile2.DIM), f.shape))
[docs]class TRRWriter(core.TrjWriter): """Write a Gromacs_ TRR trajectory. .. _Gromacs: http://www.gromacs.org """ format = "TRR" units = {'time': 'ps', 'length': 'nm', 'velocity': 'nm/ps', 'force': 'kJ/(mol*nm)'}
[docs]class TRRReader(core.TrjReader): """Read a Gromacs_ TRR trajectory. .. versionchanged:: 0.8.0 :class:`Timestep` objects returned from TRR files now have :attr:`~Timestep.has_x`, :attr:`~Timestep.has_v`, and :attr:`~Timestep.has_f` flags reflecting whether coordinates/velocities/forces were read. Attempting to access such data when the corresponding flag is set to ``False`` will raise a :exc:`NoDataError`. .. _Gromacs: http://www.gromacs.org """ format = "TRR" _Timestep = Timestep _Writer = TRRWriter units = {'time': 'ps', 'length': 'nm', 'velocity': 'nm/ps', 'force': 'kJ/(mol*nm)'}