Make skipping fixations optional when generating a data_matrix from t…

…ree sequences.
molpopgen · Apr 22, 2019 · 8a5a744 · 8a5a744
1 parent d661f16
commit 8a5a744
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/examples/tree_sequence_examples_common.cc b/examples/tree_sequence_examples_common.cc
@@ -181,7 +181,7 @@ matrix_runtime_test(const fwdpp::ts::table_collection &tables,
                     const std::vector<fwdpp::uint_t> &mcounts)
 {
     auto dm = fwdpp::ts::generate_data_matrix(tables, samples, mutations, true,
-                                              true);
+                                              true, false);
     auto rs = fwdpp::row_sums(dm);
     for (std::size_t i = 0; i < rs.first.size(); ++i)
         {
@@ -398,7 +398,7 @@ write_sfs(const options &o, const fwdpp::GSLrng_mt &rng,
                            s.data(), s.size(), sizeof(fwdpp::ts::TS_NODE_INT));
             std::iota(small_sample.begin(), small_sample.end(), 0);
             auto dm = fwdpp::ts::generate_data_matrix(tables, small_sample,
-                                                      mutations, true, false);
+                                                      mutations, true, false, true);
             auto rs = fwdpp::row_sums(dm);
             std::vector<int> sfs(small_sample.size() - 1);
             for (auto i : rs.first)

diff --git a/fwdpp/ts/generate_data_matrix.hpp b/fwdpp/ts/generate_data_matrix.hpp
@@ -20,12 +20,12 @@ namespace fwdpp
                              const std::vector<TS_NODE_INT>& samples,
                              const mcont_t& mutations,
                              const bool record_neutral,
-                             const bool record_selected, const double start,
-                             const double stop)
+                             const bool record_selected, const bool skip_fixed,
+                             const double start, const double stop)
         /// \todo Document
         /// \version 0.7.0 Added to library
         /// \version 0.7.1 Change behavior to skip sites fixed in the sample
-        /// \version 0.7.4 Add [start, stop) arguments
+        /// \version 0.7.4 Add [start, stop) arguments. Add option to skip fixed variants.
         {
             if (!(stop > start))
                 {
@@ -61,7 +61,9 @@ namespace fwdpp
                                     auto tc = tree.leaf_counts[mut->node]
                                               + tree.preserved_leaf_counts
                                                     [mut->node];
-                                    if (tc < tree.sample_size)
+                                    if (!skip_fixed
+                                        || (skip_fixed
+                                            && tc < tree.sample_size))
                                         {
                                             // Mutation leads to a polymorphism
                                             bool is_neutral
@@ -109,11 +111,11 @@ namespace fwdpp
                              const std::vector<TS_NODE_INT>& samples,
                              const mcont_t& mutations,
                              const bool record_neutral,
-                             const bool record_selected)
+                             const bool record_selected, const bool skip_fixed)
         {
-            return generate_data_matrix(tables, samples, mutations,
-                                        record_neutral, record_selected, 0.,
-                                        tables.genome_length());
+            return generate_data_matrix(
+                tables, samples, mutations, record_neutral, record_selected,
+                skip_fixed, 0., tables.genome_length());
         }
 
     } // namespace ts