Merge pull request #2786 from jeromekelleher/incremental-trees

Modular incremental trees

Author: jeromekelleher

c/tests/test_trees.c

@@ -175,6 +175,97 @@ verify_individual_nodes(tsk_treeseq_t *ts)
     }
 }
 
+static void
+verify_tree_pos(const tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *tree_parents)
+{
+    int ret;
+    const tsk_size_t N = tsk_treeseq_get_num_nodes(ts);
+    const tsk_id_t *edges_parent = ts->tables->edges.parent;
+    const tsk_id_t *edges_child = ts->tables->edges.child;
+    tsk_tree_position_t tree_pos;
+    tsk_id_t *known_parent;
+    tsk_id_t *parent = tsk_malloc(N * sizeof(*parent));
+    tsk_id_t u, index, j, e;
+    bool valid;
+
+    CU_ASSERT_FATAL(parent != NULL);
+
+    ret = tsk_tree_position_init(&tree_pos, ts, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    for (u = 0; u < (tsk_id_t) N; u++) {
+        parent[u] = TSK_NULL;
+    }
+
+    for (index = 0; index < (tsk_id_t) num_trees; index++) {
+        known_parent = tree_parents + N * (tsk_size_t) index;
+
+        valid = tsk_tree_position_next(&tree_pos);
+        CU_ASSERT_TRUE(valid);
+        CU_ASSERT_EQUAL(index, tree_pos.index);
+
+        for (j = tree_pos.out.start; j < tree_pos.out.stop; j++) {
+            e = tree_pos.out.order[j];
+            parent[edges_child[e]] = TSK_NULL;
+        }
+
+        for (j = tree_pos.in.start; j < tree_pos.in.stop; j++) {
+            e = tree_pos.in.order[j];
+            parent[edges_child[e]] = edges_parent[e];
+        }
+
+        for (u = 0; u < (tsk_id_t) N; u++) {
+            CU_ASSERT_EQUAL(parent[u], known_parent[u]);
+        }
+    }
+
+    valid = tsk_tree_position_next(&tree_pos);
+    CU_ASSERT_FALSE(valid);
+    for (j = tree_pos.out.start; j < tree_pos.out.stop; j++) {
+        e = tree_pos.out.order[j];
+        parent[edges_child[e]] = TSK_NULL;
+    }
+    for (u = 0; u < (tsk_id_t) N; u++) {
+        CU_ASSERT_EQUAL(parent[u], TSK_NULL);
+    }
+
+    for (index = (tsk_id_t) num_trees - 1; index >= 0; index--) {
+        known_parent = tree_parents + N * (tsk_size_t) index;
+
+        valid = tsk_tree_position_prev(&tree_pos);
+        CU_ASSERT_TRUE(valid);
+        CU_ASSERT_EQUAL(index, tree_pos.index);
+
+        for (j = tree_pos.out.start; j > tree_pos.out.stop; j--) {
+            e = tree_pos.out.order[j];
+            parent[edges_child[e]] = TSK_NULL;
+        }
+
+        for (j = tree_pos.in.start; j > tree_pos.in.stop; j--) {
+            CU_ASSERT_FATAL(j >= 0);
+            e = tree_pos.in.order[j];
+            parent[edges_child[e]] = edges_parent[e];
+        }
+
+        for (u = 0; u < (tsk_id_t) N; u++) {
+            CU_ASSERT_EQUAL(parent[u], known_parent[u]);
+        }
+    }
+
+    valid = tsk_tree_position_prev(&tree_pos);
+    CU_ASSERT_FALSE(valid);
+    for (j = tree_pos.out.start; j > tree_pos.out.stop; j--) {
+        e = tree_pos.out.order[j];
+        parent[edges_child[e]] = TSK_NULL;
+    }
+    for (u = 0; u < (tsk_id_t) N; u++) {
+        CU_ASSERT_EQUAL(parent[u], TSK_NULL);
+    }
+
+    tsk_tree_position_free(&tree_pos);
+    tsk_safe_free(parent);
+}
+
 static void
 verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
 {
@@ -233,6 +324,8 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
     CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(ts), breakpoints[j]);
 
     tsk_tree_free(&tree);
+
+    verify_tree_pos(ts, num_trees, parents);
 }
 
 static tsk_tree_t *
