backport fix for #2878 (pickle DCD and XDR readers at correct frame)

package/CHANGELOG

@@ -13,7 +13,7 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/20 orbeckst, VOD555, lilyminium
+??/??/20 orbeckst, VOD555, lilyminium, yuxuanzhuang
 
   * 1.0.1
 
@@ -24,6 +24,7 @@ Fixes
     density=True; the keyword was available since 0.19.0 but with incorrect
     semantics and not documented and did not produce correct results (Issue
     #2811, PR #2812)
+  * pickle correct frames of DCD, TRR, and XTC (#2878)
 
 
 06/09/20 richardjgowers, kain88-de, lilyminium, p-j-smith, bdice, joaomcteixeira,

package/MDAnalysis/lib/formats/libdcd.pyx

@@ -265,7 +265,8 @@ cdef class DCDFile:
             return
 
         current_frame = state[1]
-        self.seek(current_frame)
+        self.seek(current_frame - 1)
+        self.current_frame = current_frame
 
 
     def tell(self):

package/MDAnalysis/lib/formats/libmdaxdr.pyx

@@ -306,7 +306,8 @@ cdef class _XDRFile:
 
         # where was I
         current_frame = state[1]
-        self.seek(current_frame)
+        self.seek(current_frame - 1)
+        self.current_frame = current_frame
 
     def seek(self, frame):
         """Seek to Frame.

package/setup.py

@@ -73,7 +73,7 @@
     from commands import getoutput
 
 # NOTE: keep in sync with MDAnalysis.__version__ in version.py
-RELEASE = "1.0.0"
+RELEASE = "1.0.1-dev"
 
 is_release = 'dev' not in RELEASE
 

testsuite/MDAnalysisTests/coordinates/base.py

@@ -419,12 +419,12 @@ def test_transformations_copy(self,ref,transformed):
             ideal_coords = ref.iter_ts(i).positions + v1 + v2
             assert_array_almost_equal(ts.positions, ideal_coords, decimal = ref.prec)
 
-
     def test_add_another_transformations_raises_ValueError(self, transformed):
         # After defining the transformations, the workflow cannot be changed
         with pytest.raises(ValueError):
             transformed.add_transformations(translate([2,2,2]))
 
+
 class MultiframeReaderTest(BaseReaderTest):
     def test_last_frame(self, ref, reader):
         ts = reader[-1]
@@ -491,6 +491,7 @@ def test_iter_as_aux_lowf(self, ref, reader):
                                          decimal=ref.prec)
 
 
+
 class BaseWriterTest(object):
     @staticmethod
     @pytest.fixture()

testsuite/MDAnalysisTests/formats/test_libdcd.py

@@ -84,25 +84,48 @@ def dcd():
     with DCDFile(DCD) as dcd:
         yield dcd
 
+def _assert_compare_readers(old_reader, new_reader):
+    frame = old_reader.read()     # same as next(old_reader)
+    new_frame = new_reader.read() # same as next(new_reader)
 
-def test_pickle(dcd):
-    dcd.seek(len(dcd) - 1)
-    dump = pickle.dumps(dcd)
-    new_dcd = pickle.loads(dump)
+    assert old_reader.fname == new_reader.fname
+    assert old_reader.tell() == new_reader.tell()
+    assert_almost_equal(frame.xyz, new_frame.xyz)
+    assert_almost_equal(frame.unitcell, new_frame.unitcell)
 
-    assert dcd.fname == new_dcd.fname
-    assert dcd.tell() == new_dcd.tell()
+def test_pickle(dcd):
+    mid = len(dcd) // 2
+    dcd.seek(mid)
+    new_dcd = pickle.loads(pickle.dumps(dcd))
+    _assert_compare_readers(dcd, new_dcd)
 
+def test_pickle_last(dcd):
+    dcd.seek(len(dcd) - 1)
+    new_dcd = pickle.loads(pickle.dumps(dcd))
+    _assert_compare_readers(dcd, new_dcd)
 
 def test_pickle_closed(dcd):
     dcd.seek(len(dcd) - 1)
     dcd.close()
-    dump = pickle.dumps(dcd)
-    new_dcd = pickle.loads(dump)
+    new_dcd = pickle.loads(pickle.dumps(dcd))
 
     assert dcd.fname == new_dcd.fname
     assert dcd.tell() != new_dcd.tell()
 
+def test_pickle_after_read(dcd):
+    _ = dcd.read()
+    new_dcd = pickle.loads(pickle.dumps(dcd))
+    _assert_compare_readers(dcd, new_dcd)
+
+#@pytest.mark.xfail
+def test_pickle_immediately(dcd):
+    # do not seek before pickling: this seems to leave the DCDFile
+    # object in weird state: is this supposed to work?
+    new_dcd = pickle.loads(pickle.dumps(dcd))
+
+    assert dcd.fname == new_dcd.fname
+    assert dcd.tell() == new_dcd.tell()
+
 
 @pytest.mark.parametrize("new_frame", (10, 42, 21))
 def test_seek_normal(new_frame, dcd):

testsuite/MDAnalysisTests/formats/test_libmdaxdr.py

@@ -26,7 +26,7 @@
 import numpy as np
 
 from numpy.testing import (assert_almost_equal, assert_array_almost_equal,
-                           assert_array_equal)
+                           assert_array_equal, assert_equal)
 
 from MDAnalysis.lib.formats.libmdaxdr import TRRFile, XTCFile
 
@@ -128,20 +128,35 @@ def test_read_write_mode_file(self, xdr, tmpdir, fname):
             with pytest.raises(IOError):
                 f.read()
 
+    @staticmethod
+    def _assert_compare_readers(old_reader, new_reader):
+        frame = old_reader.read()
+        new_frame = new_reader.read()
+
+        assert old_reader.fname == new_reader.fname
+        assert old_reader.tell() == new_reader.tell()
+
+        assert_equal(old_reader.offsets, new_reader.offsets)
+        assert_almost_equal(frame.x, new_frame.x)
+        assert_almost_equal(frame.box, new_frame.box)
+        assert frame.step == new_frame.step
+        assert_almost_equal(frame.time, new_frame.time)
+
     def test_pickle(self, reader):
-        reader.seek(len(reader) - 1)
-        dump = pickle.dumps(reader)
-        new_reader = pickle.loads(dump)
+        mid = len(reader) // 2
+        reader.seek(mid)
+        new_reader = pickle.loads(pickle.dumps(reader))
+        self._assert_compare_readers(reader, new_reader)
 
-        assert reader.fname == new_reader.fname
-        assert reader.tell() == new_reader.tell()
-        assert_almost_equal(reader.offsets, new_reader.offsets)
+    def test_pickle_last_frame(self, reader):
+        reader.seek(len(reader) - 1)
+        new_reader = pickle.loads(pickle.dumps(reader))
+        self._assert_compare_readers(reader, new_reader)
 
     def test_pickle_closed(self, reader):
         reader.seek(len(reader) - 1)
         reader.close()
-        dump = pickle.dumps(reader)
-        new_reader = pickle.loads(dump)
+        new_reader = pickle.loads(pickle.dumps(reader))
 
         assert reader.fname == new_reader.fname
         assert reader.tell() != new_reader.tell()

testsuite/setup.py

@@ -88,7 +88,7 @@ def run(self):
 
 if __name__ == '__main__':
     # this must be in-sync with MDAnalysis
-    RELEASE = "1.0.0"
+    RELEASE = "1.0.1-dev"
 
     with open("README") as summary:
         LONG_DESCRIPTION = summary.read()