File: TRR.py

package info (click to toggle)
mdanalysis 2.10.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 116,696 kB
  • sloc: python: 92,135; ansic: 8,156; makefile: 215; sh: 138
file content (237 lines) | stat: -rw-r--r-- 7,853 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
# -*- 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 Lesser GNU Public Licence, v2.1 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
#
"""
TRR trajectory files --- :mod:`MDAnalysis.coordinates.TRR`
==========================================================

Read and write GROMACS TRR trajectories.

See Also
--------
MDAnalysis.coordinates.XTC: Read and write GROMACS XTC trajectory files.
MDAnalysis.coordinates.XDR: BaseReader/Writer for XDR based formats
"""
import errno
from . import base
from .XDR import XDRBaseReader, XDRBaseWriter
from ..lib.formats.libmdaxdr import TRRFile
from ..lib.mdamath import triclinic_vectors, triclinic_box


class TRRWriter(XDRBaseWriter):
    """Writer for the Gromacs TRR format.

    The Gromacs TRR trajectory format is a lossless format. The TRR format can
    store *velocities* and *forces* in addition to the coordinates. It is also
    used by other Gromacs tools to store and process other data such as modes
    from a principal component analysis.

    If the data dictionary of a :class:`Timestep` contains the key
    'lambda' the corresponding value will be used as the lambda value
    for written TRR file.  If ``None`` is found the lambda is set to 0.

    If the data dictionary of a :class:`Timestep` contains the key
    'step' the corresponding value will be used as the step value for
    the written TRR file. If the dictionary does not contain 'step', then
    the step is set to the :class:`Timestep` frame attribute.

    """

    format = "TRR"
    multiframe = True
    units = {
        "time": "ps",
        "length": "nm",
        "velocity": "nm/ps",
        "force": "kJ/(mol*nm)",
    }
    _file = TRRFile

    def _write_next_frame(self, ag):
        """Write information associated with ``ag`` at current frame into trajectory

        Parameters
        ----------
        ag : AtomGroup or Universe

        See Also
        --------
        <FormatWriter>.write(AtomGroup/Universe/TimeStep)
        The normal write() method takes a more general input


        .. versionchanged:: 1.0.0
           Renamed from `write_next_timestep` to `_write_next_frame`.
        .. versionchanged:: 2.0.0
           Deprecated support for Timestep argument has now been removed.
           Use AtomGroup or Universe as an input instead.
        .. versionchanged:: 2.1.0
           When possible, TRRWriter assigns `ts.data['step']` to `step` rather
           than `ts.frame`.
        """
        try:
            ts = ag.ts
        except AttributeError:
            try:
                # special case: can supply a Universe, too...
                ts = ag.trajectory.ts
            except AttributeError:
                errmsg = "Input obj is neither an AtomGroup or Universe"
                raise TypeError(errmsg) from None

        xyz = None
        if ts.has_positions:
            xyz = ts.positions.copy()
            if self._convert_units:
                self.convert_pos_to_native(xyz)

        velo = None
        if ts.has_velocities:
            velo = ts.velocities.copy()
            if self._convert_units:
                self.convert_velocities_to_native(velo)

        forces = None
        if ts.has_forces:
            forces = ts.forces.copy()
            if self._convert_units:
                self.convert_forces_to_native(forces)

        if self._dt is None:
            time = ts.time
        else:
            time = self._dt * ts.frame
        step = ts.data.get("step", ts.frame)

        if self._convert_units:
            dimensions = self.convert_dimensions_to_unitcell(ts, inplace=False)

        box = triclinic_vectors(dimensions)

        lmbda = 0
        if "lambda" in ts.data:
            lmbda = ts.data["lambda"]

        self._xdr.write(
            xyz, velo, forces, box, step, time, lmbda, self.n_atoms
        )


class TRRReader(XDRBaseReader):
    """Reader for the Gromacs TRR format.

    The Gromacs TRR trajectory format is a lossless format. The TRR format can
    store *velocities* and *forces* in addition to the coordinates. It is also
    used by other Gromacs tools to store and process other data such as modes
    from a principal component analysis.

    The lambda value is written in the data dictionary of the returned
    :class:`Timestep`

    Notes
    -----
    See :ref:`Notes on offsets <offsets-label>` for more information about
    offsets.

    """

    format = "TRR"
    units = {
        "time": "ps",
        "length": "nm",
        "velocity": "nm/ps",
        "force": "kJ/(mol*nm)",
    }
    _writer = TRRWriter
    _file = TRRFile

    def _read_next_timestep(self, ts=None):
        """copy next frame into timestep

        versionadded:: 2.4.0
            TRRReader implements this method so that it can use
            read_direct_xvf to read the data directly into the timestep
            rather than copying it from a temporary array.
        """
        if self._frame == self.n_frames - 1:
            raise IOError(errno.EIO, "trying to go over trajectory limit")
        if ts is None:
            ts = self.ts
        # allocate arrays to read into, will set to proper values
        # in _frame_to_ts
        ts.has_positions = True
        ts.has_velocities = True
        ts.has_forces = True
        frame = self._xdr.read_direct_xvf(
            ts.positions, ts.velocities, ts.forces
        )
        self._frame += 1
        self._frame_to_ts(frame, ts)
        return ts

    def _frame_to_ts(self, frame, ts):
        """convert a trr-frame to a mda TimeStep"""
        dt = self._kwargs["dt"]
        if dt is None:
            ts.time = frame.time
        else:
            ts.time = self._frame * dt
        ts.frame = self._frame
        ts.data["step"] = frame.step

        ts.has_positions = frame.hasx
        ts.has_velocities = frame.hasv
        ts.has_forces = frame.hasf
        ts.dimensions = triclinic_box(*frame.box)

        if self.convert_units:
            if ts.dimensions is not None:
                self.convert_pos_from_native(ts.dimensions[:3])

        if ts.has_positions:
            if self._sub is not None:
                ts.positions = frame.x[self._sub]
            else:
                ts.positions = frame.x
            if self.convert_units:
                self.convert_pos_from_native(ts.positions)

        if ts.has_velocities:
            if self._sub is not None:
                ts.velocities = frame.v[self._sub]
            else:
                ts.velocities = frame.v
            if self.convert_units:
                self.convert_velocities_from_native(ts.velocities)

        if ts.has_forces:
            if self._sub is not None:
                ts.forces = frame.f[self._sub]
            else:
                ts.forces = frame.f
            if self.convert_units:
                self.convert_forces_from_native(ts.forces)

        ts.data["lambda"] = frame.lmbda

        return ts