more

Author: molpopgen

src/sys/tree.rs

@@ -201,11 +201,17 @@ impl<'treeseq> LLTree<'treeseq> {
     }
 
     pub fn left_child_array(&self) -> &[NodeId] {
-        super::generate_slice(self.as_ll_ref().left_child, self.treeseq.num_nodes_raw() + 1)
+        super::generate_slice(
+            self.as_ll_ref().left_child,
+            self.treeseq.num_nodes_raw() + 1,
+        )
     }
 
     pub fn right_child_array(&self) -> &[NodeId] {
-        super::generate_slice(self.as_ll_ref().right_child, self.treeseq.num_nodes_raw() + 1)
+        super::generate_slice(
+            self.as_ll_ref().right_child,
+            self.treeseq.num_nodes_raw() + 1,
+        )
     }
 
     pub fn total_branch_length(&self, by_span: bool) -> Result<Time, TskitError> {
@@ -238,6 +244,34 @@ impl<'treeseq> LLTree<'treeseq> {
         let i = self.interval();
         i.1 - i.0
     }
+
+    pub fn left_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        err_if_not_tracking_samples!(
+            self.flags,
+            super::generate_slice(
+                self.as_ll_ref().left_sample,
+                self.treeseq.num_nodes_raw() + 1
+            )
+        )
+    }
+
+    pub fn right_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        err_if_not_tracking_samples!(
+            self.flags,
+            super::generate_slice(
+                self.as_ll_ref().right_sample,
+                self.treeseq.num_nodes_raw() + 1
+            )
+        )
+    }
+
+    pub fn kc_distance(&self, other: &Self, lambda: f64) -> Result<f64, TskitError> {
+        let mut kc = f64::NAN;
+        let kcp: *mut f64 = &mut kc;
+        let code =
+            unsafe { bindings::tsk_tree_kc_distance(self.as_ptr(), other.as_ptr(), lambda, kcp) };
+        handle_tsk_return_value!(code, kc)
+    }
 }
 
 // Trait defining iteration over nodes.

src/trees/tree.rs

@@ -339,17 +339,89 @@ impl<'treeseq> Tree<'treeseq> {
         X
     }
 
-    pub fn left_sample_array(&self) {
-        X
+    /// # Failing examples
+    ///
+    /// An error will be returned if ['crate::TreeFlags::SAMPLE_LISTS`] is not used:
+    ///
+    /// ```should_panic
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap(); // Error
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let l = tree.left_sample_array().unwrap();
+    ///     for _ in l {}
+    /// }
+    /// ```
+    ///
+    /// The lifetime of the slice is tied to the parent object:
+    ///
+    /// ```compile_fail
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let l = tree.left_sample_array().unwrap();
+    ///     drop(tree_iter);
+    ///     for _ in l {} // Error
+    /// }
+    /// ```
+    pub fn left_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        self.inner.left_sample_array()
     }
 
-    pub fn right_sample_array(&self) {
-        x
+    /// # Failing examples
+    ///
+    /// An error will be returned if ['crate::TreeFlags::SAMPLE_LISTS`] is not used:
+    ///
+    /// ```should_panic
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap(); // ERROR
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let r = tree.right_sample_array().unwrap();
+    ///     for _ in r {}
+    /// }
+    /// ```
+    ///
+    /// The lifetime of the slice is tied to the parent object:
+    ///
+    /// ```compile_fail
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let r = tree.right_sample_array().unwrap();
+    ///     drop(tree_iter);
+    ///     for _ in r {} // ERROR
+    /// }
+    /// ```
+    pub fn right_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        self.inner.right_sample_array()
     }
 
-    left_sample
-        right_sample
-        kc_distance
+    /// Calculate the average Kendall-Colijn (`K-C`) distance between
+    /// pairs of trees whose intervals overlap.
+    ///
+    /// # Note
+    ///
+    /// * [Citation](https://doi.org/10.1093/molbev/msw124)
+    ///
+    /// # Parameters
+    ///
+    /// * `lambda` specifies the relative weight of topology and branch length.
+    ///   If `lambda` is 0, we only consider topology.
+    ///   If `lambda` is 1, we only consider branch lengths.
+    pub fn kc_distance(&self, other: &Self, lambda: f64) -> Result<f64, TskitError> {
+        self.inner.kc_distance(&other.inner, lambda)
+    }
 }
 
 impl<'ts> streaming_iterator::StreamingIterator for Tree<'ts> {