Implement total_branch_length in C

Closes #1794

Benchmarks:

Before:
big_tree         0.26075993799986463
many_trees       0.002319710470019345

After:
big_tree         0.004731306999929075
many_trees       2.1276080024108523e-05

So 2-3 orders of magnitude (probably more compared to previous release,
since this using faster node iteration)

Author: jeromekelleher

c/tests/test_trees.c

@@ -4329,6 +4329,44 @@ test_single_tree_is_descendant(void)
     tsk_treeseq_free(&ts);
 }
 
+static void
+test_single_tree_total_branch_length(void)
+{
+    int ret;
+    tsk_treeseq_t ts;
+    tsk_tree_t tree;
+    double length;
+
+    tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, NULL,
+        NULL, NULL, NULL, 0);
+    ret = tsk_tree_init(&tree, &ts, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = tsk_tree_first(&tree);
+    CU_ASSERT_EQUAL_FATAL(ret, 1);
+
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, TSK_NULL, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 9);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, 7, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 9);
+    CU_ASSERT_EQUAL_FATAL(
+        tsk_tree_get_total_branch_length(&tree, tree.virtual_root, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 9);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, 4, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 2);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, 0, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 0);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, 5, &length), 0);
+    CU_ASSERT_EQUAL_FATAL(length, 4);
+
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_total_branch_length(&tree, -2, &length),
+        TSK_ERR_NODE_OUT_OF_BOUNDS);
+    CU_ASSERT_EQUAL_FATAL(
+        tsk_tree_get_total_branch_length(&tree, 8, &length), TSK_ERR_NODE_OUT_OF_BOUNDS);
+
+    tsk_tree_free(&tree);
+    tsk_treeseq_free(&ts);
+}
+
 static void
 test_single_tree_map_mutations(void)
 {
@@ -6605,6 +6643,7 @@ main(int argc, char **argv)
         { "test_single_tree_compute_mutation_times",
             test_single_tree_compute_mutation_times },
         { "test_single_tree_is_descendant", test_single_tree_is_descendant },
+        { "test_single_tree_total_branch_length", test_single_tree_total_branch_length },
         { "test_single_tree_map_mutations", test_single_tree_map_mutations },
         { "test_single_tree_map_mutations_internal_samples",
             test_single_tree_map_mutations_internal_samples },

c/tskit/trees.c

@@ -3631,6 +3631,40 @@ tsk_tree_get_time(const tsk_tree_t *self, tsk_id_t u, double *t)
     return ret;
 }
 
+int
+tsk_tree_get_total_branch_length(const tsk_tree_t *self, tsk_id_t node, double *ret_tbl)
+{
+    int ret = 0;
+    tsk_size_t j, num_nodes;
+    tsk_id_t u, v;
+    const tsk_id_t *restrict parent = self->parent;
+    const double *restrict time = self->tree_sequence->tables->nodes.time;
+    tsk_id_t *nodes = tsk_malloc(tsk_tree_get_size_bound(self) * sizeof(*nodes));
+    double sum = 0;
+
+    if (nodes == NULL) {
+        ret = TSK_ERR_NO_MEMORY;
+        goto out;
+    }
+    ret = tsk_tree_preorder(self, node, nodes, &num_nodes);
+    if (ret != 0) {
+        goto out;
+    }
+    /* We always skip the first node because we don't return the branch length
+     * over the input node. */
+    for (j = 1; j < num_nodes; j++) {
+        u = nodes[j];
+        v = parent[u];
+        if (v != TSK_NULL) {
+            sum += time[v] - time[u];
+        }
+    }
+    *ret_tbl = sum;
+out:
+    tsk_safe_free(nodes);
+    return ret;
+}
+
 int TSK_WARN_UNUSED
 tsk_tree_get_sites(
     const tsk_tree_t *self, const tsk_site_t **sites, tsk_size_t *sites_length)

c/tskit/trees.h

@@ -436,6 +436,31 @@ be greater than or equal to ``num_nodes``.
 */
 tsk_size_t tsk_tree_get_size_bound(const tsk_tree_t *self);
 
+/**
+@brief Returns the sum of the lengths of all branches reachable from
+    the specified node, or from all roots if node=TSK_NULL.
+
+@rst
+Return the total branch length in a particular subtree or of the
+entire tree. If the specified node is TSK_NULL (or the virtual
+root) the sum of the lengths of all branches reachable from roots
+is returned. Branch length is defined as difference between the time
+of a node and its parent. The branch length of a root is zero.
+
+Note that if the specified node is internal its branch length is
+*not* included, so that, e.g., the total branch length of a
+leaf node is zero.
+@endrst
+
+@param self A pointer to a tsk_tree_t object.
+@param node The tree node to compute branch length or TSK_NULL to return the
+    total branch length of the tree.
+@param ret_tbl A double pointer to store the returned total branch length.
+@return 0 on success or a negative value on failure.
+*/
+int tsk_tree_get_total_branch_length(
+    const tsk_tree_t *self, tsk_id_t node, double *ret_tbl);
+
 /** @} */
 
 int tsk_tree_set_root_threshold(tsk_tree_t *self, tsk_size_t root_threshold);

python/CHANGELOG.rst

@@ -54,6 +54,8 @@
   Roughly a 10X performance increase for "preorder", "postorder", "timeasc"
   and "timedesc" (:user:`jeromekelleher`, :pr:`1704`).
 
+- Substantial performance improvement for ``Tree.total_branch_length``
+  (:user:`jeromekelleher`, :issue:`1794` :pr:`1799`)
 
 --------------------
 [0.3.7] - 2021-07-08

python/_tskitmodule.c

@@ -9344,6 +9344,26 @@ Tree_get_num_edges(Tree *self)
     return ret;
 }
 
+static PyObject *
+Tree_get_total_branch_length(Tree *self)
+{
+    PyObject *ret = NULL;
+    double length;
+    int err;
+
+    if (Tree_check_state(self) != 0) {
+        goto out;
+    }
+    err = tsk_tree_get_total_branch_length(self->tree, TSK_NULL, &length);
+    if (err != 0) {
+        handle_library_error(err);
+        goto out;
+    }
+    ret = Py_BuildValue("d", length);
+out:
+    return ret;
+}
+
 static PyObject *
 Tree_get_index(Tree *self)
 {
@@ -10274,6 +10294,10 @@ static PyMethodDef Tree_methods[] = {
         .ml_meth = (PyCFunction) Tree_get_num_edges,
         .ml_flags = METH_NOARGS,
         .ml_doc = "Returns the number of branches in this tree." },
+    { .ml_name = "get_total_branch_length",
+        .ml_meth = (PyCFunction) Tree_get_total_branch_length,
+        .ml_flags = METH_NOARGS,
+        .ml_doc = "Returns the sum of the branch lengths reachable from roots" },
     { .ml_name = "get_left",
         .ml_meth = (PyCFunction) Tree_get_left,
         .ml_flags = METH_NOARGS,

python/tskit/trees.py

@@ -994,7 +994,7 @@ def total_branch_length(self):
         :return: The sum of lengths of branches in this tree.
         :rtype: float
         """
-        return sum(self.branch_length(u) for u in self.nodes())
+        return self._ll_tree.get_total_branch_length()
 
     def get_mrca(self, u, v):
         # Deprecated alias for mrca