more

Author: molpopgen

examples/haploid_wright_fisher.rs

@@ -163,7 +163,7 @@ fn stress_test_total_branch_length() {
                 if let Some(tree) = tree_iter.next() {
                     let b = tree.total_branch_length(false).unwrap();
                     let b2 = unsafe {
-                        tskit::bindings::tsk_tree_get_total_branch_length(tree.as_ptr(), -1, &mut x)
+                        tskit::bindings::tsk_tree_get_total_branch_length(tree.as_ref(), -1, &mut x)
                     };
                     assert!(b2 >= 0, "{}", b2);
                     assert!(f64::from(b) - x <= 1e-8);

src/sys/tree.rs

@@ -5,7 +5,9 @@ use super::bindings::tsk_size_t;
 use super::bindings::tsk_tree_t;
 use super::flags::TreeFlags;
 use super::newtypes::NodeId;
+use super::newtypes::Position;
 use super::newtypes::SizeType;
+use super::newtypes::Time;
 use super::tskbox::TskBox;
 use super::TreeSequence;
 use super::TskitError;
@@ -205,6 +207,37 @@ impl<'treeseq> LLTree<'treeseq> {
     pub fn right_child_array(&self) -> &[NodeId] {
         super::generate_slice(self.as_ref().right_child, self.treeseq.num_nodes_raw() + 1)
     }
+
+    pub fn total_branch_length(&self, by_span: bool) -> Result<Time, TskitError> {
+        let time: &[Time] = super::generate_slice(
+            unsafe { (*(*(*self.as_ptr()).tree_sequence).tables).nodes.time },
+            self.treeseq.num_nodes_raw() + 1,
+        );
+        let mut b = Time::from(0.);
+        for n in self.traverse_nodes(NodeTraversalOrder::Preorder) {
+            let p = self.parent(n).ok_or(TskitError::IndexError {})?;
+            if p != NodeId::NULL {
+                b += time[p.as_usize()] - time[n.as_usize()]
+            }
+        }
+
+        match by_span {
+            true => Ok(b * self.span()),
+            false => Ok(b),
+        }
+    }
+
+    pub fn interval(&self) -> (Position, Position) {
+        (
+            self.as_ref().interval.left.into(),
+            self.as_ref().interval.right.into(),
+        )
+    }
+
+    pub fn span(&self) -> Position {
+        let i = self.interval();
+        i.1 - i.0
+    }
 }
 
 // Trait defining iteration over nodes.

src/trees/tree.rs

@@ -1,6 +1,7 @@
 use crate::sys::bindings as ll_bindings;
 use crate::sys::{LLTree, TreeSequence};
 use crate::NodeId;
+use crate::Time;
 use crate::Position;
 use crate::SizeType;
 use crate::TreeFlags;
@@ -58,10 +59,20 @@ impl<'treeseq> Tree<'treeseq> {
 
     /// Return the `[left, right)` coordinates of the tree.
     pub fn interval(&self) -> (Position, Position) {
-        (
-            self.inner.as_ref().interval.left.into(),
-            self.inner.as_ref().interval.right.into(),
-        )
+        self.inner.interval()
+    }
+
+    /// Calculate the total length of the tree via a preorder traversal.
+    ///
+    /// # Parameters
+    ///
+    /// * `by_span`: if `true`, multiply the return value by [`TreeInterface::span`].
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError`] may be returned if a node index is out of range.
+    pub fn total_branch_length(&self, by_span: bool) -> Result<Time, TskitError> {
+        self.inner.total_branch_length(by_span)
     }
 
     /// # Failing examples
@@ -154,6 +165,13 @@ impl<'treeseq> Tree<'treeseq> {
         self.inner.flags()
     }
 
+    /// Return the length of the genome for which this
+    /// tree is the ancestry.
+    pub fn span(&self) -> Position {
+        let i = self.interval();
+        i.1 - i.0
+    }
+
     /// Get the parent of node `u`.
     ///
     /// Returns `None` if `u` is out of range.