diff --git a/package/MDAnalysis/coordinates/XDR.py b/package/MDAnalysis/coordinates/XDR.py index 1c9f36c6442..a022231accc 100644 --- a/package/MDAnalysis/coordinates/XDR.py +++ b/package/MDAnalysis/coordinates/XDR.py @@ -10,6 +10,8 @@ # 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 # @@ -43,6 +45,22 @@ def offsets_filename(filename, ending='npz'): ending=ending)) +def read_numpy_offsets(filename): + """read offsets into a dictionary + + Parameters + ---------- + filename : str + filename of offsets + + Returns + ------- + offsets : dict + dictionary of offsets information + """ + return {k: v for k, v in six.iteritems(np.load(filename))} + + class XDRBaseReader(base.Reader): """Base class for xdrlib file formats xtc and trr""" def __init__(self, filename, convert_units=True, sub=None, @@ -52,6 +70,12 @@ def __init__(self, filename, convert_units=True, sub=None, **kwargs) self._xdr = self._file(self.filename) + self._sub = sub + if self._sub is not None: + self.n_atoms = len(self._sub) + else: + self.n_atoms = self._xdr.n_atoms + if not refresh_offsets: self._load_offsets() else: @@ -65,12 +89,6 @@ def __init__(self, filename, convert_units=True, sub=None, except StopIteration: dt = 0 - self._sub = sub - if self._sub is not None: - self.n_atoms = len(self._sub) - else: - self.n_atoms = self._xdr.n_atoms - self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs) self._frame = 0 self._frame_to_ts(frame, self.ts) @@ -93,15 +111,15 @@ def _load_offsets(self): self._read_offsets(store=True) return - with open(fname, 'rb') as f: - data = {k: v for k, v in six.iteritems(np.load(f))} + data = read_numpy_offsets(fname) ctime_ok = getctime(self.filename) == data['ctime'] size_ok = getsize(self.filename) == data['size'] + n_atoms_ok = self._xdr.n_atoms == data['n_atoms'] - if not (ctime_ok and size_ok): + if not (ctime_ok and size_ok and n_atoms_ok): warnings.warn("Reload offsets from trajectory\n " - "ctime or size did not match") + "ctime or size or n_atoms did not match") self._read_offsets(store=True) else: self._xdr.set_offsets(data['offsets']) @@ -114,7 +132,8 @@ def _read_offsets(self, store=False): size = getsize(self.filename) try: np.savez(offsets_filename(self.filename), - offsets=offsets, size=size, ctime=ctime) + offsets=offsets, size=size, ctime=ctime, + n_atoms=self._xdr.n_atoms) except Exception as e: try: warnings.warn("Couldn't save offsets because: {}".format( @@ -122,7 +141,6 @@ def _read_offsets(self, store=False): except AttributeError: warnings.warn("Couldn't save offsets because: {}".format(e)) - def rewind(self): """Read the first frame again""" self._read_frame(0) @@ -141,9 +159,18 @@ def _reopen(self): def _read_frame(self, i): """read frame i""" - self._xdr.seek(i) self._frame = i - 1 - return self._read_next_timestep() + try: + self._xdr.seek(i) + timestep = self._read_next_timestep() + except RuntimeError: + warnings.warn('seek failed, recalculating offsets and retrying') + offsets = self._xdr.calc_offsets() + self._xdr.set_offsets(offsets) + self._read_offsets(store=True) + self._xdr.seek(i) + timestep = self._read_next_timestep() + return timestep def _read_next_timestep(self, ts=None): """copy next frame into timestep""" diff --git a/testsuite/MDAnalysisTests/coordinates/test_xdr.py b/testsuite/MDAnalysisTests/coordinates/test_xdr.py index 402f9cad762..22778331df1 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_xdr.py +++ b/testsuite/MDAnalysisTests/coordinates/test_xdr.py @@ -1,4 +1,3 @@ -import six from six.moves import zip, range import errno @@ -32,6 +31,10 @@ import MDAnalysis.core.AtomGroup from MDAnalysis.coordinates import XDR +# I want to catch all warnings in the tests. If this is not set at the start it +# could cause test that check for warnings to fail. +warnings.simplefilter('always') + class _XDRReader_Sub(TestCase): @@ -728,8 +731,7 @@ def test_offsets(self): err_msg="wrong frame offsets") outfile_offsets = XDR.offsets_filename(self.traj) - with open(outfile_offsets, 'rb') as f: - saved_offsets = {k: v for k, v in six.iteritems(np.load(f))} + saved_offsets = XDR.read_numpy_offsets(outfile_offsets) assert_array_almost_equal(self.trajectory._xdr.offsets, saved_offsets['offsets'], @@ -753,8 +755,7 @@ def test_persistent_offsets_size_mismatch(self): # check that stored offsets are not loaded when trajectory # size differs from stored size fname = XDR.offsets_filename(self.traj) - with open(fname, 'rb') as f: - saved_offsets = {k: v for k, v in six.iteritems(np.load(f))} + saved_offsets = XDR.read_numpy_offsets(fname) saved_offsets['size'] += 1 with open(fname, 'wb') as f: np.savez(f, **saved_offsets) @@ -763,15 +764,14 @@ def test_persistent_offsets_size_mismatch(self): warnings.simplefilter('always') self._reader(self.traj) assert_equal(warn[0].message.args, - ('Reload offsets from trajectory\n ctime or size did not match', )) + ('Reload offsets from trajectory\n ctime or size or n_atoms did not match', )) @dec.slow def test_persistent_offsets_ctime_mismatch(self): # check that stored offsets are not loaded when trajectory # ctime differs from stored ctime fname = XDR.offsets_filename(self.traj) - with open(fname, 'rb') as f: - saved_offsets = {k: v for k, v in six.iteritems(np.load(f))} + saved_offsets = XDR.read_numpy_offsets(fname) saved_offsets['ctime'] += 1 with open(fname, 'wb') as f: np.savez(f, **saved_offsets) @@ -780,26 +780,39 @@ def test_persistent_offsets_ctime_mismatch(self): warnings.simplefilter('always') self._reader(self.traj) assert_equal(warn[0].message.args, - ('Reload offsets from trajectory\n ctime or size did not match', )) - - # TODO: This doesn't test if the offsets work AT ALL. the old - # implementation only checked if the offsets were ok to set back to the old - # frame. But that doesn't check if any of the other offsets is potentially - # wrong. Basically the only way to check that would be to scan through the - # whole trajectory. - # @dec.slow - # def test_persistent_offsets_last_frame_wrong(self): - # # check that stored offsets are not loaded when the offsets - # # themselves appear to be wrong - # with open(XDR.offsets_filename(self.traj), 'rb') as f: - # saved_offsets = {k: v for k, v in six.iteritems(np.load(f))} - # saved_offsets['offsets'] += 1 - # with open(XDR.offsets_filename(self.traj), 'wb') as f: - # np.savez(f, **saved_offsets) - - # # with warnings.catch_warnings(): - # # u = MDAnalysis.Universe(self.top, self.traj) - # # assert_equal((u.trajectory._xdr.offsets is None), True) + ('Reload offsets from trajectory\n ctime or size or n_atoms did not match', )) + + @dec.slow + def test_persistent_offsets_natoms_mismatch(self): + # check that stored offsets are not loaded when trajectory + # ctime differs from stored ctime + fname = XDR.offsets_filename(self.traj) + saved_offsets = XDR.read_numpy_offsets(fname) + saved_offsets['n_atoms'] += 1 + np.savez(fname, **saved_offsets) + + with warnings.catch_warnings(record=True) as warn: + warnings.simplefilter('always') + self._reader(self.traj) + assert_equal(warn[0].message.args, + ('Reload offsets from trajectory\n ctime or size or n_atoms did not match', )) + + @dec.slow + def test_persistent_offsets_last_frame_wrong(self): + fname = XDR.offsets_filename(self.traj) + saved_offsets = XDR.read_numpy_offsets(fname) + + idx_frame = 3 + saved_offsets['offsets'][idx_frame] += 42 + np.savez(fname, **saved_offsets) + + with warnings.catch_warnings(record=True) as warn: + warnings.simplefilter('always') + reader = self._reader(self.traj) + reader[idx_frame] + + assert_equal(warn[0].message.args[0], + 'seek failed, recalculating offsets and retrying') @dec.slow def test_persistent_offsets_readonly(self):