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/c/tests/test_trees.c b/c/tests/test_trees.c index 5f31be9615..04b2d60dd6 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 @@ -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); @@ -1338,7 +1347,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,13 +1359,13 @@ 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); 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); @@ -3862,8 +3873,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; @@ -3877,14 +3888,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, &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); @@ -5016,6 +5027,228 @@ 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) + /* 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); + + 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. *=======================================================*/ @@ -5185,7 +5418,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 +5908,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 +6282,10 @@ 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(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); tsk_tree_free(&t); ret = tsk_tree_init(&t, &ts, 0); @@ -6054,7 +6295,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); @@ -6098,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); @@ -6397,6 +6639,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 ce129fcd80..9036c897fe 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); @@ -3452,12 +3388,13 @@ 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(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 +3410,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 +3436,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; @@ -3537,6 +3475,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) { @@ -3678,11 +3623,15 @@ 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; } @@ -3697,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; @@ -3721,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; } @@ -3760,6 +3711,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 +3728,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 +3799,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 +3824,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 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 +3863,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 +3887,84 @@ 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); + self->num_edges--; + + 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 +3972,41 @@ 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; + tsk_id_t u; + tsk_id_t path_end = TSK_NULL; + bool path_end_was_root = false; -#define IS_ROOT(U) (num_samples[U] >= root_threshold) - - 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); + self->num_edges++; + 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 +4046,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 +4190,221 @@ 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->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. + */ + 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; +} + +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) +{ + 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_tree_alloc_node_stack(self); + 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_tree_alloc_node_stack(self); + 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 bd5d9c65ba..5851b1e2ce 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 @@ -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. @@ -150,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; @@ -400,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); @@ -420,6 +456,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); @@ -427,7 +464,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/_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 ee085a6842..0f449a0b6c 100644 --- a/docs/data-model.rst +++ b/docs/data-model.rst @@ -985,8 +985,10 @@ 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. -Accessing roots -=============== +.. _sec_data_model_tree_roots: + +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, @@ -999,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 @@ -1009,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 0d6d78fcf7..8f73575162 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -17,11 +17,23 @@ :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`). + +- 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** - 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`) @@ -34,6 +46,15 @@ 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 the Tree class (:user:`jeromekelleher`, :pr:`1704`). + +- Add the ``num_edges`` property to the Tree class (:user:`jeromekelleher`, :pr:`1704`). + +- 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`). + + -------------------- [0.3.7] - 2021-07-08 -------------------- @@ -68,7 +89,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..b994838189 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; @@ -9103,7 +9089,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 +9319,40 @@ 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_num_edges(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->num_edges); +out: + return ret; +} + +static PyObject * +Tree_get_index(Tree *self) +{ + PyObject *ret = NULL; + + if (Tree_check_state(self) != 0) { + goto out; + } + ret = Py_BuildValue("n", (Py_ssize_t) self->tree->index); out: return ret; } @@ -9494,13 +9493,13 @@ Tree_get_time(Tree *self, PyObject *args) { PyObject *ret = NULL; double time; - int node, err; + 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_tree_get_time(self->tree, node_id, &time); + if (err != 0) { handle_library_error(err); goto out; } @@ -9611,18 +9610,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,18 +10050,82 @@ 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(tsk_tree_get_size_bound(self->tree) * 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 * 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 +10266,14 @@ 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_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, @@ -10329,6 +10396,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 e49da94aaf..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 @@ -297,16 +223,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/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 dde8d2be26..f56b9bf111 100644 --- a/python/tests/test_drawing.py +++ b/python/tests/test_drawing.py @@ -1173,15 +1173,15 @@ def test_draw_forky_tree(self): 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( """\ diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index 34f63bc334..7500f836e0 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), key=lambda u: (tree.time(u), u)) + + +def traversal_timedesc(tree, root=None): + yield from sorted(tree.nodes(root), key=lambda u: (tree.time(u), u), reverse=True) + + +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,407 @@ def get_samples(ts, time=None, population=None): return np.array(samples) +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): + 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, 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 + 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, 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] + 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, 4, 3, 2] + 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") == 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 + 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") == 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] + 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, 6, 5, 4] + 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") == 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 + 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") == 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] + 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, 6, 5, 4] + + 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 == list(range(15)) + assert timedesc == list(reversed(range(15))) + + # 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", [[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_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 +2880,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 +3141,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() @@ -2751,7 +3174,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 +3335,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..4dcbfec28b 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,8 @@ 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_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 @@ -2739,7 +2741,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 +2832,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 +3034,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 @@ -3073,6 +3075,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 8a802c875d..c75aa61d75 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) @@ -2610,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 @@ -6565,6 +6421,82 @@ 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 + # 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. + 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 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) + 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", + "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.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 +6537,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 +6553,12 @@ 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): + 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 @@ -6629,35 +6569,101 @@ 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 + + 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() + 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 +8080,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..8741030854 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,146 @@ 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) + self.num_edges = 0 + + 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) + self.num_edges -= 1 + + 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) + self.num_edges += 1 + + +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 +1455,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 +1481,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 +1647,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..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 @@ -1196,7 +1206,30 @@ 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. 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() + + @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): """ @@ -1229,7 +1262,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 @@ -1255,8 +1292,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. @@ -1270,6 +1309,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 @@ -2040,10 +2082,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,55 +2109,66 @@ 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 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 + # 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) + + 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 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): + # TODO add a nonrecursive version of the inorder traversal. - 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 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) - def _inorder_traversal(self, u): - # 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]) + 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 @@ -2129,25 +2180,19 @@ 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. """ - yield from sorted( - self.nodes(u, order="levelorder"), key=lambda u: (self.time(u), u) - ) + return map(int, self.timeasc(root)) - def _timedesc_traversal(self, u): + def _timedesc_traversal(self, root): """ - Sorts by decreasing time but falls back to decreasing ID for equal times. + The reverse of timeasc. """ - yield from sorted( - self.nodes(u, order="levelorder"), - key=lambda u: (self.time(u), u), - reverse=True, - ) + return map(int, self.timedesc(root)) - def _minlex_postorder_traversal(self, u): + def _minlex_postorder_traversal(self, root): """ Postorder traversal that visits leaves in minimum lexicographic order. @@ -2157,104 +2202,102 @@ 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"): """ - 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', @@ -2277,23 +2320,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): """ @@ -5559,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"``. @@ -5668,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"``.