Source code for MDAnalysis.coordinates.TNG

# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# 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
#

"""
TNG trajectory files --- :mod:`MDAnalysis.coordinates.TNG`
==========================================================


The TNG format :footcite:p:`Lundborg2014` is a format used in `GROMACS`_ for
storage of trajectory and topology information. The TNG format allows a wide
range of compression algorithms and unlike the compressed XTC format can also
store velocities and forces in addition to positions.

The classes in this module are based on the `pytng`_ package for reading TNG
files. The reader is directed to the `pytng documentation`_ for further reading
about how pytng works under the hood. Pytng is an optional dependency and must
be installed to use this Reader. Use of the reader without pytng installed will
raise an `ImportError`. The variable `HAS_PYTNG` indicates whether this module
has pytng availble.

In addition to particle-dependent trajectory information like positions,
forces and velocities, the TNG format can store trajectory metadata and
other arbitrary time dependent data. Additional information can range from
the virial and pressure components to the molecular topology of the system.
This is enabled by a block based system in which binary flags indicate the
presence or absence of various data blocks. The structure of a TNG file is
provided  in the `TNG specification`_. The TNG paper :footcite:p:`Lundborg2014` and
the `pytng documentation`_ are also good resources. The user is encouraged to
read the full list of `TNG blocks`_ to understand the full power of the TNG
format.



Current Limitations
-------------------
Currently there is only a Reader for TNG files and no Writer. This will depend
on upstream changes in pytng. Additionally, it is not currently possible to
read the molecular topology from a TNG file.

There are also some limitations to reading TNG files with pytng.
Currently we do not allow data to be read *off stride*. In essence this
means that all of the critical trajectory data (positions, box, velocities
(if present), forces (if present)) must share the same stride in trajectory
integrator steps. These critical blocks in the TNG file are henceforth called
*special blocks*. Optional blocks (all blocks that are not special blocks)
will not be read if they do not share an integrator step with, or are not
divisible by the shared integrator step of the special blocks.


References
----------

.. footbibliography::


.. Links

.. _GROMACS:
    https://www.gromacs.org/
.. _pytng:
    https://github.com/MDAnalysis/pytng
.. _pytng documentation:
    https://www.mdanalysis.org/pytng/
.. _TNG blocks:
    https://www.mdanalysis.org/pytng/documentation_pages/Blocks.html
.. _TNG specification:
    https://gitlab.com/hugomacdermott/tng/-/blob/master/Trajectoryformatspecification.mk

"""

import warnings
from typing import List, Optional

import numpy as np
from MDAnalysis.coordinates import base
from MDAnalysis.coordinates.timestep import Timestep
from MDAnalysis.lib.mdamath import triclinic_box
from MDAnalysis.lib.util import store_init_arguments

from ..due import Doi, due

try:
    import pytng
except ImportError:
    # Indicates whether pytng is found.
    HAS_PYTNG = False
else:
    # Indicates whether pytng is found.
    HAS_PYTNG = True


