File: test_mdconvert.py

package info (click to toggle)
mdtraj 1.11.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 79,324 kB
  • sloc: python: 25,217; ansic: 6,266; cpp: 5,685; xml: 1,252; makefile: 192
file content (191 lines) | stat: -rw-r--r-- 7,270 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
##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Robert McGibbon
# Contributors:
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################


"""
Code to test the mdconvert script. These tests take about two minutes to run.

This checks all pairs for formats, converting from format x -> format y. it
also tries using striding to subsample the trajectory and atom_indices, so it
does significant integration testing of the XXXTrajectoryFile modules as well.
"""

import os
import subprocess
import sys

import numpy as np
import pytest

import mdtraj as md
from mdtraj.testing import eq

on_win = sys.platform == "win32"
on_py3 = sys.version_info >= (3, 0)


def test_index(h5traj):
    # Check that extracting a specific index works
    traj, in_fn, tmpdir = h5traj
    out_fn = f"{tmpdir}/frame4.pdb"
    subprocess.check_call(["mdconvert-mdtraj", in_fn, "-i", "4", "-o", out_fn])
    frame4 = md.load(out_fn)
    eq(frame4.xyz, traj[4].xyz)


def test_slice(h5traj):
    # Check that extracting a specific slice works
    traj, in_fn, tmpdir = h5traj
    out_fn = f"{tmpdir}/frame13.pdb"
    subprocess.check_call(["mdconvert-mdtraj", in_fn, "-i", "1:5:2", "-o", out_fn])
    frame13 = md.load(out_fn)
    eq(frame13.xyz, traj[1:5:2].xyz)


extensions = [
    "xtc",
    "dcd",
    "trr",
    "nc",
    "pdb",
    "h5",
    "lh5",
    "netcdf",
]


@pytest.fixture(params=extensions, ids=lambda x: "from-" + x)
def extension(request):
    if on_win and request.param == "lh5":
        pytest.skip("No lh5 on windows py3")
    return request.param


def test_pairwise(h5traj, extension, monkeypatch):
    """ensure that the xyz coordinates are preserved by a trip
    from python -> save in format X -> mdconvert to format Y -> python
    """

    def test_base(h5traj, extension, monkeypatch):
        traj, _, tmpdir = h5traj
        ext1 = extension

        # save one copy of traj for use as a topology file
        topology_fn = f"{tmpdir}/topology.pdb"
        traj[0].save(topology_fn)

        # save a .dat file for the atom_indices so that we can test
        # mdconvert's atom_indices flag
        atom_indices = np.array([0, 3])
        atom_indices_fn = f"{tmpdir}/atom_indices.dat"
        np.savetxt(atom_indices_fn, atom_indices, fmt="%d")

        in_fn = f"{tmpdir}/traj.{ext1}"
        traj.save(in_fn)
        working_dir = f"{tmpdir}/from-{ext1}"
        os.mkdir(working_dir)

        for ext2 in extensions:
            out_fn = f"traj.{ext2}"

            command1 = ["mdconvert-mdtraj", in_fn, "-o", out_fn, "-c 6"]
            if ext2 in ["pdb", "h5", "lh5"]:
                # if we're saving a pdb or h5, we need to give it a topology too
                command1 += ["-t", topology_fn]

            # TODO: test fixture
            subprocess.check_call(command1, cwd=working_dir)

            # Use the --atom_indices flag to mdconvert
            command2 = command1 + ["-a", atom_indices_fn]
            command2[3] = "subset." + out_fn  # make sure the output goes to a different file
            subprocess.check_call(command2, cwd=working_dir)

            # Use the --stride 3 flag
            command3 = command1 + ["-s 3"]
            command3[3] = "stride." + out_fn  # change the out filename, so they don't clobbed
            subprocess.check_call(command3, cwd=working_dir)

            # ensure that the xyz coordinates are preserved by a trip
            # from python -> save in format X -> mdconvert to format Y -> python
            load_kwargs_check1 = {}
            load_kwargs_check2 = {}
            if ext2 not in ["pdb", "h5", "lh5"]:
                load_kwargs_check1["top"] = traj.topology
                load_kwargs_check2["top"] = traj.topology.subset(atom_indices)

            out1 = md.load(os.path.join(working_dir, out_fn), **load_kwargs_check1)
            out2 = md.load(
                os.path.join(working_dir, "subset." + out_fn),
                **load_kwargs_check2,
            )
            out3 = md.load(
                os.path.join(working_dir, "stride." + out_fn),
                **load_kwargs_check1,
            )

            if ext1 in ["lh5"] or ext2 in ["lh5"]:
                decimal = 3
            else:
                decimal = 6
            eq(out1.xyz, traj.xyz, decimal=decimal)
            eq(out2.xyz, traj.xyz[:, atom_indices], decimal=decimal)
            eq(out3.xyz, traj.xyz[::3], decimal=decimal)

            if ext1 not in ["lh5"] and ext2 not in ["lh5"]:
                # binpos doesn't save unitcell information
                eq(out1.unitcell_vectors, traj.unitcell_vectors, decimal=2)
                eq(out2.unitcell_vectors, traj.unitcell_vectors, decimal=2)
                eq(out3.unitcell_vectors, traj.unitcell_vectors[::3], decimal=2)

            if all(e in ["xtc", "trr", "nc", "h5"] for e in [ext1, ext2]):
                # these formats contain time information
                if all(e in ["nc"] for e in [ext1, ext2]):
                    with monkeypatch.context() as m:
                        m.setitem(sys.modules, "netCDF4", None)
                        eq(out1.time, traj.time)
                        eq(out2.time, traj.time)
                        eq(out3.time, traj.time[::3])

                eq(out1.time, traj.time)
                eq(out2.time, traj.time)
                eq(out3.time, traj.time[::3])

            if ext2 in ["pdb", "h5", "lh5"]:
                # these formats contain a topology in the file that was
                # read from disk
                eq(out1.topology, traj.topology)
                eq(out2.topology, traj.topology.subset(atom_indices))
                eq(out3.topology, traj.topology)

    if extension in ("nc"):
        # If extension is nc, we need to test with both SciPy and NetCDF4.
        # First, we use pytest.monkeypatch to remove the netCDF4 import from the environment
        # Then, we make sure the extension for both SciPy and NetCDF4 tests are different (or
        # there will be SameFileError/FileExistsError failures)
        # All these changes will be reverted outside the context manager, and the netCDF4 test will run again
        with monkeypatch.context() as m:
            m.setitem(sys.modules, "netCDF4", None)
            test_base(h5traj, ".scipy.nc", monkeypatch)

    # For testing most formats and with netCDF4 (if format is nc)
    test_base(h5traj, extension, monkeypatch)