Skip to content

Commit

Permalink
Convert kc_distance to use tree_position_t
Browse files Browse the repository at this point in the history
  • Loading branch information
@jeromekelleher
jeromekelleher committed Jul 13, 2023
1 parent 7af048d commit 3dcb4b9
Showing 1 changed file with 54 additions and 52 deletions.
106 changes: 54 additions & 52 deletions c/tskit/trees.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6104,51 +6104,55 @@ update_kc_subtree_state(
}

static int
update_kc_incremental(tsk_tree_t *self, kc_vectors *kc, tsk_edge_list_t *edges_out,
tsk_edge_list_t *edges_in, tsk_size_t *depths)
update_kc_incremental(
tsk_tree_t *tree, kc_vectors *kc, tsk_tree_position_t *tree_pos, tsk_size_t *depths)
{
int ret = 0;
tsk_edge_list_node_t *record;
tsk_edge_t *e;
tsk_id_t u;
tsk_id_t u, v, e, j;
double root_time, time;
const double *times = self->tree_sequence->tables->nodes.time;
const double *restrict times = tree->tree_sequence->tables->nodes.time;
const tsk_id_t *restrict edges_child = tree->tree_sequence->tables->edges.child;
const tsk_id_t *restrict edges_parent = tree->tree_sequence->tables->edges.parent;

tsk_bug_assert(tree_pos->index == tree->index);
tsk_bug_assert(tree_pos->interval.left == tree->interval.left);
tsk_bug_assert(tree_pos->interval.right == tree->interval.right);

/* Update state of detached subtrees */
for (record = edges_out->tail; record != NULL; record = record->prev) {
e = &record->edge;
u = e->child;
for (j = tree_pos->out.stop - 1; j >= tree_pos->out.start; j--) {
e = tree_pos->out.order[j];
u = edges_child[e];
depths[u] = 0;

if (self->parent[u] == TSK_NULL) {
root_time = times[tsk_tree_node_root(self, u)];
ret = update_kc_subtree_state(self, kc, u, depths, root_time);
if (tree->parent[u] == TSK_NULL) {
root_time = times[tsk_tree_node_root(tree, u)];
ret = update_kc_subtree_state(tree, kc, u, depths, root_time);
if (ret != 0) {
goto out;
}
}
}

/* Propagate state change down into reattached subtrees. */
for (record = edges_in->tail; record != NULL; record = record->prev) {
e = &record->edge;
u = e->child;
for (j = tree_pos->in.stop - 1; j >= tree_pos->in.start; j--) {
e = tree_pos->in.order[j];
u = edges_child[e];
v = edges_parent[e];

tsk_bug_assert(depths[e->child] == 0);
depths[u] = depths[e->parent] + 1;
tsk_bug_assert(depths[u] == 0);
depths[u] = depths[v] + 1;

root_time = times[tsk_tree_node_root(self, u)];
ret = update_kc_subtree_state(self, kc, u, depths, root_time);
root_time = times[tsk_tree_node_root(tree, u)];
ret = update_kc_subtree_state(tree, kc, u, depths, root_time);
if (ret != 0) {
goto out;
}

if (tsk_tree_is_sample(self, u)) {
time = tsk_tree_get_branch_length_unsafe(self, u);
update_kc_vectors_single_sample(self->tree_sequence, kc, u, time);
if (tsk_tree_is_sample(tree, u)) {
time = tsk_tree_get_branch_length_unsafe(tree, u);
update_kc_vectors_single_sample(tree->tree_sequence, kc, u, time);
}
}

out:
return ret;
}
Expand All @@ -6164,19 +6168,18 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
const tsk_treeseq_t *treeseqs[2] = { self, other };
tsk_tree_t trees[2];
kc_vectors kcs[2];
tsk_diff_iter_t diff_iters[2];
tsk_edge_list_t edges_out[2];
tsk_edge_list_t edges_in[2];
/* TODO the tree_pos here is redundant because we should be using this interally
* in the trees to do the advancing. Once we have converted the tree over to using
* tree_pos internally, we can get rid of these tree_pos variables and use
* the values stored in the trees themselves */
tsk_tree_position_t tree_pos[2];
tsk_size_t *depths[2];
double t0_left, t0_right, t1_left, t1_right;
int ret = 0;

for (i = 0; i < 2; i++) {
tsk_memset(&trees[i], 0, sizeof(trees[i]));
tsk_memset(&diff_iters[i], 0, sizeof(diff_iters[i]));
tsk_memset(&tree_pos[i], 0, sizeof(tree_pos[i]));
tsk_memset(&kcs[i], 0, sizeof(kcs[i]));
tsk_memset(&edges_out[i], 0, sizeof(edges_out[i]));
tsk_memset(&edges_in[i], 0, sizeof(edges_in[i]));
depths[i] = NULL;
}

Expand All @@ -6191,7 +6194,7 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
if (ret != 0) {
goto out;
}
ret = tsk_diff_iter_init_from_ts(&diff_iters[i], treeseqs[i], false);
ret = tsk_tree_position_init(&tree_pos[i], treeseqs[i], 0);
if (ret != 0) {
goto out;
}
Expand All @@ -6218,11 +6221,10 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
if (ret != 0) {
goto out;
}
ret = tsk_diff_iter_next(
&diff_iters[0], &t0_left, &t0_right, &edges_out[0], &edges_in[0]);
tsk_bug_assert(ret == TSK_TREE_OK);
ret = update_kc_incremental(
&trees[0], &kcs[0], &edges_out[0], &edges_in[0], depths[0]);
tsk_tree_position_next(&tree_pos[0]);
tsk_bug_assert(tree_pos[0].index == 0);

ret = update_kc_incremental(&trees[0], &kcs[0], &tree_pos[0], depths[0]);
if (ret != 0) {
goto out;
}
Expand All @@ -6231,37 +6233,37 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
if (ret != 0) {
goto out;
}
ret = tsk_diff_iter_next(
&diff_iters[1], &t1_left, &t1_right, &edges_out[1], &edges_in[1]);
tsk_bug_assert(ret == TSK_TREE_OK);
tsk_tree_position_next(&tree_pos[1]);
tsk_bug_assert(tree_pos[1].index != -1);

ret = update_kc_incremental(
&trees[1], &kcs[1], &edges_out[1], &edges_in[1], depths[1]);
ret = update_kc_incremental(&trees[1], &kcs[1], &tree_pos[1], depths[1]);
if (ret != 0) {
goto out;
}
while (t0_right < t1_right) {
span = t0_right - left;
tsk_bug_assert(trees[0].interval.left == tree_pos[0].interval.left);
tsk_bug_assert(trees[0].interval.right == tree_pos[0].interval.right);
tsk_bug_assert(trees[1].interval.left == tree_pos[1].interval.left);
tsk_bug_assert(trees[1].interval.right == tree_pos[1].interval.right);
while (trees[0].interval.right < trees[1].interval.right) {
span = trees[0].interval.right - left;
total += norm_kc_vectors(&kcs[0], &kcs[1], lambda_) * span;

left = t0_right;
left = trees[0].interval.right;
ret = tsk_tree_next(&trees[0]);
tsk_bug_assert(ret == TSK_TREE_OK);
ret = check_kc_distance_tree_inputs(&trees[0]);
if (ret != 0) {
goto out;
}
ret = tsk_diff_iter_next(
&diff_iters[0], &t0_left, &t0_right, &edges_out[0], &edges_in[0]);
tsk_bug_assert(ret == TSK_TREE_OK);
ret = update_kc_incremental(
&trees[0], &kcs[0], &edges_out[0], &edges_in[0], depths[0]);
tsk_tree_position_next(&tree_pos[0]);
tsk_bug_assert(tree_pos[0].index != -1);
ret = update_kc_incremental(&trees[0], &kcs[0], &tree_pos[0], depths[0]);
if (ret != 0) {
goto out;
}
}
span = t1_right - left;
left = t1_right;
span = trees[1].interval.right - left;
left = trees[1].interval.right;
total += norm_kc_vectors(&kcs[0], &kcs[1], lambda_) * span;
}
if (ret != 0) {
Expand All @@ -6272,7 +6274,7 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
out:
for (i = 0; i < 2; i++) {
tsk_tree_free(&trees[i]);
tsk_diff_iter_free(&diff_iters[i]);
tsk_tree_position_free(&tree_pos[i]);
kc_vectors_free(&kcs[i]);
tsk_safe_free(depths[i]);
}
Expand Down

0 comments on commit 3dcb4b9

Please sign in to comment.