add TableCollection and TreeSequence simplification

Author: molpopgen

src/flags.rs

@@ -0,0 +1,21 @@
+use crate::bindings as ll_bindings;
+use crate::ffi::TskitTypeAccess;
+use crate::TableCollection;
+use crate::TskitError;
+use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TSK_NULL};
+use bitflags::bitflags;
+
+bitflags! {
+    #[derive(Default)]
+    pub struct SimplificationOptions: tsk_flags_t {
+        /// Default behavior
+        const NONE = 0;
+        const FILTER_SITES = ll_bindings::TSK_FILTER_SITES;
+        const FILTER_POPULATIONS = ll_bindings::TSK_FILTER_POPULATIONS;
+        const FILTER_INDIVIDUALS = ll_bindings::TSK_FILTER_INDIVIDUALS;
+        const REDUCE_TO_SITE_TOPOLOGY = ll_bindings::TSK_REDUCE_TO_SITE_TOPOLOGY;
+        const KEEP_UNARY  = ll_bindings::TSK_KEEP_UNARY;
+        const KEEP_INPUT_ROOTS = ll_bindings::TSK_KEEP_INPUT_ROOTS;
+        const KEEP_UNARY_IN_INDIVIDUALS  = ll_bindings::TSK_KEEP_UNARY_IN_INDIVIDUALS;
+    }
+}

src/lib.rs

@@ -10,6 +10,7 @@ mod _macros; // Starts w/_ to be sorted at front by rustfmt!
 mod edge_table;
 pub mod error;
 pub mod ffi;
+mod flags;
 mod individual_table;
 pub mod metadata;
 mod migration_table;
@@ -43,6 +44,7 @@ pub const TSK_NULL: tsk_id_t = -1;
 
 pub use edge_table::{EdgeTable, EdgeTableRow};
 pub use error::TskitError;
+pub use flags:: SimplificationOptions;
 pub use individual_table::{IndividualTable, IndividualTableRow};
 pub use migration_table::{MigrationTable, MigrationTableRow};
 pub use mutation_table::{MutationTable, MutationTableRow};

src/table_collection.rs

@@ -9,10 +9,11 @@ use crate::MigrationTable;
 use crate::MutationTable;
 use crate::NodeTable;
 use crate::PopulationTable;
+use crate::SimplificationOptions;
 use crate::SiteTable;
 use crate::TableAccess;
 use crate::TskReturnValue;
-use crate::{tsk_flags_t, tsk_id_t, tsk_size_t};
+use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TSK_NULL};
 use ll_bindings::tsk_table_collection_free;
 
 /// A table collection.
@@ -487,6 +488,42 @@ impl TableCollection {
     pub fn tree_sequence(self) -> Result<crate::TreeSequence, TskitError> {
         crate::TreeSequence::new(self)
     }
+
+    pub fn simplify(
+        &mut self,
+        samples: &[tsk_id_t],
+        options: SimplificationOptions,
+        idmap: bool,
+    ) -> Result<Option<Vec<tsk_id_t>>, TskitError> {
+        match idmap {
+            true => {
+                let mut output_node_map: Vec<tsk_id_t> =
+                    vec![TSK_NULL; self.nodes().num_rows() as usize];
+                let rv = unsafe {
+                    ll_bindings::tsk_table_collection_simplify(
+                        self.as_mut_ptr(),
+                        samples.as_ptr(),
+                        samples.len() as tsk_size_t,
+                        options.bits(),
+                        output_node_map.as_mut_ptr(),
+                    )
+                };
+                handle_tsk_return_value!(rv, Some(output_node_map))
+            }
+            false => {
+                let rv = unsafe {
+                    ll_bindings::tsk_table_collection_simplify(
+                        self.as_mut_ptr(),
+                        samples.as_ptr(),
+                        samples.len() as tsk_size_t,
+                        options.bits(),
+                        std::ptr::null_mut(),
+                    )
+                };
+                handle_tsk_return_value!(rv, None)
+            }
+        }
+    }
 }
 
 impl TableAccess for TableCollection {

src/trees.rs

@@ -7,6 +7,7 @@ use crate::MigrationTable;
 use crate::MutationTable;
 use crate::NodeTable;
 use crate::PopulationTable;
+use crate::SimplificationOptions;
 use crate::SiteTable;
 use crate::TableAccess;
 use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TableCollection, TSK_NULL};
@@ -831,6 +832,47 @@ impl TreeSequence {
     pub fn num_samples(&self) -> tsk_size_t {
         unsafe { ll_bindings::tsk_treeseq_get_num_samples(self.as_ptr()) }
     }
+
+    pub fn simplify(
+        &mut self,
+        samples: &[tsk_id_t],
+        options: SimplificationOptions,
+        idmap: bool,
+    ) -> Result<(Self, Option<Vec<tsk_id_t>>), TskitError> {
+        let tables = TableCollection::new(unsafe { (*self.inner.tables).sequence_length })?;
+        let mut ts = tables.tree_sequence()?;
+
+        match idmap {
+            true => {
+                let mut output_node_map: Vec<tsk_id_t> =
+                    vec![TSK_NULL; unsafe { (*self.inner.tables).nodes.num_rows } as usize];
+                let rv = unsafe {
+                    ll_bindings::tsk_treeseq_simplify(
+                        self.as_mut_ptr(),
+                        samples.as_ptr(),
+                        samples.len() as tsk_size_t,
+                        options.bits(),
+                        ts.as_mut_ptr(),
+                        output_node_map.as_mut_ptr(),
+                    )
+                };
+                handle_tsk_return_value!(rv, (ts, Some(output_node_map)))
+            }
+            false => {
+                let rv = unsafe {
+                    ll_bindings::tsk_treeseq_simplify(
+                        self.as_mut_ptr(),
+                        samples.as_ptr(),
+                        samples.len() as tsk_size_t,
+                        options.bits(),
+                        ts.as_mut_ptr(),
+                        std::ptr::null_mut(),
+                    )
+                };
+                handle_tsk_return_value!(rv, (ts, None))
+            }
+        }
+    }
 }
 
 impl TableAccess for TreeSequence {