Use toplevel metadata to adjust times when importing a tree sequence. (…

…#1129)

* The "generation" field is used to convert backwards times to
  forward time.
* This adjustment is applied to node times and to mutation
  origin times.
* Add test of complex evolve/export/import/evolve work flows

Closes #1114
Closes #1111

cpp/fwdpy11_types/DiploidPopulation.cc

@@ -43,7 +43,8 @@ PYBIND11_MAKE_OPAQUE(std::vector<fwdpy11::DiploidMetadata>);
 PYBIND11_MAKE_OPAQUE(fwdpy11::DiploidPopulation::genome_container);
 PYBIND11_MAKE_OPAQUE(fwdpy11::DiploidPopulation::mutation_container);
 
-fwdpy11::DiploidPopulation create_DiploidPopulation_from_tree_sequence(py::object ts);
+fwdpy11::DiploidPopulation
+create_DiploidPopulation_from_tree_sequence(py::object ts, std::uint32_t generation);
 
 namespace
 {
@@ -309,7 +310,7 @@ init_DiploidPopulation(py::module& m)
                     = load(f).cast<decltype(rv.ancient_sample_genetic_value_matrix)>();
                 return rv;
             })
-        .def_static("_create_from_tskit", [](py::object ts) {
-            return create_DiploidPopulation_from_tree_sequence(ts);
+        .def_static("_create_from_tskit", [](py::object ts, std::uint32_t generation) {
+            return create_DiploidPopulation_from_tree_sequence(ts, generation);
         });
 }

cpp/fwdpy11_types/ts_from_tskit.cc

@@ -1,3 +1,4 @@
+#include "fwdpp/ts/generate_offspring.hpp"
 #include <vector>
 #include <utility>
 #include <cstdint>
@@ -14,8 +15,9 @@
 namespace py = pybind11;
 
 // TODO put this in a header in case anyone else finds it useful?
-std::tuple<fwdpp::ts::std_table_collection::node_table, fwdpp::ts::std_table_collection::edge_table, int, double>
-convert_tables_from_tskit(py::object ts)
+std::tuple<fwdpp::ts::std_table_collection::node_table,
+           fwdpp::ts::std_table_collection::edge_table, int, double>
+convert_tables_from_tskit(py::object ts, std::uint32_t generation)
 {
     py::object tstables = ts.attr("tables");
     py::object tsnodes = tstables.attr("nodes");
@@ -36,14 +38,12 @@ convert_tables_from_tskit(py::object ts)
     edges.reserve(left_.shape(0));
     for (py::ssize_t i = 0; i < left_.shape(0); ++i)
         {
-            edges.push_back(
-                fwdpp::ts::edge{ left_(i), right_(i), parent_(i), child_(i) });
+            edges.push_back(fwdpp::ts::edge{left_(i), right_(i), parent_(i), child_(i)});
         }
 
     // same treatment for nodes
     auto time = tsnodes.attr("time").cast<py::array_t<double>>();
-    auto population
-        = tsnodes.attr("population").cast<py::array_t<std::int32_t>>();
+    auto population = tsnodes.attr("population").cast<py::array_t<std::int32_t>>();
     auto time_ = time.unchecked<1>();
     auto population_ = population.unchecked<1>();
 
@@ -62,33 +62,34 @@ convert_tables_from_tskit(py::object ts)
                 {
                     ++ntips;
                 }
-            nodes.push_back(fwdpp::ts::node{ population_(i), t });
+            t += generation;
+            nodes.push_back(fwdpp::ts::node{population_(i), t});
         }
 
     double l = ts.attr("get_sequence_length")().cast<double>();
     return std::make_tuple(std::move(nodes), std::move(edges), ntips, l);
 }
 
 fwdpy11::DiploidPopulation
-create_DiploidPopulation_from_tree_sequence(py::object ts)
+create_DiploidPopulation_from_tree_sequence(py::object ts, std::uint32_t generation)
 {
-    auto t = convert_tables_from_tskit(ts);
+    auto t = convert_tables_from_tskit(ts, generation);
     auto nodes(std::move(std::get<0>(t)));
     auto edges(std::move(std::get<1>(t)));
     auto twoN = std::get<2>(t);
     auto l = std::get<3>(t);
     if (twoN % 2 != 0.0)
         {
-            throw std::invalid_argument(
-                "tree sequence has odd number of tips");
+            throw std::invalid_argument("tree sequence has odd number of tips");
         }
 
     fwdpy11::DiploidPopulation pop(twoN / 2, l);
     pop.tables->nodes.swap(nodes);
     pop.tables->edges.swap(edges);
     // NOTE: this may be an issue when we allow variable survival probabilitites!
-    pop.tables->edge_offset = static_cast<fwdpp::ts::table_index_t>(pop.tables->num_edges());
-    if(!fwdpp::ts::edge_table_minimally_sorted(*pop.tables))
+    pop.tables->edge_offset
+        = static_cast<fwdpp::ts::table_index_t>(pop.tables->num_edges());
+    if (!fwdpp::ts::edge_table_minimally_sorted(*pop.tables))
         {
             throw std::runtime_error("edge table is not sorted");
         }
@@ -102,15 +103,14 @@ create_DiploidPopulation_from_tree_sequence(py::object ts)
         {
             pop.diploid_metadata[i].nodes[0] = 2 * i;
             pop.diploid_metadata[i].nodes[1] = 2 * i + 1;
-            if (pop.tables->nodes[2*i].deme 
-                != pop.tables->nodes[2*i+1].deme)
+            if (pop.tables->nodes[2 * i].deme != pop.tables->nodes[2 * i + 1].deme)
                 {
                     throw std::invalid_argument(
                         "inconsistent deme fields for nodes in the same "
                         "individual");
                 }
-            pop.diploid_metadata[i].deme
-                = pop.tables->nodes[2*i].deme;
+            pop.diploid_metadata[i].deme = pop.tables->nodes[2 * i].deme;
         }
+    pop.generation = generation;
     return pop;
 }

