# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*- # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8 # # 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 # import os import platform import sys import warnings import MDAnalysis as mda import numpy as np import pytest from MDAnalysis.transformations import translate from numpy.testing import assert_almost_equal, assert_equal from MDAnalysisTests.datafiles import ( CRD, DCD, GRO, PDB, PSF, TRR, XTC, LAMMPS_chain, LAMMPSDUMP_chain1, LAMMPSDUMP_chain2, PDB_closed, PDB_small, ) from MDAnalysisTests.util import no_warning class TestChainReader(object): prec = 3 @pytest.fixture() def universe(self): return mda.Universe(PSF, [DCD, CRD, DCD, CRD, DCD, CRD, CRD]) @pytest.fixture() def transformed(ref): return mda.Universe( PSF, [DCD, CRD, DCD, CRD, DCD, CRD, CRD], transformations=[translate([10, 10, 10])], ) def test_regular_repr(self): u = mda.Universe(PSF, [DCD, CRD, DCD]) assert_equal( "", u.trajectory.__repr__(), ) def test_truncated_repr(self, universe): assert_equal( "", universe.trajectory.__repr__(), ) def test_next_trajectory(self, universe): universe.trajectory.rewind() universe.trajectory.next() assert_equal(universe.trajectory.ts.frame, 1, "loading frame 2") def test_n_atoms(self, universe): assert_equal( universe.trajectory.n_atoms, 3341, "wrong number of atoms" ) def test_n_frames(self, universe): assert_equal( universe.trajectory.n_frames, 3 * 98 + 4, "wrong number of frames in chained dcd", ) def test_iteration(self, universe): for ts in universe.trajectory: pass # just forward to last frame assert_equal( universe.trajectory.n_frames - 1, ts.frame, "iteration yielded wrong number of frames ({0:d}), " "should be {1:d}".format(ts.frame, universe.trajectory.n_frames), ) def test_jump_lastframe_trajectory(self, universe): universe.trajectory[-1] assert_equal( universe.trajectory.ts.frame + 1, universe.trajectory.n_frames, "indexing last frame with trajectory[-1]", ) def test_slice_trajectory(self, universe): frames = [ts.frame for ts in universe.trajectory[5:17:3]] assert_equal(frames, [5, 8, 11, 14], "slicing dcd [5:17:3]") def test_full_slice(self, universe): trj_iter = universe.trajectory[:] frames = [ts.frame for ts in trj_iter] assert_equal(frames, np.arange(universe.trajectory.n_frames)) def test_frame_numbering(self, universe): universe.trajectory[98] # index is 0-based and frames are 0-based assert_equal(universe.trajectory.frame, 98, "wrong frame number") def test_next_after_frame_numbering(self, universe): universe.trajectory[98] # index is 0-based and frames are 0-based universe.trajectory.next() assert_equal(universe.trajectory.frame, 99, "wrong frame number") def test_frame(self, universe): universe.trajectory[0] coord0 = universe.atoms.positions.copy() # forward to frame where we repeat original dcd again: # dcd:0..97 crd:98 dcd:99..196 universe.trajectory[99] assert_equal( universe.atoms.positions, coord0, "coordinates at frame 1 and 100 should be the same!", ) def test_time(self, universe): universe.trajectory[30] # index and frames 0-based assert_almost_equal( universe.trajectory.time, 30.0, 5, err_msg="Wrong time of frame" ) def test_write_dcd(self, universe, tmpdir): """test that ChainReader written dcd (containing crds) is correct (Issue 81)""" outfile = str(tmpdir) + "chain-reader.dcd" with mda.Writer(outfile, universe.atoms.n_atoms) as W: for ts in universe.trajectory: W.write(universe) universe.trajectory.rewind() u = mda.Universe(PSF, outfile) for ts_orig, ts_new in zip(universe.trajectory, u.trajectory): assert_almost_equal( ts_orig._pos, ts_new._pos, self.prec, err_msg="Coordinates disagree at frame {0:d}".format( ts_orig.frame ), ) def test_transform_iteration(self, universe, transformed): vector = np.float32([10, 10, 10]) # # Are the transformations applied and # are the coordinates "overtransformed"? # iterate once: for ts in transformed.trajectory: frame = ts.frame ref = universe.trajectory[frame].positions + vector assert_almost_equal(ts.positions, ref, decimal=6) # iterate again: for ts in transformed.trajectory: frame = ts.frame ref = universe.trajectory[frame].positions + vector assert_almost_equal(ts.positions, ref, decimal=6) def test_transform_slice(self, universe, transformed): vector = np.float32([10, 10, 10]) # what happens when we slice the trajectory? for ts in transformed.trajectory[5:17:3]: frame = ts.frame ref = universe.trajectory[frame].positions + vector assert_almost_equal(ts.positions, ref, decimal=6) def test_transform_switch(self, universe, transformed): vector = np.float32([10, 10, 10]) # grab a frame: ref = universe.trajectory[2].positions + vector assert_almost_equal( transformed.trajectory[2].positions, ref, decimal=6 ) # now switch to another frame newref = universe.trajectory[10].positions + vector assert_almost_equal( transformed.trajectory[10].positions, newref, decimal=6 ) # what happens when we comeback to the previous frame? assert_almost_equal( transformed.trajectory[2].positions, ref, decimal=6 ) def test_transfrom_rewind(self, universe, transformed): vector = np.float32([10, 10, 10]) ref = universe.trajectory[0].positions + vector transformed.trajectory.rewind() assert_almost_equal( transformed.trajectory.ts.positions, ref, decimal=6 ) class TestChainReaderCommonDt(object): common_dt = 100.0 prec = 3 @pytest.fixture() def trajectory(self): universe = mda.Universe( PSF, [DCD, CRD, DCD, CRD, DCD, CRD, CRD], dt=self.common_dt ) return universe.trajectory def test_time(self, trajectory): # We test this for the beginning, middle and end of the trajectory. for frame_n in (0, trajectory.n_frames // 2, -1): trajectory[frame_n] assert_almost_equal( trajectory.time, trajectory.frame * self.common_dt, 5, err_msg="Wrong time for frame {0:d}".format(frame_n), ) class TestChainReaderFormats(object): """Test of ChainReader with explicit formats (Issue 76).""" def test_set_all_format_tuples(self): universe = mda.Universe( GRO, [(PDB, "pdb"), (XTC, "xtc"), (TRR, "trr")] ) assert universe.trajectory.n_frames == 21 assert_equal(universe.trajectory.filenames, [PDB, XTC, TRR]) def test_set_all_format_lammps(self): universe = mda.Universe( LAMMPS_chain, [LAMMPSDUMP_chain1, LAMMPSDUMP_chain2], format="LAMMPSDUMP", continuous=True, ) assert universe.trajectory.n_frames == 11 # Test whether the amount of time between frames is consistently = 1 time_values = np.array([ts.time for ts in universe.trajectory]) dt_array = time_values[1:] - time_values[:-1] assert np.unique(dt_array) == 1 assert_equal(time_values, np.arange(11)) def test_set_format_tuples_and_format(self): universe = mda.Universe( GRO, [(PDB, "pdb"), GRO, GRO, (XTC, "xtc"), (TRR, "trr")], format="gro", ) assert universe.trajectory.n_frames == 23 assert_equal(universe.trajectory.filenames, [PDB, GRO, GRO, XTC, TRR]) with pytest.raises(TypeError) as errinfo: mda.Universe( GRO, [(PDB, "pdb"), GRO, GRO, (XTC, "xtc"), (TRR, "trr")], format="pdb", ) assert "Unable to read" in str(errinfo.value) def test_set_one_format_tuple(self): universe = mda.Universe(PSF, [(PDB_small, "pdb"), DCD]) assert universe.trajectory.n_frames == 99 def test_set_all_formats(self): with pytest.raises(TypeError) as errinfo: mda.Universe(PDB, [PDB, GRO], format="gro") assert "Unable to read" in str(errinfo.value) universe = mda.Universe(GRO, [PDB, PDB, PDB], format="pdb") assert_equal(universe.trajectory.filenames, [PDB, PDB, PDB]) def build_trajectories(folder, sequences, fmt="xtc"): """ A scenario is given as a series of time sequences. The result is returned as a list of time and origin of the frame. Each element of that result list correspond to a frame of the conatenated trajectory, The first element of each tuple is the time it corresponds to, the second element is the index of the input the frame comes from. See original gist and PR 1728 for comparison with gromacs trjcat. https://gist.github.com/jbarnoud/cacd0957d3df01d1577f640b20e86039 """ template = "trjcat_test_{{}}.{}".format(fmt) template = os.path.join(folder, template) # Use an empty universe to have a topology utop = mda.Universe.empty(1, trajectory=True) # Create synthetic trajectories. The times come from the user input, # the coordinates indicate the index of the input. fnames = [] for index, subseq in enumerate(sequences): coords = np.zeros((len(subseq), 1, 3), dtype=np.float32) + index u = mda.Universe(utop._topology, coords, to_guess=()) out_traj = mda.Writer(template.format(index), n_atoms=len(u.atoms)) fnames.append(out_traj.filename) with out_traj: for ts, time in zip(u.trajectory, subseq): # The time step needs a box to avoid a warning ts.dimensions = [10, 10, 10, 90, 90, 90] # The time is set from the user input ts.time = time out_traj.write(u) return utop, fnames class TestChainReaderContinuous(object): class SequenceInfo(object): def __init__(self, seq, n_frames, order): self.seq = seq self.n_frames = n_frames self.order = order @pytest.mark.parametrize("fmt", ("xtc", "trr")) @pytest.mark.parametrize( "seq_info", [ SequenceInfo( seq=([0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7]), n_frames=8, order=[0, 0, 1, 1, 2, 2, 2, 2], ), SequenceInfo( seq=([0, 1, 2, 4], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), n_frames=10, order=np.ones(10), ), SequenceInfo( seq=([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3]), n_frames=4, order=np.ones(4), ), SequenceInfo( seq=([5, 6, 7, 8, 9], [2, 3, 4, 5, 6], [0, 1, 2, 3]), n_frames=10, order=[2, 2, 1, 1, 1, 0, 0, 0, 0, 0], ), SequenceInfo(seq=([0, 1, 2],) * 3, n_frames=3, order=[2, 2, 2]), SequenceInfo( seq=( [0, 1, 2, 3], [3, 4], [4, 5, 6, 7], ), n_frames=8, order=[0, 0, 0, 1, 2, 2, 2, 2], ), SequenceInfo( seq=([5, 6, 7, 8, 9], [2, 3, 4, 5, 6], [0, 1, 2, 3]), n_frames=10, order=[2, 2, 1, 1, 1, 0, 0, 0, 0, 0], ), SequenceInfo( seq=[ list(range(0, 6)), list(range(2, 5)), list(range(2, 5)), list(range(2, 5)), list(range(3, 8)), ], n_frames=8, order=[0, 0, 3, 4, 4, 4, 4, 4], ), ], ) def test_order(self, seq_info, tmpdir, fmt): folder = str(tmpdir) utop, fnames = build_trajectories( folder, sequences=seq_info.seq, fmt=fmt ) u = mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) assert u.trajectory.n_frames == seq_info.n_frames for i, ts in enumerate(u.trajectory): assert_almost_equal(i, ts.time, decimal=4) # check we have used the right trajectory assert seq_info.order[i] == int(ts.positions[0, 0]) def test_start_frames(self, tmpdir): folder = str(tmpdir) sequences = ([0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7]) utop, fnames = build_trajectories( folder, sequences=sequences, ) u = mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) assert_equal(u.trajectory._start_frames, [0, 2, 4]) def test_missing(self, tmpdir): folder = str(tmpdir) sequences = ([0, 1, 2, 3], [5, 6, 7, 8, 9]) utop, fnames = build_trajectories( folder, sequences=sequences, ) u = mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) assert u.trajectory.n_frames == 9 def test_warning(self, tmpdir): folder = str(tmpdir) # this sequence *should* trigger a warning sequences = ([0, 1, 2, 3], [5, 6, 7]) utop, fnames = build_trajectories( folder, sequences=sequences, ) with pytest.warns(UserWarning): mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) def test_interleaving_error(self, tmpdir): folder = str(tmpdir) # interleaving is not supported by chainreader sequences = ([0, 2, 4, 6], [1, 3, 5, 7]) utop, fnames = build_trajectories( folder, sequences=sequences, ) with pytest.raises(RuntimeError): mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) def test_easy_trigger_warning(self, tmpdir): folder = str(tmpdir) # this sequence shouldn't trigger a warning sequences = ([0, 1, 2, 3], [2, 3, 4, 5], [4, 5, 6, 7]) utop, fnames = build_trajectories( folder, sequences=sequences, ) with no_warning(UserWarning): with warnings.catch_warnings(): # for windows Python 3.10 ignore: # ImportWarning('_SixMetaPathImporter.find_spec() not found # TODO: remove when we no longer have a dependency # that still imports six if sys.version_info >= (3, 10): warnings.filterwarnings( action="ignore", category=ImportWarning ) mda.Universe( utop._topology, fnames, continuous=True, to_guess=() ) def test_single_frames(self, tmpdir): folder = str(tmpdir) sequences = ( [0, 1, 2, 3], [ 5, ], ) utop, fnames = build_trajectories( folder, sequences=sequences, ) with pytest.raises(RuntimeError): mda.Universe(utop._topology, fnames, continuous=True, to_guess=()) def test_mixed_filetypes(self): with pytest.raises(ValueError): mda.Universe(PDB, [XTC, TRR], continuous=True) def test_unsupported_filetypes(self): with pytest.raises(NotImplementedError): mda.Universe(PSF, [DCD, DCD], continuous=True) # see issue 2353. The PDB reader has multiple format endings. To ensure # the not implemented error is thrown we do a check here. A more # careful test in the future would be a dummy reader with multiple # formats, just in case PDB will allow continuous reading in the future. with pytest.raises(ValueError): mda.Universe(PDB, [PDB, XTC], continuous=True) @pytest.mark.parametrize( "l, ref", ( [((0, 3), (3, 3), (4, 7)), (0, 1, 2)], [((0, 9), (0, 4)), (0, 1)], [((0, 3), (2, 2), (3, 3), (2, 6), (5, 9)), (0, 1, 3, 2, 4)], [((0, 2), (4, 9), (0, 4), (7, 9)), (0, 2, 1, 3)], [((0, 5), (2, 4), (2, 4), (2, 4), (3, 7)), (0, 1, 2, 3, 4)], ), ) def test_multilevel_arg_sort(l, ref): indices = mda.coordinates.chain.multi_level_argsort(l) assert_equal(indices, ref) @pytest.mark.parametrize( "l, ref", ( [((0, 4), (3, 6), (6, 9)), (0, 1, 2)], [((0, 3), (3, 4), (4, 7)), (0, 1, 2)], [((0, 3), (3, 5), (4, 7)), (0, 1, 2)], [((0, 3), (0, 4), (3, 5), (4, 7)), (1, 2, 3)], [((0, 3), (0, 4)), (1,)], [((0, 3), (0, 3)), (1,)], [((1, 3), (0, 4)), (0,)], [((0, 3),), (0,)], [((0, 3), (5, 9)), (0, 1)], [((0, 3), (0, 3), (5, 9)), (1, 2)], [((0, 3), (0, 3), (0, 3), (5, 9)), (2, 3)], [((0, 5), (2, 4), (2, 4), (2, 4), (3, 7)), (0, 3, 4)], [((0, 3), (2, 4), (2, 4), (2, 4), (3, 7)), (0, 3, 4)], [((0, 3), (2, 4), (4, 7), (4, 7), (4, 7), (6, 9)), (0, 1, 4, 5)], [((0, 6), (2, 5), (2, 5), (2, 5), (3, 8)), (0, 3, 4)], ), ) def test_filter_times(l, ref): indices = mda.coordinates.chain.filter_times(l, dt=1) assert_equal(indices, ref) def test_issue_4008(): # test for transformation on a chainreader trajectory # in issue #4008 it was observed that the transformation did not work on # the frames which were "borders" of the chained trajectories def shift_by_group(ag): # this just shifts all atoms in the box by their center of mass, # i.e. afterwards the center of mass of all atoms should be [0,0,0] def wrapped(ts): shift = ag.center_of_mass() ts.positions[:] -= shift return ts return wrapped u = mda.Universe(GRO, [XTC, XTC]) trafos = (shift_by_group(u.atoms),) u.trajectory.add_transformations(*trafos) com = np.zeros((len(u.trajectory), 3)) for i, ts in enumerate(u.trajectory): com[i] = u.atoms.center_of_mass() # see issue for details, but in general the com shouldn't be far from zero # in the issue report values of ~25 were seen before the fix # with the fix all values should be <1e-6 assert np.abs(com).max() < 1e-6 def test_issue_3349(): # test for repr of a chainreader when the subreaders are memoryreaders cdx = np.random.random(3341 * 3 * 10) cdx = cdx.reshape(-1, 3341, 3).astype(np.float32) * 10 u = mda.Universe(PSF, (cdx, DCD)) u2 = mda.Universe(PSF, (cdx, cdx, cdx, DCD)) u_expected_filenames = np.array([None, str(DCD)]) u2_expected_filenames = np.array([None, None, None, str(DCD)]) assert_equal( "", u.trajectory.__repr__(), ) assert_equal(u_expected_filenames, u.trajectory.filenames) assert_equal( "", u2.trajectory.__repr__(), ) assert_equal(u2_expected_filenames, u2.trajectory.filenames)