Set time_units to generations when exporting to tskit. (#1242)

molpopgen · Dec 5, 2023 · cacaabf · cacaabf
1 parent 5742492
commit cacaabf
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/fwdpy11/tskit_tools/_dump_tables_to_tskit.py b/fwdpy11/tskit_tools/_dump_tables_to_tskit.py
@@ -162,6 +162,7 @@ def _dump_tables_to_tskit(
     edge_view = np.array(self.tables.edges, copy=False)
 
     tc = tskit.TableCollection(self.tables.genome_length)
+    tc.time_units = "generations"
 
     tc.metadata_schema = fwdpy11.tskit_tools.metadata_schema.TopLevelMetadata
 

diff --git a/tests/test_tree_sequences.py b/tests/test_tree_sequences.py
@@ -186,6 +186,7 @@ def setUpClass(self):
             max(self.pop.mcounts) == 2 * self.pop.N
         ), "Nothing fixed, so test case is not helpful"
         self.dumped_ts = self.pop.dump_tables_to_tskit()
+        assert self.dumped_ts.time_units == "generations"
 
     def test_mut_lookup(self):
         self.assertEqual(len(self.pop.mut_lookup),
@@ -853,6 +854,7 @@ def test_mutation_table_contents(self):
 
     def test_dump_to_tskit(self):
         ts = self.pop.dump_tables_to_tskit()
+        assert ts.time_units == "generations"
         self.assertEqual(mcounts_comparison(self.pop, ts), True)
 
     def test_binary_round_trip(self):
@@ -949,6 +951,7 @@ def test_mutation_table_contents(self):
 
     def test_dump_to_tskit(self):
         ts = self.pop.dump_tables_to_tskit()
+        assert ts.time_units == "generations"
         self.assertEqual(mcounts_comparison(self.pop, ts), True)
 
     def test_binary_round_trip(self):
@@ -1479,6 +1482,7 @@ def setUpClass(self):
         )
         self.amd = np.array(self.pop.ancient_sample_metadata, copy=False)
         self.tskit_ts = self.pop.dump_tables_to_tskit()
+        assert self.tskit_ts.time_units == "generations"
 
     def test_ancient_sample_records(self):
         self.assertEqual(len(self.pop.ancient_sample_metadata), self.pop.N)