fwdpy11/_types/diploid_population.py

@@ -109,7 +109,11 @@ def create_from_tskit(cls, ts: tskit.TreeSequence,
               of fwdpy11.
 
         """
-        ll = ll_DiploidPopulation._create_from_tskit(ts)
+        generation = 0
+        if len(ts.metadata) > 0:
+            if "generation" in ts.metadata:
+                generation = ts.metadata["generation"]
+        ll = ll_DiploidPopulation._create_from_tskit(ts, generation)
         if import_mutations is True:
             import fwdpy11.tskit_tools
             for mutation in ts.mutations():
@@ -145,7 +149,7 @@ def create_from_tskit(cls, ts: tskit.TreeSequence,
                 if i is not None:
                     mvec.append(i)
             ll._set_mutations(mvec, mutation_nodes,
-                              [int(i.time) for i in ts.mutations()])
+                              [int(i.time) - ll._generation for i in ts.mutations()])
         return cls(0, 0.0, ll_pop=ll)
 
     @ classmethod

tests/test_recreate_pop_from_own_treeseq.py

@@ -1,4 +1,3 @@
-#
 # Copyright (C) 2023 Kevin Thornton <krthornt@uci.edu>
 #
 # This file is part of fwdpy11.
@@ -90,32 +89,33 @@ def pop():
 
     fwdpy11.evolvets(rng, pop, params, 100, suppress_table_indexing=True)
 
-    assert pop.generation == 3
-
     return pop
 
 
 # Related to GitHub issue 1109.
 def test_recreate_pop_from_own_treeseq(pop):
+    gen = pop.generation
     ts = pop.dump_tables_to_tskit()
     nm = ts.num_mutations
     assert nm > 0
 
     pop = fwdpy11.DiploidPopulation.create_from_tskit(
         ts, import_mutations=True)
+    assert pop.generation == gen
     assert len(pop.tables.mutations) == nm
-    assert pop.generation == 0
     validate_mutation_table(pop)
+    assert(all([pop.mutations[i.key].g >= 0 for i in pop.tables.mutations]))
 
 
 def test_multiple_export_recreate_evolve_steps():
     pop, rng, params = pop_rng_params()
     for i in range(5):
         fwdpy11.evolvets(rng, pop, params, 100, suppress_table_indexing=True)
-        assert pop.generation == 3
         ts = pop.dump_tables_to_tskit()
         pop = fwdpy11.DiploidPopulation.create_from_tskit(
             ts, import_mutations=True)
+        assert(all([pop.mutations[i.key].g >= 0 for i in pop.tables.mutations]))
+    assert pop.generation == 15
 
 
 def test_start_pop_from_treeseq_without_mutations(pop):
@@ -157,3 +157,55 @@ def test_start_pop_from_treeseq_without_mutations(pop):
     ts = tables.tree_sequence()
     pop = fwdpy11.DiploidPopulation.create_from_tskit(
         ts, import_mutations=True)
+
+
+def test_complex_start_stop_workflow():
+    yaml = """
+description: split
+time_units: generations
+demes:
+ - name: A
+   epochs:
+    - start_size: 10
+      end_time: 10
+ - name: B
+   start_time: 10
+   ancestors: [A]
+   epochs:
+    - start_size: 10
+      end_size: 50
+ - name: C
+   start_time: 10
+   ancestors: [A]
+   epochs:
+    - start_size: 10
+      end_size: 50
+"""
+    demography = fwdpy11.ForwardDemesGraph.from_demes(
+        yaml, burnin=10, burnin_is_exact=True)
+    pdict = {'recregions': [],
+             'nregions': [],
+             'sregions': [],
+             'rates': (0, 0, 0),
+             'gvalue': fwdpy11.Multiplicative(2.),
+             'simlen': 11,
+             'demography': demography
+             }
+    params = fwdpy11.ModelParams(**pdict)
+    pop = fwdpy11.DiploidPopulation(demography.initial_sizes, 1.0)
+    rng = fwdpy11.GSLrng(10)
+    fwdpy11.evolvets(rng, pop, params, 10)
+    assert len(pop.deme_sizes()) == 2
+    assert pop.generation == 11
+    ts = pop.dump_tables_to_tskit()
+    restored = fwdpy11.DiploidPopulation.create_from_tskit(ts)
+    pdict['simlen'] = 100
+    params = fwdpy11.ModelParams(**pdict)
+    # Evolve rest of model
+    fwdpy11.evolvets(rng, restored, params, 10)
+    model_time = restored.generation
+    assert model_time == 20
+    ts = restored.dump_tables_to_tskit()
+    restored2 = fwdpy11.DiploidPopulation.create_from_tskit(ts)
+    fwdpy11.evolvets(rng, restored2, params, 10)
+    assert model_time == restored2.generation