Skip to content

Commit f93e792

Browse files
jeromekelleherjeromekelleher
committed
Initial version of tsk_tree_position_t and tests
1 parent 67f2ef3 commit f93e792

File tree

5 files changed

+1007
-0
lines changed

5 files changed

+1007
-0
lines changed

‎c/tests/test_trees.c‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,97 @@ verify_individual_nodes(tsk_treeseq_t *ts)
175175
}
176176
}
177177

178+
static void
179+
verify_tree_pos(const tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *tree_parents)
180+
{
181+
int ret;
182+
const tsk_size_t N = tsk_treeseq_get_num_nodes(ts);
183+
const tsk_id_t *edges_parent = ts->tables->edges.parent;
184+
const tsk_id_t *edges_child = ts->tables->edges.child;
185+
tsk_tree_position_t tree_pos;
186+
tsk_id_t *known_parent;
187+
tsk_id_t *parent = tsk_malloc(N * sizeof(*parent));
188+
tsk_id_t u, index, j, e;
189+
bool valid;
190+
191+
CU_ASSERT_FATAL(parent != NULL);
192+
193+
ret = tsk_tree_position_init(&tree_pos, ts, 0);
194+
CU_ASSERT_EQUAL_FATAL(ret, 0);
195+
196+
for (u = 0; u < (tsk_id_t) N; u++) {
197+
parent[u] = TSK_NULL;
198+
}
199+
200+
for (index = 0; index < (tsk_id_t) num_trees; index++) {
201+
known_parent = tree_parents + N * (tsk_size_t) index;
202+
203+
valid = tsk_tree_position_next(&tree_pos);
204+
CU_ASSERT_TRUE(valid);
205+
CU_ASSERT_EQUAL(index, tree_pos.index);
206+
207+
for (j = tree_pos.out.start; j < tree_pos.out.stop; j++) {
208+
e = tree_pos.out.order[j];
209+
parent[edges_child[e]] = TSK_NULL;
210+
}
211+
212+
for (j = tree_pos.in.start; j < tree_pos.in.stop; j++) {
213+
e = tree_pos.in.order[j];
214+
parent[edges_child[e]] = edges_parent[e];
215+
}
216+
217+
for (u = 0; u < (tsk_id_t) N; u++) {
218+
CU_ASSERT_EQUAL(parent[u], known_parent[u]);
219+
}
220+
}
221+
222+
valid = tsk_tree_position_next(&tree_pos);
223+
CU_ASSERT_FALSE(valid);
224+
for (j = tree_pos.out.start; j < tree_pos.out.stop; j++) {
225+
e = tree_pos.out.order[j];
226+
parent[edges_child[e]] = TSK_NULL;
227+
}
228+
for (u = 0; u < (tsk_id_t) N; u++) {
229+
CU_ASSERT_EQUAL(parent[u], TSK_NULL);
230+
}
231+
232+
for (index = (tsk_id_t) num_trees - 1; index >= 0; index--) {
233+
known_parent = tree_parents + N * (tsk_size_t) index;
234+
235+
valid = tsk_tree_position_prev(&tree_pos);
236+
CU_ASSERT_TRUE(valid);
237+
CU_ASSERT_EQUAL(index, tree_pos.index);
238+
239+
for (j = tree_pos.out.start; j > tree_pos.out.stop; j--) {
240+
e = tree_pos.out.order[j];
241+
parent[edges_child[e]] = TSK_NULL;
242+
}
243+
244+
for (j = tree_pos.in.start; j > tree_pos.in.stop; j--) {
245+
CU_ASSERT_FATAL(j >= 0);
246+
e = tree_pos.in.order[j];
247+
parent[edges_child[e]] = edges_parent[e];
248+
}
249+
250+
for (u = 0; u < (tsk_id_t) N; u++) {
251+
CU_ASSERT_EQUAL(parent[u], known_parent[u]);
252+
}
253+
}
254+
255+
valid = tsk_tree_position_prev(&tree_pos);
256+
CU_ASSERT_FALSE(valid);
257+
for (j = tree_pos.out.start; j > tree_pos.out.stop; j--) {
258+
e = tree_pos.out.order[j];
259+
parent[edges_child[e]] = TSK_NULL;
260+
}
261+
for (u = 0; u < (tsk_id_t) N; u++) {
262+
CU_ASSERT_EQUAL(parent[u], TSK_NULL);
263+
}
264+
265+
tsk_tree_position_free(&tree_pos);
266+
tsk_safe_free(parent);
267+
}
268+
178269
static void
179270
verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
180271
{
@@ -233,6 +324,8 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
233324
CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(ts), breakpoints[j]);
234325

