From 176218e58c948ed5fcfaf33ff9584c561950de29 Mon Sep 17 00:00:00 2001 From: Daniel Goldstein Date: Thu, 23 Apr 2020 11:28:05 -0400 Subject: [PATCH] add depth method to tree class --- c/tests/test_trees.c | 46 ++++++++++++++++++++++++++++++++++ c/tskit/trees.c | 22 ++++++++++++++++ c/tskit/trees.h | 1 + python/_tskitmodule.c | 22 ++++++++++++++++ python/tests/test_highlevel.py | 10 ++++++++ python/tests/test_lowlevel.py | 1 + python/tskit/trees.py | 12 +++++++++ 7 files changed, 114 insertions(+) diff --git a/c/tests/test_trees.c b/c/tests/test_trees.c index 0ece01fc32..d49e0ca93f 100644 --- a/c/tests/test_trees.c +++ b/c/tests/test_trees.c @@ -3507,6 +3507,51 @@ test_single_tree_iter_times(void) tsk_treeseq_free(&ts); } +static void +test_single_tree_iter_depths(void) +{ + int ret = 0; + const char *nodes = "1 0 0\n" + "1 0 0\n" + "1 0 0\n" + "1 0 0\n" + "0 1 0\n" + "0 2 0\n" + "0 3 0\n"; + 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 }; + unsigned int depth; + tsk_treeseq_t ts; + tsk_tree_t tree; + tsk_id_t u; + uint32_t num_nodes = 7; + + tsk_treeseq_from_text(&ts, 6, nodes, edges, NULL, NULL, NULL, NULL, NULL); + ret = tsk_tree_init(&tree, &ts, 0); + CU_ASSERT_EQUAL(ret, 0); + ret = tsk_tree_first(&tree); + CU_ASSERT_EQUAL(ret, 1); + 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); + CU_ASSERT_EQUAL(ret, 0); + CU_ASSERT_EQUAL(depth, depths[u]); + } + + ret = tsk_tree_depth(&tree, (tsk_id_t) num_nodes, &depth); + CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + ret = tsk_tree_depth(&tree, TSK_NULL, &depth); + CU_ASSERT_EQUAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); + + ret = tsk_tree_next(&tree); + CU_ASSERT_EQUAL(ret, 0); + + tsk_tree_free(&tree); + tsk_treeseq_free(&ts); +} static void test_single_tree_simplify(void) { @@ -5287,6 +5332,7 @@ main(int argc, char **argv) test_single_tree_general_samples_iter }, { "test_single_nonbinary_tree_iter", test_single_nonbinary_tree_iter }, { "test_single_tree_iter_times", test_single_tree_iter_times }, + { "test_single_tree_iter_depths", test_single_tree_iter_depths }, { "test_single_tree_simplify", test_single_tree_simplify }, { "test_single_tree_simplify_no_sample_nodes", test_single_tree_simplify_no_sample_nodes }, diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 881bd249d0..80b4f03d1c 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3565,6 +3565,28 @@ tsk_tree_get_sites(tsk_tree_t *self, tsk_site_t **sites, tsk_size_t *sites_lengt return 0; } +int TSK_WARN_UNUSED +tsk_tree_depth(tsk_tree_t *self, tsk_id_t u, tsk_size_t *depth_ret) +{ + tsk_id_t v; + tsk_size_t depth; + int ret = 0; + + ret = tsk_tree_check_node(self, u); + if (ret != 0) { + goto out; + } + + depth = 0; + for (v = self->parent[u]; v != TSK_NULL; v = self->parent[v]) { + depth++; + } + + *depth_ret = depth; +out: + return ret; +} + static void tsk_tree_check_state(tsk_tree_t *self) { diff --git a/c/tskit/trees.h b/c/tskit/trees.h index d5eb133007..d72e5e1b0b 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -381,6 +381,7 @@ int tsk_tree_get_num_samples(tsk_tree_t *self, tsk_id_t u, size_t *num_samples); int tsk_tree_get_num_tracked_samples( tsk_tree_t *self, tsk_id_t u, size_t *num_tracked_samples); int tsk_tree_get_sites(tsk_tree_t *self, tsk_site_t **sites, tsk_size_t *sites_length); +int tsk_tree_depth(tsk_tree_t *self, tsk_id_t u, tsk_size_t *depth); typedef struct { tsk_id_t node; diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 5468b0c798..74b84a16ef 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -8765,6 +8765,26 @@ Tree_get_children(Tree *self, PyObject *args) return ret; } +static PyObject * +Tree_depth(Tree *self, PyObject *args) +{ + PyObject *ret = NULL; + tsk_size_t depth; + int node, err; + + if (Tree_get_node_argument(self, args, &node) != 0) { + goto out; + } + err = tsk_tree_depth(self->tree, node, &depth); + if (ret != 0) { + handle_library_error(err); + goto out; + } + ret = Py_BuildValue("I", (unsigned int) depth); +out: + return ret; +} + static bool Tree_check_sample_list(Tree *self) { @@ -9196,6 +9216,8 @@ static PyMethodDef Tree_methods[] = { "Returns True if the specified node is a sample." }, {"is_descendant", (PyCFunction) Tree_is_descendant, METH_VARARGS, "Returns True if u is a descendant of v." }, + {"depth", (PyCFunction) Tree_depth, METH_VARARGS, + "Returns the depth of node u" }, {"get_parent", (PyCFunction) Tree_get_parent, METH_VARARGS, "Returns the parent of node u" }, {"get_time", (PyCFunction) Tree_get_time, METH_VARARGS, diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py index 146f76166e..7022b343f5 100644 --- a/python/tests/test_highlevel.py +++ b/python/tests/test_highlevel.py @@ -434,10 +434,20 @@ def verify_tree_structure(self, st): self.assertEqual(st.num_samples(), len(samples)) self.assertEqual(sorted(st.samples()), sorted(samples)) + def verify_tree_depths(self, st): + for root in st.roots: + stack = [(root, 0)] + while len(stack) > 0: + u, depth = stack.pop() + self.assertEqual(st.depth(u), depth) + for c in st.children(u): + stack.append((c, depth + 1)) + def verify_tree(self, st): self.verify_tree_mrcas(st) self.verify_tree_branch_lengths(st) self.verify_tree_structure(st) + self.verify_tree_depths(st) def verify_trees(self, ts): pts = tests.PythonTreeSequence(ts) diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index f008c2194c..5a188442be 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -1730,6 +1730,7 @@ def test_bounds_checking(self): self.assertRaises(ValueError, st.get_right_sample, v) self.assertRaises(ValueError, st.is_descendant, v, 0) self.assertRaises(ValueError, st.is_descendant, 0, v) + self.assertRaises(ValueError, st.depth, v) n = ts.get_num_samples() for v in [-100, -1, n + 1, n + 100, n * 100]: self.assertRaises(ValueError, st.get_next_sample, v) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 56fd05052f..66abde0c54 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -952,6 +952,18 @@ def time(self, u): """ return self._ll_tree.get_time(u) + def depth(self, u): + """ + Returns the number of nodes on the path from ``u`` to a + root, not including ``u``. Thus, the depth of a root is + zero. + + :param int u: The node of interest. + :return: The depth of u. + :rtype: int + """ + return self._ll_tree.depth(u) + def get_population(self, u): # Deprecated alias for self.population return self.population(u)