Skip to content

Commit

Permalink
table collection API update (#351)
Browse files Browse the repository at this point in the history 
* Track if a table_collection has up-to-date indexes.

* Add ts::count_trees.

* Add test of number of trees vs tskit
  • Loading branch information
@molpopgen
molpopgen authored Jul 7, 2021
1 parent 2b3ed31 commit 2a3a951
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 4 deletions.
50 changes: 50 additions & 0 deletions fwdpp/ts/table_collection_functions.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,56 @@ namespace fwdpp
sort_edge_table(edge_table_offset, tables);
sort_mutation_table_and_rebuild_site_table(tables);
}

template <typename TableCollectionType>
std::size_t
count_trees(const TableCollectionType& tables)
/// Count the number of trees.
/// Tables must be indexed.
///
/// Implementation ported from forrustts.
{
if (!tables.indexed())
{
throw tables_error("tables are not indexed");
}
std::size_t num_trees = 0;

auto input_edge = begin(tables.input_left);
auto output_edge = begin(tables.output_right);
const auto input_end = end(tables.input_left);
const auto output_end = end(tables.output_right);

double tree_left = 0.0;

while (input_edge < input_end || tree_left < tables.genome_length())
{
while (output_edge < output_end
&& tables.edges[*output_edge].right == tree_left)
{
++output_edge;
}
while (input_edge < input_end
&& tables.edges[*input_edge].left == tree_left)
{
++input_edge;
}
auto tree_right = tables.genome_length();
if (input_edge < input_end)
{
tree_right
= std::min(tree_right, tables.edges[*input_edge].left);
}
if (output_edge < output_end)
{
tree_right
= std::min(tree_right, tables.edges[*output_edge].right);
}
tree_left = tree_right;
++num_trees;
}
return num_trees;
}
}
}

Expand Down
21 changes: 17 additions & 4 deletions fwdpp/ts/types/table_collection.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ namespace fwdpp
private:
/// Length of the genomic region.
double L;
bool indexed_;

public:
static constexpr SignedInteger null = generate_null_id<SignedInteger>();
Expand Down Expand Up @@ -76,8 +77,8 @@ namespace fwdpp
std::ptrdiff_t edge_offset;

explicit table_collection(const double maxpos)
: L{maxpos}, nodes{}, edges{}, mutations{}, sites{}, input_left{},
output_right{}, edge_offset{0}
: L{maxpos}, indexed_{false}, nodes{}, edges{}, mutations{}, sites{},
input_left{}, output_right{}, edge_offset{0}
{
if (maxpos <= 0 || !std::isfinite(maxpos))
{
Expand All @@ -88,8 +89,8 @@ namespace fwdpp
table_collection(const id_type num_initial_nodes,
const double initial_time, id_type pop,
const double maxpos)
: L{maxpos}, nodes{}, edges{}, mutations{}, sites{}, input_left{},
output_right{}, edge_offset{0}
: L{maxpos}, indexed_{false}, nodes{}, edges{}, mutations{}, sites{},
input_left{}, output_right{}, edge_offset{0}
{
if (maxpos <= 0 || !std::isfinite(maxpos))
{
Expand All @@ -114,6 +115,7 @@ namespace fwdpp
id_type
push_back_node(double time, id_type pop, std::uint32_t flags)
{
indexed_ = false;
nodes.push_back(node{pop, time, flags});
return static_cast<id_type>(nodes.size() - 1);
}
Expand All @@ -122,13 +124,15 @@ namespace fwdpp
id_type
emplace_back_node(args&&... Args)
{
indexed_ = false;
nodes.emplace_back(node{std::forward<args>(Args)...});
return static_cast<id_type>(nodes.size() - 1);
}

std::size_t
push_back_edge(double l, double r, id_type parent, id_type child)
{
indexed_ = false;
edges.push_back(edge{l, r, parent, child});
return edges.size();
}
Expand All @@ -137,6 +141,7 @@ namespace fwdpp
std::size_t
emplace_back_edge(args&&... Args)
{
indexed_ = false;
edges.emplace_back(edge{std::forward<args>(Args)...});
return edges.size();
}
Expand Down Expand Up @@ -182,6 +187,7 @@ namespace fwdpp
output_right.clear();
input_left.resize(edges.size());
output_right.resize(edges.size());
indexed_ = false;
std::iota(begin(input_left), end(input_left), 0);
std::iota(begin(output_right), end(output_right), 0);
std::sort(begin(input_left), end(input_left),
Expand All @@ -202,6 +208,7 @@ namespace fwdpp
}
return edges[i].right < edges[j].right;
});
indexed_ = true;
}

std::size_t
Expand All @@ -222,6 +229,12 @@ namespace fwdpp
return L;
}

inline bool
indexed() const
{
return indexed_;
}

inline bool
operator==(const table_collection& b) const
{
Expand Down
13 changes: 13 additions & 0 deletions testsuite/tree_sequences/test_edge_buffering_std_table_collection.cc
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include <iostream>
#include <boost/test/unit_test.hpp>
#include <fwdpp/ts/table_collection.hpp>
#include <fwdpp/ts/table_collection_functions.hpp>
#include "wfevolve_table_collection.hpp"
#include "tskit_utils.hpp"

Expand Down Expand Up @@ -104,9 +105,21 @@ BOOST_AUTO_TEST_CASE(test_kc_distance)
ll.tables_buffering, ll.nsteps, ll.samples_buffering);
auto tsk_tables_sort = dump_table_collection_to_tskit(ll.tables_buffering, ll.nsteps,
ll.samples_sorting);
ll.tables_buffering.build_indexes();
ll.tables_sorting.build_indexes();
tsk_treeseq_wrapper treeseq_buffer(tsk_tables_buffer.get());
tsk_treeseq_wrapper treeseq_sort(tsk_tables_buffer.get());
double kc_distance = std::numeric_limits<double>::quiet_NaN();
auto ntrees_buffer = fwdpp::ts::count_trees(ll.tables_buffering);
BOOST_REQUIRE(ntrees_buffer > 0);
auto ntrees_sort = fwdpp::ts::count_trees(ll.tables_sorting);
BOOST_REQUIRE(ntrees_sort > 0);
BOOST_REQUIRE_EQUAL(ntrees_buffer, ntrees_sort);
BOOST_REQUIRE_EQUAL(static_cast<tsk_size_t>(ntrees_buffer),
tsk_treeseq_get_num_trees(treeseq_buffer.get()));
BOOST_REQUIRE_EQUAL(static_cast<tsk_size_t>(ntrees_buffer),
tsk_treeseq_get_num_trees(treeseq_sort.get()));

BOOST_REQUIRE_NO_THROW({
auto rv = tsk_treeseq_kc_distance(treeseq_buffer.get(), treeseq_sort.get(), 0.,
&kc_distance);
Expand Down

0 comments on commit 2a3a951

Please sign in to comment.