235326
tsk_tree_free(&tree);
327+
328+
verify_tree_pos(ts, num_trees, parents);
236329
}
237330

238331
static tsk_tree_t *
@@ -5233,6 +5326,65 @@ test_single_tree_tracked_samples(void)
52335326
tsk_tree_free(&tree);
52345327
}
52355328

5329+
static void
5330+
test_single_tree_tree_pos(void)
5331+
{
5332+
tsk_treeseq_t ts;
5333+
tsk_tree_position_t tree_pos;
5334+
bool valid;
5335+
int ret;
5336+
5337+
tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, NULL,
5338+
NULL, NULL, NULL, 0);
5339+
5340+
ret = tsk_tree_position_init(&tree_pos, &ts, 0);
5341+
CU_ASSERT_EQUAL_FATAL(ret, 0);
5342+
valid = tsk_tree_position_next(&tree_pos);
5343+
CU_ASSERT_FATAL(valid);
5344+
5345+
CU_ASSERT_EQUAL_FATAL(tree_pos.interval.left, 0);
5346+
CU_ASSERT_EQUAL_FATAL(tree_pos.interval.right, 1);
5347+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.start, 0);
5348+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.stop, 6);
5349+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.order, ts.tables->indexes.edge_insertion_order);
5350+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 0);
5351+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 0);
5352+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_removal_order);
5353+
5354+
valid = tsk_tree_position_next(&tree_pos);
5355+
CU_ASSERT_FATAL(!valid);
5356+
5357+
tsk_tree_position_print_state(&tree_pos, _devnull);
5358+
5359+
CU_ASSERT_EQUAL_FATAL(tree_pos.index, -1);
5360+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 0);
5361+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 6);
5362+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_removal_order);
5363+
5364+
valid = tsk_tree_position_prev(&tree_pos);
5365+
CU_ASSERT_FATAL(valid);
5366+
5367+
CU_ASSERT_EQUAL_FATAL(tree_pos.interval.left, 0);
5368+
CU_ASSERT_EQUAL_FATAL(tree_pos.interval.right, 1);
5369+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.start, 5);
5370+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.stop, -1);
5371+
CU_ASSERT_EQUAL_FATAL(tree_pos.in.order, ts.tables->indexes.edge_removal_order);
5372+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 5);
5373+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 5);
5374+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_insertion_order);
5375+
5376+
valid = tsk_tree_position_prev(&tree_pos);
5377+
CU_ASSERT_FATAL(!valid);
5378+
5379+
CU_ASSERT_EQUAL_FATAL(tree_pos.index, -1);
5380+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 5);
5381+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, -1);
5382+
CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_insertion_order);
5383+
5384+
tsk_tree_position_free(&tree_pos);
5385+
tsk_treeseq_free(&ts);
5386+
}
5387+
52365388
/*=======================================================
52375389
* Multi tree tests.
52385390
*======================================================*/
@@ -8185,6 +8337,7 @@ main(int argc, char **argv)
81858337
{ "test_single_tree_map_mutations_internal_samples",
81868338
test_single_tree_map_mutations_internal_samples },
81878339
{ "test_single_tree_tracked_samples", test_single_tree_tracked_samples },
8340+
{ "test_single_tree_tree_pos", test_single_tree_tree_pos },
81888341

81898342
/* Multi tree tests */
81908343
{ "test_simple_multi_tree", test_simple_multi_tree },

‎c/tskit/trees.c‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3553,6 +3553,164 @@ tsk_treeseq_split_edges(const tsk_treeseq_t *self, double time, tsk_flags_t flag
35533553
return ret;
35543554
}
35553555