@@ -5233,6 +5326,65 @@ test_single_tree_tracked_samples(void)
     tsk_tree_free(&tree);
 }
 
+static void
+test_single_tree_tree_pos(void)
+{
+    tsk_treeseq_t ts;
+    tsk_tree_position_t tree_pos;
+    bool valid;
+    int ret;
+
+    tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, NULL,
+        NULL, NULL, NULL, 0);
+
+    ret = tsk_tree_position_init(&tree_pos, &ts, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    valid = tsk_tree_position_next(&tree_pos);
+    CU_ASSERT_FATAL(valid);
+
+    CU_ASSERT_EQUAL_FATAL(tree_pos.interval.left, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.interval.right, 1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.start, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.stop, 6);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.order, ts.tables->indexes.edge_insertion_order);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_removal_order);
+
+    valid = tsk_tree_position_next(&tree_pos);
+    CU_ASSERT_FATAL(!valid);
+
+    tsk_tree_position_print_state(&tree_pos, _devnull);
+
+    CU_ASSERT_EQUAL_FATAL(tree_pos.index, -1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 6);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_removal_order);
+
+    valid = tsk_tree_position_prev(&tree_pos);
+    CU_ASSERT_FATAL(valid);
+
+    CU_ASSERT_EQUAL_FATAL(tree_pos.interval.left, 0);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.interval.right, 1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.start, 5);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.stop, -1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.in.order, ts.tables->indexes.edge_removal_order);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 5);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, 5);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_insertion_order);
+
+    valid = tsk_tree_position_prev(&tree_pos);
+    CU_ASSERT_FATAL(!valid);
+
+    CU_ASSERT_EQUAL_FATAL(tree_pos.index, -1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.start, 5);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.stop, -1);
+    CU_ASSERT_EQUAL_FATAL(tree_pos.out.order, ts.tables->indexes.edge_insertion_order);
+
+    tsk_tree_position_free(&tree_pos);
+    tsk_treeseq_free(&ts);
+}
+
 /*=======================================================
  * Multi tree tests.
  *======================================================*/
@@ -8185,6 +8337,7 @@ main(int argc, char **argv)
         { "test_single_tree_map_mutations_internal_samples",
             test_single_tree_map_mutations_internal_samples },
         { "test_single_tree_tracked_samples", test_single_tree_tracked_samples },
+        { "test_single_tree_tree_pos", test_single_tree_tree_pos },
 
         /* Multi tree tests */
         { "test_simple_multi_tree", test_simple_multi_tree },

c/tskit/haplotype_matching.c

