diff --git a/c/tskit/trees.c b/c/tskit/trees.c index c2feda4b2d..70dcc6846c 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -9929,7 +9929,7 @@ tsk_matvec_calculator_print_state(const tsk_matvec_calculator_t *self, FILE *out fprintf(out, "options = %d\n", self->options); fprintf(out, "position = %f\n", self->position); fprintf(out, "tree_pos:\n"); - tsk_tree_position_print_state(&self->tree_pos); + tsk_tree_position_print_state(&self->tree_pos, out); fprintf(out, "samples = %lld: [", (long long) num_samples); fprintf(out, "]\n"); fprintf(out, "node\tparent\tx\tv\tw"); @@ -9964,7 +9964,7 @@ tsk_matvec_calculator_init(tsk_matvec_calculator_t *self, const tsk_treeseq_t *t self->options = options; self->result = result; self->num_nodes = num_nodes; - self->position = 0.0; + self->position = windows[0]; self->parent = tsk_malloc(num_nodes * sizeof(*self->parent)); self->x = tsk_calloc(num_nodes, sizeof(*self->x)); @@ -10183,81 +10183,82 @@ tsk_matvec_calculator_run(tsk_matvec_calculator_t *self) tsk_size_t j, k, m; tsk_id_t e, p, c; tsk_size_t n = tsk_treeseq_get_num_samples(self->ts); - double tree_right; - const double sequence_length = self->ts->tables->sequence_length; const tsk_size_t num_edges = self->ts->tables->edges.num_rows; - const tsk_id_t *restrict I = self->ts->tables->indexes.edge_insertion_order; - const tsk_id_t *restrict O = self->ts->tables->indexes.edge_removal_order; const double *restrict edge_right = self->ts->tables->edges.right; const double *restrict edge_left = self->ts->tables->edges.left; const tsk_id_t *restrict edge_child = self->ts->tables->edges.child; const tsk_id_t *restrict edge_parent = self->ts->tables->edges.parent; - const tsk_size_t num_trees = self->tree_sequence->num_trees; + const tsk_size_t num_trees = self->ts->num_trees; + const double *restrict windows = self->windows; double *restrict out; - const double *restrict breakpoints = self->tree_sequence->breakpoints; + const double *restrict breakpoints = self->ts->breakpoints; + tsk_tree_position_t tree_pos = self->tree_pos; tsk_id_t index; bool valid; + double next_position; - j = 0; - k = 0; m = 0; - tree_right = sequence_length; + self->position = windows[0]; /* seek to the first window */ index = (tsk_id_t) tsk_search_sorted(breakpoints, num_trees + 1, windows[0]); - if (breakpoints[index] > x) { + if (breakpoints[index] > windows[0]) { index--; } - ret = tsk_tree_position_seek_forward(&self->tree_pos, index); + ret = tsk_tree_position_seek_forward(&tree_pos, index); if (ret != 0) { goto out; } - self->position = windows[0]; - for (j = tree_pos.in.start; j != tree_pos.in.stop; j++) { + for (j = (tsk_size_t) tree_pos.in.start; j != (tsk_size_t) tree_pos.in.stop; j++) { e = tree_pos.in.order[j]; if (edge_left[e] <= self->position && self->position < edge_right[e]) { - tsk_matvec_calculator_insert_edge(self, edge_parent[e], edge_child[e], e); + p = edge_parent[e]; + c = edge_child[e]; + tsk_matvec_calculator_insert_edge(self, p, c); } } - valid = true; - while (valid && m < self->num_windows) { - if (self->position == self->tree_pos.left) { - for (k = tree_pos.out.start; k != tree_pos.out.stop; k++) { + valid = tsk_tree_position_next(&tree_pos); + j = (tsk_size_t) tree_pos.in.start - 1; + k = (tsk_size_t) tree_pos.out.start - 1; + while (m < self->num_windows) { + if (valid && self->position == tree_pos.interval.left) { + for (k = (tsk_size_t) tree_pos.out.start; + k != (tsk_size_t) tree_pos.out.stop; k++) { e = tree_pos.out.order[k]; p = edge_parent[e]; c = edge_child[e]; tsk_matvec_calculator_remove_edge(self, p, c); } - for (j = tree_pos.in.start; j != tree_pos.in.stop; j++) { + for (j = (tsk_size_t) tree_pos.in.start; j != (tsk_size_t) tree_pos.in.stop; + j++) { e = tree_pos.in.order[j]; p = edge_parent[e]; c = edge_child[e]; tsk_matvec_calculator_insert_edge(self, p, c); } + valid = tsk_tree_position_next(&tree_pos); } - next_position = self->windows[m + 1]; - if (j < num_edges) { - next_position = TSK_MIN(next_position, edge_left[tree_pos.in.order[j]]); + next_position = windows[m + 1]; + if (j + 1 < num_edges) { + next_position = TSK_MIN(next_position, edge_left[tree_pos.in.order[j + 1]]); } - if (k < num_edges) { - next_position = TSK_MIN(next_position, edge_right[tree_pos.out.order[k]]); + if (k + 1 < num_edges) { + next_position + = TSK_MIN(next_position, edge_right[tree_pos.out.order[k + 1]]); } self->position = next_position; - if (self->position == self->windows[m + 1]) { + if (self->position == windows[m + 1]) { out = GET_2D_ROW(self->result, self->num_weights * n, m); tsk_matvec_calculator_write_output(self, out); m += 1; } - if (self->position == self->tree_pos.right) { - valid = tsk_tree_position_next(&tree_pos); - } if (self->options & TSK_DEBUG) { tsk_matvec_calculator_print_state(self, tsk_get_debug_stream()); } } - /* out: */ +out: return ret; } diff --git a/python/tests/test_relatedness_vector.py b/python/tests/test_relatedness_vector.py index fa588df538..6537a6df68 100644 --- a/python/tests/test_relatedness_vector.py +++ b/python/tests/test_relatedness_vector.py @@ -248,8 +248,6 @@ def run(self): valid = tree_pos.next() j = tree_pos.in_range.start - 1 k = tree_pos.out_range.start - 1 - print(tree_pos) - print("windows", self.windows) while m < num_windows: if valid and self.position == tree_pos.interval.left: for k in range(tree_pos.out_range.start, tree_pos.out_range.stop, 1):