Skip to content

Commit

Permalink
Merge pull request #656 from kain88-de/xdrlib-offset-checking
Browse files Browse the repository at this point in the history
Xdrlib offset checking, issue (631)
  • Loading branch information
mnmelo committed Feb 7, 2016
2 parents 73d6f59 + d6ca51b commit 795e573
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 42 deletions.
55 changes: 41 additions & 14 deletions package/MDAnalysis/coordinates/XDR.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
#
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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'])
Expand All @@ -114,15 +132,15 @@ 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(
e.message))
except AttributeError:
warnings.warn("Couldn't save offsets because: {}".format(e))


def rewind(self):
"""Read the first frame again"""
self._read_frame(0)
Expand All @@ -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"""
Expand Down
69 changes: 41 additions & 28 deletions testsuite/MDAnalysisTests/coordinates/test_xdr.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import six
from six.moves import zip, range

import errno
Expand Down Expand Up @@ -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):

Expand Down Expand Up @@ -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'],
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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):
Expand Down

0 comments on commit 795e573

Please sign in to comment.