3556+
/* ======================================================== *
3557+
* tree_position
3558+
* ======================================================== */
3559+
3560+
static void
3561+
tsk_tree_position_set_null(tsk_tree_position_t *self)
3562+
{
3563+
self->index = -1;
3564+
self->interval.left = 0;
3565+
self->interval.right = 0;
3566+
}
3567+
3568+
int
3569+
tsk_tree_position_init(tsk_tree_position_t *self, const tsk_treeseq_t *tree_sequence,
3570+
tsk_flags_t TSK_UNUSED(options))
3571+
{
3572+
memset(self, 0, sizeof(*self));
3573+
self->tree_sequence = tree_sequence;
3574+
tsk_tree_position_set_null(self);
3575+
return 0;
3576+
}
3577+
3578+
int
3579+
tsk_tree_position_free(tsk_tree_position_t *TSK_UNUSED(self))
3580+
{
3581+
return 0;
3582+
}
3583+
3584+
int
3585+
tsk_tree_position_print_state(const tsk_tree_position_t *self, FILE *out)
3586+
{
3587+
fprintf(out, "Tree position state\n");
3588+
fprintf(out, "index = %d\n", (int) self->index);
3589+
fprintf(
3590+
out, "out = start=%d\tstop=%d\n", (int) self->out.start, (int) self->out.stop);
3591+
fprintf(
3592+
out, "in = start=%d\tstop=%d\n", (int) self->in.start, (int) self->in.stop);
3593+
return 0;
3594+
}
3595+
3596+
bool
3597+
tsk_tree_position_next(tsk_tree_position_t *self)
3598+
{
3599+
const tsk_table_collection_t *tables = self->tree_sequence->tables;
3600+
const tsk_id_t M = (tsk_id_t) tables->edges.num_rows;
3601+
const tsk_id_t num_trees = (tsk_id_t) self->tree_sequence->num_trees;
3602+
const double *restrict left_coords = tables->edges.left;
3603+
const tsk_id_t *restrict left_order = tables->indexes.edge_insertion_order;
3604+
const double *restrict right_coords = tables->edges.right;
3605+
const tsk_id_t *restrict right_order = tables->indexes.edge_removal_order;
3606+
const double *restrict breakpoints = self->tree_sequence->breakpoints;
3607+
tsk_id_t j, left_current_index, right_current_index;
3608+
double left;
3609+
3610+
if (self->index == -1) {
3611+
self->interval.right = 0;
3612+
self->in.stop = 0;
3613+
self->out.stop = 0;
3614+
self->direction = TSK_DIR_FORWARD;
3615+
}
3616+
3617+
if (self->direction == TSK_DIR_FORWARD) {
3618+
left_current_index = self->in.stop;
3619+
right_current_index = self->out.stop;
3620+
} else {
3621+
left_current_index = self->out.stop + 1;
3622+
right_current_index = self->in.stop + 1;
3623+
}
3624+
3625+
left = self->interval.right;
3626+
3627+
j = right_current_index;
3628+
self->out.start = j;
3629+
while (j < M && right_coords[right_order[j]] == left) {
3630+
j++;
3631+
}
3632+
self->out.stop = j;
3633+
self->out.order = right_order;
3634+
3635+
j = left_current_index;
3636+
self->in.start = j;
3637+
while (j < M && left_coords[left_order[j]] == left) {
3638+
j++;
3639+
}
3640+
self->in.stop = j;
3641+
self->in.order = left_order;
3642+
3643+
self->direction = TSK_DIR_FORWARD;
3644+
self->index++;
3645+
if (self->index == num_trees) {
3646+
tsk_tree_position_set_null(self);
3647+
} else {
3648+
self->interval.left = left;
3649+
self->interval.right = breakpoints[self->index + 1];
3650+
}
3651+
return self->index != -1;
3652+
}
3653+
3654+
bool
3655+
tsk_tree_position_prev(tsk_tree_position_t *self)
3656+
{
3657+
const tsk_table_collection_t *tables = self->tree_sequence->tables;
3658+
const tsk_id_t M = (tsk_id_t) tables->edges.num_rows;
3659+
const double sequence_length = tables->sequence_length;
3660+
const tsk_id_t num_trees = (tsk_id_t) self->tree_sequence->num_trees;
3661+
const double *restrict left_coords = tables->edges.left;
3662+
const tsk_id_t *restrict left_order = tables->indexes.edge_insertion_order;
3663+
const double *restrict right_coords = tables->edges.right;
3664+
const tsk_id_t *restrict right_order = tables->indexes.edge_removal_order;
3665+
const double *restrict breakpoints = self->tree_sequence->breakpoints;
3666+
tsk_id_t j, left_current_index, right_current_index;
3667+
double right;
3668+
3669+
if (self->index == -1) {
3670+
self->index = num_trees;
3671+
self->interval.left = sequence_length;
3672+
self->in.stop = M - 1;
3673+
self->out.stop = M - 1;
3674+
self->direction = TSK_DIR_REVERSE;
3675+
}
3676+
3677+
if (self->direction == TSK_DIR_REVERSE) {
3678+
left_current_index = self->out.stop;
3679+
right_current_index = self->in.stop;
3680+
} else {
3681+
left_current_index = self->in.stop - 1;
3682+
right_current_index = self->out.stop - 1;
3683+
}
3684+
3685+
right = self->interval.left;
3686+
3687+
j = left_current_index;
3688+
self->out.start = j;
3689+
while (j >= 0 && left_coords[left_order[j]] == right) {
3690+
j--;
3691+
}
3692+
self->out.stop = j;
3693+
self->out.order = left_order;
3694+
3695+
j = right_current_index;
3696+
self->in.start = j;
3697+
while (j >= 0 && right_coords[right_order[j]] == right) {
3698+
j--;
3699+
}
3700+
self->in.stop = j;
3701+
self->in.order = right_order;
3702+
3703+
self->index--;
3704+
self->direction = TSK_DIR_REVERSE;
3705+
if (self->index == -1) {
3706+
tsk_tree_position_set_null(self);
3707+
} else {
3708+
self->interval.left = breakpoints[self->index];
3709+
self->interval.right = right;
3710+
}
3711+
return self->index != -1;
3712+
}
3713+
35563714
/* ======================================================== *
35573715
* Tree
35583716
* ======================================================== */

