C version passes

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;
 }
 

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):