Minimal support for ancestry iteration using ForwardDemesGraph.

molpopgen · Oct 14, 2022 · b7112ed · b7112ed
1 parent a3c4da8
commit b7112ed
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 2 deletions.
diff --git a/cpptests/test_forward_demes_graph.cc b/cpptests/test_forward_demes_graph.cc
@@ -2,7 +2,9 @@
 
 #include <boost/test/unit_test_suite.hpp>
 #include <fwdpy11/types/DiploidPopulation.hpp>
+#include <fwdpy11/discrete_demography/exceptions.hpp>
 #include <core/demes/forward_graph.hpp>
+#include <numeric>
 #include <stdexcept>
 
 #include "forward_demes_graph_fixtures.hpp"
@@ -45,6 +47,14 @@ BOOST_FIXTURE_TEST_CASE(single_deme_model_with_burn_in, SingleDemeModel)
             for (auto p : g.offspring_deme_sizes())
                 {
                     BOOST_REQUIRE_EQUAL(p, 100.0);
+                    if (p > 0.0)
+                        {
+                            auto ancestry = g.offspring_ancestry_proportions(num);
+                            auto sum_proportions = std::accumulate(
+                                std::begin(ancestry), std::end(ancestry), 0.0);
+                            BOOST_REQUIRE(std::fabs(sum_proportions - 1.0)
+                                          <= std::numeric_limits<double>::epsilon());
+                        }
                     num += 1;
                 }
             BOOST_REQUIRE_EQUAL(num, g.number_of_demes());
@@ -62,10 +72,34 @@ BOOST_FIXTURE_TEST_CASE(single_deme_model_with_burn_in, SingleDemeModel)
                     num += 1;
                 }
             BOOST_REQUIRE_EQUAL(num, g.number_of_demes());
+
             pop.generation += 1;
             g.iterate_state();
         }
     BOOST_REQUIRE_EQUAL(pop.generation, 10);
 }
 
+BOOST_FIXTURE_TEST_CASE(single_deme_model_with_burn_in_invalid_ancestry_request,
+                        SingleDemeModel)
+{
+    fwdpy11_core::ForwardDemesGraph g(yaml, 10);
+    BOOST_REQUIRE_EQUAL(g.number_of_demes(), 1);
+    fwdpy11::DiploidPopulation pop(100, 1.0);
+    g.initialize_model(pop.generation);
+    auto end_time = g.model_end_time();
+    BOOST_REQUIRE_EQUAL(end_time, 11);
+    while (g.iterating_model())
+        {
+            // this is an invalid index
+            BOOST_REQUIRE_THROW(
+                {
+                    // We are expecting (?) the API to set the error code here.
+                    /* auto ancestry = */ g.offspring_ancestry_proportions(
+                        g.number_of_demes());
+                },
+                fwdpy11::discrete_demography::DemographyError);
+            g.iterate_state();
+        }
+}
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/lib/core/demes/forward_graph.hpp b/lib/core/demes/forward_graph.hpp
@@ -6,8 +6,7 @@
 
 namespace fwdpy11_core
 {
-    template<typename T>
-    struct ForwardDemesGraphDataIterator
+    template <typename T> struct ForwardDemesGraphDataIterator
     {
         const T *first;
         const T *last;
@@ -48,5 +47,7 @@ namespace fwdpy11_core
         ForwardDemesGraphDataIterator<double> offspring_deme_sizes() const;
         ForwardDemesGraphDataIterator<double> offspring_selfing_rates() const;
         ForwardDemesGraphDataIterator<double> offspring_cloning_rates() const;
+        ForwardDemesGraphDataIterator<double>
+        offspring_ancestry_proportions(std::size_t offspring_deme) const;
     };
 }
diff --git a/lib/demes/forward_graph.cc b/lib/demes/forward_graph.cc
@@ -216,4 +216,17 @@ namespace fwdpy11_core
         throw_if_null(begin, __FILE__, __LINE__);
         return ForwardDemesGraphDataIterator<double>{begin, begin + number_of_demes()};
     }
+
+    ForwardDemesGraphDataIterator<double>
+    ForwardDemesGraph::offspring_ancestry_proportions(std::size_t offspring_deme) const
+    {
+        std::int32_t status;
+        auto begin = demes_forward_graph_ancestry_proportions(offspring_deme, &status,
+                                                              pimpl->graph.get());
+        pimpl->handle_error_code(status);
+        // NOTE: this might be overly strict,
+        // but we keep it for now.
+        throw_if_null(begin, __FILE__, __LINE__);
+        return ForwardDemesGraphDataIterator<double>{begin, begin + number_of_demes()};
+    }
 }