Add depth method to tree class

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

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

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;

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,

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)

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)

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)