‎c/tskit/trees.h‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1739,6 +1739,40 @@ bool tsk_tree_equals(const tsk_tree_t *self, const tsk_tree_t *other);
17391739
int tsk_diff_iter_init_from_ts(
17401740
tsk_diff_iter_t *self, const tsk_treeseq_t *tree_sequence, tsk_flags_t options);
17411741

1742+
/* Temporarily putting this here to avoid problems with doxygen. Will need to
1743+
* move up the file later when it gets incorporated into the tsk_tree_t object.
1744+
*/
1745+
typedef struct {
1746+
tsk_id_t index;
1747+
struct {
1748+
double left;
1749+
double right;
1750+
} interval;
1751+
struct {
1752+
tsk_id_t start;
1753+
tsk_id_t stop;
1754+
const tsk_id_t *order;
1755+
} in;
1756+
struct {
1757+
tsk_id_t start;
1758+
tsk_id_t stop;
1759+
const tsk_id_t *order;
1760+
} out;
1761+
tsk_id_t left_current_index;
1762+
tsk_id_t right_current_index;
1763+
int direction;
1764+
const tsk_treeseq_t *tree_sequence;
1765+
} tsk_tree_position_t;
1766+
1767+
int tsk_tree_position_init(
1768+
tsk_tree_position_t *self, const tsk_treeseq_t *tree_sequence, tsk_flags_t options);
1769+
int tsk_tree_position_free(tsk_tree_position_t *self);
1770+
int tsk_tree_position_print_state(const tsk_tree_position_t *self, FILE *out);
1771+
bool tsk_tree_position_next(tsk_tree_position_t *self);
1772+
bool tsk_tree_position_prev(tsk_tree_position_t *self);
1773+
int tsk_tree_position_seek_forward(tsk_tree_position_t *self, tsk_id_t index);
1774+
int tsk_tree_position_seek_backward(tsk_tree_position_t *self, tsk_id_t index);
1775+
17421776
#ifdef __cplusplus
17431777
}
17441778
#endif

0 commit comments

Comments
 (0)