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://www.MDAnalysis.org
# 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 as np
import errno

from . import statno
from . import core
from . import libxdrfile2
from MDAnalysis import NoDataError


[docs]class Timestep(core.Timestep): """Timestep for a Gromacs_ TRR trajectory. TRR Timestep always has positions, velocities and forces allocated, meaning that the `_pos`, `_velocities` and `_forces` attributes will always exist. Whether or not this data is valid for the current frame is determined by the flags `has_positions`, `has_velocities` and `has_forces`. These are controlled by the entries in the data dictionary: `position_source`, `velocity_source` and `force_source`. When accessing the `.positions` attribute, this will only be returned if `position_source` matches the current `frame`, otherwise a :class:`~MDAnalysis.NoDataError` will be raised. This scheme applies to both velocities and forces. When writing to the Timestep, it is important that the :attr:`~Timestep.frame` attribute is updated first! This is because the source of the data is taken as the current frame of the Timestep. For example:: ts = MDAnalysis.coordinates.TRR.Timestep(N) # N being the number of atoms ts.frame = 0 # The first frame ts.velocities = vel_array # Where vel_array is an existing array of shape (N, DIM) # This will also automatically set the `has_velocities` flag # And the velocity source will be "0" Attempting to populate the array before setting the frame instead will cause the velocity data to be considered "out of date":: ts = MDAnalysis.coordinates.TRR.Timestep(N) # N being the number of atoms ts.velocities = vel_array ts.frame = 0 # This frame number doesn't match the frame number of when # velocities was populated .. versionchanged:: 0.8.0 TRR :class:`Timestep` objects are now fully aware of the existence or not of coordinate/velocity/force information in frames. .. versionchanged:: 0.11.0 Velocities and forces can now be directly accessed via the `velocities` and `forces` attributes, with a :class:`~MDAnalysis.NoDataError` being raised if no data was given for that frame. Management redone to use `has_positions`, `has_velocities` and `has_forces`. Setting these will update the data dictionary entries which track the data. Timestep lmbda now stored in the data dictionary .. _Gromacs: http://www.gromacs.org """ # The exception error _nodataerr = ("You are accessing the {attrs} of a Timestep but there are none in" " this frame. It might be the case that your trajectory doesn't" " have {attrs}, or not every frame. In the latter case you can " "(1) get rid of {attr}-less frames before passing the trajectory" " to MDAnalysis, or " "(2) reflow your code to not access {attrs} when they're not there," " by making use of the '{flag}' flag of Timestep objects.") def __init__(self, n_atoms, **kwargs): kwargs.update(positions=True, velocities=True, forces=True) super(Timestep, self).__init__(n_atoms, **kwargs) # Set initial sources to None, so they are never valid # _source is compared against current frame when accessing # if they do not match, then the Timestep returns _nodataerr self.data['position_source'] = None self.data['velocity_source'] = None self.data['force_source'] = None # TRR always has pos vel & force allocated self._pos = np.zeros((self.n_atoms, 3), dtype=np.float32, order=self.order) self._velocities = np.zeros((self.n_atoms, 3), dtype=np.float32, order=self.order) self._forces = np.zeros((self.n_atoms, 3), dtype=np.float32, order=self.order) self.data['lmbda'] = 0 @property def has_positions(self): # When setting position the frame is recorded # if frame gets changed (by reading next frame) # then this data is "out of date", so has_positions # is False return self.data['position_source'] == self.frame @has_positions.setter def has_positions(self, val): # If val evaluates to True, say that position data # comes from this frame # If val evaluates to False, say that position data # comes from -1 (ie. never) self.data['position_source'] = self.frame if val else None @property def positions(self): # If the position info came from this frame, yield it # else play dumb if self.has_positions: return self._pos else: raise NoDataError(self._nodataerr.format( attr='position', attrs='positions', flag='has_positions')) @positions.setter def positions(self, new): self._pos[:] = new self.has_positions = True @property def has_velocities(self): return self.data['velocity_source'] == self.frame @has_velocities.setter def has_velocities(self, val): self.data['velocity_source'] = self.frame if val else None @property def velocities(self): if self.has_velocities: return self._velocities else: raise NoDataError(self._nodataerr.format( attr='velocity', attrs='velocities', flag='has_velocities')) @velocities.setter def velocities(self, new): self._velocities[:] = new self.has_velocities = True @property def has_forces(self): return self.data['force_source'] == self.frame @has_forces.setter def has_forces(self, val): self.data['force_source'] = self.frame if val else None @property def forces(self): if self.has_forces: return self._forces else: raise NoDataError(self._nodataerr.format( attr='force', attrs='forces', flag='has_forces')) @forces.setter def forces(self, new): self._forces[:] = new self.has_forces = True
[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_velocities`, and :attr:`~Timestep.has_forces` 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`. .. versionchanged:: 0.11.0 Returned Timesteps will always be a :class:`Timestep` instance. Attributes lmbda and status now stored in the TS.data dictionary .. _Gromacs: http://www.gromacs.org """ format = "TRR" _Timestep = Timestep _Writer = TRRWriter units = {'time': 'ps', 'length': 'nm', 'velocity': 'nm/ps', 'force': 'kJ/(mol*nm)'} def _allocate_sub(self, DIM): self._pos_buf = np.zeros((self._trr_n_atoms, DIM), dtype=np.float32, order='C') self._velocities_buf = np.zeros((self._trr_n_atoms, DIM), dtype=np.float32, order='C') self._forces_buf = np.zeros((self._trr_n_atoms, DIM), dtype=np.float32, order='C') def _read_trj_natoms(self, filename): return libxdrfile2.read_trr_natoms(filename) def _read_trj_n_frames(self, filename): self._n_frames, self._offsets = libxdrfile2.read_trr_n_frames(filename) self._store_offsets() def _read_next_timestep(self, ts=None): if ts is None: ts = self.ts elif not isinstance(ts, Timestep): # Requires using a TRR Timestep ts = Timestep.from_timestep(ts) if self.xdrfile is None: self.open_trajectory() if self._sub is None: ts.data['status'], ts._frame, ts.time, ts.data['lmbda'],\ has_x, has_v, has_f = libxdrfile2.read_trr( self.xdrfile, ts._unitcell, ts._pos, ts._velocities, ts._forces) else: ts.data['status'], ts._frame, ts.time, ts.data['lmbda'],\ has_x, has_v, has_f = libxdrfile2.read_trr( self.xdrfile, ts._unitcell, self._pos_buf, self._velocities_buf, self._forces_buf) ts._pos[:] = self._pos_buf[self._sub] ts._velocities[:] = self._velocities_buf[self._sub] ts._forces[:] = self._forces_buf[self._sub] # Updated frame number first! ts.frame += 1 # Update the sources of data if has_x: ts.has_positions = True if has_v: ts.has_velocities = True if has_f: ts.has_forces = True if ((ts.data['status'] == libxdrfile2.exdrENDOFFILE) or (ts.data['status'] == libxdrfile2.exdrINT)): # seems that trr files can get a exdrINT when reaching EOF (??) raise IOError(errno.EIO, "End of file reached for {0} file".format(self.format), self.filename) elif not ts.data['status'] == libxdrfile2.exdrOK: raise IOError(errno.EBADF,("Problem with {0} file, status {1}" "".format((self.format, statno.ERRORCODE[ts.data['status']]), self.filename))) if self.convert_units: # TRRs have the annoying possibility of frames without coordinates/velocities/forces... if ts.has_positions: self.convert_pos_from_native(ts._pos) # in-place ! self.convert_pos_from_native(ts._unitcell) # in-place ! (note: trr contain unit vecs!) ts.time = self.convert_time_from_native(ts.time) # in-place does not work with scalars if ts.has_velocities: self.convert_velocities_from_native(ts._velocities) # in-place if ts.has_forces: self.convert_forces_from_native(ts._forces) # in-place return ts