functions for getting sample/node lists out of table collections

Author: molpopgen

src/table_collection.rs

@@ -415,6 +415,85 @@ impl TableCollection {
         handle_tsk_return_value!(rv)
     }
 
+    /// Obtain a vector containing the indexes ("ids")
+    /// of all nodes for which [`crate::TSK_NODE_IS_SAMPLE`]
+    /// is `true`.
+    pub fn samples_as_vector(&self) -> Vec<tsk_id_t> {
+        let mut samples: Vec<tsk_id_t> = vec![];
+        for (i, row) in self.nodes_iter(false).enumerate() {
+            if row.flags & crate::TSK_NODE_IS_SAMPLE > 0 {
+                samples.push(i as tsk_id_t);
+            }
+        }
+        samples
+    }
+
+    /// Obtain a vector containing the indexes ("ids") of all nodes
+    /// satisfying a certain criterion.
+    ///
+    /// # Parameters
+    ///
+    /// * `f`: a function.  The function is passed the current table
+    ///    collection and each [`crate::node_table::NodeTableRow`].
+    ///    If `f` returns `true`, the index of that row is included
+    ///    in the return value.
+    ///
+    /// # Examples
+    ///
+    /// Get all nodes with time > 0.0:
+    ///
+    /// ```
+    /// use tskit::TSK_NULL;
+    /// use tskit::tsk_id_t;
+    /// use tskit::TableAccess;
+    ///
+    /// let mut tables = tskit::TableCollection::new(100.).unwrap();
+    /// tables
+    ///     .add_node(tskit::TSK_NODE_IS_SAMPLE, 0.0, TSK_NULL, TSK_NULL)
+    ///     .unwrap();
+    /// tables
+    ///     .add_node(tskit::TSK_NODE_IS_SAMPLE, 1.0, TSK_NULL, TSK_NULL)
+    ///     .unwrap();
+    /// let samples = tables.create_node_id_vector(
+    ///     |row: &tskit::NodeTableRow, _: tsk_id_t| row.time > 0.,
+    /// );
+    /// assert_eq!(samples[0], 1);
+    ///
+    /// // Get all nodes that have a mutation:
+    ///
+    /// fn node_has_mutation(
+    ///     tables: &tskit::TableCollection,
+    ///     row: &tskit::NodeTableRow,
+    ///     i: tsk_id_t,
+    /// ) -> bool {
+    ///     for mrow in tables.mutations_iter(false) {
+    ///         if mrow.node == i {
+    ///             return true;
+    ///         }
+    ///     }
+    ///     false
+    /// }
+    ///
+    /// // Get all nodes that have a mutation:
+    ///
+    /// tables.add_mutation(0, 0, TSK_NULL, 0.0, None).unwrap();
+    /// let samples_with_mut = tables.create_node_id_vector(
+    ///     |row: &tskit::NodeTableRow, i: tsk_id_t| node_has_mutation(&tables, row, i));
+    /// assert_eq!(samples_with_mut[0], 0);
+    /// ```
+    pub fn create_node_id_vector(
+        &self,
+        mut f: impl FnMut(&crate::NodeTableRow, tsk_id_t) -> bool,
+    ) -> Vec<tsk_id_t> {
+        let mut samples: Vec<tsk_id_t> = vec![];
+        for (i, row) in self.nodes_iter(true).enumerate() {
+            if f(&row, i as tsk_id_t) {
+                samples.push(i as tsk_id_t);
+            }
+        }
+        samples
+    }
+
     /// Sort the tables.  
     /// The [``bookmark``](crate::types::Bookmark) can
     /// be used to affect where sorting starts from for each table.