[docs] class TNGReader(base.ReaderBase): r"""Reader for the TNG format The TNG format :footcite:p:`Lundborg2014` is a format used in `GROMACS`_ for storage of trajectory and topology information. This reader is are based on the `pytng`_ package for reading TNG files. The reader is directed to the `pytng documentation`_ and `TNG specification`_ for further reading. The TNG format allows a wide range of compression algorithms and unlike the compressed XTC format can also store velocities and forces in addition to positions. The contents of the *special blocks* (positions, box, velocities, forces) are read into the timestep if present. Additional blocks are read into the `ts.data` dictionary if they are available at the current frame. .. versionadded:: 2.4.0 """ format = "TNG" # NOTE: Time units are in seconds unlike other GROMACS formats units = { "time": "second", "length": "nm", "velocity": "nm/ps", "force": "kJ/(mol*nm)", } _box_blockname = "TNG_TRAJ_BOX_SHAPE" _positions_blockname = "TNG_TRAJ_POSITIONS" _velocities_blockname = "TNG_TRAJ_VELOCITIES" _forces_blockname = "TNG_TRAJ_FORCES" _special_blocks = [ _box_blockname, _positions_blockname, _velocities_blockname, _forces_blockname, ] @due.dcite(Doi("10.1002/jcc.23495"), description="The TNG paper", path=__name__) @store_init_arguments def __init__(self, filename: str, convert_units: bool = True, **kwargs): """Initialize a TNG trajectory Parameters ---------- filename : str filename of the trajectory convert_units : bool (optional) convert into MDAnalysis units """ if not HAS_PYTNG: raise ImportError("TNGReader: To read TNG files please install pytng") super(TNGReader, self).__init__(filename, **kwargs) self.filename = filename self.convert_units = convert_units self._file_iterator = pytng.TNGFileIterator(self.filename, "r") self.n_atoms = self._file_iterator.n_atoms self._n_steps = self._file_iterator.n_steps # names of the blocks self._block_names = list(self._file_iterator.block_ids.keys()) # block ids, dict of C long long self._block_dictionary = self._file_iterator.block_ids self._block_strides = self._file_iterator.block_strides self._data_frames = self._file_iterator.n_data_frames # init all special blocks to not present self._special_block_present = {k: False for k in self._special_blocks} self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs) # check for all the special blocks self._has_box = self._box_blockname in self._block_names if self._has_box: self._special_block_present[self._box_blockname] = True self.ts.dimensions = np.zeros(6, dtype=np.float32) self._has_positions = self._positions_blockname in self._block_names if self._has_positions: self._special_block_present[self._positions_blockname] = True self.ts.positions = self._file_iterator.make_ndarray_for_block_from_name( self._positions_blockname ) self._has_velocities = self._velocities_blockname in self._block_names if self._has_velocities: self._special_block_present[self._velocities_blockname] = True self.ts.velocities = self._file_iterator.make_ndarray_for_block_from_name( self._velocities_blockname ) self._has_forces = self._forces_blockname in self._block_names if self._has_forces: self._special_block_present[self._forces_blockname] = True self.ts.forces = self._file_iterator.make_ndarray_for_block_from_name( self._forces_blockname ) # check for any additional blocks that will be read into ts.data self._additional_blocks = [ block for block in self._block_names if block not in self._special_blocks ] self._check_strides_and_frames() self._frame = 0 # box needs a temporary place to be stored as ts.dimensions is # wrong shape initially self._box_temp = self._file_iterator.make_ndarray_for_block_from_name( self._box_blockname ) self._frame = -1 self._read_next_timestep() def _check_strides_and_frames(self): """ Check that the strides and frame numbers of the blocks in the TNG file match up so that the file can be iterated over retrieving data at each integrator step """ strides = [] n_data_frames = [] for block in self.special_blocks: stride = self._block_strides[block] strides.append(stride) n_data_frame = self._data_frames[block] n_data_frames.append(n_data_frame) strides_eq = all(v == strides[0] for v in strides) frames_eq = all(v == n_data_frames[0] for v in n_data_frames) if (not strides_eq) or (not frames_eq): raise IOError( "Strides of TNG special blocks not equal, file cannot be read" ) self._global_stride = strides[0] # NOTE frame number is 0 indexed so increment self._n_frames = n_data_frames[0] + 1 self._additional_blocks_to_read = [] for block in self._additional_blocks: stride_add = self._block_strides[block] if stride_add != self._global_stride: if stride_add % self._global_stride: # pragma: no cover warnings.warn( f"TNG additional block {block} does not match" " strides of other blocks and is not" " divisible by the global stride_length." " It will not be read" ) else: self._additional_blocks_to_read.append(block) # pragma: no cover else: self._additional_blocks_to_read.append(block)
[docs] def close(self): """close the reader""" self._file_iterator._close()
[docs] @staticmethod def parse_n_atoms(filename: str) -> int: """parse the number of atoms in the TNG trajectory file Parameters ---------- filename : str The name of the TNG file Returns ------- n_atoms : int The number of atoms in the TNG file """ if not HAS_PYTNG: raise ImportError("TNGReader: To read TNG files please install pytng") with pytng.TNGFileIterator(filename, "r") as tng: n_atoms = tng.n_atoms return n_atoms
@property def n_frames(self) -> int: """number of frames in trajectory Returns ------- n_frames : int The number of data containing frames in the trajectory """ return self._n_frames @property def blocks(self) -> List[str]: """list of the blocks that are in the file Returns ------- blocks : list The block names present in the TNG trajectory """ return self._block_names @property def special_blocks(self) -> List[str]: """list of the special blocks that are in the file Returns ------- special_blocks : list The special block names (positions, box, velocities, forces) present in the TNG trajectory """ return [k for k, v in self._special_block_present.items() if v] @property def additional_blocks(self) -> List[str]: """list of the additional (non-special) blocks that are being read from the trajectory. This may be exclude some blocks present in the file if they do not fall on the same trajectory stride as the positions and velocities. Returns ------- additional_blocks : list The additional block names in the TNG trajectory """ return self._additional_blocks_to_read def _reopen(self): """reopen the trajectory""" self.ts.frame = 0 self._frame = -1 self._file_iterator._close() self._file_iterator._open(self.filename, "r") def _frame_to_step(self, frame: int) -> int: """Convert a frame number to an integrator step Parameters ---------- frame : int The frame number Returns ------- integrators_step : int The integrator step of the frame in the TNG file """ return frame * self._global_stride def _read_frame(self, i: int) -> Timestep: """read frame i Parameters ---------- i : int The trajectory frame to be read Returns ------- ts : Timestep Data from frame i encapsulated in an MDA `:class:Timestep` """ self._frame = i - 1 ts = self._read_next_timestep() return ts def _read_next_timestep(self, ts: Optional[Timestep] = None) -> Timestep: """Read next frame into a timestep Parameters ---------- ts : Timestep The timestep to read the data into Returns ------- ts : Timestep The timestep filled with data from the next step """ if self._frame == self.n_frames - 1: raise IOError("trying to go over trajectory limit") if ts is None: ts = self.ts # convert from frames to integrator steps step = self._frame_to_step(self._frame + 1) iterator_step = self._file_iterator.read_step(step) self._frame += 1 ts = self._frame_to_ts(iterator_step, ts) return ts def _frame_to_ts( self, curr_step: "pytng.TNGCurrentIntegratorStep", ts: Timestep ) -> Timestep: """convert a TNGCurrentIteratorStep to an MDA Timestep Parameters ---------- curr_step : pytng.TNGCurrentIntegratorStep The current timestep in the TNG trajectory ts : Timestep The current timestep in the MDA trajectory Returns ------- ts: Timestep The current timestep in the MDA trajectory with data from TNG trajectory integrator step Raises ------ IOError The reading of one of the attributes from the TNG file failed """ ts.frame = self._frame time = curr_step.get_time() if self.convert_units: time = self.convert_time_from_native(time) ts.time = time ts.data["step"] = curr_step.step if self._has_box: curr_step.get_box(self._box_temp) ts.dimensions = triclinic_box(*self._box_temp.reshape(3, 3)) if not curr_step.read_success: raise IOError("Failed to read box from TNG file") if self.convert_units: self.convert_pos_from_native(ts.dimensions[:3]) if self._has_positions: curr_step.get_positions(ts.positions) if not curr_step.read_success: raise IOError("Failed to read positions from TNG file") if self.convert_units: self.convert_pos_from_native(ts.positions) if self._has_velocities: curr_step.get_velocities(ts.velocities) if not curr_step.read_success: raise IOError("Failed to read velocities from TNG file") if self.convert_units: self.convert_velocities_from_native(ts.velocities) if self._has_forces: curr_step.get_forces(ts.forces) if not curr_step.read_success: raise IOError("Failed to read forces from TNG file") if self.convert_units: self.convert_forces_from_native(self.ts.forces) for block in self._additional_blocks_to_read: add_block_stride = self._block_strides[block] # check we are on stride for our block if not (add_block_stride % self._global_stride): block_data = self._file_iterator.make_ndarray_for_block_from_name(block) # additional blocks read into ts.data dictionary ts.data[block] = curr_step.get_blockid( self._block_dictionary[block], block_data ) if not curr_step.read_success: raise IOError( f"Failed to read additional block {block} from TNG file" ) return ts def __getstate__(self): """Make a dictionary of the class state to pickle Reader instance. Must be done manually as pytng uses a non-trivial`__cinit__`. """ state = self.__dict__.copy() # cant have PyTNG file iterator in as it is non-pickable del state["_file_iterator"] return state def __setstate__(self, state): """Restore class from `state` dictionary in unpickling of Reader instance """ self.__dict__ = state # reconstruct file iterator self._file_iterator = pytng.TNGFileIterator(self.filename, "r") # unlike self._read_frame(self._frame), # the following lines update the state of the C-level file iterator # without updating the ts object. # This is necessary to preserve the modification, # e.g. changing coordinates, in the ts object. # see PR #3722 for more details. step = self._frame_to_step(self._frame) _ = self._file_iterator.read_step(step)
[docs] def Writer(self): """Writer for TNG files Raises ------ NotImplementedError Currently there is not writer for TNG files pending implementation upstream in pytng. """ raise NotImplementedError("There is currently no writer for TNG files")