Check that we can pickle ReplicaExchange

Author: benmwebb

test/test.py

@@ -6,16 +6,51 @@
 import sys
 import subprocess
 import ihm.reader
+import pickle
+import IMP
+
 
 TOPDIR = os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), '..'))
 
+
 class Tests(unittest.TestCase):
     def test_simple(self):
         """Test model building"""
         os.chdir(os.path.join(TOPDIR, 'production_scripts'))
         p = subprocess.check_call(["python", "sample.py", "--test"])
         # todo: assert outputs, run analysis
 
+    def test_pickle(self):
+        """Test that pickled ReplicaExchange object works"""
+        # Set up modeling but don't run sampling
+        os.chdir(os.path.join(TOPDIR, 'production_scripts'))
+        with open('sample.py') as fh:
+            contents = fh.read().replace('rex.execute_macro()', '')
+        sys.argv = [sys.argv[0], '--test']
+        g = {}
+        exec(contents, g)
+        rex = g['rex']
+        del g
+        rex.vars['number_of_frames'] = 2
+
+        dump = pickle.dumps((rex.model, rex))
+
+        # Run the original ReplicaExchange and get the final score
+        IMP.random_number_generator.seed(99)
+        rex.execute_macro()
+        rs = IMP.pmi.tools.get_restraint_set(rex.model)
+        original_score = rs.evaluate(False)
+        del rex, rs
+
+        # With the same random seed, we should get the exact same trajectory
+        # with the pickled object
+        newm, newrex = pickle.loads(dump)
+        IMP.random_number_generator.seed(99)
+        newrex.execute_macro()
+        rs = IMP.pmi.tools.get_restraint_set(newrex.model)
+        new_score = rs.evaluate(False)
+        self.assertAlmostEqual(original_score, new_score, delta=1e-4)
+
     def test_mmcif(self):
         """Test generation of mmCIF output"""
         os.chdir(os.path.join(TOPDIR, 'production_scripts'))