Fix errors in pickling and binary serialization. Fixes #389.

molpopgen · Feb 25, 2020 · 260402f · 260402f
1 parent 75ca78a
commit 260402f
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 1 deletion.
diff --git a/fwdpy11/headers/fwdpy11/serialization.hpp b/fwdpy11/headers/fwdpy11/serialization.hpp
@@ -92,7 +92,7 @@ namespace fwdpy11
             w(buffer, &msize);
             if (msize > 0)
                 {
-                    w(buffer, pop->genetic_value_matrix.data(), msize);
+                    w(buffer, pop->ancient_sample_genetic_value_matrix.data(), msize);
                 }
 
             return buffer;

diff --git a/fwdpy11/src/fwdpy11_types/DiploidPopulation.cc b/fwdpy11/src/fwdpy11_types/DiploidPopulation.cc
@@ -327,6 +327,8 @@ init_DiploidPopulation(py::module& m)
                         dump(s, f);
                     }
                 dump(self.tables.preserved_nodes, f);
+                dump(self.genetic_value_matrix, f);
+                dump(self.ancient_sample_genetic_value_matrix, f);
             },
             R"delim(
              Pickle the population to an open file.
@@ -443,6 +445,12 @@ init_DiploidPopulation(py::module& m)
                     = load(f).cast<decltype(rv.tables.preserved_nodes)>();
                 rv.tables.build_indexes();
                 rv.rebuild_mutation_lookup(false);
+                rv.genetic_value_matrix
+                    = load(f).cast<decltype(rv.genetic_value_matrix)>();
+                rv.ancient_sample_genetic_value_matrix
+                    = load(f)
+                          .cast<decltype(
+                              rv.ancient_sample_genetic_value_matrix)>();
                 return rv;
             },
             R"delim(

diff --git a/tests/test_record_genetic_value_matrix.py b/tests/test_record_genetic_value_matrix.py
@@ -25,7 +25,9 @@
 """
 
 import fwdpy11
+import pickle
 import unittest
+import tempfile
 import numpy as np
 
 
@@ -142,6 +144,28 @@ def test_ancient_sample_genetic_value_reconstruction(self):
             gv = np.sum(gv, axis=1)
             self.assertTrue(np.allclose(gv, gvslice))
 
+    def test_pickling(self):
+        pp = pickle.dumps(self.pop, -1)
+        up = pickle.loads(pp)
+        gv = self.pop.genetic_values
+        upgv = up.genetic_values
+        self.assertTrue(np.array_equal(gv, upgv))
+        gv = self.pop.ancient_sample_genetic_values
+        upgv = up.ancient_sample_genetic_values
+        self.assertTrue(np.array_equal(gv, upgv))
+
+    def test_pickle_to_file(self):
+        with tempfile.NamedTemporaryFile(delete=False) as tfile:
+            self.pop.pickle_to_file(tfile)
+            tfile.seek(0)
+            pop2 = fwdpy11.DiploidPopulation.load_from_pickle_file(tfile)
+        gv = self.pop.genetic_values
+        gv2 = pop2.genetic_values
+        self.assertTrue(np.array_equal(gv, gv2))
+        gv = self.pop.ancient_sample_genetic_values
+        gv2 = pop2.ancient_sample_genetic_values
+        self.assertTrue(np.array_equal(gv, gv2))
+
 
 class TestTwoTraitsIsotropy(unittest.TestCase):
     @classmethod
@@ -232,6 +256,28 @@ def test_ancient_sample_genetic_value_reconstruction(self):
                 d1 = np.power(j - 0.0, 2.0)
                 self.assertTrue(np.isclose(np.exp(-(d0+d1)/(2.*self.VS)), w))
 
+    def test_pickling(self):
+        pp = pickle.dumps(self.pop, -1)
+        up = pickle.loads(pp)
+        gv = self.pop.genetic_values
+        upgv = up.genetic_values
+        self.assertTrue(np.array_equal(gv, upgv))
+        gv = self.pop.ancient_sample_genetic_values
+        upgv = up.ancient_sample_genetic_values
+        self.assertTrue(np.array_equal(gv, upgv))
+
+    def test_pickle_to_file(self):
+        with tempfile.NamedTemporaryFile(delete=False) as tfile:
+            self.pop.pickle_to_file(tfile)
+            tfile.seek(0)
+            pop2 = fwdpy11.DiploidPopulation.load_from_pickle_file(tfile)
+        gv = self.pop.genetic_values
+        gv2 = pop2.genetic_values
+        self.assertTrue(np.array_equal(gv, gv2))
+        gv = self.pop.ancient_sample_genetic_values
+        gv2 = pop2.ancient_sample_genetic_values
+        self.assertTrue(np.array_equal(gv, gv2))
+
 
 if __name__ == "__main__":
     unittest.main()