From 62e829ad60f503ba3ee95bbb0adf5c3da3571344 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 15 Sep 2021 13:49:50 +0100 Subject: [PATCH 1/8] Add virtual_root to tree arrays Closes #1670 --- Minimal C changes to implement virtual root support Includes the minimal changes to get the Python tests passing as well. Skipping viz tests for now, as the rendering is arbitrary. Add virtual roots to the quintuply linked tree arrays. Test properties of the virtual_root --- c/tests/test_trees.c | 20 +- c/tskit/trees.c | 413 ++++++++++++++++----------------- c/tskit/trees.h | 5 + python/CHANGELOG.rst | 9 +- python/_tskitmodule.c | 32 +-- python/tests/__init__.py | 5 +- python/tests/test_drawing.py | 3 + python/tests/test_highlevel.py | 18 +- python/tests/test_lowlevel.py | 19 +- python/tests/test_topology.py | 209 ++++++++++++++--- python/tests/tsutil.py | 169 ++++++++++++-- python/tskit/trees.py | 24 +- 12 files changed, 609 insertions(+), 317 deletions(-) diff --git a/c/tests/test_trees.c b/c/tests/test_trees.c index 5f31be9615..ffde0cd9ec 100644 --- a/c/tests/test_trees.c +++ b/c/tests/test_trees.c @@ -1338,7 +1338,9 @@ test_simplest_multi_root_tree(void) CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&ts), 1); ret = tsk_tree_init(&t, &ts, 0); + tsk_tree_print_state(&t, _devnull); + /* Make sure the initial roots are set correctly */ CU_ASSERT_EQUAL(t.left_root, 0); CU_ASSERT_EQUAL(t.left_sib[0], TSK_NULL); @@ -1348,7 +1350,6 @@ test_simplest_multi_root_tree(void) CU_ASSERT_EQUAL(t.left_sib[2], 1); CU_ASSERT_EQUAL(t.right_sib[2], TSK_NULL); - CU_ASSERT_EQUAL(ret, 0); CU_ASSERT_EQUAL(ret, 0); ret = tsk_tree_first(&t); CU_ASSERT_EQUAL(ret, 1); @@ -3882,7 +3883,7 @@ test_single_tree_iter_depths(void) CU_ASSERT_EQUAL(depth, depths[u]); } - ret = tsk_tree_depth(&tree, (tsk_id_t) num_nodes, &depth); + ret = tsk_tree_depth(&tree, (tsk_id_t) num_nodes + 1, &depth); CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); ret = tsk_tree_depth(&tree, TSK_NULL, &depth); CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); @@ -5185,7 +5186,12 @@ test_empty_tree_kc(void) CU_ASSERT_EQUAL_FATAL(t.left_root, TSK_NULL); CU_ASSERT_EQUAL_FATAL(t.left, 0); CU_ASSERT_EQUAL_FATAL(t.right, 1); - CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 0, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); + CU_ASSERT_EQUAL_FATAL(t.parent[0], TSK_NULL); + CU_ASSERT_EQUAL_FATAL(t.left_child[0], TSK_NULL); + CU_ASSERT_EQUAL_FATAL(t.right_child[0], TSK_NULL); + CU_ASSERT_EQUAL_FATAL(t.left_sib[0], TSK_NULL); + CU_ASSERT_EQUAL_FATAL(t.right_sib[0], TSK_NULL); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 1, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); ret = tsk_tree_kc_distance(&t, &t, 0, &result); CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MULTIPLE_ROOTS); @@ -5670,7 +5676,7 @@ test_tree_errors(void) tsk_node_t node; tsk_treeseq_t ts, other_ts; tsk_tree_t t, other_t; - tsk_id_t bad_nodes[] = { num_nodes, num_nodes + 1, -1 }; + tsk_id_t bad_nodes[] = { num_nodes + 1, num_nodes + 2, -1 }; tsk_id_t tracked_samples[] = { 0, 0, 0 }; tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, NULL, NULL, @@ -6044,7 +6050,9 @@ test_empty_tree_sequence(void) CU_ASSERT_EQUAL_FATAL(t.left_root, TSK_NULL); CU_ASSERT_EQUAL_FATAL(t.left, 0); CU_ASSERT_EQUAL_FATAL(t.right, 1); - CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 0, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 0, &v), 0); + CU_ASSERT_EQUAL_FATAL(v, TSK_NULL); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 1, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); tsk_tree_free(&t); ret = tsk_tree_init(&t, &ts, 0); @@ -6054,7 +6062,7 @@ test_empty_tree_sequence(void) CU_ASSERT_EQUAL_FATAL(t.left_root, TSK_NULL); CU_ASSERT_EQUAL_FATAL(t.left, 0); CU_ASSERT_EQUAL_FATAL(t.right, 1); - CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 0, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 1, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); tsk_tree_free(&t); tsk_treeseq_free(&ts); diff --git a/c/tskit/trees.c b/c/tskit/trees.c index ce129fcd80..2881684680 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3154,81 +3154,11 @@ tsk_treeseq_simplify(const tsk_treeseq_t *self, const tsk_id_t *samples, * Tree * ======================================================== */ -int TSK_WARN_UNUSED -tsk_tree_clear(tsk_tree_t *self) -{ - int ret = 0; - tsk_size_t j; - tsk_id_t u; - const tsk_size_t N = self->num_nodes; - const tsk_size_t num_samples = self->tree_sequence->num_samples; - const bool sample_counts = !(self->options & TSK_NO_SAMPLE_COUNTS); - const bool sample_lists = !!(self->options & TSK_SAMPLE_LISTS); - - self->left = 0; - self->right = 0; - self->index = -1; - /* TODO we should profile this method to see if just doing a single loop over - * the nodes would be more efficient than multiple memsets. - */ - tsk_memset(self->parent, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->left_child, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->right_child, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->left_sib, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->right_sib, 0xff, N * sizeof(tsk_id_t)); - - if (sample_counts) { - tsk_memset(self->num_samples, 0, N * sizeof(*self->num_samples)); - tsk_memset(self->marked, 0, N * sizeof(*self->marked)); - /* We can't reset the tracked samples via tsk_memset because we don't - * know where the tracked samples are. - */ - for (j = 0; j < self->num_nodes; j++) { - if (!tsk_treeseq_is_sample(self->tree_sequence, (tsk_id_t) j)) { - self->num_tracked_samples[j] = 0; - } - } - } - if (sample_lists) { - tsk_memset(self->left_sample, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->right_sample, 0xff, N * sizeof(tsk_id_t)); - tsk_memset(self->next_sample, 0xff, num_samples * sizeof(tsk_id_t)); - } - /* Set the sample attributes */ - for (j = 0; j < num_samples; j++) { - u = self->samples[j]; - if (sample_counts) { - self->num_samples[u] = 1; - } - if (sample_lists) { - /* We are mapping to *indexes* into the list of samples here */ - self->left_sample[u] = (tsk_id_t) j; - self->right_sample[u] = (tsk_id_t) j; - } - } - self->left_root = TSK_NULL; - if (sample_counts && self->root_threshold == 1 && num_samples > 0) { - self->left_root = self->samples[0]; - for (j = 0; j < num_samples; j++) { - /* Set initial roots */ - u = self->samples[j]; - if (j < num_samples - 1) { - self->right_sib[u] = self->samples[j + 1]; - } - if (j > 0) { - self->left_sib[u] = self->samples[j - 1]; - } - } - } - return ret; -} - int TSK_WARN_UNUSED tsk_tree_init(tsk_tree_t *self, const tsk_treeseq_t *tree_sequence, tsk_flags_t options) { int ret = TSK_ERR_NO_MEMORY; - tsk_size_t num_samples; - tsk_size_t num_nodes; + tsk_size_t num_samples, num_nodes, N; tsk_memset(self, 0, sizeof(tsk_tree_t)); if (tree_sequence == NULL) { @@ -3238,32 +3168,35 @@ tsk_tree_init(tsk_tree_t *self, const tsk_treeseq_t *tree_sequence, tsk_flags_t num_nodes = tree_sequence->tables->nodes.num_rows; num_samples = tree_sequence->num_samples; self->num_nodes = num_nodes; + self->virtual_root = (tsk_id_t) num_nodes; self->tree_sequence = tree_sequence; self->samples = tree_sequence->samples; self->options = options; self->root_threshold = 1; - self->parent = tsk_malloc(num_nodes * sizeof(tsk_id_t)); - self->left_child = tsk_malloc(num_nodes * sizeof(tsk_id_t)); - self->right_child = tsk_malloc(num_nodes * sizeof(tsk_id_t)); - self->left_sib = tsk_malloc(num_nodes * sizeof(tsk_id_t)); - self->right_sib = tsk_malloc(num_nodes * sizeof(tsk_id_t)); + + /* Allocate space in the quintuply linked tree for the virtual root */ + N = num_nodes + 1; + self->parent = tsk_malloc(N * sizeof(*self->parent)); + self->left_child = tsk_malloc(N * sizeof(*self->left_child)); + self->right_child = tsk_malloc(N * sizeof(*self->right_child)); + self->left_sib = tsk_malloc(N * sizeof(*self->left_sib)); + self->right_sib = tsk_malloc(N * sizeof(*self->right_sib)); if (self->parent == NULL || self->left_child == NULL || self->right_child == NULL || self->left_sib == NULL || self->right_sib == NULL) { goto out; } if (!(self->options & TSK_NO_SAMPLE_COUNTS)) { - self->num_samples = tsk_calloc(num_nodes, sizeof(*self->num_samples)); - self->num_tracked_samples - = tsk_calloc(num_nodes, sizeof(*self->num_tracked_samples)); - self->marked = tsk_calloc(num_nodes, sizeof(uint8_t)); + self->num_samples = tsk_calloc(N, sizeof(*self->num_samples)); + self->num_tracked_samples = tsk_calloc(N, sizeof(*self->num_tracked_samples)); + self->marked = tsk_calloc(N, sizeof(*self->marked)); if (self->num_samples == NULL || self->num_tracked_samples == NULL || self->marked == NULL) { goto out; } } if (self->options & TSK_SAMPLE_LISTS) { - self->left_sample = tsk_malloc(num_nodes * sizeof(*self->left_sample)); - self->right_sample = tsk_malloc(num_nodes * sizeof(*self->right_sample)); + self->left_sample = tsk_malloc(N * sizeof(*self->left_sample)); + self->right_sample = tsk_malloc(N * sizeof(*self->right_sample)); self->next_sample = tsk_malloc(num_samples * sizeof(*self->next_sample)); if (self->left_sample == NULL || self->right_sample == NULL || self->next_sample == NULL) { @@ -3291,6 +3224,8 @@ tsk_tree_set_root_threshold(tsk_tree_t *self, tsk_size_t root_threshold) goto out; } self->root_threshold = root_threshold; + /* Reset the roots */ + ret = tsk_tree_clear(self); out: return ret; } @@ -3360,6 +3295,7 @@ tsk_tree_set_tracked_samples( if (ret != 0) { goto out; } + self->num_tracked_samples[self->virtual_root] = num_tracked_samples; for (j = 0; j < num_tracked_samples; j++) { u = tracked_samples[j]; if (u < 0 || u >= (tsk_id_t) self->num_nodes) { @@ -3430,7 +3366,7 @@ int TSK_WARN_UNUSED tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options) { int ret = TSK_ERR_GENERIC; - tsk_size_t N = self->num_nodes; + tsk_size_t N = self->num_nodes + 1; if (!(options & TSK_NO_INIT)) { ret = tsk_tree_init(dest, self->tree_sequence, options); @@ -3453,11 +3389,11 @@ tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options) dest->sites_length = self->sites_length; dest->root_threshold = self->root_threshold; - tsk_memcpy(dest->parent, self->parent, N * sizeof(tsk_id_t)); - tsk_memcpy(dest->left_child, self->left_child, N * sizeof(tsk_id_t)); - tsk_memcpy(dest->right_child, self->right_child, N * sizeof(tsk_id_t)); - tsk_memcpy(dest->left_sib, self->left_sib, N * sizeof(tsk_id_t)); - tsk_memcpy(dest->right_sib, self->right_sib, N * sizeof(tsk_id_t)); + tsk_memcpy(dest->parent, self->parent, N * sizeof(*self->parent)); + tsk_memcpy(dest->left_child, self->left_child, N * sizeof(*self->left_child)); + tsk_memcpy(dest->right_child, self->right_child, N * sizeof(*self->right_child)); + tsk_memcpy(dest->left_sib, self->left_sib, N * sizeof(*self->left_sib)); + tsk_memcpy(dest->right_sib, self->right_sib, N * sizeof(*self->right_sib)); if (!(dest->options & TSK_NO_SAMPLE_COUNTS)) { if (self->options & TSK_NO_SAMPLE_COUNTS) { ret = TSK_ERR_UNSUPPORTED_OPERATION; @@ -3473,10 +3409,11 @@ tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options) ret = TSK_ERR_UNSUPPORTED_OPERATION; goto out; } - tsk_memcpy(dest->left_sample, self->left_sample, N * sizeof(tsk_id_t)); - tsk_memcpy(dest->right_sample, self->right_sample, N * sizeof(tsk_id_t)); + tsk_memcpy(dest->left_sample, self->left_sample, N * sizeof(*self->left_sample)); + tsk_memcpy( + dest->right_sample, self->right_sample, N * sizeof(*self->right_sample)); tsk_memcpy(dest->next_sample, self->next_sample, - self->tree_sequence->num_samples * sizeof(tsk_id_t)); + self->tree_sequence->num_samples * sizeof(*self->next_sample)); } ret = 0; out: @@ -3498,7 +3435,7 @@ static int tsk_tree_check_node(const tsk_tree_t *self, tsk_id_t u) { int ret = 0; - if (u < 0 || u >= (tsk_id_t) self->num_nodes) { + if (u < 0 || u > (tsk_id_t) self->num_nodes) { ret = TSK_ERR_NODE_OUT_OF_BOUNDS; } return ret; @@ -3760,6 +3697,11 @@ tsk_tree_check_state(const tsk_tree_t *self) tsk_bug_assert(children != NULL); + /* Check the virtual root properties */ + tsk_bug_assert(self->parent[self->virtual_root] == TSK_NULL); + tsk_bug_assert(self->left_sib[self->virtual_root] == TSK_NULL); + tsk_bug_assert(self->right_sib[self->virtual_root] == TSK_NULL); + for (j = 0; j < self->tree_sequence->num_samples; j++) { u = self->samples[j]; while (self->parent[u] != TSK_NULL) { @@ -3772,6 +3714,7 @@ tsk_tree_check_state(const tsk_tree_t *self) } else { tsk_bug_assert(self->left_sib[self->left_root] == TSK_NULL); } + /* Iterate over the roots and make sure they are set */ for (u = self->left_root; u != TSK_NULL; u = self->right_sib[u]) { tsk_bug_assert(is_root[u]); @@ -3842,7 +3785,7 @@ tsk_tree_print_state(const tsk_tree_t *self, FILE *out) } fprintf(out, "\n"); - for (j = 0; j < self->num_nodes; j++) { + for (j = 0; j < self->num_nodes + 1; j++) { fprintf(out, "%lld\t%lld\t%lld\t%lld\t%lld\t%lld", (long long) j, (long long) self->parent[j], (long long) self->left_child[j], (long long) self->right_child[j], (long long) self->left_sib[j], @@ -3867,15 +3810,17 @@ tsk_tree_print_state(const tsk_tree_t *self, FILE *out) /* Methods for positioning the tree along the sequence */ -/* parent, left_child and right_sib are restrict pointers in the calling function, - * so we pass these as parameters to ensure the relationships are clear to the - * compiler. */ +/* The following methods are very performance sensitive and so we use a + * lot of restrict pointers. Because we are saying that we don't have + * any aliases to these pointers, we pass around the reference to parent + * since it's used in all the functions. */ static inline void -tsk_tree_update_sample_lists(tsk_tree_t *self, const tsk_id_t *restrict parent, - const tsk_id_t *restrict left_child, const tsk_id_t *restrict right_sib, - tsk_id_t node) +tsk_tree_update_sample_lists( + tsk_tree_t *self, tsk_id_t node, const tsk_id_t *restrict parent) { tsk_id_t u, v, sample_index; + tsk_id_t *restrict left_child = self->left_child; + tsk_id_t *restrict right_sib = self->right_sib; tsk_id_t *restrict left = self->left_sample; tsk_id_t *restrict right = self->right_sample; tsk_id_t *restrict next = self->next_sample; @@ -3904,26 +3849,17 @@ tsk_tree_update_sample_lists(tsk_tree_t *self, const tsk_id_t *restrict parent, } } -static void -tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) +static inline void +tsk_tree_remove_branch( + tsk_tree_t *self, tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent) { - tsk_id_t *restrict parent = self->parent; tsk_id_t *restrict left_child = self->left_child; tsk_id_t *restrict right_child = self->right_child; tsk_id_t *restrict left_sib = self->left_sib; tsk_id_t *restrict right_sib = self->right_sib; - tsk_size_t *restrict num_samples = self->num_samples; - tsk_size_t *restrict num_tracked_samples = self->num_tracked_samples; - uint8_t *restrict marked = self->marked; - const uint8_t mark = self->mark; - const tsk_size_t root_threshold = self->root_threshold; - tsk_id_t lsib, rsib, u, path_end, lroot, rroot; - bool path_end_was_root; - -#define IS_ROOT(U) (num_samples[U] >= root_threshold) + tsk_id_t lsib = left_sib[c]; + tsk_id_t rsib = right_sib[c]; - lsib = left_sib[c]; - rsib = right_sib[c]; if (lsib == TSK_NULL) { left_child[p] = rsib; } else { @@ -3937,53 +3873,83 @@ tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) parent[c] = TSK_NULL; left_sib[c] = TSK_NULL; right_sib[c] = TSK_NULL; +} - if (!(self->options & TSK_NO_SAMPLE_COUNTS)) { - /* keep the compiler happy */ - path_end_was_root = false; - path_end = TSK_NULL; +static inline void +tsk_tree_insert_branch( + tsk_tree_t *self, tsk_id_t p, tsk_id_t c, tsk_id_t *restrict parent) +{ + tsk_id_t *restrict left_child = self->left_child; + tsk_id_t *restrict right_child = self->right_child; + tsk_id_t *restrict left_sib = self->left_sib; + tsk_id_t *restrict right_sib = self->right_sib; + tsk_id_t u; + parent[c] = p; + u = right_child[p]; + if (u == TSK_NULL) { + left_child[p] = c; + left_sib[c] = TSK_NULL; + right_sib[c] = TSK_NULL; + } else { + right_sib[u] = c; + left_sib[c] = u; + right_sib[c] = TSK_NULL; + } + right_child[p] = c; +} + +static inline void +tsk_tree_insert_root(tsk_tree_t *self, tsk_id_t root, tsk_id_t *restrict parent) +{ + tsk_tree_insert_branch(self, self->virtual_root, root, parent); + parent[root] = TSK_NULL; +} + +static inline void +tsk_tree_remove_root(tsk_tree_t *self, tsk_id_t root, tsk_id_t *restrict parent) +{ + tsk_tree_remove_branch(self, self->virtual_root, root, parent); +} + +static void +tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) +{ + tsk_id_t *restrict parent = self->parent; + tsk_size_t *restrict num_samples = self->num_samples; + tsk_size_t *restrict num_tracked_samples = self->num_tracked_samples; + uint8_t *restrict marked = self->marked; + const uint8_t mark = self->mark; + const tsk_size_t root_threshold = self->root_threshold; + tsk_id_t u; + tsk_id_t path_end = TSK_NULL; + bool path_end_was_root = false; + +#define POTENTIAL_ROOT(U) (num_samples[U] >= root_threshold) + + tsk_tree_remove_branch(self, p, c, parent); + + if (!(self->options & TSK_NO_SAMPLE_COUNTS)) { u = p; while (u != TSK_NULL) { path_end = u; - path_end_was_root = IS_ROOT(u); + path_end_was_root = POTENTIAL_ROOT(u); num_samples[u] -= num_samples[c]; num_tracked_samples[u] -= num_tracked_samples[c]; marked[u] = mark; u = parent[u]; } - if (path_end_was_root && !IS_ROOT(path_end)) { - /* remove path_end from the list of roots */ - lroot = left_sib[path_end]; - rroot = right_sib[path_end]; - self->left_root = TSK_NULL; - if (lroot != TSK_NULL) { - right_sib[lroot] = rroot; - self->left_root = lroot; - } - if (rroot != TSK_NULL) { - left_sib[rroot] = lroot; - self->left_root = rroot; - } - left_sib[path_end] = TSK_NULL; - right_sib[path_end] = TSK_NULL; - } - if (IS_ROOT(c)) { - if (self->left_root != TSK_NULL) { - lroot = left_sib[self->left_root]; - if (lroot != TSK_NULL) { - right_sib[lroot] = c; - } - left_sib[c] = lroot; - left_sib[self->left_root] = c; - } - right_sib[c] = self->left_root; - self->left_root = c; + + if (path_end_was_root && !POTENTIAL_ROOT(path_end)) { + tsk_tree_remove_root(self, path_end, parent); + } + if (POTENTIAL_ROOT(c)) { + tsk_tree_insert_root(self, c, parent); } } if (self->options & TSK_SAMPLE_LISTS) { - tsk_tree_update_sample_lists(self, parent, left_child, right_sib, p); + tsk_tree_update_sample_lists(self, p, parent); } } @@ -3991,93 +3957,40 @@ static void tsk_tree_insert_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) { tsk_id_t *restrict parent = self->parent; - tsk_id_t *restrict left_child = self->left_child; - tsk_id_t *restrict right_child = self->right_child; - tsk_id_t *restrict left_sib = self->left_sib; - tsk_id_t *restrict right_sib = self->right_sib; tsk_size_t *restrict num_samples = self->num_samples; tsk_size_t *restrict num_tracked_samples = self->num_tracked_samples; uint8_t *restrict marked = self->marked; const uint8_t mark = self->mark; const tsk_size_t root_threshold = self->root_threshold; - tsk_id_t lsib, rsib, u, path_end, lroot; - bool path_end_was_root; - -#define IS_ROOT(U) (num_samples[U] >= root_threshold) + tsk_id_t u; + tsk_id_t path_end = TSK_NULL; + bool path_end_was_root = false; - parent[c] = p; - u = right_child[p]; - lsib = left_sib[c]; - rsib = right_sib[c]; - if (u == TSK_NULL) { - left_child[p] = c; - left_sib[c] = TSK_NULL; - right_sib[c] = TSK_NULL; - } else { - right_sib[u] = c; - left_sib[c] = u; - right_sib[c] = TSK_NULL; - } - right_child[p] = c; +#define POTENTIAL_ROOT(U) (num_samples[U] >= root_threshold) if (!(self->options & TSK_NO_SAMPLE_COUNTS)) { - - /* keep compiler happy */ - path_end = TSK_NULL; - path_end_was_root = false; - u = p; while (u != TSK_NULL) { path_end = u; - path_end_was_root = IS_ROOT(u); + path_end_was_root = POTENTIAL_ROOT(u); num_samples[u] += num_samples[c]; num_tracked_samples[u] += num_tracked_samples[c]; marked[u] = mark; u = parent[u]; } - if (IS_ROOT(c)) { - if (path_end_was_root) { - /* Remove c from the root list */ - self->left_root = TSK_NULL; - if (lsib != TSK_NULL) { - right_sib[lsib] = rsib; - self->left_root = lsib; - } - if (rsib != TSK_NULL) { - left_sib[rsib] = lsib; - self->left_root = rsib; - } - } else { - /* Replace c with path_end in root list */ - if (lsib != TSK_NULL) { - right_sib[lsib] = path_end; - } - if (rsib != TSK_NULL) { - left_sib[rsib] = path_end; - } - left_sib[path_end] = lsib; - right_sib[path_end] = rsib; - self->left_root = path_end; - } - } else { - if (IS_ROOT(path_end) && !path_end_was_root) { - /* Add a path_end as new root */ - if (self->left_root != TSK_NULL) { - lroot = left_sib[self->left_root]; - if (lroot != TSK_NULL) { - right_sib[lroot] = path_end; - } - left_sib[path_end] = lroot; - left_sib[self->left_root] = path_end; - } - right_sib[path_end] = self->left_root; - self->left_root = path_end; - } + if (POTENTIAL_ROOT(c)) { + tsk_tree_remove_root(self, c, parent); + } + if (POTENTIAL_ROOT(path_end) && !path_end_was_root) { + tsk_tree_insert_root(self, path_end, parent); } } + + tsk_tree_insert_branch(self, p, c, parent); + if (self->options & TSK_SAMPLE_LISTS) { - tsk_tree_update_sample_lists(self, parent, left_child, right_sib, p); + tsk_tree_update_sample_lists(self, p, parent); } } @@ -4117,12 +4030,7 @@ tsk_tree_advance(tsk_tree_t *self, int direction, const double *restrict out_bre tsk_tree_insert_edge(self, edge_parent[k], edge_child[k]); } - if (self->left_root != TSK_NULL) { - /* Ensure that left_root is the left-most root */ - while (self->left_sib[self->left_root] != TSK_NULL) { - self->left_root = self->left_sib[self->left_root]; - } - } + self->left_root = self->left_child[self->virtual_root]; self->direction = direction; self->index = self->index + direction; @@ -4266,6 +4174,73 @@ tsk_tree_prev(tsk_tree_t *self) return ret; } +int TSK_WARN_UNUSED +tsk_tree_clear(tsk_tree_t *self) +{ + int ret = 0; + tsk_size_t j; + tsk_id_t u; + const tsk_size_t N = self->num_nodes + 1; + const tsk_size_t num_samples = self->tree_sequence->num_samples; + const bool sample_counts = !(self->options & TSK_NO_SAMPLE_COUNTS); + const bool sample_lists = !!(self->options & TSK_SAMPLE_LISTS); + + self->left = 0; + self->right = 0; + self->index = -1; + /* TODO we should profile this method to see if just doing a single loop over + * the nodes would be more efficient than multiple memsets. + */ + tsk_memset(self->parent, 0xff, N * sizeof(*self->parent)); + tsk_memset(self->left_child, 0xff, N * sizeof(*self->left_child)); + tsk_memset(self->right_child, 0xff, N * sizeof(*self->right_child)); + tsk_memset(self->left_sib, 0xff, N * sizeof(*self->left_sib)); + tsk_memset(self->right_sib, 0xff, N * sizeof(*self->right_sib)); + + if (sample_counts) { + tsk_memset(self->num_samples, 0, N * sizeof(*self->num_samples)); + tsk_memset(self->marked, 0, N * sizeof(*self->marked)); + /* We can't reset the tracked samples via memset because we don't + * know where the tracked samples are. + */ + for (j = 0; j < self->num_nodes; j++) { + if (!tsk_treeseq_is_sample(self->tree_sequence, (tsk_id_t) j)) { + self->num_tracked_samples[j] = 0; + } + } + self->num_tracked_samples[self->virtual_root] = 0; + self->num_samples[self->virtual_root] = num_samples; + } + if (sample_lists) { + tsk_memset(self->left_sample, 0xff, N * sizeof(tsk_id_t)); + tsk_memset(self->right_sample, 0xff, N * sizeof(tsk_id_t)); + tsk_memset(self->next_sample, 0xff, num_samples * sizeof(tsk_id_t)); + } + /* Set the sample attributes */ + for (j = 0; j < num_samples; j++) { + u = self->samples[j]; + if (sample_counts) { + self->num_samples[u] = 1; + } + if (sample_lists) { + /* We are mapping to *indexes* into the list of samples here */ + self->left_sample[u] = (tsk_id_t) j; + self->right_sample[u] = (tsk_id_t) j; + } + } + self->left_root = TSK_NULL; + if (sample_counts && self->root_threshold == 1 && num_samples > 0) { + for (j = 0; j < num_samples; j++) { + /* Set initial roots */ + if (self->root_threshold == 1) { + tsk_tree_insert_root(self, self->samples[j], self->parent); + } + } + self->left_root = self->left_child[self->virtual_root]; + } + return ret; +} + /* Parsimony methods */ static inline uint64_t diff --git a/c/tskit/trees.h b/c/tskit/trees.h index bd5d9c65ba..9730353350 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -125,6 +125,11 @@ typedef struct { * other roots can be found using right_sib. */ tsk_id_t left_root; + /** + @brief The ID of the "virtual root" whose children are the roots of the + tree. + */ + tsk_id_t virtual_root; /** @brief The parent of node u is parent[u]. Equal to TSK_NULL if node u is a root or is not a node in the current tree. diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index 0d6d78fcf7..857535be93 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -17,6 +17,13 @@ :issue:`1554`, :issue:`1573`, :issue:`1589`,:issue:`1598`,:issue:`1628`, :pr:`1571`, :pr:`1579`, :pr:`1585`, :pr:`1590`, :pr:`1602`, :pr:`1618`, :pr:`1620`, :pr:`1652`). +- The Tree class now conceptually has an extra node, the "virtual root" whose + children are the roots of the tree. The quintuply linked tree arrays + (parent_array, left_child_array, right_child_array, left_sib_array and right_sib_array) + all have one extra element. + (:user:`jeromekelleher`, :issue:`1691`, :pr:`1704`). + + **Features** - Add ``__setitem__`` to all tables allowing single rows to be updated. For example @@ -68,7 +75,7 @@ - SVG visualization of a tree sequence can be restricted to displaying between left and right genomic coordinates using the ``x_lim`` parameter. The default settings - now mean that if the left or right flanks of a tree sequence are entirely empty, + now mean that if the left or right flanks of a tree sequence are entirely empty, these regions will not be plotted in the SVG (:user:`hyanwong`, :pr:`1288`). - SVG visualization of a single tree allows all mutations on an edge to be plotted diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index d932819295..460c1d2c03 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -9103,7 +9103,7 @@ static int Tree_check_bounds(Tree *self, int node) { int ret = 0; - if (node < 0 || node >= (int) self->tree->num_nodes) { + if (node < 0 || node > (int) self->tree->num_nodes) { PyErr_SetString(PyExc_ValueError, "Node index out of bounds"); ret = -1; } @@ -9333,27 +9333,27 @@ Tree_get_num_roots(Tree *self) } static PyObject * -Tree_get_index(Tree *self) +Tree_get_virtual_root(Tree *self) { PyObject *ret = NULL; if (Tree_check_state(self) != 0) { goto out; } - ret = Py_BuildValue("n", (Py_ssize_t) self->tree->index); + ret = Py_BuildValue("n", (Py_ssize_t) self->tree->virtual_root); out: return ret; } static PyObject * -Tree_get_left_root(Tree *self) +Tree_get_index(Tree *self) { PyObject *ret = NULL; if (Tree_check_state(self) != 0) { goto out; } - ret = Py_BuildValue("i", (int) self->tree->left_root); + ret = Py_BuildValue("n", (Py_ssize_t) self->tree->index); out: return ret; } @@ -9493,18 +9493,18 @@ static PyObject * Tree_get_time(Tree *self, PyObject *args) { PyObject *ret = NULL; - double time; - int node, err; + tsk_node_t node; + int node_id, err; - if (Tree_get_node_argument(self, args, &node) != 0) { + if (Tree_get_node_argument(self, args, &node_id) != 0) { goto out; } - err = tsk_tree_get_time(self->tree, node, &time); - if (ret != 0) { + err = tsk_treeseq_get_node(self->tree->tree_sequence, node_id, &node); + if (err != 0) { handle_library_error(err); goto out; } - ret = Py_BuildValue("d", time); + ret = Py_BuildValue("d", node.time); out: return ret; } @@ -10055,14 +10055,14 @@ Tree_set_root_threshold(Tree *self, PyObject *args) * corresponding arrays in the tsk_tree object. We use properties and * return a new array each time rather than trying to create a single array * at Tree initialisation time to avoid a circular reference counting loop, - * which (it seems) the even cyclic garbage collection support can't resolve. + * which (it seems) even cyclic garbage collection support can't resolve. */ static PyObject * Tree_make_array(Tree *self, int dtype, void *data) { PyObject *ret = NULL; PyArrayObject *array = NULL; - npy_intp dims = self->tree->num_nodes; + npy_intp dims = self->tree->num_nodes + 1; array = (PyArrayObject *) PyArray_SimpleNewFromData(1, &dims, dtype, data); if (array == NULL) { @@ -10203,10 +10203,10 @@ static PyMethodDef Tree_methods[] = { .ml_meth = (PyCFunction) Tree_get_index, .ml_flags = METH_NOARGS, .ml_doc = "Returns the index this tree occupies within the tree sequence." }, - { .ml_name = "get_left_root", - .ml_meth = (PyCFunction) Tree_get_left_root, + { .ml_name = "get_virtual_root", + .ml_meth = (PyCFunction) Tree_get_virtual_root, .ml_flags = METH_NOARGS, - .ml_doc = "Returns the root of the tree." }, + .ml_doc = "Returns the virtual root of the tree." }, { .ml_name = "get_left", .ml_meth = (PyCFunction) Tree_get_left, .ml_flags = METH_NOARGS, diff --git a/python/tests/__init__.py b/python/tests/__init__.py index e49da94aaf..cb23fa5836 100644 --- a/python/tests/__init__.py +++ b/python/tests/__init__.py @@ -297,16 +297,15 @@ def edge_diffs(self): left = right def trees(self): - rtt = tsutil.RootThresholdTree(self._tree_sequence) pt = PythonTree(self._tree_sequence.get_num_nodes()) pt.index = 0 - for left, right in rtt.iterate(): + for (left, right), rtt in tsutil.algorithm_R(self._tree_sequence): pt.parent[:] = rtt.parent pt.left_child[:] = rtt.left_child pt.right_child[:] = rtt.right_child pt.left_sib[:] = rtt.left_sib pt.right_sib[:] = rtt.right_sib - pt.left_root = rtt.left_root + pt.left_root = rtt.left_child[-1] pt.left = left pt.right = right # Add in all the sites diff --git a/python/tests/test_drawing.py b/python/tests/test_drawing.py index dde8d2be26..9a7923fbab 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -1171,6 +1171,7 @@ def test_draw_forky_tree(self): ) self.verify_text_rendering(t.draw_text(order="minlex"), tree) + @pytest.mark.skip("FIXME rendering differently now") def test_draw_multiroot_forky_tree(self): tree = ( " 13 \n" @@ -2353,6 +2354,7 @@ def test_known_svg_ts_no_axes(self, overwrite_viz, draw_plotbox): svg, "ts_no_axes.svg", overwrite_viz, width=200 * ts.num_trees ) + @pytest.mark.skip("FIXME rendering differently now") def test_known_svg_ts_internal_sample(self, overwrite_viz, draw_plotbox): ts = tsutil.jiggle_samples(self.get_simple_ts()) svg = ts.draw_svg( @@ -2537,6 +2539,7 @@ def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox, capl width=200 * ts.num_trees, ) + @pytest.mark.skip("FIXME rendering differently now") def test_known_svg_ts_multiroot(self, overwrite_viz, draw_plotbox, caplog): tables = wf.wf_sim( 6, diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index 34f63bc334..201a2fd340 100644 --- a/python/tests/test_highlevel.py +++ b/python/tests/test_highlevel.py @@ -2751,7 +2751,10 @@ def is_descendant(tree, u, v): tree = self.get_tree() for u, v in itertools.product(range(tree.num_nodes), repeat=2): assert is_descendant(tree, u, v) == tree.is_descendant(u, v) - for bad_node in [-1, -2, tree.num_nodes, tree.num_nodes + 1]: + # All nodes are descendents of themselves + for u in range(tree.num_nodes + 1): + assert tree.is_descendant(u, u) + for bad_node in [-1, -2, tree.num_nodes + 1]: with pytest.raises(ValueError): tree.is_descendant(0, bad_node) with pytest.raises(ValueError): @@ -2909,12 +2912,13 @@ def test_interval(self): def verify_tree_arrays(self, tree): ts = tree.tree_sequence - assert tree.parent_array.shape == (ts.num_nodes,) - assert tree.left_child_array.shape == (ts.num_nodes,) - assert tree.right_child_array.shape == (ts.num_nodes,) - assert tree.left_sib_array.shape == (ts.num_nodes,) - assert tree.right_sib_array.shape == (ts.num_nodes,) - for u in range(ts.num_nodes): + N = ts.num_nodes + 1 + assert tree.parent_array.shape == (N,) + assert tree.left_child_array.shape == (N,) + assert tree.right_child_array.shape == (N,) + assert tree.left_sib_array.shape == (N,) + assert tree.right_sib_array.shape == (N,) + for u in range(N): assert tree.parent(u) == tree.parent_array[u] assert tree.left_child(u) == tree.left_child_array[u] assert tree.right_child(u) == tree.right_child_array[u] diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 17eee43d72..972f8c0815 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -2511,7 +2511,9 @@ def test_constructor(self): st = _tskit.Tree(ts) assert st.get_num_nodes() == ts.get_num_nodes() # An uninitialised tree should always be zero. - assert st.get_left_root() == 0 + samples = ts.get_samples() + assert st.get_left_child(st.get_virtual_root()) == samples[0] + assert st.get_right_child(st.get_virtual_root()) == samples[-1] assert st.get_left() == 0 assert st.get_right() == 0 for j in range(ts.get_num_samples()): @@ -2656,7 +2658,7 @@ def test_mrca_interface(self): for ts in self.get_example_tree_sequences(): num_nodes = ts.get_num_nodes() st = _tskit.Tree(ts) - for v in [num_nodes, 10 ** 6, _tskit.NULL]: + for v in [num_nodes + 1, 10 ** 6, _tskit.NULL]: with pytest.raises(ValueError): st.get_mrca(v, v) with pytest.raises(ValueError): @@ -2694,9 +2696,8 @@ def get_times(tree): with pytest.raises(ValueError): st.get_newick(root=0, precision=100) for precision in range(17): - tree = st.get_newick( - root=st.get_left_root(), precision=precision - ).decode() + root = st.get_left_child(st.get_virtual_root()) + tree = st.get_newick(root=root, precision=precision).decode() times = get_times(tree) assert len(times) > ts.get_num_samples() for t in times: @@ -2712,7 +2713,7 @@ def test_cleared_tree(self): def check_tree(tree): assert tree.get_index() == -1 - assert tree.get_left_root() == samples[0] + assert tree.get_left_child(tree.get_virtual_root()) == samples[0] assert tree.get_mrca(0, 1) == _tskit.NULL for u in range(ts.get_num_nodes()): assert tree.get_parent(u) == _tskit.NULL @@ -2739,7 +2740,7 @@ def test_newick_interface(self): with pytest.raises(TypeError): st.get_newick(ts, buffer_size=bad_type) while st.next(): - u = st.get_left_root() + u = st.get_left_child(st.get_virtual_root()) newick = st.get_newick(u) assert newick.endswith(b";") with pytest.raises(ValueError): @@ -2830,7 +2831,7 @@ def test_sample_list(self): assert t.get_left_sample(j) == _tskit.NULL assert t.get_right_sample(j) == _tskit.NULL # The roots should have all samples. - u = t.get_left_root() + u = t.get_left_child(t.get_virtual_root()) samples = [] while u != _tskit.NULL: sample = t.get_left_sample(u) @@ -3032,7 +3033,7 @@ def test_array_properties(self, array): t1.first() a = getattr(t1, array + "_array") assert a.dtype == np.int32 - assert a.shape == (ts1.get_num_nodes(),) + assert a.shape == (ts1.get_num_nodes() + 1,) assert a.base == t1 assert not a.flags.writeable assert a.flags.aligned diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index 8a802c875d..a2c72d7669 100644 --- a/python/tests/test_topology.py +++ b/python/tests/test_topology.py @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2018-2020 Tskit Developers +# Copyright (c) 2018-2021 Tskit Developers # Copyright (c) 2016-2017 University of Oxford # # Permission is hereby granted, free of charge, to any person obtaining a copy @@ -323,6 +323,12 @@ def test_multiroot_tree(self): ts = tsutil.decapitate(ts, ts.num_edges // 2) self.verify(ts) + def test_all_missing_data(self): + tables = tskit.TableCollection(1) + for _ in range(10): + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0) + self.verify(tables.tree_sequence()) + class TestKCMetric(unittest.TestCase): """ @@ -1775,12 +1781,15 @@ def test_zero_nodes(self): assert t.roots == [] assert t.root == tskit.NULL assert t.parent_dict == {} + assert t.virtual_root == 0 + assert t.left_child(t.virtual_root) == -1 + assert t.right_child(t.virtual_root) == -1 assert list(t.nodes()) == [] assert list(ts.haplotypes()) == [] assert list(ts.variants()) == [] methods = [t.parent, t.left_child, t.right_child, t.left_sib, t.right_sib] for method in methods: - for u in [-1, 0, 1, 100]: + for u in [-1, 1, 100]: with pytest.raises(ValueError): method(u) tsp = ts.simplify() @@ -1805,6 +1814,7 @@ def test_one_node_zero_samples(self): assert t.interval == (0, 1) assert t.roots == [] assert t.root == tskit.NULL + assert t.virtual_root == 1 assert t.parent_dict == {} assert list(t.nodes()) == [] assert list(ts.haplotypes()) == [] @@ -1812,7 +1822,7 @@ def test_one_node_zero_samples(self): methods = [t.parent, t.left_child, t.right_child, t.left_sib, t.right_sib] for method in methods: assert method(0) == tskit.NULL - for u in [-1, 1, 100]: + for u in [-1, 2, 100]: with pytest.raises(ValueError): method(u) @@ -1859,6 +1869,7 @@ def test_one_node_one_sample(self): assert t.interval == (0, 1) assert t.roots == [0] assert t.root == 0 + assert t.virtual_root == 1 assert t.parent_dict == {} assert list(t.nodes()) == [0] assert list(ts.haplotypes(isolated_as_missing=False)) == [""] @@ -1866,7 +1877,7 @@ def test_one_node_one_sample(self): methods = [t.parent, t.left_child, t.right_child, t.left_sib, t.right_sib] for method in methods: assert method(0) == tskit.NULL - for u in [-1, 1, 100]: + for u in [-1, 2, 100]: with pytest.raises(ValueError): method(u) tsp = ts.simplify() @@ -1892,6 +1903,7 @@ def test_one_node_one_sample_sites(self): assert t.interval == (0, 1) assert t.roots == [0] assert t.root == 0 + assert t.virtual_root == 1 assert t.parent_dict == {} assert list(t.nodes()) == [0] assert list(ts.haplotypes(isolated_as_missing=False)) == ["1"] @@ -1899,7 +1911,7 @@ def test_one_node_one_sample_sites(self): methods = [t.parent, t.left_child, t.right_child, t.left_sib, t.right_sib] for method in methods: assert method(0) == tskit.NULL - for u in [-1, 1, 100]: + for u in [-1, 2, 100]: with pytest.raises(ValueError): method(u) tsp = ts.simplify(filter_sites=False) @@ -6565,6 +6577,81 @@ def test_coalescent_trees(self): next(new_trees) +class TestVirtualRootAPIs(ExampleTopologyMixin): + """ + Tests the APIs based on getting roots. + """ + + def verify(self, ts): + for tree in ts.trees(): + left_child = tree.left_child_array + right_child = tree.right_child_array + assert tree.virtual_root == ts.num_nodes + assert tree.left_root == tree.left_child(tree.virtual_root) + assert tree.right_root == tree.right_child(tree.virtual_root) + assert tree.left_root == left_child[-1] + assert tree.right_root == right_child[-1] + assert tree.parent(tree.virtual_root) == tskit.NULL + assert tree.left_sib(tree.virtual_root) == tskit.NULL + assert tree.right_sib(tree.virtual_root) == tskit.NULL + + u = tree.left_root + roots = [] + while u != tskit.NULL: + roots.append(u) + u = tree.right_sib(u) + assert roots == list(tree.roots) + + # The branch_length for roots is defined as 0, and it's consistent + # to have the same for the virtual root. + assert tree.branch_length(tree.virtual_root) == 0 + assert tree.depth(tree.virtual_root) == 0 + assert tree.num_children(tree.virtual_root) == tree.num_roots + assert tree.num_samples(tree.virtual_root) == tree.num_samples() + # We're not using tracked samples here. + assert tree.num_tracked_samples(tree.virtual_root) == 0 + # The virtual_root is internal because it has children (the roots) + assert tree.is_internal(tree.virtual_root) + assert not tree.is_leaf(tree.virtual_root) + assert not tree.is_sample(tree.virtual_root) + # The mrca of the virtual_root and anything except itself is -1 + assert tree.mrca(0, tree.virtual_root) == tskit.NULL + assert tree.mrca(tree.virtual_root, tree.virtual_root) == tree.virtual_root + # The virtual_root is a descendant of nothing other than itself + assert not tree.is_descendant(0, tree.virtual_root) + assert tree.is_descendant(tree.virtual_root, tree.virtual_root) + + assert list(tree.leaves(tree.virtual_root)) == list(tree.leaves()) + assert list(tree.samples(tree.virtual_root)) == list(tree.samples()) + + orders = [ + "preorder", + "inorder", + "levelorder", + "breadthfirst", + # https://github.com/tskit-dev/tskit/issues/1725 + # "postorder", + # "timeasc", + # "timedesc", + # "minlex_postorder", + ] + for order in orders: + l_vr = list(tree.nodes(tree.virtual_root, order=order)) + l_standard = list(tree.nodes(order=order)) + assert len(l_vr) == 1 + len(l_standard) + assert tree.virtual_root in l_vr + + # For pre-order, virtual_root should be first node visited: + assert next(tree.nodes(tree.virtual_root)) == tree.virtual_root + + # Methods that imply looking up tree sequence properties of the + # node raise an error + # Some methods don't apply + for method in [tree.time, tree.population]: + with pytest.raises(tskit.LibraryError, match="Node out of bounds"): + method(tree.virtual_root) + + class TestSampleLists(ExampleTopologyMixin): """ Tests for the sample lists algorithm. @@ -6605,6 +6692,12 @@ def verify(self, ts): break index = tree1.next_sample[index] assert samples1 == samples2 + # The python implementation here doesn't maintain roots + np.testing.assert_array_equal(tree1.parent, tree2.parent_array[:-1]) + np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array[:-1]) + np.testing.assert_array_equal( + tree1.right_child, tree2.right_child_array[:-1] + ) assert right == ts.sequence_length @@ -6615,10 +6708,9 @@ class TestOneSampleRoot(ExampleTopologyMixin): """ def verify(self, ts): - tree1 = tsutil.RootThresholdTree(ts, root_threshold=1) tree2 = tskit.Tree(ts) tree2.first() - for interval in tree1.iterate(): + for interval, tree1 in tsutil.algorithm_R(ts, root_threshold=1): assert interval == tree2.interval assert tree1.roots() == tree2.roots # Definition here is the set unique path ends from samples @@ -6629,35 +6721,96 @@ def verify(self, ts): u = tree2.parent(u) roots.add(path_end) assert set(tree1.roots()) == roots + np.testing.assert_array_equal(tree1.parent, tree2.parent_array) + np.testing.assert_array_equal(tree1.left_child, tree2.left_child_array) + np.testing.assert_array_equal(tree1.right_child, tree2.right_child_array) + np.testing.assert_array_equal(tree1.left_sib, tree2.left_sib_array) + np.testing.assert_array_equal(tree1.right_sib, tree2.right_sib_array) tree2.next() assert tree2.index == -1 -class TestKSamplesRoot(ExampleTopologyMixin): +class RootThreshold(ExampleTopologyMixin): """ Tests for the root criteria of subtending at least k samples. """ def verify(self, ts): - for k in range(1, 5): - tree1 = tsutil.RootThresholdTree(ts, root_threshold=k) - tree2 = tskit.Tree(ts, root_threshold=k) - tree2.first() - for interval in tree1.iterate(): - assert interval == tree2.interval - # Definition here is the set unique path ends from samples - # that subtend at least k samples - roots = set() - for u in ts.samples(): - while u != tskit.NULL: - path_end = u - u = tree2.parent(u) - if tree2.num_samples(path_end) >= k: - roots.add(path_end) - assert set(tree1.roots()) == roots - assert tree1.roots() == tree2.roots - tree2.next() - assert tree2.index == -1 + k = self.root_threshold + trees_py = tsutil.algorithm_R(ts, root_threshold=k) + tree_lib = tskit.Tree(ts, root_threshold=k) + tree_lib.first() + tree_leg = tsutil.LegacyRootThresholdTree(ts, root_threshold=k) + for (interval_py, tree_py), interval_leg in itertools.zip_longest( + trees_py, tree_leg.iterate() + ): + assert interval_py == tree_lib.interval + assert interval_leg == tree_lib.interval + + # Definition here is the set unique path ends from samples + # that subtend at least k samples + roots = set() + for u in ts.samples(): + while u != tskit.NULL: + path_end = u + u = tree_lib.parent(u) + if tree_lib.num_samples(path_end) >= k: + roots.add(path_end) + assert set(tree_py.roots()) == roots + assert set(tree_lib.roots) == roots + assert set(tree_leg.roots()) == roots + assert len(tree_leg.roots()) == tree_lib.num_roots + assert tree_py.roots() == tree_lib.roots + + # # The python class has identical behaviour to the lib version + assert tree_py.left_child[-1] == tree_lib.left_root + np.testing.assert_array_equal(tree_py.parent, tree_lib.parent_array) + np.testing.assert_array_equal(tree_py.left_child, tree_lib.left_child_array) + np.testing.assert_array_equal( + tree_py.right_child, tree_lib.right_child_array + ) + np.testing.assert_array_equal(tree_py.left_sib, tree_lib.left_sib_array) + np.testing.assert_array_equal(tree_py.right_sib, tree_lib.right_sib_array) + + # NOTE: the legacy left_root value is *not* necessarily the same as the + # new left_root. + # assert tree_leg.left_root == tree_py.left_child[-1] + + # The virtual root version is identical to the legacy tree + # except for the extra node and the details of the sib arrays. + np.testing.assert_array_equal(tree_py.parent[:-1], tree_leg.parent) + np.testing.assert_array_equal(tree_py.left_child[:-1], tree_leg.left_child) + np.testing.assert_array_equal( + tree_py.right_child[:-1], tree_leg.right_child + ) + # The sib arrays are identical except for root nodes. + for u in range(ts.num_nodes): + if u not in roots: + assert tree_py.left_sib[u] == tree_leg.left_sib[u] + assert tree_py.right_sib[u] == tree_leg.right_sib[u] + + tree_lib.next() + assert tree_lib.index == -1 + + +class TestRootThreshold1(RootThreshold): + root_threshold = 1 + + +class TestRootThreshold2(RootThreshold): + root_threshold = 2 + + +class TestRootThreshold3(RootThreshold): + root_threshold = 3 + + +class TestRootThreshold4(RootThreshold): + root_threshold = 4 + + +class TestRootThreshold10(RootThreshold): + root_threshold = 10 class TestSquashEdges: @@ -8074,7 +8227,7 @@ def test_is_isolated_bad(self): with pytest.raises(ValueError): tree.is_isolated(tskit.NULL) with pytest.raises(ValueError): - tree.is_isolated(ts.num_nodes) + tree.is_isolated(ts.num_nodes + 1) with pytest.raises(ValueError): tree.is_isolated(-2) with pytest.raises(TypeError): diff --git a/python/tests/tsutil.py b/python/tests/tsutil.py index d1bb725ee0..9e34844150 100644 --- a/python/tests/tsutil.py +++ b/python/tests/tsutil.py @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2018-2019 Tskit Developers +# Copyright (c) 2018-2021 Tskit Developers # Copyright (C) 2017 University of Oxford # # Permission is hereby granted, free of charge, to any person obtaining a copy @@ -1175,6 +1175,143 @@ def algorithm_T(ts): left = right +class QuintuplyLinkedTree: + def __init__(self, n, root_threshold=1): + self.root_threshold = root_threshold + self.parent = np.zeros(n + 1, dtype=np.int32) - 1 + self.left_child = np.zeros(n + 1, dtype=np.int32) - 1 + self.right_child = np.zeros(n + 1, dtype=np.int32) - 1 + self.left_sib = np.zeros(n + 1, dtype=np.int32) - 1 + self.right_sib = np.zeros(n + 1, dtype=np.int32) - 1 + self.num_samples = np.zeros(n + 1, dtype=np.int32) + + def __str__(self): + s = "id\tparent\tlchild\trchild\tlsib\trsib\tnsamp\n" + for j in range(len(self.parent)): + s += ( + f"{j}\t{self.parent[j]}\t" + f"{self.left_child[j]}\t{self.right_child[j]}\t" + f"{self.left_sib[j]}\t{self.right_sib[j]}\t" + f"{self.num_samples[j]}\n" + ) + return s + + def roots(self): + roots = [] + u = self.left_child[-1] + while u != -1: + roots.append(u) + u = self.right_sib[u] + return roots + + def remove_branch(self, p, c): + lsib = self.left_sib[c] + rsib = self.right_sib[c] + if lsib == -1: + self.left_child[p] = rsib + else: + self.right_sib[lsib] = rsib + if rsib == -1: + self.right_child[p] = lsib + else: + self.left_sib[rsib] = lsib + self.parent[c] = -1 + self.left_sib[c] = -1 + self.right_sib[c] = -1 + + def insert_branch(self, p, c): + assert self.parent[c] == -1, "contradictory edges" + self.parent[c] = p + u = self.right_child[p] + if u == -1: + self.left_child[p] = c + self.left_sib[c] = -1 + self.right_sib[c] = -1 + else: + self.right_sib[u] = c + self.left_sib[c] = u + self.right_sib[c] = -1 + self.right_child[p] = c + + def is_potential_root(self, u): + return self.num_samples[u] >= self.root_threshold + + # Note we cheat a bit here and use the -1 == last element semantics from Python. + # We could use self.insert_branch(N, root) and then set self.parent[root] = -1. + def insert_root(self, root): + self.insert_branch(-1, root) + + def remove_root(self, root): + self.remove_branch(-1, root) + + def remove_edge(self, edge): + self.remove_branch(edge.parent, edge.child) + + u = edge.parent + while u != -1: + path_end = u + path_end_was_root = self.is_potential_root(u) + self.num_samples[u] -= self.num_samples[edge.child] + u = self.parent[u] + + if path_end_was_root and not self.is_potential_root(path_end): + self.remove_root(path_end) + if self.is_potential_root(edge.child): + self.insert_root(edge.child) + + def insert_edge(self, edge): + u = edge.parent + while u != -1: + path_end = u + path_end_was_root = self.is_potential_root(u) + self.num_samples[u] += self.num_samples[edge.child] + u = self.parent[u] + + if self.is_potential_root(edge.child): + self.remove_root(edge.child) + if self.is_potential_root(path_end) and not path_end_was_root: + self.insert_root(path_end) + + self.insert_branch(edge.parent, edge.child) + + +def algorithm_R(ts, root_threshold=1): + """ + Quintuply linked tree with root tracking. + """ + sequence_length = ts.sequence_length + N = ts.num_nodes + M = ts.num_edges + tree = QuintuplyLinkedTree(N, root_threshold=root_threshold) + edges = list(ts.edges()) + in_order = ts.tables.indexes.edge_insertion_order + out_order = ts.tables.indexes.edge_removal_order + + # Initialise the tree + for u in ts.samples(): + tree.num_samples[u] = 1 + if tree.is_potential_root(u): + tree.insert_root(u) + + j = 0 + k = 0 + left = 0 + while j < M or left < sequence_length: + while k < M and edges[out_order[k]].right == left: + tree.remove_edge(edges[out_order[k]]) + k += 1 + while j < M and edges[in_order[j]].left == left: + tree.insert_edge(edges[in_order[j]]) + j += 1 + right = sequence_length + if j < M: + right = min(right, edges[in_order[j]].left) + if k < M: + right = min(right, edges[out_order[k]].right) + yield (left, right), tree + left = right + + class SampleListTree: """ Straightforward implementation of the quintuply linked tree for developing @@ -1315,15 +1452,8 @@ def sample_lists(self): sequence_length = ts.sequence_length edges = list(ts.edges()) M = len(edges) - time = [ts.node(edge.parent).time for edge in edges] - in_order = sorted( - range(M), - key=lambda j: (edges[j].left, time[j], edges[j].parent, edges[j].child), - ) - out_order = sorted( - range(M), - key=lambda j: (edges[j].right, -time[j], -edges[j].parent, -edges[j].child), - ) + in_order = ts.tables.indexes.edge_insertion_order + out_order = ts.tables.indexes.edge_removal_order j = 0 k = 0 left = 0 @@ -1348,10 +1478,12 @@ def sample_lists(self): left = right -class RootThresholdTree: +class LegacyRootThresholdTree: """ - Straightforward implementation of the quintuply linked tree for developing - and testing the root_threshold feature. + Implementation of the quintuply linked tree with root tracking using the + pre C 1.0/Python 0.4.0 algorithm. We keep this version around to make sure + that we can be clear what the differences in the semantics of the new + and old versions are. NOTE: The interface is pretty awkward; it's not intended for anything other than testing. @@ -1512,15 +1644,8 @@ def iterate(self): sequence_length = ts.sequence_length edges = list(ts.edges()) M = len(edges) - time = [ts.node(edge.parent).time for edge in edges] - in_order = sorted( - range(M), - key=lambda j: (edges[j].left, time[j], edges[j].parent, edges[j].child), - ) - out_order = sorted( - range(M), - key=lambda j: (edges[j].right, -time[j], -edges[j].parent, -edges[j].child), - ) + in_order = ts.tables.indexes.edge_insertion_order + out_order = ts.tables.indexes.edge_removal_order j = 0 k = 0 left = 0 diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 7627f4f9f7..ef17713d8d 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -1196,7 +1196,13 @@ def right_sample(self, u): def next_sample(self, u): return self._ll_tree.get_next_sample(u) - # TODO do we also have right_root? + @property + def virtual_root(self): + """ + The ID of the virtual root in this tree. + """ + return self._ll_tree.get_virtual_root() + @property def left_root(self): """ @@ -1229,7 +1235,11 @@ def left_root(self): :rtype: int """ - return self._ll_tree.get_left_root() + return self.left_child(self.virtual_root) + + @property + def right_root(self): + return self.right_child(self.virtual_root) def get_children(self, u): # Deprecated alias for self.children @@ -2040,10 +2050,8 @@ def num_samples(self, u=None): :return: The number of samples in the subtree rooted at u. :rtype: int """ - if u is None: - return sum(self._ll_tree.get_num_samples(u) for u in self.roots) - else: - return self._ll_tree.get_num_samples(u) + u = self.virtual_root if u is None else u + return self._ll_tree.get_num_samples(u) def get_num_tracked_leaves(self, u): # Deprecated alias for num_tracked_samples. The method name is inaccurate @@ -2069,6 +2077,10 @@ def num_tracked_samples(self, u=None): the subtree rooted at u. :rtype: int """ + # This should work, there's a but somethings wrong somewhere + # https://github.com/tskit-dev/tskit/issues/1724 + # u = self.virtual_root if u is None else u + # return self._ll_tree.get_num_tracked_samples(u) roots = [u] if u is None: roots = self.roots From 0a9886672a122d5e524916e023ac61006db4048c Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Sat, 9 Oct 2021 11:11:16 +0100 Subject: [PATCH 2/8] Update drawing samples with arbitrary node reordering --- c/tskit/trees.c | 2 +- python/tests/data/svg/internal_sample_ts.svg | 44 ++++++++++---------- python/tests/data/svg/ts_multiroot.svg | 32 +++++++------- python/tests/test_drawing.py | 21 ++++------ 4 files changed, 48 insertions(+), 51 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 2881684680..9dc00d677e 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3810,7 +3810,7 @@ tsk_tree_print_state(const tsk_tree_t *self, FILE *out) /* Methods for positioning the tree along the sequence */ -/* The following methods are very performance sensitive and so we use a +/* The following methods are performance sensitive and so we use a * lot of restrict pointers. Because we are saying that we don't have * any aliases to these pointers, we pass around the reference to parent * since it's used in all the functions. */ diff --git a/python/tests/data/svg/internal_sample_ts.svg b/python/tests/data/svg/internal_sample_ts.svg index e51d7f5e8c..897dd2afa2 100644 --- a/python/tests/data/svg/internal_sample_ts.svg +++ b/python/tests/data/svg/internal_sample_ts.svg @@ -99,7 +99,17 @@ - + + + + 7 + + + + + 8 + + @@ -149,16 +159,6 @@ 9 - - - - 7 - - - - - 8 - @@ -227,7 +227,17 @@ - + + + + 7 + + + + + 8 + + @@ -262,16 +272,6 @@ 6 - - - - 7 - - - - - 8 - diff --git a/python/tests/data/svg/ts_multiroot.svg b/python/tests/data/svg/ts_multiroot.svg index b3b8510b2c..19e4e0974d 100644 --- a/python/tests/data/svg/ts_multiroot.svg +++ b/python/tests/data/svg/ts_multiroot.svg @@ -466,21 +466,7 @@ - - - - - 2 - - - - - 4 - - - 6 - - + @@ -494,7 +480,7 @@ 7 - + @@ -508,6 +494,20 @@ 13 + + + + + 2 + + + + + 4 + + + 6 + diff --git a/python/tests/test_drawing.py b/python/tests/test_drawing.py index 9a7923fbab..f56b9bf111 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -1171,18 +1171,17 @@ def test_draw_forky_tree(self): ) self.verify_text_rendering(t.draw_text(order="minlex"), tree) - @pytest.mark.skip("FIXME rendering differently now") def test_draw_multiroot_forky_tree(self): tree = ( - " 13 \n" - "┏━┳━┳━╋━┳━━┓ \n" - "┃ ┃ ┃ ┃ ┃ 12 \n" - "┃ ┃ ┃ ┃ ┃ ┏┻┓ \n" - "┃ ┃ ┃ ┃ ┃ ┃ ┃ 11 \n" - "┃ ┃ ┃ ┃ ┃ ┃ ┃ ┏━┻┓ \n" - "┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ 10 \n" - "┃ ┃ ┃ ┃ ┃ ┃ ┃ ┃ ┏┻┓ \n" - "2 4 5 6 9 1 7 8 0 3 \n" + " 13 \n" + " ┏━┳━┳━╋━┳━━┓ \n" + " ┃ ┃ ┃ ┃ ┃ 12 \n" + " ┃ ┃ ┃ ┃ ┃ ┏┻┓\n" + " 11 ┃ ┃ ┃ ┃ ┃ ┃ ┃\n" + "┏━┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n" + "┃ 10 ┃ ┃ ┃ ┃ ┃ ┃ ┃\n" + "┃ ┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ ┃\n" + "8 0 3 2 4 5 6 9 1 7\n" ) nodes = io.StringIO( """\ @@ -2354,7 +2353,6 @@ def test_known_svg_ts_no_axes(self, overwrite_viz, draw_plotbox): svg, "ts_no_axes.svg", overwrite_viz, width=200 * ts.num_trees ) - @pytest.mark.skip("FIXME rendering differently now") def test_known_svg_ts_internal_sample(self, overwrite_viz, draw_plotbox): ts = tsutil.jiggle_samples(self.get_simple_ts()) svg = ts.draw_svg( @@ -2539,7 +2537,6 @@ def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox, capl width=200 * ts.num_trees, ) - @pytest.mark.skip("FIXME rendering differently now") def test_known_svg_ts_multiroot(self, overwrite_viz, draw_plotbox, caplog): tables = wf.wf_sim( 6, From a0af17b2f59763b04a7d98973420d02c26dcb76c Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Sat, 9 Oct 2021 14:06:20 +0100 Subject: [PATCH 3/8] Refactor node traversal tests Change timeasc and timedesc to keep the tree ordering within a timeslice rather than sorting by node ID (potentially breaking) Closes #1776 Closes #1725 --- c/tests/test_trees.c | 235 +++++++++++++- c/tests/testlib.c | 28 +- c/tskit/trees.c | 149 ++++++++- c/tskit/trees.h | 9 +- docs/data-model.rst | 4 + python/_tskitmodule.c | 98 ++++-- python/tests/__init__.py | 74 ----- python/tests/test_highlevel.py | 545 ++++++++++++++++++++++++++++----- python/tests/test_lowlevel.py | 19 ++ python/tests/test_topology.py | 170 +--------- python/tskit/trees.py | 171 +++++------ 11 files changed, 1044 insertions(+), 458 deletions(-) diff --git a/c/tests/test_trees.c b/c/tests/test_trees.c index ffde0cd9ec..22de5ea100 100644 --- a/c/tests/test_trees.c +++ b/c/tests/test_trees.c @@ -3863,8 +3863,8 @@ test_single_tree_iter_depths(void) const char *edges = "0 6 4 0,1\n" "0 6 5 2,3\n" "0 6 6 4,5\n"; - unsigned int depths[] = { 2, 2, 2, 2, 1, 1, 0 }; - tsk_size_t depth; + int depths[] = { 2, 2, 2, 2, 1, 1, 0 }; + int depth; tsk_treeseq_t ts; tsk_tree_t tree; tsk_id_t u; @@ -3878,14 +3878,14 @@ test_single_tree_iter_depths(void) CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&ts), num_nodes); for (u = 0; u < (tsk_id_t) num_nodes; u++) { - ret = tsk_tree_depth(&tree, u, &depth); + ret = tsk_tree_get_depth(&tree, u, &depth); CU_ASSERT_EQUAL(ret, 0); CU_ASSERT_EQUAL(depth, depths[u]); } - ret = tsk_tree_depth(&tree, (tsk_id_t) num_nodes + 1, &depth); + ret = tsk_tree_get_depth(&tree, (tsk_id_t) num_nodes + 1, &depth); CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); - ret = tsk_tree_depth(&tree, TSK_NULL, &depth); + ret = tsk_tree_get_depth(&tree, TSK_NULL, &depth); CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); ret = tsk_tree_next(&tree); @@ -5017,6 +5017,225 @@ test_non_sample_leaf_sample_lists(void) tsk_tree_free(&t); } +static void +test_virtual_root_properties(void) +{ + int ret; + tsk_treeseq_t ts; + tsk_tree_t t; + int depth; + double time; + tsk_id_t node; + + tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, NULL, + NULL, NULL, NULL, 0); + + ret = tsk_tree_init(&t, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsk_tree_first(&t); + CU_ASSERT_EQUAL_FATAL(ret, 1); + + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_depth(&t, t.virtual_root, &depth), 0) + CU_ASSERT_EQUAL_FATAL(depth, -1); + + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_time(&t, t.virtual_root, &time), 0) + CU_ASSERT_TRUE(isinf(time)); + + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&t, t.virtual_root, 0, &node), 0) + CU_ASSERT_EQUAL(node, t.virtual_root); + + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&t, 0, t.virtual_root, &node), 0) + CU_ASSERT_EQUAL(node, t.virtual_root); + + CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, t.virtual_root, &node), 0) + CU_ASSERT_EQUAL(node, TSK_NULL); + + /* The definition of "descendant" is that node v is on the path from + * u to a root. Since there is no parent link from roots to the + * virtual_root, it's consistent with this definition to return false + * for every node. */ + CU_ASSERT_FALSE(tsk_tree_is_descendant(&t, 0, t.virtual_root)); + CU_ASSERT_FALSE( + tsk_tree_is_descendant(&t, t.left_child[t.virtual_root], t.virtual_root)); + CU_ASSERT_FALSE(tsk_tree_is_descendant(&t, t.virtual_root, 0)); + /* The virtual_root *is* a descendent of itself, though. This is + * consistent with other nodes that are not "in" the tree being + * descendents of themselves, despite not being roots in the tree. */ + CU_ASSERT_TRUE(tsk_tree_is_descendant(&t, t.virtual_root, t.virtual_root)); + + CU_ASSERT_FALSE(tsk_tree_is_sample(&t, t.virtual_root)); + + tsk_tree_free(&t); + tsk_treeseq_free(&ts); +} + +/*======================================================= + * Tree traversals + *=======================================================*/ + +static void +verify_node_lists(tsk_size_t n, tsk_id_t *l1, tsk_id_t *l2) +{ + tsk_size_t j; + + for (j = 0; j < n; j++) { + /* printf("%d %d\n", l1[j], l2[j]); */ + CU_ASSERT_EQUAL(l1[j], l2[j]); + } +} + +static void +test_single_tree_traversal(void) +{ + int ret; + tsk_treeseq_t ts; + tsk_tree_t t; + tsk_size_t num_nodes = 7; + tsk_id_t preorder[] = { 6, 4, 0, 1, 5, 2, 3 }; + tsk_id_t preorder_vr[] = { 7, 6, 4, 0, 1, 5, 2, 3 }; + tsk_id_t postorder[] = { 0, 1, 4, 2, 3, 5, 6 }; + tsk_id_t postorder_vr[] = { 0, 1, 4, 2, 3, 5, 6, 7 }; + tsk_id_t nodes[num_nodes + 1]; + tsk_size_t n; + + tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, NULL, + NULL, NULL, NULL, 0); + + ret = tsk_tree_init(&t, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsk_tree_first(&t); + CU_ASSERT_EQUAL_FATAL(ret, 1); + + ret = tsk_tree_preorder(&t, -1, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, num_nodes); + verify_node_lists(n, nodes, preorder); + + ret = tsk_tree_preorder(&t, t.virtual_root, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, num_nodes + 1); + verify_node_lists(n, nodes, preorder_vr); + + ret = tsk_tree_preorder(&t, 5, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 3); + verify_node_lists(n, nodes, preorder + 4); + + ret = tsk_tree_postorder(&t, -1, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, num_nodes); + verify_node_lists(n, nodes, postorder); + + ret = tsk_tree_postorder(&t, t.virtual_root, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, num_nodes + 1); + verify_node_lists(n, nodes, postorder_vr); + + ret = tsk_tree_postorder(&t, 4, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 3); + verify_node_lists(n, nodes, postorder); + + /* Check errors */ + ret = tsk_tree_preorder(&t, -2, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + ret = tsk_tree_preorder(&t, 8, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + + ret = tsk_tree_postorder(&t, -2, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + ret = tsk_tree_postorder(&t, 8, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + + tsk_tree_free(&t); + tsk_treeseq_free(&ts); +} + +/* printed out in tree order. +0.90┊ ┊ 11 ┊ ┊ + ┊ ┊ ┏┻┓ ┊ ┊ +0.80┊ 10 ┊ ┃ ┃ ┊ ┊ + ┊ ┏┻┓ ┊ ┃ ┃ ┊ ┊ +0.40┊ 9 ┃ ┃ ┊ 9 ┃ ┃ ┊ 9 ┊ + ┊ ┏━┻┓ ┃ ┃ ┊ ┏━┻━┓ ┃ ┃ ┊ ┏━┻━━┓ ┊ +0.30┊ ┃ ┃ ┃ ┃ ┊ ┃ 8 ┃ ┃ ┊ ┃ 8 ┊ + ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊ ┃ ┏┻┓ ┊ +0.20┊ ┃ 7 ┃ ┃ ┊ 7 ┃ ┃ ┃ ┃ ┊ 7 ┃ ┃ ┊ + ┊ ┃ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┃ ┃ ┊ ┏━┻┓ ┃ ┃ ┊ +0.10┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ ┃ 6 ┃ ┃ ┊ + ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊ +0.00┊ 5 2 3 4 0 1 ┊ 3 4 1 2 0 5 ┊ 4 0 3 1 2 5 ┊ + 0 4 8 10 +*/ + +static void +test_multiroot_tree_traversal(void) +{ + int ret; + tsk_treeseq_t ts; + + tsk_tree_t t; + tsk_id_t preorder[] = { 5, 9, 2, 7, 3, 4, 10, 0, 1 }; + tsk_id_t preorder_vr[] = { 12, 5, 9, 2, 7, 3, 4, 10, 0, 1 }; + tsk_id_t postorder[] = { 5, 2, 3, 4, 7, 9, 0, 1, 10 }; + tsk_id_t postorder_vr[] = { 5, 2, 3, 4, 7, 9, 0, 1, 10, 12 }; + tsk_id_t nodes[13]; + tsk_size_t n; + + tsk_treeseq_from_text(&ts, 10, multiroot_ex_nodes, multiroot_ex_edges, NULL, NULL, + NULL, NULL, NULL, 0); + + ret = tsk_tree_init(&t, &ts, 0); + CU_ASSERT_EQUAL_FATAL(ret, 0); + ret = tsk_tree_first(&t); + CU_ASSERT_EQUAL_FATAL(ret, 1); + + ret = tsk_tree_preorder(&t, -1, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 9); + verify_node_lists(n, nodes, preorder); + + ret = tsk_tree_preorder(&t, t.virtual_root, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 10); + verify_node_lists(n, nodes, preorder_vr); + + ret = tsk_tree_preorder(&t, 10, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 3); + verify_node_lists(n, nodes, preorder + 6); + + ret = tsk_tree_postorder(&t, -1, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 9); + verify_node_lists(n, nodes, postorder); + + ret = tsk_tree_postorder(&t, t.virtual_root, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 10); + verify_node_lists(n, nodes, postorder_vr); + + ret = tsk_tree_postorder(&t, 10, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 3); + verify_node_lists(n, nodes, postorder + 6); + + /* Nodes that aren't "in" the tree have singleton traversal lists */ + + ret = tsk_tree_preorder(&t, 11, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 1); + CU_ASSERT_EQUAL_FATAL(nodes[0], 11); + + ret = tsk_tree_postorder(&t, 11, nodes, &n); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(n, 1); + CU_ASSERT_EQUAL_FATAL(nodes[0], 11); + + tsk_tree_free(&t); + tsk_treeseq_free(&ts); +} + /*======================================================= * KC Distance tests. *=======================================================*/ @@ -6405,6 +6624,12 @@ main(int argc, char **argv) { "test_internal_sample_sample_sets", test_internal_sample_sample_sets }, { "test_non_sample_leaf_sample_lists", test_non_sample_leaf_sample_lists }, + { "test_virtual_root_properties", test_virtual_root_properties }, + + /* tree traversal orders */ + { "test_single_tree_traversal", test_single_tree_traversal }, + { "test_multiroot_tree_traversal", test_multiroot_tree_traversal }, + /* KC distance tests */ { "test_single_tree_kc", test_single_tree_kc }, { "test_isolated_node_kc", test_isolated_node_kc }, diff --git a/c/tests/testlib.c b/c/tests/testlib.c index ae39068ef5..d966b3cfd6 100644 --- a/c/tests/testlib.c +++ b/c/tests/testlib.c @@ -309,20 +309,20 @@ const char *internal_sample_ex_mutations = "0 2 1\n" /*** An example of a tree sequence with multiple roots. ***/ /* -0.90┊ ┊ 11 ┊ ┊ - ┊ ┊ ┏┻┓ ┊ ┊ -0.80┊ 10 ┊ ┃ ┃ ┊ ┊ - ┊ ┏┻┓ ┊ ┃ ┃ ┊ ┊ -0.40┊ ┃ ┃ 9 ┊ ┃ ┃ 9 ┊ 9 ┊ - ┊ ┃ ┃ ┏━┻┓ ┊ ┃ ┃ ┏━┻━┓ ┊ ┏━┻━━┓ ┊ -0.30┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃ 8 ┃ ┊ ┃ 8 ┊ - ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┏┻┓ ┃ ┊ ┃ ┏┻┓ ┊ -0.20┊ ┃ ┃ ┃ 7 ┊ ┃ ┃ ┃ ┃ 7 ┊ 7 ┃ ┃ ┊ - ┊ ┃ ┃ ┃ ┏┻┓ ┊ ┃ ┃ ┃ ┃ ┏┻┓ ┊ ┏┻━┓ ┃ ┃ ┊ -0.10┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ 6 ┃ ┃ ┃ ┊ - ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┃ ┊ -0.00┊ 0 1 2 3 4 5 ┊ 0 5 1 2 3 4 ┊ 0 3 4 1 2 5 ┊ - 0.00 4.00 8.00 10.00 +0.90┊ ┊ 11 ┊ ┊ + ┊ ┊ ┏┻┓ ┊ ┊ +0.80┊ 10 ┊ ┃ ┃ ┊ ┊ + ┊ ┏┻┓ ┊ ┃ ┃ ┊ ┊ +0.40┊ 9 ┃ ┃ ┊ 9 ┃ ┃ ┊ 9 ┊ + ┊ ┏━┻┓ ┃ ┃ ┊ ┏━┻━┓ ┃ ┃ ┊ ┏━┻━━┓ ┊ +0.30┊ ┃ ┃ ┃ ┃ ┊ ┃ 8 ┃ ┃ ┊ ┃ 8 ┊ + ┊ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊ ┃ ┏┻┓ ┊ +0.20┊ ┃ 7 ┃ ┃ ┊ 7 ┃ ┃ ┃ ┃ ┊ 7 ┃ ┃ ┊ + ┊ ┃ ┏┻┓ ┃ ┃ ┊ ┏┻┓ ┃ ┃ ┃ ┃ ┊ ┏━┻┓ ┃ ┃ ┊ +0.10┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ ┃ 6 ┃ ┃ ┊ + ┊ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┃ ┃ ┃ ┃ ┊ ┃ ┏┻┓ ┃ ┃ ┊ +0.00┊ 5 2 3 4 0 1 ┊ 3 4 1 2 0 5 ┊ 4 0 3 1 2 5 ┊ + 0 4 8 10 */ const char *multiroot_ex_nodes = "1 0.0 0 -1\n" "1 0.0 0 -1\n" diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 9dc00d677e..e8943bd2dc 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3474,6 +3474,13 @@ tsk_tree_get_mrca(const tsk_tree_t *self, tsk_id_t u, tsk_id_t v, tsk_id_t *mrca goto out; } + /* Simplest to make the virtual_root a special case here to avoid + * doing the time lookup. */ + if (u == self->virtual_root || v == self->virtual_root) { + *mrca = self->virtual_root; + return 0; + } + tu = time[u]; tv = time[v]; while (u != v) { @@ -3615,11 +3622,16 @@ tsk_tree_get_time(const tsk_tree_t *self, tsk_id_t u, double *t) int ret = 0; tsk_node_t node; - ret = tsk_treeseq_get_node(self->tree_sequence, u, &node); - if (ret != 0) { - goto out; + if (u == self->virtual_root) { + *t = INFINITY; + } else { + + ret = tsk_treeseq_get_node(self->tree_sequence, u, &node); + if (ret != 0) { + goto out; + } + *t = node.time; } - *t = node.time; out: return ret; } @@ -3634,22 +3646,24 @@ tsk_tree_get_sites( } /* u must be a valid node in the tree. For internal use */ -static tsk_size_t -tsk_tree_get_depth(const tsk_tree_t *self, tsk_id_t u) +static int +tsk_tree_get_depth_unsafe(const tsk_tree_t *self, tsk_id_t u) { tsk_id_t v; - tsk_size_t depth = 0; + int depth = 0; + if (u == self->virtual_root) { + return -1; + } for (v = self->parent[u]; v != TSK_NULL; v = self->parent[v]) { depth++; } - return depth; } int TSK_WARN_UNUSED -tsk_tree_depth(const tsk_tree_t *self, tsk_id_t u, tsk_size_t *depth_ret) +tsk_tree_get_depth(const tsk_tree_t *self, tsk_id_t u, int *depth_ret) { int ret = 0; @@ -3658,7 +3672,7 @@ tsk_tree_depth(const tsk_tree_t *self, tsk_id_t u, tsk_size_t *depth_ret) goto out; } - *depth_ret = tsk_tree_get_depth(self, u); + *depth_ret = tsk_tree_get_depth_unsafe(self, u); out: return ret; } @@ -4241,6 +4255,121 @@ tsk_tree_clear(tsk_tree_t *self) return ret; } +/* Traversal orders */ + +int +tsk_tree_preorder( + const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes_ret) +{ + int ret = 0; + const tsk_id_t *restrict right_child = self->right_child; + const tsk_id_t *restrict left_sib = self->left_sib; + tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack)); + tsk_size_t num_nodes = 0; + tsk_id_t u, v; + int stack_top; + + if (stack == NULL) { + ret = TSK_ERR_NO_MEMORY; + goto out; + } + + if (root == -1) { + stack_top = -1; + for (u = right_child[self->virtual_root]; u != TSK_NULL; u = left_sib[u]) { + stack_top++; + stack[stack_top] = u; + } + } else { + ret = tsk_tree_check_node(self, root); + if (ret != 0) { + goto out; + } + stack_top = 0; + stack[stack_top] = root; + } + + while (stack_top >= 0) { + u = stack[stack_top]; + stack_top--; + nodes[num_nodes] = u; + num_nodes++; + for (v = right_child[u]; v != TSK_NULL; v = left_sib[v]) { + stack_top++; + stack[stack_top] = v; + } + } + *num_nodes_ret = num_nodes; +out: + /* can't use tsk_safe_free because this is a restrict pointer */ + if (stack != NULL) { + free(stack); + } + return ret; +} + +int +tsk_tree_postorder( + const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes_ret) +{ + int ret = 0; + const tsk_id_t *restrict right_child = self->right_child; + const tsk_id_t *restrict left_sib = self->left_sib; + const tsk_id_t *restrict parent = self->parent; + tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack)); + tsk_size_t num_nodes = 0; + tsk_id_t u, v, postorder_parent; + int stack_top; + bool is_virtual_root = root == self->virtual_root; + + if (stack == NULL) { + ret = TSK_ERR_NO_MEMORY; + goto out; + } + + if (root == -1 || is_virtual_root) { + stack_top = -1; + for (u = right_child[self->virtual_root]; u != TSK_NULL; u = left_sib[u]) { + stack_top++; + stack[stack_top] = u; + } + } else { + ret = tsk_tree_check_node(self, root); + if (ret != 0) { + goto out; + } + stack_top = 0; + stack[stack_top] = root; + } + + postorder_parent = TSK_NULL; + while (stack_top >= 0) { + u = stack[stack_top]; + if (right_child[u] != TSK_NULL && u != postorder_parent) { + for (v = right_child[u]; v != TSK_NULL; v = left_sib[v]) { + stack_top++; + stack[stack_top] = v; + } + } else { + stack_top--; + postorder_parent = parent[u]; + nodes[num_nodes] = u; + num_nodes++; + } + } + if (is_virtual_root) { + nodes[num_nodes] = root; + num_nodes++; + } + *num_nodes_ret = num_nodes; +out: + /* can't use tsk_safe_free because this is a restrict pointer */ + if (stack != NULL) { + free(stack); + } + return ret; +} + /* Parsimony methods */ static inline uint64_t diff --git a/c/tskit/trees.h b/c/tskit/trees.h index 9730353350..69ed9c6397 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -1,7 +1,7 @@ /* * MIT License * - * Copyright (c) 2019 Tskit Developers + * Copyright (c) 2019-2021 Tskit Developers * Copyright (c) 2015-2018 University of Oxford * * Permission is hereby granted, free of charge, to any person obtaining a copy @@ -425,6 +425,7 @@ int tsk_tree_set_tracked_samples_from_sample_list( int tsk_tree_get_parent(const tsk_tree_t *self, tsk_id_t u, tsk_id_t *parent); int tsk_tree_get_time(const tsk_tree_t *self, tsk_id_t u, double *t); +int tsk_tree_get_depth(const tsk_tree_t *self, tsk_id_t u, int *depth); int tsk_tree_get_mrca(const tsk_tree_t *self, tsk_id_t u, tsk_id_t v, tsk_id_t *mrca); int tsk_tree_get_num_samples( const tsk_tree_t *self, tsk_id_t u, tsk_size_t *num_samples); @@ -432,7 +433,11 @@ int tsk_tree_get_num_tracked_samples( const tsk_tree_t *self, tsk_id_t u, tsk_size_t *num_tracked_samples); int tsk_tree_get_sites( const tsk_tree_t *self, const tsk_site_t **sites, tsk_size_t *sites_length); -int tsk_tree_depth(const tsk_tree_t *self, tsk_id_t u, tsk_size_t *depth); + +int tsk_tree_preorder( + const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes); +int tsk_tree_postorder( + const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes); typedef struct { tsk_id_t node; diff --git a/docs/data-model.rst b/docs/data-model.rst index ee085a6842..55ae632e57 100644 --- a/docs/data-model.rst +++ b/docs/data-model.rst @@ -985,9 +985,13 @@ details of how to use the quintuply linked structure in the C API. ordering of the children of a node should therefore not be depended on. +.. _sec_data_model_tree_roots: + Accessing roots =============== +.. todo:: Update this with a discussion of the virtual root + The roots of a tree are defined as the unique endpoints of upward paths starting from sample nodes (if no path leads upward from a sample node, that node is also a root). Thus, trees can have multiple roots in ``tskit``. diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 460c1d2c03..f3bf52c452 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -772,20 +772,6 @@ table_get_offset_array(tsk_size_t num_rows, tsk_size_t *data) PyArrayObject *array; npy_intp dims = (npy_intp) num_rows + 1; - /* npy_intp j; */ - /* uint32_t *offset32; */ - /* /1* TODO make this conditional on the size of the data *1/ */ - /* /1* Actually, this should probably always return a uint64. Let's get */ - /* * other stuff working first and see how much it breaks if we always */ - /* * return 64 bit. *1/ */ - /* array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_UINT32, 0); */ - /* if (array == NULL) { */ - /* goto out; */ - /* } */ - /* offset32 = PyArray_DATA(array); */ - /* for (j = 0; j < dims; j++) { */ - /* offset32[j] = (uint32_t) data[j]; */ - /* } */ array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_UINT64, 0); if (array == NULL) { goto out; @@ -9493,18 +9479,18 @@ static PyObject * Tree_get_time(Tree *self, PyObject *args) { PyObject *ret = NULL; - tsk_node_t node; + double time; int node_id, err; if (Tree_get_node_argument(self, args, &node_id) != 0) { goto out; } - err = tsk_treeseq_get_node(self->tree->tree_sequence, node_id, &node); + err = tsk_tree_get_time(self->tree, node_id, &time); if (err != 0) { handle_library_error(err); goto out; } - ret = Py_BuildValue("d", node.time); + ret = Py_BuildValue("d", time); out: return ret; } @@ -9611,18 +9597,18 @@ static PyObject * Tree_depth(Tree *self, PyObject *args) { PyObject *ret = NULL; - tsk_size_t depth; + int depth; int node, err; if (Tree_get_node_argument(self, args, &node) != 0) { goto out; } - err = tsk_tree_depth(self->tree, node, &depth); + err = tsk_tree_get_depth(self->tree, node, &depth); if (ret != 0) { handle_library_error(err); goto out; } - ret = Py_BuildValue("I", (unsigned int) depth); + ret = Py_BuildValue("i", depth); out: return ret; } @@ -10051,6 +10037,70 @@ Tree_set_root_threshold(Tree *self, PyObject *args) return ret; } +typedef int tsk_traversal_func( + const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes); + +static PyObject * +Tree_get_traversal_array(Tree *self, PyObject *args, tsk_traversal_func *func) +{ + PyObject *ret = NULL; + PyArrayObject *array = NULL; + int32_t *data = NULL; + int root = TSK_NULL; + npy_intp dims; + tsk_size_t length; + int err; + + if (Tree_check_state(self) != 0) { + goto out; + } + if (!PyArg_ParseTuple(args, "i", &root)) { + goto out; + } + data = PyDataMem_NEW((self->tree->num_nodes + 1) * sizeof(*data)); + if (data == NULL) { + ret = PyErr_NoMemory(); + goto out; + } + err = func(self->tree, (tsk_id_t) root, data, &length); + if (err != 0) { + handle_library_error(err); + goto out; + } + dims = (npy_intp) length; + array = (PyArrayObject *) PyArray_SimpleNewFromData(1, &dims, NPY_INT32, data); + if (array == NULL) { + goto out; + } + /* Set the OWNDATA flag on that the data will be freed with the array */ + PyArray_ENABLEFLAGS(array, NPY_ARRAY_OWNDATA); + /* Not strictly necessary since we're creating a new array, but let's + * keep the door open to future optimisations. */ + PyArray_CLEARFLAGS(array, NPY_ARRAY_WRITEABLE); + + ret = (PyObject *) array; + data = NULL; + array = NULL; +out: + Py_XDECREF(array); + if (data != NULL) { + PyDataMem_FREE(data); + } + return ret; +} + +static PyObject * +Tree_get_preorder(Tree *self, PyObject *args) +{ + return Tree_get_traversal_array(self, args, tsk_tree_preorder); +} + +static PyObject * +Tree_get_postorder(Tree *self, PyObject *args) +{ + return Tree_get_traversal_array(self, args, tsk_tree_postorder); +} + /* The x_array properties are the high-performance zero-copy interface to the * corresponding arrays in the tsk_tree object. We use properties and * return a new array each time rather than trying to create a single array @@ -10329,6 +10379,14 @@ static PyMethodDef Tree_methods[] = { .ml_meth = (PyCFunction) Tree_get_root_threshold, .ml_flags = METH_NOARGS, .ml_doc = "Returns the root threshold for this tree." }, + { .ml_name = "get_preorder", + .ml_meth = (PyCFunction) Tree_get_preorder, + .ml_flags = METH_VARARGS, + .ml_doc = "Returns the nodes in this tree in preorder." }, + { .ml_name = "get_postorder", + .ml_meth = (PyCFunction) Tree_get_postorder, + .ml_flags = METH_VARARGS, + .ml_doc = "Returns the nodes in this tree in postorder." }, { NULL } /* Sentinel */ }; diff --git a/python/tests/__init__.py b/python/tests/__init__.py index cb23fa5836..ef16e9c0ca 100644 --- a/python/tests/__init__.py +++ b/python/tests/__init__.py @@ -83,80 +83,6 @@ def children(self, u): v = self.right_sib[v] return ret - def _preorder_nodes(self, u, node_ist): - node_ist.append(u) - for c in self.children(u): - self._preorder_nodes(c, node_ist) - - def _postorder_nodes(self, u, node_ist): - for c in self.children(u): - self._postorder_nodes(c, node_ist) - node_ist.append(u) - - def _inorder_nodes(self, u, node_ist): - children = self.children(u) - if len(children) > 0: - mid = len(children) // 2 - for v in children[:mid]: - self._inorder_nodes(v, node_ist) - node_ist.append(u) - for v in children[mid:]: - self._inorder_nodes(v, node_ist) - else: - node_ist.append(u) - - def _levelorder_nodes(self, u, node_ist, level): - node_ist[level].append(u) if level < len(node_ist) else node_ist.append([u]) - for c in self.children(u): - self._levelorder_nodes(c, node_ist, level + 1) - - def _minlex_postorder_nodes(self, u, node_ist): - node_ist.extend(self._minlex_postorder_nodes_helper(u)[1]) - - def _minlex_postorder_nodes_helper(self, u): - """ - For a given input ID u, this function returns a tuple whose first value - is the minimum leaf node ID under node u, and whose second value is - a list containing the minlex postorder for the subtree rooted at node u. - The first value is needed for sorting, and the second value is what - finally gets returned. - """ - children = self.children(u) - if len(children) > 0: - children_return = [self._minlex_postorder_nodes_helper(c) for c in children] - # sorts by first value, which is the minimum leaf node ID - children_return.sort() - minlex_postorder = [] - for _, child_minlex_postorder in children_return: - minlex_postorder.extend(child_minlex_postorder) - minlex_postorder.extend([u]) - return (children_return[0][0], minlex_postorder) - else: - return (u, [u]) - - def nodes(self, root=None, order="preorder"): - roots = [root] - if root is None: - roots = self.roots - for u in roots: - node_list = [] - if order == "preorder": - self._preorder_nodes(u, node_list) - elif order == "inorder": - self._inorder_nodes(u, node_list) - elif order == "postorder": - self._postorder_nodes(u, node_list) - elif order == "levelorder" or order == "breadthfirst": - # Returns nodes in their respective levels - # Nested list comprehension flattens node_list in order - self._levelorder_nodes(u, node_list, 0) - node_list = iter([i for level in node_list for i in level]) - elif order == "minlex_postorder": - self._minlex_postorder_nodes(u, node_list) - else: - raise ValueError("order not supported") - yield from node_list - def get_interval(self): return self.left, self.right diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index 201a2fd340..f3d94f15b7 100644 --- a/python/tests/test_highlevel.py +++ b/python/tests/test_highlevel.py @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2018-2020 Tskit Developers +# Copyright (c) 2018-2021 Tskit Developers # Copyright (c) 2015-2018 University of Oxford # # Permission is hereby granted, free of charge, to any person obtaining a copy @@ -58,6 +58,87 @@ from tskit import UNKNOWN_TIME +def traversal_preorder(tree, root=None): + roots = tree.roots if root is None else [root] + for node in roots: + yield node + for child in tree.children(node): + yield from traversal_preorder(tree, child) + + +def traversal_postorder(tree, root=None): + roots = tree.roots if root is None else [root] + for node in roots: + for child in tree.children(node): + yield from traversal_postorder(tree, child) + yield node + + +def traversal_inorder(tree, root=None): + roots = tree.roots if root is None else [root] + for node in roots: + children = list(tree.children(node)) + half = len(children) // 2 + for child in children[:half]: + yield from traversal_inorder(tree, child) + yield node + for child in children[half:]: + yield from traversal_inorder(tree, child) + + +def traversal_levelorder(tree, root=None): + yield from sorted(list(tree.nodes(root)), key=lambda u: tree.depth(u)) + + +def _traversal_minlex_postorder(tree, u): + """ + For a given input ID u, this function returns a tuple whose first value + is the minimum leaf node ID under node u, and whose second value is + a list containing the minlex postorder for the subtree rooted at node u. + The first value is needed for sorting, and the second value is what + finally gets returned. + """ + children = tree.children(u) + if len(children) > 0: + children_return = [_traversal_minlex_postorder(tree, c) for c in children] + # sorts by first value, which is the minimum leaf node ID + children_return.sort(key=lambda x: x[0]) + minlex_postorder = [] + for _, child_minlex_postorder in children_return: + minlex_postorder.extend(child_minlex_postorder) + minlex_postorder.extend([u]) + return (children_return[0][0], minlex_postorder) + else: + return (u, [u]) + + +def traversal_minlex_postorder(tree, root=None): + roots = tree.roots if root is None else [root] + root_lists = [_traversal_minlex_postorder(tree, node) for node in roots] + for _, node_list in sorted(root_lists, key=lambda x: x[0]): + yield from node_list + + +def traversal_timeasc(tree, root=None): + yield from sorted(tree.nodes(root, order="postorder"), key=lambda u: tree.time(u)) + + +def traversal_timedesc(tree, root=None): + yield from sorted(tree.nodes(root, order="preorder"), key=lambda u: -tree.time(u)) + + +traversal_map = { + "preorder": traversal_preorder, + "postorder": traversal_postorder, + "inorder": traversal_inorder, + "levelorder": traversal_levelorder, + "breadthfirst": traversal_levelorder, + "minlex_postorder": traversal_minlex_postorder, + "timeasc": traversal_timeasc, + "timedesc": traversal_timedesc, +} + + def insert_uniform_mutations(tables, num_mutations, nodes): """ Returns n evenly mutations over the specified list of nodes. @@ -363,6 +444,389 @@ def get_samples(ts, time=None, population=None): return np.array(samples) +class TestTreeTraversals: + @pytest.mark.parametrize("ts", get_example_tree_sequences()) + @pytest.mark.parametrize("order", list(traversal_map.keys())) + def test_traversals_virtual_root(self, ts, order): + tree = ts.first() + node_list2 = list(traversal_map[order](tree, tree.virtual_root)) + node_list1 = list(tree.nodes(tree.virtual_root, order=order)) + assert tree.virtual_root in node_list1 + assert node_list1 == node_list2 + + @pytest.mark.parametrize("ts", get_example_tree_sequences()) + @pytest.mark.parametrize("order", list(traversal_map.keys())) + def test_traversals(self, ts, order): + tree = next(ts.trees()) + traverser = traversal_map[order] + node_list1 = list(tree.nodes(order=order)) + node_list2 = list(traverser(tree)) + assert node_list1 == node_list2 + + def test_binary_example(self): + t = tskit.Tree.generate_balanced(5) + # 8 + # ┏━━┻━┓ + # ┃ 7 + # ┃ ┏━┻┓ + # 5 ┃ 6 + # ┏┻┓ ┃ ┏┻┓ + # 0 1 2 3 4 + + def f(node=None, order=None): + return list(t.nodes(node, order)) + + assert f(order="preorder") == [8, 5, 0, 1, 7, 2, 6, 3, 4] + assert f(order="postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8] + assert f(order="inorder") == [0, 5, 1, 8, 2, 7, 3, 6, 4] + assert f(order="levelorder") == [8, 5, 7, 0, 1, 2, 6, 3, 4] + assert f(order="breadthfirst") == [8, 5, 7, 0, 1, 2, 6, 3, 4] + assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8] + assert f(order="timedesc") == [8, 7, 5, 6, 0, 1, 2, 3, 4] + assert f(order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8] + + q = t.virtual_root + assert f(q, order="preorder") == [q, 8, 5, 0, 1, 7, 2, 6, 3, 4] + assert f(q, order="postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8, q] + assert f(q, order="inorder") == [q, 0, 5, 1, 8, 2, 7, 3, 6, 4] + assert f(q, order="levelorder") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4] + assert f(q, order="breadthfirst") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4] + assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, q] + assert f(q, order="timedesc") == [q, 8, 7, 5, 6, 0, 1, 2, 3, 4] + assert f(q, order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8, q] + + assert f(7, order="preorder") == [7, 2, 6, 3, 4] + assert f(7, order="postorder") == [2, 3, 4, 6, 7] + assert f(7, order="inorder") == [2, 7, 3, 6, 4] + assert f(7, order="levelorder") == [7, 2, 6, 3, 4] + assert f(7, order="breadthfirst") == [7, 2, 6, 3, 4] + assert f(7, order="timeasc") == [2, 3, 4, 6, 7] + assert f(7, order="timedesc") == [7, 6, 2, 3, 4] + assert f(7, order="minlex_postorder") == [2, 3, 4, 6, 7] + + def test_ternary_example(self): + t = tskit.Tree.generate_balanced(7, arity=3) + # 10 + # ┏━━━┳┻━━━┓ + # 7 8 9 + # ┏┻┓ ┏┻┓ ┏━╋━┓ + # 0 1 2 3 4 5 6 + + def f(node=None, order=None): + return list(t.nodes(node, order)) + + assert f(order="preorder") == [10, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6] + assert f(order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10] + assert f(order="inorder") == [0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6] + assert f(order="levelorder") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="breadthfirst") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + assert f(order="timedesc") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10] + + q = t.virtual_root + assert f(q, order="preorder") == [q, 10, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6] + assert f(q, order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10, q] + assert f(q, order="inorder") == [q, 0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6] + assert f(q, order="levelorder") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="breadthfirst") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, q] + assert f(q, order="timedesc") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10, q] + + assert f(9, order="preorder") == [9, 4, 5, 6] + assert f(9, order="postorder") == [4, 5, 6, 9] + assert f(9, order="inorder") == [4, 9, 5, 6] + assert f(9, order="levelorder") == [9, 4, 5, 6] + assert f(9, order="breadthfirst") == [9, 4, 5, 6] + assert f(9, order="timeasc") == [4, 5, 6, 9] + assert f(9, order="timedesc") == [9, 4, 5, 6] + assert f(9, order="minlex_postorder") == [4, 5, 6, 9] + + def test_multiroot_example(self): + tables = tskit.Tree.generate_balanced(7, arity=3).tree_sequence.dump_tables() + tables.edges.truncate(len(tables.edges) - 3) + t = tables.tree_sequence().first() + + # 7 8 9 + # ┏┻┓ ┏┻┓ ┏━╋━┓ + # 0 1 2 3 4 5 6 + def f(node=None, order=None): + return list(t.nodes(node, order)) + + assert f(order="preorder") == [7, 0, 1, 8, 2, 3, 9, 4, 5, 6] + assert f(order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9] + assert f(order="inorder") == [0, 7, 1, 2, 8, 3, 4, 9, 5, 6] + assert f(order="levelorder") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="breadthfirst") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + assert f(order="timedesc") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9] + + q = t.virtual_root + assert f(q, order="preorder") == [q, 7, 0, 1, 8, 2, 3, 9, 4, 5, 6] + assert f(q, order="postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, q] + assert f(q, order="inorder") == [0, 7, 1, q, 2, 8, 3, 4, 9, 5, 6] + assert f(q, order="levelorder") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="breadthfirst") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, q] + assert f(q, order="timedesc") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, q] + + assert f(9, order="preorder") == [9, 4, 5, 6] + assert f(9, order="postorder") == [4, 5, 6, 9] + assert f(9, order="inorder") == [4, 9, 5, 6] + assert f(9, order="levelorder") == [9, 4, 5, 6] + assert f(9, order="breadthfirst") == [9, 4, 5, 6] + assert f(9, order="minlex_postorder") == [4, 5, 6, 9] + assert f(9, order="timeasc") == [4, 5, 6, 9] + assert f(9, order="timedesc") == [9, 4, 5, 6] + + def test_multiroot_non_lexical_example(self): + nodes = io.StringIO( + """\ + id is_sample time population individual metadata + 0 1 0.000000 0 -1 b'' + 1 1 0.000000 0 -1 b'' + 2 1 0.000000 0 -1 b'' + 3 1 0.000000 0 -1 b'' + 4 1 0.000000 0 -1 b'' + 5 1 0.000000 0 -1 b'' + 6 1 0.000000 0 -1 b'' + 7 1 0.000000 0 -1 b'' + 8 1 0.000000 0 -1 b'' + 9 1 0.000000 0 -1 b'' + 10 0 0.047734 0 -1 b'' + 11 0 0.061603 0 -1 b'' + 12 0 0.189503 0 -1 b'' + 13 0 0.275885 0 -1 b'' + 14 0 0.518301 0 -1 b'' + """ + ) + edges = io.StringIO( + """\ + left right parent child + 0.000000 10000.000000 10 0 + 0.000000 10000.000000 10 2 + 0.000000 10000.000000 11 9 + 0.000000 10000.000000 11 10 + 0.000000 10000.000000 12 3 + 0.000000 10000.000000 12 7 + 0.000000 10000.000000 13 5 + 0.000000 10000.000000 13 11 + 0.000000 10000.000000 14 1 + 0.000000 10000.000000 14 8 + """ + ) + ts = tskit.load_text( + nodes, edges, sequence_length=10000, strict=False, base64_metadata=False + ) + t = ts.first() + + # Note: this is drawn out in "tree" order. + # 14 + # ┏┻┓ + # 13 ┃ ┃ + # ┏━┻━┓ ┃ ┃ + # 12 ┃ ┃ ┃ ┃ + # ┏┻┓ ┃ ┃ ┃ ┃ + # ┃ ┃ ┃ 11 ┃ ┃ + # ┃ ┃ ┃ ┏━┻┓ ┃ ┃ + # ┃ ┃ ┃ ┃ 10 ┃ ┃ + # ┃ ┃ ┃ ┃ ┏┻┓ ┃ ┃ + # 4 6 3 7 5 9 0 2 1 8 + + def f(node=None, order=None): + return list(t.nodes(node, order)) + + pre = f(order="preorder") + post = f(order="postorder") + inord = f(order="inorder") + level = f(order="levelorder") + breadth = f(order="breadthfirst") + timeasc = f(order="timeasc") + timedesc = f(order="timedesc") + minlex = f(order="minlex_postorder") + assert pre == [4, 6, 12, 3, 7, 13, 5, 11, 9, 10, 0, 2, 14, 1, 8] + assert post == [4, 6, 3, 7, 12, 5, 9, 0, 2, 10, 11, 13, 1, 8, 14] + assert inord == [4, 6, 3, 12, 7, 5, 13, 9, 11, 0, 10, 2, 1, 14, 8] + assert level == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2] + assert breadth == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2] + assert timeasc == [4, 6, 3, 7, 5, 9, 0, 2, 1, 8, 10, 11, 12, 13, 14] + assert timedesc == [14, 13, 12, 11, 10, 4, 6, 3, 7, 5, 9, 0, 2, 1, 8] + + # And the minlex tree: + # 14 + # ┏┻┓ + # 13 ┃ ┃ + # ┏━┻━┓ ┃ ┃ + # ┃ ┃ ┃ ┃ 12 + # ┃ ┃ ┃ ┃ ┏┻┓ + # 11 ┃ ┃ ┃ ┃ ┃ + # ┏┻━┓ ┃ ┃ ┃ ┃ ┃ + # 10 ┃ ┃ ┃ ┃ ┃ ┃ + # ┏┻┓ ┃ ┃ ┃ ┃ ┃ ┃ + # 0 2 9 5 1 8 3 7 4 6 + assert minlex == [0, 2, 10, 9, 11, 5, 13, 1, 8, 14, 3, 7, 12, 4, 6] + + @pytest.mark.parametrize( + ["order", "expected"], + [ + ("preorder", [[9, 6, 2, 3, 7, 4, 5, 0, 1], [10, 4, 8, 5, 0, 1, 6, 2, 3]]), + ("inorder", [[2, 6, 3, 9, 4, 7, 0, 5, 1], [4, 10, 0, 5, 1, 8, 2, 6, 3]]), + ("postorder", [[2, 3, 6, 4, 0, 1, 5, 7, 9], [4, 0, 1, 5, 2, 3, 6, 8, 10]]), + ("levelorder", [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]]), + ( + "breadthfirst", + [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]], + ), + ("timeasc", [[2, 3, 4, 0, 1, 5, 6, 7, 9], [4, 0, 1, 2, 3, 5, 6, 8, 10]]), + ("timedesc", [[9, 7, 6, 5, 2, 3, 4, 0, 1], [10, 8, 6, 5, 4, 0, 1, 2, 3]]), + ( + "minlex_postorder", + [[0, 1, 5, 4, 7, 2, 3, 6, 9], [0, 1, 5, 2, 3, 6, 8, 4, 10]], + ), + ], + ) + def test_ts_example(self, order, expected): + # 1.20┊ ┊ 10 ┊ + # ┊ ┊ ┏━┻━━┓ ┊ + # 0.90┊ 9 ┊ ┃ ┃ ┊ + # ┊ ┏━━┻━┓ ┊ ┃ ┃ ┊ + # 0.60┊ ┃ ┃ ┊ ┃ 8 ┊ + # ┊ ┃ ┃ ┊ ┃ ┏━┻━┓ ┊ + # 0.44┊ ┃ 7 ┊ ┃ ┃ ┃ ┊ + # ┊ ┃ ┏━┻┓ ┊ ┃ ┃ ┃ ┊ + # 0.21┊ 6 ┃ ┃ ┊ ┃ ┃ 6 ┊ + # ┊ ┏┻┓ ┃ ┃ ┊ ┃ ┃ ┏┻┓ ┊ + # 0.15┊ ┃ ┃ ┃ 5 ┊ ┃ 5 ┃ ┃ ┊ + # ┊ ┃ ┃ ┃ ┏┻┓ ┊ ┃ ┏┻┓ ┃ ┃ ┊ + # 0.00┊ 2 3 4 0 1 ┊ 4 0 1 2 3 ┊ + # 0.00 0.50 1.00 + nodes = """\ + id is_sample population time + 0 1 0 0.00000000000000 + 1 1 0 0.00000000000000 + 2 1 0 0.00000000000000 + 3 1 0 0.00000000000000 + 4 1 0 0.00000000000000 + 5 0 0 0.14567111023387 + 6 0 0 0.21385545626353 + 7 0 0 0.43508024345063 + 8 0 0 0.60156352971203 + 9 0 0 0.90000000000000 + 10 0 0 1.20000000000000 + """ + edges = """\ + id left right parent child + 0 0.00000000 1.00000000 5 0,1 + 1 0.00000000 1.00000000 6 2,3 + 2 0.00000000 0.50000000 7 4,5 + 3 0.50000000 1.00000000 8 5,6 + 4 0.00000000 0.50000000 9 6,7 + 5 0.50000000 1.00000000 10 4,8 + """ + ts = tskit.load_text( + nodes=io.StringIO(nodes), edges=io.StringIO(edges), strict=False + ) + tree_orders = [list(tree.nodes(order=order)) for tree in ts.trees()] + assert tree_orders == expected + + def test_polytomy_inorder(self): + """ + If there are N children, current inorder traversal first visits + floor(N/2) children, then the parent, then the remaining children. + Here we explicitly test that behaviour. + """ + # + # __4__ + # / / \ \ + # 0 1 2 3 + # + nodes_polytomy_4 = """\ + id is_sample population time + 0 1 0 0.00000000000000 + 1 1 0 0.00000000000000 + 2 1 0 0.00000000000000 + 3 1 0 0.00000000000000 + 4 0 0 1.00000000000000 + """ + edges_polytomy_4 = """\ + id left right parent child + 0 0.00000000 1.00000000 4 0,1,2,3 + """ + # + # __5__ + # / /|\ \ + # 0 1 2 3 4 + # + nodes_polytomy_5 = """\ + id is_sample population time + 0 1 0 0.00000000000000 + 1 1 0 0.00000000000000 + 2 1 0 0.00000000000000 + 3 1 0 0.00000000000000 + 4 1 0 0.00000000000000 + 5 0 0 1.00000000000000 + """ + edges_polytomy_5 = """\ + id left right parent child + 0 0.00000000 1.00000000 5 0,1,2,3,4 + """ + for nodes_string, edges_string, expected_result in [ + [nodes_polytomy_4, edges_polytomy_4, [[0, 1, 4, 2, 3]]], + [nodes_polytomy_5, edges_polytomy_5, [[0, 1, 5, 2, 3, 4]]], + ]: + ts = tskit.load_text( + nodes=io.StringIO(nodes_string), + edges=io.StringIO(edges_string), + strict=False, + ) + tree_orders = [] + for tree in ts.trees(): + tree_orders.append(list(tree.nodes(order="inorder"))) + assert tree_orders == expected_result + + def test_minlex_postorder_multiple_roots(self): + # + # 10 8 9 11 + # / \ / \ / \ / \ + # 5 3 2 4 6 7 1 0 + # + nodes_string = """\ + id is_sample population time + 0 1 0 0.00000000000000 + 1 1 0 0.00000000000000 + 2 1 0 0.00000000000000 + 3 1 0 0.00000000000000 + 4 1 0 0.00000000000000 + 5 1 0 0.00000000000000 + 6 1 0 0.00000000000000 + 7 1 0 0.00000000000000 + 8 0 0 1.00000000000000 + 9 0 0 1.00000000000000 + 10 0 0 1.00000000000000 + 11 0 0 1.00000000000000 + """ + edges_string = """\ + id left right parent child + 0 0.00000000 1.00000000 8 2,4 + 1 0.00000000 1.00000000 9 6,7 + 2 0.00000000 1.00000000 10 5,3 + 3 0.00000000 1.00000000 11 1,0 + """ + expected_result = [[0, 1, 11, 2, 4, 8, 3, 5, 10, 6, 7, 9]] + ts = tskit.load_text( + nodes=io.StringIO(nodes_string), + edges=io.StringIO(edges_string), + strict=False, + ) + tree_orders = [] + for tree in ts.trees(): + tree_orders.append(list(tree.nodes(order="minlex_postorder"))) + assert tree_orders == expected_result + + class TestMRCACalculator: """ Class to test the Schieber-Vishkin algorithm. @@ -2398,6 +2862,16 @@ def test_num_children(self): for u in tree.nodes(): assert tree.num_children(u) == len(tree.children(u)) + def test_virtual_root_semantics(self): + for ts in get_example_tree_sequences(): + for tree in ts.trees(): + assert math.isinf(tree.time(tree.virtual_root)) + assert tree.depth(tree.virtual_root) == -1 + assert tree.parent(tree.virtual_root) == -1 + assert list(tree.children(tree.virtual_root)) == tree.roots + with pytest.raises(tskit.LibraryError, match="bounds"): + tree.population(tree.virtual_root) + def test_root_properties(self): tested = set() for ts in get_example_tree_sequences(): @@ -2649,75 +3123,6 @@ def verify_nx_nearest_neighbor_search(self): nearest_neighbor_of = [min(dist_dod[u], key=dist_dod[u].get) for u in range(3)] assert [2, 2, 1] == [nearest_neighbor_of[u] for u in range(3)] - def test_traversals(self): - for ts in get_example_tree_sequences(): - tree = next(ts.trees()) - self.verify_traversals(tree) - - # Verify time-ordered traversals separately, because the PythonTree - # class does not contain time information at the moment - for root in tree.roots: - time_ordered = tree.nodes(root, order="timeasc") - t = tree.time(next(time_ordered)) - for u in time_ordered: - next_t = tree.time(u) - assert next_t >= t - t = next_t - time_ordered = tree.nodes(root, order="timedesc") - t = tree.time(next(time_ordered)) - for u in time_ordered: - next_t = tree.time(u) - assert next_t <= t - t = next_t - - def verify_traversals(self, tree): - t1 = tree - t2 = tests.PythonTree.from_tree(t1) - assert list(t1.nodes()) == list(t2.nodes()) - orders = [ - "inorder", - "postorder", - "levelorder", - "breadthfirst", - "minlex_postorder", - ] - if tree.num_roots == 1: - with pytest.raises(ValueError): - list(t1.nodes(order="bad order")) - assert list(t1.nodes()) == list(t1.nodes(t1.get_root())) - assert list(t1.nodes()) == list(t1.nodes(t1.get_root(), "preorder")) - for u in t1.nodes(): - assert list(t1.nodes(u)) == list(t2.nodes(u)) - for test_order in orders: - assert sorted(list(t1.nodes())) == sorted( - list(t1.nodes(order=test_order)) - ) - assert list(t1.nodes(order=test_order)) == list( - t1.nodes(t1.get_root(), order=test_order) - ) - assert list(t1.nodes(order=test_order)) == list( - t1.nodes(t1.get_root(), test_order) - ) - assert list(t1.nodes(order=test_order)) == list( - t2.nodes(order=test_order) - ) - for u in t1.nodes(): - assert list(t1.nodes(u, test_order)) == list( - t2.nodes(u, test_order) - ) - else: - for test_order in orders: - all_nodes = [] - for root in t1.roots: - assert list(t1.nodes(root, order=test_order)) == list( - t2.nodes(root, order=test_order) - ) - all_nodes.extend(t1.nodes(root, order=test_order)) - # minlex_postorder reorders the roots, so this last test is - # not appropriate - if test_order != "minlex_postorder": - assert all_nodes == list(t1.nodes(order=test_order)) - def test_total_branch_length(self): # Note: this definition works when we have no non-sample branches. t1 = self.get_tree() diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 972f8c0815..8270e2022e 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -3074,6 +3074,25 @@ def test_array_lifetime(self, array): a2[:] = 0 assert a3 is not a2 + @pytest.mark.parametrize("ordering", ["preorder", "postorder"]) + def test_traversal_arrays(self, ordering): + ts = self.get_example_tree_sequence(10) + tree = _tskit.Tree(ts) + tree.first() + method = getattr(tree, "get_" + ordering) + for bad_type in [None, {}]: + with pytest.raises(TypeError): + method(bad_type) + for bad_node in [-2, 10 ** 6]: + with pytest.raises(_tskit.LibraryError, match="out of bounds"): + method(bad_node) + a = method(tree.get_virtual_root()) + assert a.dtype == np.int32 + assert not a.flags.writeable + assert a.flags.aligned + assert a.flags.c_contiguous + assert a.flags.owndata + class TestTableMetadataSchema(MetadataTestMixin): def test_metadata_schema_attribute(self): diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index a2c72d7669..cb479c8ee2 100644 --- a/python/tests/test_topology.py +++ b/python/tests/test_topology.py @@ -2622,162 +2622,6 @@ def test_nonbinary_tree_sequence_permuted_nodes(self): self.verify_permuted_nodes(ts) -class TestTraversalOrder: - """ - Tests node traversal orders. - """ - - # - # 9 10 - # / \ / \ - # / \ / 8 - # / \ / / \ - # 7 \ / / \ - # / \ 6 / / 6 - # / 5 / \ / 5 / \ - # / / \ / \ / / \ / \ - # 4 0 1 2 3 4 0 1 2 3 - # - # 0 ------------------ 0.5 ------------------ 1.0 - nodes = """\ - id is_sample population time - 0 1 0 0.00000000000000 - 1 1 0 0.00000000000000 - 2 1 0 0.00000000000000 - 3 1 0 0.00000000000000 - 4 1 0 0.00000000000000 - 5 0 0 0.14567111023387 - 6 0 0 0.21385545626353 - 7 0 0 0.43508024345063 - 8 0 0 0.60156352971203 - 9 0 0 0.90000000000000 - 10 0 0 1.20000000000000 - """ - edges = """\ - id left right parent child - 0 0.00000000 1.00000000 5 0,1 - 1 0.00000000 1.00000000 6 2,3 - 2 0.00000000 0.50000000 7 4,5 - 3 0.50000000 1.00000000 8 5,6 - 4 0.00000000 0.50000000 9 6,7 - 5 0.50000000 1.00000000 10 4,8 - """ - node_order_results = { - "preorder": [[9, 6, 2, 3, 7, 4, 5, 0, 1], [10, 4, 8, 5, 0, 1, 6, 2, 3]], - "inorder": [[2, 6, 3, 9, 4, 7, 0, 5, 1], [4, 10, 0, 5, 1, 8, 2, 6, 3]], - "postorder": [[2, 3, 6, 4, 0, 1, 5, 7, 9], [4, 0, 1, 5, 2, 3, 6, 8, 10]], - "levelorder": [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]], - "breadthfirst": [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]], - "timeasc": [[0, 1, 2, 3, 4, 5, 6, 7, 9], [0, 1, 2, 3, 4, 5, 6, 8, 10]], - "timedesc": [[9, 7, 6, 5, 4, 3, 2, 1, 0], [10, 8, 6, 5, 4, 3, 2, 1, 0]], - "minlex_postorder": [[0, 1, 5, 4, 7, 2, 3, 6, 9], [0, 1, 5, 2, 3, 6, 8, 4, 10]], - } - - def test_traversal_order(self): - ts = tskit.load_text( - nodes=io.StringIO(self.nodes), edges=io.StringIO(self.edges), strict=False - ) - for test_order, expected_result in self.node_order_results.items(): - tree_orders = [] - for tree in ts.trees(): - tree_orders.append(list(tree.nodes(order=test_order))) - assert tree_orders == expected_result - - def test_polytomy_inorder(self): - """ - If there are N children, current inorder traversal first visits - floor(N/2) children, then the parent, then the remaining children. - Here we explicitly test that behaviour. - """ - # - # __4__ - # / / \ \ - # 0 1 2 3 - # - nodes_polytomy_4 = """\ - id is_sample population time - 0 1 0 0.00000000000000 - 1 1 0 0.00000000000000 - 2 1 0 0.00000000000000 - 3 1 0 0.00000000000000 - 4 0 0 1.00000000000000 - """ - edges_polytomy_4 = """\ - id left right parent child - 0 0.00000000 1.00000000 4 0,1,2,3 - """ - # - # __5__ - # / /|\ \ - # 0 1 2 3 4 - # - nodes_polytomy_5 = """\ - id is_sample population time - 0 1 0 0.00000000000000 - 1 1 0 0.00000000000000 - 2 1 0 0.00000000000000 - 3 1 0 0.00000000000000 - 4 1 0 0.00000000000000 - 5 0 0 1.00000000000000 - """ - edges_polytomy_5 = """\ - id left right parent child - 0 0.00000000 1.00000000 5 0,1,2,3,4 - """ - for nodes_string, edges_string, expected_result in [ - [nodes_polytomy_4, edges_polytomy_4, [[0, 1, 4, 2, 3]]], - [nodes_polytomy_5, edges_polytomy_5, [[0, 1, 5, 2, 3, 4]]], - ]: - ts = tskit.load_text( - nodes=io.StringIO(nodes_string), - edges=io.StringIO(edges_string), - strict=False, - ) - tree_orders = [] - for tree in ts.trees(): - tree_orders.append(list(tree.nodes(order="inorder"))) - assert tree_orders == expected_result - - def test_minlex_postorder_multiple_roots(self): - # - # 10 8 9 11 - # / \ / \ / \ / \ - # 5 3 2 4 6 7 1 0 - # - nodes_string = """\ - id is_sample population time - 0 1 0 0.00000000000000 - 1 1 0 0.00000000000000 - 2 1 0 0.00000000000000 - 3 1 0 0.00000000000000 - 4 1 0 0.00000000000000 - 5 1 0 0.00000000000000 - 6 1 0 0.00000000000000 - 7 1 0 0.00000000000000 - 8 0 0 1.00000000000000 - 9 0 0 1.00000000000000 - 10 0 0 1.00000000000000 - 11 0 0 1.00000000000000 - """ - edges_string = """\ - id left right parent child - 0 0.00000000 1.00000000 8 2,4 - 1 0.00000000 1.00000000 9 6,7 - 2 0.00000000 1.00000000 10 5,3 - 3 0.00000000 1.00000000 11 1,0 - """ - expected_result = [[0, 1, 11, 2, 4, 8, 3, 5, 10, 6, 7, 9]] - ts = tskit.load_text( - nodes=io.StringIO(nodes_string), - edges=io.StringIO(edges_string), - strict=False, - ) - tree_orders = [] - for tree in ts.trees(): - tree_orders.append(list(tree.nodes(order="minlex_postorder"))) - assert tree_orders == expected_result - - class TestSimplifyExamples(TopologyTestCase): """ Tests for simplify where we write out the input and expected output @@ -6605,7 +6449,8 @@ def verify(self, ts): # The branch_length for roots is defined as 0, and it's consistent # to have the same for the virtual root. assert tree.branch_length(tree.virtual_root) == 0 - assert tree.depth(tree.virtual_root) == 0 + # The virtual root has depth -1 from the root + assert tree.depth(tree.virtual_root) == -1 assert tree.num_children(tree.virtual_root) == tree.num_roots assert tree.num_samples(tree.virtual_root) == tree.num_samples() # We're not using tracked samples here. @@ -6629,11 +6474,10 @@ def verify(self, ts): "inorder", "levelorder", "breadthfirst", - # https://github.com/tskit-dev/tskit/issues/1725 - # "postorder", - # "timeasc", - # "timedesc", - # "minlex_postorder", + "postorder", + "timeasc", + "timedesc", + "minlex_postorder", ] for order in orders: l_vr = list(tree.nodes(tree.virtual_root, order=order)) @@ -6647,7 +6491,7 @@ def verify(self, ts): # Methods that imply looking up tree sequence properties of the # node raise an error # Some methods don't apply - for method in [tree.time, tree.population]: + for method in [tree.population]: with pytest.raises(tskit.LibraryError, match="Node out of bounds"): method(tree.virtual_root) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index ef17713d8d..1e2f905d44 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -1265,8 +1265,10 @@ def get_time(self, u): def time(self, u): """ - Returns the time of the specified node. - Equivalent to ``tree.tree_sequence.node(u).time``. + Returns the time of the specified node. This is equivalently + to ``tree.tree_sequence.node(u).time`` except for the special + case of the tree's :ref:`virtual root `, + which is defined as positive infinity. :param int u: The node of interest. :return: The time of u. @@ -1280,6 +1282,9 @@ def depth(self, u): root, not including ``u``. Thus, the depth of a root is zero. + As a special case, the depth of the :ref:`virtual root + ` is defined as -1. + :param int u: The node of interest. :return: The depth of u. :rtype: int @@ -2086,50 +2091,37 @@ def num_tracked_samples(self, u=None): roots = self.roots return sum(self._ll_tree.get_num_tracked_samples(root) for root in roots) - def _preorder_traversal(self, u): - stack = collections.deque([u]) - # For perf we store these to avoid lookups in the tight loop - pop = stack.pop - extend = stack.extend - get_children = self.children - # Note: the usual style is to be explicit about what we're testing - # and use while len(stack) > 0, but this form is slightly faster. - while stack: - v = pop() - extend(reversed(get_children(v))) - yield v + def preorder(self, u=NULL): + return self._ll_tree.get_preorder(u) - def _postorder_traversal(self, u): - stack = collections.deque([u]) - parent = NULL - # For perf we store these to avoid lookups in the tight loop - pop = stack.pop - extend = stack.extend - get_children = self.children - get_parent = self.get_parent - # Note: the usual style is to be explicit about what we're testing - # and use while len(stack) > 0, but this form is slightly faster. - while stack: - v = stack[-1] - children = [] if v == parent else get_children(v) - if children: - extend(reversed(children)) - else: - parent = get_parent(v) - yield pop() + def postorder(self, u=NULL): + return self._ll_tree.get_postorder(u) + + def _preorder_traversal(self, root): + return map(int, self.preorder(root)) + + def _postorder_traversal(self, root): + return map(int, self.postorder(root)) - def _inorder_traversal(self, u): + def _inorder_traversal(self, root): # TODO add a nonrecursive version of the inorder traversal. - children = self.get_children(u) - mid = len(children) // 2 - for c in children[:mid]: - yield from self._inorder_traversal(c) - yield u - for c in children[mid:]: - yield from self._inorder_traversal(c) - - def _levelorder_traversal(self, u): - queue = collections.deque([u]) + + def traverse(u): + children = self.get_children(u) + mid = len(children) // 2 + for c in children[:mid]: + yield from traverse(c) + yield u + for c in children[mid:]: + yield from traverse(c) + + roots = self.roots if root == NULL else [root] + for root in roots: + yield from traverse(root) + + def _levelorder_traversal(self, root): + roots = self.roots if root == NULL else [root] + queue = collections.deque(roots) # For perf we store these to avoid lookups in the tight loop pop = queue.popleft extend = queue.extend @@ -2141,25 +2133,26 @@ def _levelorder_traversal(self, u): extend(children(v)) yield v - def _timeasc_traversal(self, u): + def _timeasc_traversal(self, root): """ - Sorts by increasing time but falls back to increasing ID for equal times. + A stricter version of postorder where nodes are visiting in time-increasing + order first, and in preorder within equal time-slices. """ + # TODO implement with numpy? yield from sorted( - self.nodes(u, order="levelorder"), key=lambda u: (self.time(u), u) + self.nodes(root, order="postorder"), key=lambda u: self.time(u) ) - def _timedesc_traversal(self, u): + def _timedesc_traversal(self, root): """ - Sorts by decreasing time but falls back to decreasing ID for equal times. + A stricter version of preorder where nodes are visiting in time-descreasing + order first, and in preorder within equal time-slices. """ yield from sorted( - self.nodes(u, order="levelorder"), - key=lambda u: (self.time(u), u), - reverse=True, + self.nodes(root, order="preorder"), key=lambda u: -self.time(u) ) - def _minlex_postorder_traversal(self, u): + def _minlex_postorder_traversal(self, root): """ Postorder traversal that visits leaves in minimum lexicographic order. @@ -2169,55 +2162,46 @@ def _minlex_postorder_traversal(self, u): multiple Trees of a TreeSequence, as it leads to more consistency between adjacent Trees. """ - # We skip perf optimisations here (compared to _preorder_traversal and - # _postorder_traversal) as this ordering is unlikely to be used in perf - # sensitive applications - stack = collections.deque([u]) - parent = NULL # We compute a dictionary mapping from internal node ID to min leaf ID # under the node, using a first postorder traversal - min_leaf_dict = {} - while len(stack) > 0: - v = stack[-1] - children = [] if v == parent else self.children(v) - if children: - # The first time visiting a node, we push its children onto the stack. - # reversed is not strictly necessary, but it gives the postorder - # we would intuitively expect. - stack.extend(reversed(children)) + min_leaf = {} + for u in self.nodes(root, order="postorder"): + if self.is_leaf(u): + min_leaf[u] = u else: - # The second time visiting a node, we record its min leaf ID - # underneath, pop it, and update the parent variable - if v != parent: - # at a leaf node - min_leaf_dict[v] = v - else: - # at a parent after finishing all its children - min_leaf_dict[v] = min([min_leaf_dict[c] for c in self.children(v)]) - parent = self.get_parent(v) - stack.pop() + min_leaf[u] = min(min_leaf[v] for v in self.children(u)) + + stack = [] - # Now we do a second postorder traversal - stack.clear() - stack.extend([u]) + def push(nodes): + stack.extend(sorted(nodes, key=lambda u: min_leaf[u], reverse=True)) + + # The postorder traversal isn't robust to using virtual_root directly + # as a node because we depend on tree.parent() returning the last + # node we visiting on the path from "root". So, we treat this as a + # special case. + is_virtual_root = root == self.virtual_root + roots = self.roots if root == -1 or is_virtual_root else [root] + + push(roots) parent = NULL while len(stack) > 0: v = stack[-1] children = [] if v == parent else self.children(v) - if children: + if len(children) > 0: # The first time visiting a node, we push onto the stack its children # in order of reverse min leaf ID under each child. This guarantees # that the earlier children visited have smaller min leaf ID, # which is equivalent to the minlex condition. - stack.extend( - sorted(children, key=lambda u: min_leaf_dict[u], reverse=True) - ) + push(children) else: # The second time visiting a node, we pop and yield it, and # we update the parent variable parent = self.get_parent(v) yield stack.pop() + if is_virtual_root: + yield self.virtual_root def nodes(self, root=None, order="preorder"): """ @@ -2254,6 +2238,7 @@ def nodes(self, root=None, order="preorder"): `Wikipedia `__. + - FIXME BEFORE MERGING timeas and timedesc - 'timeasc': visits the nodes in order of increasing time, falling back to increasing ID if times are equal. - 'timedesc': visits the nodes in order of decreasing time, falling back to @@ -2289,23 +2274,9 @@ def nodes(self, root=None, order="preorder"): iterator = methods[order] except KeyError: raise ValueError(f"Traversal ordering '{order}' not supported") - roots = [root] - if root is None: - roots = self.roots - if order == "minlex_postorder" and len(roots) > 1: - # we need to visit the roots in minlex order as well - # we first visit all the roots and then sort by the min value - root_values = [] - for u in roots: - root_minlex_postorder = list(iterator(u)) - min_value = root_minlex_postorder[0] - root_values.append([min_value, root_minlex_postorder]) - root_values.sort() - for _, nodes_for_root in root_values: - yield from nodes_for_root - else: - for u in roots: - yield from iterator(u) + + root = -1 if root is None else root + return iterator(root) def _node_edges(self): """ From 364a009ebd3f1fb20a2840aa7608c7a69c3b2d94 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 12 Oct 2021 14:36:09 +0100 Subject: [PATCH 4/8] Revert changes to timeasc ordering. --- python/tests/test_highlevel.py | 38 +++++++++++++++++----------------- python/tskit/trees.py | 13 ++++-------- 2 files changed, 23 insertions(+), 28 deletions(-) diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index f3d94f15b7..4264ef0241 100644 --- a/python/tests/test_highlevel.py +++ b/python/tests/test_highlevel.py @@ -120,11 +120,11 @@ def traversal_minlex_postorder(tree, root=None): def traversal_timeasc(tree, root=None): - yield from sorted(tree.nodes(root, order="postorder"), key=lambda u: tree.time(u)) + yield from sorted(tree.nodes(root), key=lambda u: (tree.time(u), u)) def traversal_timedesc(tree, root=None): - yield from sorted(tree.nodes(root, order="preorder"), key=lambda u: -tree.time(u)) + yield from sorted(tree.nodes(root), key=lambda u: (tree.time(u), u), reverse=True) traversal_map = { @@ -482,7 +482,7 @@ def f(node=None, order=None): assert f(order="levelorder") == [8, 5, 7, 0, 1, 2, 6, 3, 4] assert f(order="breadthfirst") == [8, 5, 7, 0, 1, 2, 6, 3, 4] assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8] - assert f(order="timedesc") == [8, 7, 5, 6, 0, 1, 2, 3, 4] + assert f(order="timedesc") == [8, 7, 6, 5, 4, 3, 2, 1, 0] assert f(order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8] q = t.virtual_root @@ -492,7 +492,7 @@ def f(node=None, order=None): assert f(q, order="levelorder") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4] assert f(q, order="breadthfirst") == [q, 8, 5, 7, 0, 1, 2, 6, 3, 4] assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, q] - assert f(q, order="timedesc") == [q, 8, 7, 5, 6, 0, 1, 2, 3, 4] + assert f(q, order="timedesc") == [q, 8, 7, 6, 5, 4, 3, 2, 1, 0] assert f(q, order="minlex_postorder") == [0, 1, 5, 2, 3, 4, 6, 7, 8, q] assert f(7, order="preorder") == [7, 2, 6, 3, 4] @@ -501,7 +501,7 @@ def f(node=None, order=None): assert f(7, order="levelorder") == [7, 2, 6, 3, 4] assert f(7, order="breadthfirst") == [7, 2, 6, 3, 4] assert f(7, order="timeasc") == [2, 3, 4, 6, 7] - assert f(7, order="timedesc") == [7, 6, 2, 3, 4] + assert f(7, order="timedesc") == [7, 6, 4, 3, 2] assert f(7, order="minlex_postorder") == [2, 3, 4, 6, 7] def test_ternary_example(self): @@ -520,8 +520,8 @@ def f(node=None, order=None): assert f(order="inorder") == [0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6] assert f(order="levelorder") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] assert f(order="breadthfirst") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] - assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - assert f(order="timedesc") == [10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="timeasc") == list(range(11)) + assert f(order="timedesc") == list(reversed(range(11))) assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10] q = t.virtual_root @@ -530,8 +530,8 @@ def f(node=None, order=None): assert f(q, order="inorder") == [q, 0, 7, 1, 10, 2, 8, 3, 4, 9, 5, 6] assert f(q, order="levelorder") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] assert f(q, order="breadthfirst") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] - assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, q] - assert f(q, order="timedesc") == [q, 10, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="timeasc") == list(range(12)) + assert f(q, order="timedesc") == list(reversed(range(12))) assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, 10, q] assert f(9, order="preorder") == [9, 4, 5, 6] @@ -540,7 +540,7 @@ def f(node=None, order=None): assert f(9, order="levelorder") == [9, 4, 5, 6] assert f(9, order="breadthfirst") == [9, 4, 5, 6] assert f(9, order="timeasc") == [4, 5, 6, 9] - assert f(9, order="timedesc") == [9, 4, 5, 6] + assert f(9, order="timedesc") == [9, 6, 5, 4] assert f(9, order="minlex_postorder") == [4, 5, 6, 9] def test_multiroot_example(self): @@ -559,8 +559,8 @@ def f(node=None, order=None): assert f(order="inorder") == [0, 7, 1, 2, 8, 3, 4, 9, 5, 6] assert f(order="levelorder") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] assert f(order="breadthfirst") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] - assert f(order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] - assert f(order="timedesc") == [7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(order="timeasc") == list(range(10)) + assert f(order="timedesc") == list(reversed(range(10))) assert f(order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9] q = t.virtual_root @@ -569,8 +569,8 @@ def f(node=None, order=None): assert f(q, order="inorder") == [0, 7, 1, q, 2, 8, 3, 4, 9, 5, 6] assert f(q, order="levelorder") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] assert f(q, order="breadthfirst") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] - assert f(q, order="timeasc") == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, q] - assert f(q, order="timedesc") == [q, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6] + assert f(q, order="timeasc") == list(range(10)) + [q] + assert f(q, order="timedesc") == [q] + list(reversed(range(10))) assert f(q, order="minlex_postorder") == [0, 1, 7, 2, 3, 8, 4, 5, 6, 9, q] assert f(9, order="preorder") == [9, 4, 5, 6] @@ -580,7 +580,7 @@ def f(node=None, order=None): assert f(9, order="breadthfirst") == [9, 4, 5, 6] assert f(9, order="minlex_postorder") == [4, 5, 6, 9] assert f(9, order="timeasc") == [4, 5, 6, 9] - assert f(9, order="timedesc") == [9, 4, 5, 6] + assert f(9, order="timedesc") == [9, 6, 5, 4] def test_multiroot_non_lexical_example(self): nodes = io.StringIO( @@ -652,8 +652,8 @@ def f(node=None, order=None): assert inord == [4, 6, 3, 12, 7, 5, 13, 9, 11, 0, 10, 2, 1, 14, 8] assert level == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2] assert breadth == [4, 6, 12, 13, 14, 3, 7, 5, 11, 1, 8, 9, 10, 0, 2] - assert timeasc == [4, 6, 3, 7, 5, 9, 0, 2, 1, 8, 10, 11, 12, 13, 14] - assert timedesc == [14, 13, 12, 11, 10, 4, 6, 3, 7, 5, 9, 0, 2, 1, 8] + assert timeasc == list(range(15)) + assert timedesc == list(reversed(range(15))) # And the minlex tree: # 14 @@ -680,8 +680,8 @@ def f(node=None, order=None): "breadthfirst", [[9, 6, 7, 2, 3, 4, 5, 0, 1], [10, 4, 8, 5, 6, 0, 1, 2, 3]], ), - ("timeasc", [[2, 3, 4, 0, 1, 5, 6, 7, 9], [4, 0, 1, 2, 3, 5, 6, 8, 10]]), - ("timedesc", [[9, 7, 6, 5, 2, 3, 4, 0, 1], [10, 8, 6, 5, 4, 0, 1, 2, 3]]), + ("timeasc", [[0, 1, 2, 3, 4, 5, 6, 7, 9], [0, 1, 2, 3, 4, 5, 6, 8, 10]]), + ("timedesc", [[9, 7, 6, 5, 4, 3, 2, 1, 0], [10, 8, 6, 5, 4, 3, 2, 1, 0]]), ( "minlex_postorder", [[0, 1, 5, 4, 7, 2, 3, 6, 9], [0, 1, 5, 2, 3, 6, 8, 4, 10]], diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 1e2f905d44..b578470df8 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -2135,21 +2135,17 @@ def _levelorder_traversal(self, root): def _timeasc_traversal(self, root): """ - A stricter version of postorder where nodes are visiting in time-increasing - order first, and in preorder within equal time-slices. + Sorts by increasing time but falls back to increasing ID for equal times. """ # TODO implement with numpy? - yield from sorted( - self.nodes(root, order="postorder"), key=lambda u: self.time(u) - ) + yield from sorted(self.nodes(root), key=lambda u: (self.time(u), u)) def _timedesc_traversal(self, root): """ - A stricter version of preorder where nodes are visiting in time-descreasing - order first, and in preorder within equal time-slices. + The reverse of timeasc. """ yield from sorted( - self.nodes(root, order="preorder"), key=lambda u: -self.time(u) + self.nodes(root), key=lambda u: (self.time(u), u), reverse=True ) def _minlex_postorder_traversal(self, root): @@ -2238,7 +2234,6 @@ def nodes(self, root=None, order="preorder"): `Wikipedia `__. - - FIXME BEFORE MERGING timeas and timedesc - 'timeasc': visits the nodes in order of increasing time, falling back to increasing ID if times are equal. - 'timedesc': visits the nodes in order of decreasing time, falling back to From 07faaf225919beb23886ea2e246a9aa6ae93817e Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 12 Oct 2021 14:52:46 +0100 Subject: [PATCH 5/8] Numpy versions of timeasc and timedesc --- c/tests/test_trees.c | 2 +- python/tests/test_topology.py | 5 +++-- python/tskit/trees.py | 23 ++++++++++++++++++----- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/c/tests/test_trees.c b/c/tests/test_trees.c index 22de5ea100..d8fe481895 100644 --- a/c/tests/test_trees.c +++ b/c/tests/test_trees.c @@ -1,7 +1,7 @@ /* * MIT License * - * Copyright (c) 2019 Tskit Developers + * Copyright (c) 2019-2021 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 diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index cb479c8ee2..3f6737d2aa 100644 --- a/python/tests/test_topology.py +++ b/python/tests/test_topology.py @@ -6459,8 +6459,9 @@ def verify(self, ts): assert tree.is_internal(tree.virtual_root) assert not tree.is_leaf(tree.virtual_root) assert not tree.is_sample(tree.virtual_root) - # The mrca of the virtual_root and anything except itself is -1 - assert tree.mrca(0, tree.virtual_root) == tskit.NULL + # The mrca of the virtual_root and anything is itself + assert tree.mrca(0, tree.virtual_root) == tree.virtual_root + assert tree.mrca(tree.virtual_root, 0) == tree.virtual_root assert tree.mrca(tree.virtual_root, tree.virtual_root) == tree.virtual_root # The virtual_root is a descendant of nothing other than itself assert not tree.is_descendant(0, tree.virtual_root) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index b578470df8..a4e9e2fb01 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -2097,6 +2097,22 @@ def preorder(self, u=NULL): def postorder(self, u=NULL): return self._ll_tree.get_postorder(u) + def timeasc(self, u=NULL): + nodes = self.preorder(u) + is_virtual_root = u == self.virtual_root + time = self.tree_sequence.tables.nodes.time + if is_virtual_root: + # We could avoid creating this array if we wanted to, but + # it's not that often people will be using this with the + # virtual_root as an argument, so doesn't seem worth + # the complexity + time = np.append(time, [np.inf]) + order = np.lexsort([nodes, time[nodes]]) + return nodes[order] + + def timedesc(self, u=NULL): + return self.timeasc(u)[::-1] + def _preorder_traversal(self, root): return map(int, self.preorder(root)) @@ -2137,16 +2153,13 @@ def _timeasc_traversal(self, root): """ Sorts by increasing time but falls back to increasing ID for equal times. """ - # TODO implement with numpy? - yield from sorted(self.nodes(root), key=lambda u: (self.time(u), u)) + yield from self.timeasc(root) def _timedesc_traversal(self, root): """ The reverse of timeasc. """ - yield from sorted( - self.nodes(root), key=lambda u: (self.time(u), u), reverse=True - ) + yield from self.timedesc(root) def _minlex_postorder_traversal(self, root): """ From d670511e7c54a0b9fa2d74af9dad362fbbeb488e Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Tue, 12 Oct 2021 20:04:21 +0100 Subject: [PATCH 6/8] Count edges in the tree. Keep a track of the number of edges that are used to build the topology of the tree and document in Python and C interfaces. --- c/tests/test_trees.c | 17 ++++++++++++++++- c/tskit/trees.c | 4 ++++ c/tskit/trees.h | 6 ++++++ python/CHANGELOG.rst | 9 +++++++++ python/_tskitmodule.c | 17 +++++++++++++++++ python/tests/test_lowlevel.py | 1 + python/tests/test_topology.py | 8 ++++++++ python/tests/tsutil.py | 3 +++ python/tskit/trees.py | 13 +++++++++++++ 9 files changed, 77 insertions(+), 1 deletion(-) diff --git a/c/tests/test_trees.c b/c/tests/test_trees.c index d8fe481895..04b2d60dd6 100644 --- a/c/tests/test_trees.c +++ b/c/tests/test_trees.c @@ -47,6 +47,7 @@ check_trees_equal(tsk_tree_t *self, tsk_tree_t *other) CU_ASSERT_FATAL(self->sites_length == other->sites_length); CU_ASSERT_FATAL(self->sites == other->sites); CU_ASSERT_FATAL(self->samples == other->samples); + CU_ASSERT_FATAL(self->num_edges == other->num_edges); CU_ASSERT_FATAL(tsk_memcmp(self->parent, other->parent, N * sizeof(tsk_id_t)) == 0); CU_ASSERT_FATAL(tsk_tree_equals(self, other)); } @@ -184,6 +185,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents) tsk_size_t k, l, tree_sites_length; const tsk_site_t *sites = NULL; tsk_tree_t tree; + tsk_size_t num_edges; tsk_size_t num_nodes = tsk_treeseq_get_num_nodes(ts); tsk_size_t num_sites = tsk_treeseq_get_num_sites(ts); tsk_size_t num_mutations = tsk_treeseq_get_num_mutations(ts); @@ -202,11 +204,16 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents) tsk_tree_print_state(&tree, _devnull); /* tsk_tree_print_state(&tree, stdout); */ CU_ASSERT_EQUAL(tree.left, breakpoints[j]); + num_edges = 0; for (u = 0; u < (tsk_id_t) num_nodes; u++) { ret = tsk_tree_get_parent(&tree, u, &v); CU_ASSERT_EQUAL(ret, 0); CU_ASSERT_EQUAL(v, parents[j * (tsk_id_t) num_nodes + u]); + if (v != TSK_NULL) { + num_edges++; + } } + CU_ASSERT_EQUAL(num_edges, tree.num_edges); ret = tsk_tree_get_sites(&tree, &sites, &tree_sites_length); CU_ASSERT_EQUAL(ret, 0); for (k = 0; k < tree_sites_length; k++) { @@ -1222,6 +1229,7 @@ test_simplest_degenerate_multiple_root_records(void) CU_ASSERT_EQUAL(ret, 1); CU_ASSERT_EQUAL(tsk_tree_get_num_roots(&t), 2); CU_ASSERT_EQUAL(t.left_root, 2); + CU_ASSERT_EQUAL(t.num_edges, 2); CU_ASSERT_EQUAL(t.right_sib[2], 3); CU_ASSERT_EQUAL(t.right_sib[3], TSK_NULL); @@ -1311,6 +1319,7 @@ test_simplest_zero_root_tree(void) ret = tsk_tree_first(&t); CU_ASSERT_EQUAL(ret, 1); CU_ASSERT_EQUAL(tsk_tree_get_num_roots(&t), 0); + CU_ASSERT_EQUAL(t.num_edges, 4); CU_ASSERT_EQUAL(t.left_root, TSK_NULL); CU_ASSERT_EQUAL(t.right_sib[2], 3); CU_ASSERT_EQUAL(t.right_sib[3], TSK_NULL); @@ -1356,6 +1365,7 @@ test_simplest_multi_root_tree(void) CU_ASSERT_EQUAL(tsk_tree_get_num_roots(&t), 2); CU_ASSERT_EQUAL(t.left_root, 0); CU_ASSERT_EQUAL(t.right_sib[0], 3); + CU_ASSERT_EQUAL(t.num_edges, 2); tsk_tree_print_state(&t, _devnull); @@ -5039,7 +5049,10 @@ test_virtual_root_properties(void) CU_ASSERT_EQUAL_FATAL(depth, -1); CU_ASSERT_EQUAL_FATAL(tsk_tree_get_time(&t, t.virtual_root, &time), 0) - CU_ASSERT_TRUE(isinf(time)); + /* Workaround problems in IEEE floating point macros. We may want to + * add tsk_isinf (like tsk_isnan) at some point, but not worth it just + * for this test case */ + CU_ASSERT_TRUE(isinf((float) time)); CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&t, t.virtual_root, 0, &node), 0) CU_ASSERT_EQUAL(node, t.virtual_root); @@ -6269,6 +6282,7 @@ test_empty_tree_sequence(void) CU_ASSERT_EQUAL_FATAL(t.left_root, TSK_NULL); CU_ASSERT_EQUAL_FATAL(t.left, 0); CU_ASSERT_EQUAL_FATAL(t.right, 1); + CU_ASSERT_EQUAL_FATAL(t.num_edges, 0); CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 0, &v), 0); CU_ASSERT_EQUAL_FATAL(v, TSK_NULL); CU_ASSERT_EQUAL_FATAL(tsk_tree_get_parent(&t, 1, &v), TSK_ERR_NODE_OUT_OF_BOUNDS); @@ -6325,6 +6339,7 @@ test_zero_edges(void) CU_ASSERT_EQUAL_FATAL(ret, 1); CU_ASSERT_EQUAL(t.left, 0); CU_ASSERT_EQUAL(t.right, 2); + CU_ASSERT_EQUAL(t.num_edges, 0); CU_ASSERT_EQUAL(t.parent[0], TSK_NULL); CU_ASSERT_EQUAL(t.parent[1], TSK_NULL); CU_ASSERT_EQUAL(t.left_root, 0); diff --git a/c/tskit/trees.c b/c/tskit/trees.c index e8943bd2dc..67af3f0a43 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3388,6 +3388,7 @@ tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options) dest->sites = self->sites; dest->sites_length = self->sites_length; dest->root_threshold = self->root_threshold; + dest->num_edges = self->num_edges; tsk_memcpy(dest->parent, self->parent, N * sizeof(*self->parent)); tsk_memcpy(dest->left_child, self->left_child, N * sizeof(*self->left_child)); @@ -3942,6 +3943,7 @@ tsk_tree_remove_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) #define POTENTIAL_ROOT(U) (num_samples[U] >= root_threshold) tsk_tree_remove_branch(self, p, c, parent); + self->num_edges--; if (!(self->options & TSK_NO_SAMPLE_COUNTS)) { u = p; @@ -4002,6 +4004,7 @@ tsk_tree_insert_edge(tsk_tree_t *self, tsk_id_t p, tsk_id_t c) } tsk_tree_insert_branch(self, p, c, parent); + self->num_edges++; if (self->options & TSK_SAMPLE_LISTS) { tsk_tree_update_sample_lists(self, p, parent); @@ -4201,6 +4204,7 @@ tsk_tree_clear(tsk_tree_t *self) self->left = 0; self->right = 0; + self->num_edges = 0; self->index = -1; /* TODO we should profile this method to see if just doing a single loop over * the nodes would be more efficient than multiple memsets. diff --git a/c/tskit/trees.h b/c/tskit/trees.h index 69ed9c6397..8906aae11e 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -155,6 +155,12 @@ typedef struct { if node u has no siblings to its right. */ tsk_id_t *right_sib; + /** + @brief The total number of edges defining the topology of this tree. + This is equal to the number of tree sequence edges that intersect with + the tree's genomic interval. + */ + tsk_size_t num_edges; tsk_size_t num_nodes; tsk_flags_t options; diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index 857535be93..e38f73fca2 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -29,6 +29,7 @@ - Add ``__setitem__`` to all tables allowing single rows to be updated. For example ``tables.nodes[0] = tables.nodes[0].replace(flags=tskit.NODE_IS_SAMPLE)`` (:user:`jeromekelleher`, :user:`benjeffery`, :issue:`1545`, :pr:`1600`). + - Added a new parameter ``time`` to ``TreeSequence.samples()`` allowing to select samples at a specific time point or time interval. (:user:`mufernando`, :user:`petrelharp`, :issue:`1692`, :pr:`1700`) @@ -41,6 +42,14 @@ tree sequence. This is then used to generate an error if ``time_units`` is ``uncalibrated`` when using the branch lengths in statistics. (:user:`benjeffery`, :issue:`1644`, :pr:`1760`) +- Add the ``virtual_root`` property to Trees COMPLETEME + +- Add the ``num_edges`` property to Trees COMPLETEME + +- Improved performance for tree traversal methods (and library methods that are based + on tree traversals). TODO + + -------------------- [0.3.7] - 2021-07-08 -------------------- diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index f3bf52c452..7daf3f874a 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -9331,6 +9331,19 @@ Tree_get_virtual_root(Tree *self) return ret; } +static PyObject * +Tree_get_num_edges(Tree *self) +{ + PyObject *ret = NULL; + + if (Tree_check_state(self) != 0) { + goto out; + } + ret = Py_BuildValue("n", (Py_ssize_t) self->tree->num_edges); +out: + return ret; +} + static PyObject * Tree_get_index(Tree *self) { @@ -10257,6 +10270,10 @@ static PyMethodDef Tree_methods[] = { .ml_meth = (PyCFunction) Tree_get_virtual_root, .ml_flags = METH_NOARGS, .ml_doc = "Returns the virtual root of the tree." }, + { .ml_name = "get_num_edges", + .ml_meth = (PyCFunction) Tree_get_num_edges, + .ml_flags = METH_NOARGS, + .ml_doc = "Returns the number of branches in this tree." }, { .ml_name = "get_left", .ml_meth = (PyCFunction) Tree_get_left, .ml_flags = METH_NOARGS, diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 8270e2022e..544c622803 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -2714,6 +2714,7 @@ def test_cleared_tree(self): def check_tree(tree): assert tree.get_index() == -1 assert tree.get_left_child(tree.get_virtual_root()) == samples[0] + assert tree.get_num_edge() == 0 assert tree.get_mrca(0, 1) == _tskit.NULL for u in range(ts.get_num_nodes()): assert tree.get_parent(u) == _tskit.NULL diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py index 3f6737d2aa..c75aa61d75 100644 --- a/python/tests/test_topology.py +++ b/python/tests/test_topology.py @@ -6556,6 +6556,9 @@ def verify(self, ts): tree2 = tskit.Tree(ts) tree2.first() for interval, tree1 in tsutil.algorithm_R(ts, root_threshold=1): + root_reachable_nodes = len(tree2.preorder()) + size_bound = tree1.num_edges + ts.num_samples + assert size_bound >= root_reachable_nodes assert interval == tree2.interval assert tree1.roots() == tree2.roots # Definition here is the set unique path ends from samples @@ -6592,6 +6595,11 @@ def verify(self, ts): assert interval_py == tree_lib.interval assert interval_leg == tree_lib.interval + root_reachable_nodes = len(tree_lib.preorder()) + size_bound = tree_py.num_edges + ts.num_samples + assert size_bound >= root_reachable_nodes + assert tree_py.num_edges == tree_lib.num_edges + # Definition here is the set unique path ends from samples # that subtend at least k samples roots = set() diff --git a/python/tests/tsutil.py b/python/tests/tsutil.py index 9e34844150..8741030854 100644 --- a/python/tests/tsutil.py +++ b/python/tests/tsutil.py @@ -1184,6 +1184,7 @@ def __init__(self, n, root_threshold=1): self.left_sib = np.zeros(n + 1, dtype=np.int32) - 1 self.right_sib = np.zeros(n + 1, dtype=np.int32) - 1 self.num_samples = np.zeros(n + 1, dtype=np.int32) + self.num_edges = 0 def __str__(self): s = "id\tparent\tlchild\trchild\tlsib\trsib\tnsamp\n" @@ -1246,6 +1247,7 @@ def remove_root(self, root): def remove_edge(self, edge): self.remove_branch(edge.parent, edge.child) + self.num_edges -= 1 u = edge.parent while u != -1: @@ -1273,6 +1275,7 @@ def insert_edge(self, edge): self.insert_root(path_end) self.insert_branch(edge.parent, edge.child) + self.num_edges += 1 def algorithm_R(ts, root_threshold=1): diff --git a/python/tskit/trees.py b/python/tskit/trees.py index a4e9e2fb01..380385f3e9 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -1203,6 +1203,19 @@ def virtual_root(self): """ return self._ll_tree.get_virtual_root() + @property + def num_edges(self): + """ + The total number of edges in this tree. This is equal to the + number of tree sequence edges that intersect with this tree's + genomic interval. + + Note that this may be greater than the number of branches that + are reachable from the tree's roots, since we can have topology + that is not associated with any samples. + """ + return self._ll_tree.get_num_edges() + @property def left_root(self): """ From edbaaca4c3c0c09ed270a81e0818138db776101d Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Wed, 13 Oct 2021 12:45:19 +0100 Subject: [PATCH 7/8] Add tsk_tree_size_bound to enable safe allocations Closes #1725 --- c/tskit/trees.c | 37 +++++++++++++++++++++++++++++++--- c/tskit/trees.h | 25 +++++++++++++++++++++++ python/_tskitmodule.c | 2 +- python/tests/test_highlevel.py | 18 +++++++++++++++++ python/tests/test_lowlevel.py | 2 +- python/tskit/trees.py | 6 ++++-- 6 files changed, 83 insertions(+), 7 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 67af3f0a43..9036c897fe 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3626,7 +3626,6 @@ tsk_tree_get_time(const tsk_tree_t *self, tsk_id_t u, double *t) if (u == self->virtual_root) { *t = INFINITY; } else { - ret = tsk_treeseq_get_node(self->tree_sequence, u, &node); if (ret != 0) { goto out; @@ -4259,8 +4258,40 @@ tsk_tree_clear(tsk_tree_t *self) return ret; } +tsk_size_t +tsk_tree_get_size_bound(const tsk_tree_t *self) +{ + tsk_size_t bound = 0; + + if (self->tree_sequence != NULL) { + /* This is a safe upper bound which can be computed cheaply. + * We have at most n roots and each edge adds at most one new + * node to the tree. We also allow space for the virtual root, + * to simplify client code. + * + * In the common case of a binary tree with a single root, we have + * 2n - 1 nodes in total, and 2n - 2 edges. Therefore, we return + * 3n - 1, which is an over-estimate of 1/2 and we allocate + * 1.5 times as much memory as we need. + * + * Since tracking the exact number of nodes in the tree would require + * storing the number of nodes beneath every node and complicate + * the tree transition method, this seems like a good compromise + * and will result in less memory usage overall in nearly all cases. + */ + bound = 1 + self->tree_sequence->num_samples + self->num_edges; + } + return bound; +} + /* Traversal orders */ +static tsk_id_t * +tsk_tree_alloc_node_stack(const tsk_tree_t *self) +{ + return tsk_malloc(tsk_tree_get_size_bound(self) * sizeof(tsk_id_t)); +} + int tsk_tree_preorder( const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes_ret) @@ -4268,7 +4299,7 @@ tsk_tree_preorder( int ret = 0; const tsk_id_t *restrict right_child = self->right_child; const tsk_id_t *restrict left_sib = self->left_sib; - tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack)); + tsk_id_t *restrict stack = tsk_tree_alloc_node_stack(self); tsk_size_t num_nodes = 0; tsk_id_t u, v; int stack_top; @@ -4320,7 +4351,7 @@ tsk_tree_postorder( const tsk_id_t *restrict right_child = self->right_child; const tsk_id_t *restrict left_sib = self->left_sib; const tsk_id_t *restrict parent = self->parent; - tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack)); + tsk_id_t *restrict stack = tsk_tree_alloc_node_stack(self); tsk_size_t num_nodes = 0; tsk_id_t u, v, postorder_parent; int stack_top; diff --git a/c/tskit/trees.h b/c/tskit/trees.h index 8906aae11e..5851b1e2ce 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -411,6 +411,31 @@ int tsk_tree_prev(tsk_tree_t *self); int tsk_tree_clear(tsk_tree_t *self); void tsk_tree_print_state(const tsk_tree_t *self, FILE *out); + +/** +@brief Return an upper bound on the number of nodes reachable + from the roots of this tree. + +@rst +This function provides an upper bound on the number of nodes that +can be reached in tree traversals, and is intended to be used +for memory allocation purposes. If ``num_nodes`` is the number +of nodes visited in a tree traversal from the virtual root +(e.g., ``tsk_tree_preorder(tree, tree->virtual_root, nodes, +&num_nodes)``), the bound ``N`` returned here is guaranteed to +be greater than or equal to ``num_nodes``. + +.. warning:: The precise value returned is not defined and should + not be depended on, as it may change from version-to-version. + +@endrst + +@param self A pointer to a tsk_tree_t object. +@return An upper bound on the number nodes reachable from the roots + of this tree, or zero if this tree has not been initialised. +*/ +tsk_size_t tsk_tree_get_size_bound(const tsk_tree_t *self); + /** @} */ int tsk_tree_set_root_threshold(tsk_tree_t *self, tsk_size_t root_threshold); diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 7daf3f874a..b994838189 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -10070,7 +10070,7 @@ Tree_get_traversal_array(Tree *self, PyObject *args, tsk_traversal_func *func) if (!PyArg_ParseTuple(args, "i", &root)) { goto out; } - data = PyDataMem_NEW((self->tree->num_nodes + 1) * sizeof(*data)); + data = PyDataMem_NEW(tsk_tree_get_size_bound(self->tree) * sizeof(*data)); if (data == NULL) { ret = PyErr_NoMemory(); goto out; diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index 4264ef0241..7500f836e0 100644 --- a/python/tests/test_highlevel.py +++ b/python/tests/test_highlevel.py @@ -445,6 +445,24 @@ def get_samples(ts, time=None, population=None): class TestTreeTraversals: + def test_bad_traversal_order(self): + ts = msprime.sim_ancestry(2, random_seed=234) + tree = ts.first() + for bad_order in ["pre", "post", "preorderorder", ("x",), b"preorder"]: + with pytest.raises(ValueError, match="Traversal order"): + tree.nodes(order=bad_order) + + @pytest.mark.parametrize("order", list(traversal_map.keys())) + def test_returned_types(self, order): + ts = msprime.sim_ancestry(2, random_seed=234) + tree = ts.first() + iterator = tree.nodes(order=order) + assert isinstance(iterator, collections.abc.Iterable) + lst = list(iterator) + assert len(lst) > 0 + for u in lst: + assert isinstance(u, int) + @pytest.mark.parametrize("ts", get_example_tree_sequences()) @pytest.mark.parametrize("order", list(traversal_map.keys())) def test_traversals_virtual_root(self, ts, order): diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 544c622803..4dcbfec28b 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -2714,7 +2714,7 @@ def test_cleared_tree(self): def check_tree(tree): assert tree.get_index() == -1 assert tree.get_left_child(tree.get_virtual_root()) == samples[0] - assert tree.get_num_edge() == 0 + assert tree.get_num_edges() == 0 assert tree.get_mrca(0, 1) == _tskit.NULL for u in range(ts.get_num_nodes()): assert tree.get_parent(u) == _tskit.NULL diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 380385f3e9..eb8f3388f5 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -2127,9 +2127,11 @@ def timedesc(self, u=NULL): return self.timeasc(u)[::-1] def _preorder_traversal(self, root): + # Return Python integers for compatibility return map(int, self.preorder(root)) def _postorder_traversal(self, root): + # Return Python integers for compatibility return map(int, self.postorder(root)) def _inorder_traversal(self, root): @@ -2166,13 +2168,13 @@ def _timeasc_traversal(self, root): """ Sorts by increasing time but falls back to increasing ID for equal times. """ - yield from self.timeasc(root) + return map(int, self.timeasc(root)) def _timedesc_traversal(self, root): """ The reverse of timeasc. """ - yield from self.timedesc(root) + return map(int, self.timedesc(root)) def _minlex_postorder_traversal(self, root): """ From a882a5d68a1d855cb89892abb22ffeb845eadf46 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Thu, 14 Oct 2021 12:22:32 +0100 Subject: [PATCH 8/8] Document orderings Closes #1691 Closes #1706 --- c/CHANGELOG.rst | 3 + docs/_static/different_time_samples.svg | 60 ++++++------ docs/_static/tree_structure1.svg | 85 +++++++++-------- docs/_static/tree_structure2.svg | 80 ++++++++-------- docs/data-model.rst | 48 +++++++--- docs/examples.py | 8 +- docs/python-api.rst | 6 ++ python/CHANGELOG.rst | 13 ++- python/tskit/trees.py | 119 ++++++++++++++---------- 9 files changed, 237 insertions(+), 185 deletions(-) diff --git a/c/CHANGELOG.rst b/c/CHANGELOG.rst index 942c8184f0..e34d09d0d5 100644 --- a/c/CHANGELOG.rst +++ b/c/CHANGELOG.rst @@ -14,6 +14,7 @@ - The previously deprecated option ``TSK_SAMPLE_COUNTS`` has been removed. (:user:`benjeffery`, :issue:`1744`, :pr:`1761`). +- FIXME breaking changes for tree API and virtual root **Features** @@ -37,6 +38,8 @@ tree sequence. This is then used to generate an error if ``time_units`` is ``uncalibrated`` when using the branch lengths in statistics. (:user:`benjeffery`, :issue:`1644`, :pr:`1760`) +- FIXME add features for virtual root, num_edges, stack allocation size etc + **Fixes** ---------------------- diff --git a/docs/_static/different_time_samples.svg b/docs/_static/different_time_samples.svg index b93b0f4dca..cc312dc703 100644 --- a/docs/_static/different_time_samples.svg +++ b/docs/_static/different_time_samples.svg @@ -1,39 +1,33 @@ - - - - - - - - - - - - - - - - - - - - - - - 3 + + + + + + + + + + + 2 - - 4 - 2 - 0 - 1 + + + + + 0 + + + + + 1 + + + + 3 - - - - - + + 4 diff --git a/docs/_static/tree_structure1.svg b/docs/_static/tree_structure1.svg index 2c6689f9d5..7a54088e70 100644 --- a/docs/_static/tree_structure1.svg +++ b/docs/_static/tree_structure1.svg @@ -1,49 +1,48 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 6 + + + + + + + + + + + + 0 + + + + + 1 + + + + + 2 + + + + 5 - - 7 - 0 - 1 - 2 - 3 - 4 + + + + + 3 + + + + + 4 + + + + 6 - - 5 - - - - - + + 7 diff --git a/docs/_static/tree_structure2.svg b/docs/_static/tree_structure2.svg index 52882ed7fe..406623b7ef 100644 --- a/docs/_static/tree_structure2.svg +++ b/docs/_static/tree_structure2.svg @@ -1,47 +1,47 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 7 - 0 - 1 - 2 - 6 - 3 - 4 + + + + + + + + + + + 3 - - 5 + + + + 4 + + 6 - - - + + + + + + 0 + + + + + 1 + + + + + 2 + + + + 5 + + + 7 diff --git a/docs/data-model.rst b/docs/data-model.rst index 55ae632e57..0f449a0b6c 100644 --- a/docs/data-model.rst +++ b/docs/data-model.rst @@ -987,10 +987,8 @@ details of how to use the quintuply linked structure in the C API. .. _sec_data_model_tree_roots: -Accessing roots -=============== - -.. todo:: Update this with a discussion of the virtual root +Roots +===== The roots of a tree are defined as the unique endpoints of upward paths starting from sample nodes (if no path leads upward from a sample node, @@ -1003,6 +1001,10 @@ example, we get a tree with two roots: :width: 200px :alt: An example tree with multiple roots +We keep track of roots in tskit by using a special additional node +called the **virtual root**, whose children are the roots. In the +quintuply linked tree encoding this is an extra element at the end +of each of the tree arrays, as shown here: =========== =========== =========== =========== =========== =========== node parent left_child right_child left_sib right_sib @@ -1013,17 +1015,37 @@ node parent left_child right_child left_sib right_sib 3 6 -1 -1 -1 4 4 6 -1 -1 3 -1 5 7 0 2 -1 -1 -6 -1 3 4 7 -1 -7 -1 5 5 -1 6 +6 -1 3 4 -1 7 +7 -1 5 5 6 -1 +**8** **-1** **6** **7** **-1** **-1** =========== =========== =========== =========== =========== =========== -To gain efficient access to the roots in the quintuply linked encoding we keep -one extra piece of information: the ``left_root``. In this example -the leftmost root is ``7``. Roots are considered siblings, and so -once we have one root we can find all the other roots efficiently using -the ``left_sib`` and ``right_sib`` arrays. For example, we can see here -that the right sibling of ``7`` is ``6``, and the left sibling of ``6`` -is ``7``. +In this example, node 8 is the virtual root; its left child is 6 +and its right child is 7. +Importantly, though, this is an asymmetric +relationship, since the parent of the "real" roots 6 and 7 is null +(-1) and *not* the virtual root. To emphasise that this is not a "real" +node, we've shown the values for the virtual root here in bold. + +The main function of the virtual root is to efficiently keep track of +tree roots in the internal library algorithms, and is usually not +something we need to think about unless working directly with +the quintuply linked tree structure. However, the virtual root can be +useful in some algorithms and so it can optionally be returned in traversal +orders (see :meth:`.Tree.nodes`). The virtual root has the following +properties: + +- Its ID is always equal to the number of nodes in the tree sequence (i.e., + the length of the node table). However, there is **no corresponding row** + in the node table, and any attempts to access information about the + virtual root via either the tree sequence or tables APIs will fail with + an out-of-bounds error. +- The parent and siblings of the virtual root are null. +- The time of the virtual root is defined as positive infinity (if + accessed via :meth:`.Tree.time`). This is useful in defining the + time-based node traversal orderings. +- The virtual root is the parent of no other node---roots do **not** + have parent pointers to the virtual root. .. _sec_data_model_missing_data: diff --git a/docs/examples.py b/docs/examples.py index dee4a55c58..95aa651907 100644 --- a/docs/examples.py +++ b/docs/examples.py @@ -257,7 +257,7 @@ def stats(): def tree_structure(): - def write_table(tree): + def write_table(tree, show_virtual_root=False): fmt = "{:<12}" heading = [ "node", @@ -273,7 +273,7 @@ def write_table(tree): print(line) print(col_def) - for u in range(ts.num_nodes): + for u in range(ts.num_nodes + int(show_virtual_root)): line = "".join( fmt.format(v) for v in [ @@ -325,7 +325,7 @@ def write_table(tree): ) tree = ts.first() - write_table(tree) + write_table(tree, show_virtual_root=True) print(tree.draw_text()) tree.draw_svg("_static/tree_structure2.svg", time_scale="rank") @@ -404,6 +404,6 @@ def finding_nearest_neighbors(): # allele_frequency_spectra() # missing_data() # stats() -# tree_structure() +tree_structure() tree_traversal() finding_nearest_neighbors() diff --git a/docs/python-api.rst b/docs/python-api.rst index 35198bf329..a8e7264d6f 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -6,6 +6,12 @@ compare arrays representing different trees along the sequence, you must take **copies** of the arrays. +.. |virtual_root_array_note| replace:: The length of these arrays is + equal to the number of nodes in the tree sequence plus 1, with the + final element corresponding to the tree's :meth:`~.Tree.virtual_root`. + Please see the :ref:`tree roots ` section + for more details. + .. currentmodule:: tskit .. _sec_python_api: diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index e38f73fca2..8f73575162 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -23,6 +23,10 @@ all have one extra element. (:user:`jeromekelleher`, :issue:`1691`, :pr:`1704`). +- Tree traversal orders returned by the ``nodes`` method have changed when there + are multiple roots. Previously orders were defined locally for each root, but + are now globally across all roots. (:user:`jeromekelleher`, :pr:`1704`). + **Features** @@ -42,12 +46,13 @@ tree sequence. This is then used to generate an error if ``time_units`` is ``uncalibrated`` when using the branch lengths in statistics. (:user:`benjeffery`, :issue:`1644`, :pr:`1760`) -- Add the ``virtual_root`` property to Trees COMPLETEME +- Add the ``virtual_root`` property to the Tree class (:user:`jeromekelleher`, :pr:`1704`). -- Add the ``num_edges`` property to Trees COMPLETEME +- Add the ``num_edges`` property to the Tree class (:user:`jeromekelleher`, :pr:`1704`). -- Improved performance for tree traversal methods (and library methods that are based - on tree traversals). TODO +- Improved performance for tree traversal methods in the ``nodes`` iterator. + Roughly a 10X performance increase for "preorder", "postorder", "timeasc" + and "timedesc" (:user:`jeromekelleher`, :pr:`1704`). -------------------- diff --git a/python/tskit/trees.py b/python/tskit/trees.py index eb8f3388f5..485b164106 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -1050,11 +1050,13 @@ def parent_array(self): """ A numpy array (dtype=np.int32) encoding the parent of each node in this tree, such that ``tree.parent_array[u] == tree.parent(u)`` - for all ``0 <= u < ts.num_nodes``. See the :meth:`~.parent` + for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.parent` method for details on the semantics of tree parents and the :ref:`sec_data_model_tree_structure` section for information on the quintuply linked tree encoding. + .. note:: |virtual_root_array_note| + .. warning:: |tree_array_warning| """ return self._parent_array @@ -1085,11 +1087,13 @@ def left_child_array(self): """ A numpy array (dtype=np.int32) encoding the left child of each node in this tree, such that ``tree.left_child_array[u] == tree.left_child(u)`` - for all ``0 <= u < ts.num_nodes``. See the :meth:`~.left_child` + for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.left_child` method for details on the semantics of tree left_child and the :ref:`sec_data_model_tree_structure` section for information on the quintuply linked tree encoding. + .. note:: |virtual_root_array_note| + .. warning:: |tree_array_warning| """ return self._left_child_array @@ -1118,11 +1122,13 @@ def right_child_array(self): """ A numpy array (dtype=np.int32) encoding the right child of each node in this tree, such that ``tree.right_child_array[u] == tree.right_child(u)`` - for all ``0 <= u < ts.num_nodes``. See the :meth:`~.right_child` + for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.right_child` method for details on the semantics of tree right_child and the :ref:`sec_data_model_tree_structure` section for information on the quintuply linked tree encoding. + .. note:: |virtual_root_array_note| + .. warning:: |tree_array_warning| """ return self._right_child_array @@ -1147,11 +1153,13 @@ def left_sib_array(self): """ A numpy array (dtype=np.int32) encoding the left sib of each node in this tree, such that ``tree.left_sib_array[u] == tree.left_sib(u)`` - for all ``0 <= u < ts.num_nodes``. See the :meth:`~.left_sib` + for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.left_sib` method for details on the semantics of tree left_sib and the :ref:`sec_data_model_tree_structure` section for information on the quintuply linked tree encoding. + .. note:: |virtual_root_array_note| + .. warning:: |tree_array_warning| """ return self._left_sib_array @@ -1176,11 +1184,13 @@ def right_sib_array(self): """ A numpy array (dtype=np.int32) encoding the right sib of each node in this tree, such that ``tree.right_sib_array[u] == tree.right_sib(u)`` - for all ``0 <= u < ts.num_nodes``. See the :meth:`~.right_sib` + for all ``0 <= u <= ts.num_nodes``. See the :meth:`~.right_sib` method for details on the semantics of tree right_sib and the :ref:`sec_data_model_tree_structure` section for information on the quintuply linked tree encoding. + .. note:: |virtual_root_array_note| + .. warning:: |tree_array_warning| """ return self._right_sib_array @@ -1199,7 +1209,11 @@ def next_sample(self, u): @property def virtual_root(self): """ - The ID of the virtual root in this tree. + The ID of the virtual root in this tree. This is equal to + :attr:`TreeSequence.num_nodes`. + + Please see the :ref:`tree roots ` + section for more details. """ return self._ll_tree.get_virtual_root() @@ -2104,6 +2118,8 @@ def num_tracked_samples(self, u=None): roots = self.roots return sum(self._ll_tree.get_num_tracked_samples(root) for root in roots) + # TODO document these traversal arrays + # https://github.com/tskit-dev/tskit/issues/1788 def preorder(self, u=NULL): return self._ll_tree.get_preorder(u) @@ -2229,52 +2245,59 @@ def push(nodes): def nodes(self, root=None, order="preorder"): """ - Returns an iterator over the node IDs reachable from the root(s) in this tree. - See :meth:`Tree.roots` for which nodes are considered roots. If the root - parameter is provided, only the subtree rooted at this single root node - will be iterated over. If this parameter is None, iterate over the node IDs in - the subtrees rooted at each root node in turn. If the order parameter is - provided, iterate over the nodes in each subtree in the specified tree traversal - order. + Returns an iterator over the node IDs reachable from the root(s) in this + tree in the specified traversal order. .. note:: Unlike the :meth:`TreeSequence.nodes` method, this iterator produces integer node IDs, not :class:`Node` objects. - The currently implemented traversal orders are: - - - 'preorder': starting at root, yield the current node, then recurse - and do a preorder on each child of the current node. See also `Wikipedia - `__. - - 'inorder': starting at root, assuming binary trees, recurse and do - an inorder on the first child, then yield the current node, then - recurse and do an inorder on the second child. In the case of ``n`` - child nodes (not necessarily 2), the first ``n // 2`` children are - visited in the first stage, and the remaining ``n - n // 2`` children - are visited in the second stage. See also `Wikipedia - `__. - - 'postorder': starting at root, recurse and do a postorder on each - child of the current node, then yield the current node. See also - `Wikipedia - `__. - - 'levelorder' ('breadthfirst'): visit the nodes under root (including - the root) in increasing order of their depth from root. See also - `Wikipedia - ` section for more + information on the virtual root.) + + The ``order`` parameter defines the order in which tree nodes + are visited in the iteration. The available orders are: + + - 'preorder': starting at root, yield the current node, then recurse + and do a preorder on each child of the current node. See also `Wikipedia + `__. + - 'inorder': starting at root, assuming binary trees, recurse and do + an inorder on the first child, then yield the current node, then + recurse and do an inorder on the second child. In the case of ``n`` + child nodes (not necessarily 2), the first ``n // 2`` children are + visited in the first stage, and the remaining ``n - n // 2`` children + are visited in the second stage. See also `Wikipedia + `__. + - 'postorder': starting at root, recurse and do a postorder on each + child of the current node, then yield the current node. See also + `Wikipedia + `__. + - 'levelorder' ('breadthfirst'): visit the nodes under root (including + the root) in increasing order of their depth from root. See also + `Wikipedia + `__. - - 'timeasc': visits the nodes in order of increasing time, falling back to - increasing ID if times are equal. - - 'timedesc': visits the nodes in order of decreasing time, falling back to - decreasing ID if times are equal. - - 'minlex_postorder': a usual postorder has ambiguity in the order in - which children of a node are visited. We constrain this by outputting - a postorder such that the leaves visited, when their IDs are - listed out, have minimum `lexicographic order - `__ out of all valid - traversals. This traversal is useful for drawing multiple trees of - a ``TreeSequence``, as it leads to more consistency between adjacent - trees. Note that internal non-leaf nodes are not counted in - assessing the lexicographic order. + - 'timeasc': visits the nodes in order of increasing time, falling back to + increasing ID if times are equal. + - 'timedesc': visits the nodes in order of decreasing time, falling back to + decreasing ID if times are equal. + - 'minlex_postorder': a usual postorder has ambiguity in the order in + which children of a node are visited. We constrain this by outputting + a postorder such that the leaves visited, when their IDs are + listed out, have minimum `lexicographic order + `__ out of all valid + traversals. This traversal is useful for drawing multiple trees of + a ``TreeSequence``, as it leads to more consistency between adjacent + trees. Note that internal non-leaf nodes are not counted in + assessing the lexicographic order. :param int root: The root of the subtree we are traversing. :param str order: The traversal ordering. Currently 'preorder', @@ -5565,7 +5588,7 @@ def draw_svg( :param str order: The left-to-right ordering of child nodes in each drawn tree. This can be either: ``"minlex"``, which minimises the differences between adjacent trees (see also the ``"minlex_postorder"`` traversal - order for the :meth:`.nodes` method); or ``"tree"`` which draws trees + order for the :meth:`.Tree.nodes` method); or ``"tree"`` which draws trees in the left-to-right order defined by the :ref:`quintuply linked tree structure `. If not specified or None, this defaults to ``"minlex"``. @@ -5674,7 +5697,7 @@ def draw_text( :param str order: The left-to-right ordering of child nodes in the drawn tree. This can be either: ``"minlex"``, which minimises the differences between adjacent trees (see also the ``"minlex_postorder"`` traversal - order for the :meth:`.nodes` method); or ``"tree"`` which draws trees + order for the :meth:`.Tree.nodes` method); or ``"tree"`` which draws trees in the left-to-right order defined by the :ref:`quintuply linked tree structure `. If not specified or None, this defaults to ``"minlex"``.