@@ -209,7 +209,7 @@ int
 tsk_ls_hmm_free(tsk_ls_hmm_t *self)
 {
     tsk_tree_free(&self->tree);
-    tsk_diff_iter_free(&self->diffs);
+    tsk_tree_position_free(&self->tree_pos);
     tsk_safe_free(self->recombination_rate);
     tsk_safe_free(self->mutation_rate);
     tsk_safe_free(self->recombination_rate);
@@ -248,9 +248,8 @@ tsk_ls_hmm_reset(tsk_ls_hmm_t *self)
     tsk_memset(self->transition_parent, 0xff,
         self->max_transitions * sizeof(*self->transition_parent));
 
-    /* This is safe because we've already zero'd out the memory. */
-    tsk_diff_iter_free(&self->diffs);
-    ret = tsk_diff_iter_init_from_ts(&self->diffs, self->tree_sequence, false);
+    tsk_tree_position_free(&self->tree_pos);
+    ret = tsk_tree_position_init(&self->tree_pos, self->tree_sequence, 0);
     if (ret != 0) {
         goto out;
     }
@@ -306,47 +305,48 @@ tsk_ls_hmm_update_tree(tsk_ls_hmm_t *self)
     int ret = 0;
     tsk_id_t *restrict parent = self->parent;
     tsk_id_t *restrict T_index = self->transition_index;
+    const tsk_id_t *restrict edges_child = self->tree_sequence->tables->edges.child;
+    const tsk_id_t *restrict edges_parent = self->tree_sequence->tables->edges.parent;
     tsk_value_transition_t *restrict T = self->transitions;
-    tsk_edge_list_node_t *record;
-    tsk_edge_list_t records_out, records_in;
-    tsk_edge_t edge;
-    double left, right;
-    tsk_id_t u;
+    tsk_id_t u, c, p, j, e;
     tsk_value_transition_t *vt;
 
-    ret = tsk_diff_iter_next(&self->diffs, &left, &right, &records_out, &records_in);
-    if (ret < 0) {
-        goto out;
-    }
+    tsk_tree_position_next(&self->tree_pos);
+    tsk_bug_assert(self->tree_pos.index != -1);
+    tsk_bug_assert(self->tree_pos.index == self->tree.index);
 
-    for (record = records_out.head; record != NULL; record = record->next) {
-        u = record->edge.child;
+    for (j = self->tree_pos.out.start; j < self->tree_pos.out.stop; j++) {
+        e = self->tree_pos.out.order[j];
+        c = edges_child[e];
+        u = c;
         if (T_index[u] == TSK_NULL) {
             /* Ensure the subtree we're detaching has a transition at the root */
             while (T_index[u] == TSK_NULL) {
                 u = parent[u];
                 tsk_bug_assert(u != TSK_NULL);
             }
             tsk_bug_assert(self->num_transitions < self->max_transitions);
-            T_index[record->edge.child] = (tsk_id_t) self->num_transitions;
-            T[self->num_transitions].tree_node = record->edge.child;
+            T_index[c] = (tsk_id_t) self->num_transitions;
+            T[self->num_transitions].tree_node = c;
             T[self->num_transitions].value = T[T_index[u]].value;
             self->num_transitions++;
         }
-        parent[record->edge.child] = TSK_NULL;
+        parent[c] = TSK_NULL;
     }
 
-    for (record = records_in.head; record != NULL; record = record->next) {
-        edge = record->edge;
-        parent[edge.child] = edge.parent;
-        u = edge.parent;
-        if (parent[edge.parent] == TSK_NULL) {
+    for (j = self->tree_pos.in.start; j < self->tree_pos.in.stop; j++) {
+        e = self->tree_pos.in.order[j];
+        c = edges_child[e];
+        p = edges_parent[e];
+        parent[c] = p;
+        u = p;
+        if (parent[p] == TSK_NULL) {
             /* Grafting onto a new root. */
-            if (T_index[record->edge.parent] == TSK_NULL) {
-                T_index[edge.parent] = (tsk_id_t) self->num_transitions;
+            if (T_index[p] == TSK_NULL) {
+                T_index[p] = (tsk_id_t) self->num_transitions;
                 tsk_bug_assert(self->num_transitions < self->max_transitions);
-                T[self->num_transitions].tree_node = edge.parent;
-                T[self->num_transitions].value = T[T_index[edge.child]].value;
+                T[self->num_transitions].tree_node = p;
+                T[self->num_transitions].value = T[T_index[c]].value;
                 self->num_transitions++;
             }
         } else {
@@ -356,18 +356,17 @@ tsk_ls_hmm_update_tree(tsk_ls_hmm_t *self)
             }
             tsk_bug_assert(u != TSK_NULL);
         }
-        tsk_bug_assert(T_index[u] != -1 && T_index[edge.child] != -1);
-        if (T[T_index[u]].value == T[T_index[edge.child]].value) {
-            vt = &T[T_index[edge.child]];
+        tsk_bug_assert(T_index[u] != -1 && T_index[c] != -1);
+        if (T[T_index[u]].value == T[T_index[c]].value) {
+            vt = &T[T_index[c]];
             /* Mark the value transition as unusued */
             vt->value = -1;
             vt->tree_node = TSK_NULL;
-            T_index[edge.child] = TSK_NULL;
+            T_index[c] = TSK_NULL;
         }
     }
 
     ret = tsk_ls_hmm_remove_dead_roots(self);
-out:
     return ret;
 }
 

c/tskit/haplotype_matching.h

@@ -1,7 +1,7 @@
 /*
  * MIT License
  *
- * Copyright (c) 2019-2022 Tskit Developers
+ * Copyright (c) 2019-2023 Tskit Developers
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -98,7 +98,10 @@ typedef struct _tsk_ls_hmm_t {
     tsk_size_t num_nodes;
     /* state */
     tsk_tree_t tree;
-    tsk_diff_iter_t diffs;
+    /* NOTE: this tree_position will be redundant once we integrate the top-level
+     * tree class with this.
+     */
+    tsk_tree_position_t tree_pos;
     tsk_id_t *parent;
     /* The probability value transitions on the tree */
     tsk_value_transition_t *transitions;