Merge pull request #58 from molpopgen/validate_tables_simplification

Closes #59. Closes #60.

Author: molpopgen

forrustts_examples/neutral_wf.rs

@@ -1,5 +1,6 @@
 use clap::{value_t, value_t_or_exit, App, Arg};
 use forrustts::wright_fisher::*;
+use forrustts::SimplificationFlags;
 
 fn main() {
     let matches = App::new("neutral_wf")
@@ -58,6 +59,13 @@ fn main() {
                 .help("Survival probability")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name("validate_tables")
+                .short("v")
+                .long("validate_tables")
+                .help("Validate all tables prior to simplification")
+                .takes_value(false),
+        )
         .get_matches();
 
     // TODO: default params
@@ -69,6 +77,7 @@ fn main() {
     let psurvival = value_t!(matches.value_of("psurvival"), f64).unwrap_or(0.0);
     let seed = value_t_or_exit!(matches.value_of("seed"), usize);
     let outfile = value_t_or_exit!(matches.value_of("outfile"), String);
+    let validate_tables =  matches.is_present("validate_tables");
 
     // TODO: parameter validation..
 
@@ -86,6 +95,12 @@ fn main() {
         SimulationFlags::USE_STATE
     };
 
+    let mut simplification_flags = SimplificationFlags::empty();
+
+    if validate_tables {
+        simplification_flags |= SimplificationFlags::VALIDATE_ALL;
+    }
+
     let (mut tables, is_sample) = neutral_wf(
         PopulationParams {
             size: popsize,
@@ -98,6 +113,7 @@ fn main() {
             seed,
             nsteps: g,
             flags,
+            simplification_flags,
         },
     )
     .unwrap();

src/error.rs

@@ -6,7 +6,7 @@ use thiserror::Error;
 ///
 /// Some members of this enum implement ``From``
 /// in order to redirect other error types.
-#[derive(Error, Debug)]
+#[derive(Error, Debug, PartialEq)]
 pub enum ForrusttsError {
     /// An error that occurs during simplification.
     #[error("{value:?}")]

src/nested_forward_list.rs

@@ -11,7 +11,7 @@
 use thiserror::Error;
 
 /// Errror type for [``NestedForwardList``] operations.
-#[derive(Error, Debug)]
+#[derive(Error, Debug, PartialEq)]
 pub enum NestedForwardListError {
     /// Tail of a list is unexpectedly null.
     #[error("Tail is null")]

src/simplification_common.rs

@@ -1,6 +1,7 @@
 /// Common functions to reuse in various "simplify tables"
 /// functions
 use crate::simplification_logic;
+use crate::validate_edge_table;
 use crate::ForrusttsError;
 use crate::SamplesInfo;
 use crate::SimplificationBuffers;
@@ -9,6 +10,16 @@ use crate::SimplificationOutput;
 use crate::{IdType, NULL_ID};
 use crate::{Node, TableCollection};
 
+pub fn validate_tables(
+    tables: &TableCollection,
+    flags: &SimplificationFlags,
+) -> Result<(), ForrusttsError> {
+    if flags.contains(SimplificationFlags::VALIDATE_EDGES) {
+        validate_edge_table(tables.genome_length(), tables.edges(), tables.nodes())?;
+    }
+    Ok(())
+}
+
 fn setup_idmap(nodes: &[Node], idmap: &mut Vec<IdType>) {
     idmap.resize(nodes.len(), NULL_ID);
     idmap.iter_mut().for_each(|x| *x = NULL_ID);
@@ -27,12 +38,7 @@ pub fn setup_simplification(
         });
     }
 
-    if flags.bits() != 0 {
-        return Err(ForrusttsError::SimplificationError {
-            value: "SimplificationFlags must be zero".to_string(),
-        });
-    }
-
+    validate_tables(tables, &flags)?;
     setup_idmap(&tables.nodes_, &mut output.idmap);
 
     state.clear();

src/simplification_flags.rs

@@ -4,19 +4,20 @@ bitflags! {
     /// Boolean flags affecting simplification
     /// behavior.
     ///
-    /// Currently, this is unused, and exists
-    /// as a placeholder for the future.
-    ///
     /// # Example
     ///
     /// ```
     /// let e = forrustts::SimplificationFlags::empty();
-    /// assert!(e.contains(forrustts::SimplificationFlags::NONE));
+    /// assert_eq!(e.bits(), 0);
     /// ```
     #[derive(Default)]
     pub struct SimplificationFlags: u32 {
-        /// Placeholder
-        const NONE = 0;
+        /// Validate that input edges are sorted
+        const VALIDATE_EDGES = 1 << 0;
+        /// Validate that input mutations are sorted
+        const VALIDATE_MUTATIONS = 1 << 1;
+        /// Validate all tables.
+        const VALIDATE_ALL = Self::VALIDATE_EDGES.bits | Self::VALIDATE_MUTATIONS.bits;
     }
 }
 
@@ -27,7 +28,6 @@ mod test {
     #[test]
     fn test_empty() {
         let e = SimplificationFlags::empty();
-        assert!(e.contains(SimplificationFlags::NONE));
         assert_eq!(e.bits(), 0);
     }
 }

src/simplify_tables.rs

@@ -91,3 +91,43 @@ pub fn simplify_tables(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // TODO: we need lots more tests of these validations!
+
+    #[test]
+    fn test_simplify_tables_unsorted_edges() {
+        let mut tables = TableCollection::new(1000).unwrap();
+
+        tables.add_node(0, 0).unwrap(); // parent
+        tables.add_node(1, 0).unwrap(); // child
+        tables.add_edge(100, tables.genome_length(), 0, 1).unwrap();
+        tables.add_edge(0, 100, 0, 1).unwrap();
+
+        let mut output = SimplificationOutput::new();
+
+        let mut samples = SamplesInfo::new();
+        samples.samples.push(1);
+
+        let _ = simplify_tables_without_state(
+            &samples,
+            SimplificationFlags::VALIDATE_ALL,
+            &mut tables,
+            &mut output,
+        )
+        .map_or_else(
+            |x: ForrusttsError| {
+                assert_eq!(
+                    x,
+                    ForrusttsError::TablesError {
+                        value: TablesError::EdgesNotSortedByLeft
+                    }
+                )
+            },
+            |_| panic!(),
+        );
+    }
+}

src/tables.rs

@@ -1,4 +1,5 @@
 use crate::tsdef::{IdType, Position, Time, NULL_ID};
+use bitflags::bitflags;
 use std::cmp::Ordering;
 use thiserror::Error;
 
@@ -269,6 +270,40 @@ fn sort_mutation_table(sites: &[Site], mutations: &mut MutationTable) {
     });
 }
 
+bitflags! {
+    /// Modifies behavior of
+    /// [``TableCollection::validate_tables``]
+    ///
+    /// ```
+    /// let f = forrustts::TableValidationFlags::empty();
+    /// assert_eq!(f.contains(forrustts::TableValidationFlags::VALIDATE_ALL), true);
+    /// ```
+    #[derive(Default)]
+    pub struct TableValidationFlags: u32 {
+        /// Validate all tables.
+        /// This is also the "default"/empty.
+        const VALIDATE_ALL = 0;
+    }
+}
+
+bitflags! {
+    /// Modifies behavior of
+    /// [``TableCollection::sort_tables``]
+    ///
+    /// ```
+    /// let f = forrustts::TableSortingFlags::empty();
+    /// assert_eq!(f.contains(forrustts::TableSortingFlags::SORT_ALL), true);
+    /// ```
+    #[derive(Default)]
+    pub struct TableSortingFlags: u32 {
+        /// Sort all tables.
+        /// This is also the "default"/empty.
+        const SORT_ALL = 0;
+        /// Do not sort the edge table.
+        const SKIP_EDGE_TABLE = 1 << 0;
+    }
+}
+
 /// Perform a data integrity check on an [``EdgeTable``].
 ///
 /// This checks, amongst other things, the sorting order
@@ -653,10 +688,26 @@ impl TableCollection {
     }
 
     /// Sort all tables for simplification.
+    #[deprecated(since = "0.1.0", note = "use sort_tables instead")]
     pub fn sort_tables_for_simplification(&mut self) {
-        sort_edge_table(&self.nodes_, &mut self.edges_);
+        self.sort_tables(TableSortingFlags::empty());
+    }
+
+    /// Sort all tables for simplification.
+    pub fn sort_tables(&mut self, flags: TableSortingFlags) {
+        if !flags.contains(TableSortingFlags::SKIP_EDGE_TABLE) {
+            sort_edge_table(&self.nodes_, &mut self.edges_);
+        }
         sort_mutation_table(&self.sites_, &mut self.mutations_);
     }
+
+    /// Run a validation check on the tables.
+    pub fn validate(&self, flags: TableValidationFlags) -> TablesResult<bool> {
+        if flags.contains(TableValidationFlags::VALIDATE_ALL) {
+            validate_edge_table(self.genome_length(), &self.edges_, &self.nodes_)?;
+        }
+        Ok(true)
+    }
 }
 
 #[cfg(test)]

src/test_simplify_tables.rs

@@ -38,10 +38,10 @@ mod test {
         genome_length: Position,
         psurvival: f64,
         seed: usize,
+        // None here means "never simplify".
         simplification_interval: Option<Time>,
         flags: SimulationFlags,
     ) -> Result<(TableCollection, Vec<i32>), ForrusttsError> {
-        // None here means "never simplify".
         neutral_wf(
             PopulationParams {
                 size: 250,
@@ -54,6 +54,7 @@ mod test {
                 seed,
                 nsteps: num_generations,
                 flags,
+                simplification_flags: SimplificationFlags::VALIDATE_ALL,
             },
         )
     }
@@ -89,7 +90,7 @@ mod test {
         );
 
         // Now, sort and simplify the tables we got from the sim:
-        tables.sort_tables_for_simplification();
+        tables.sort_tables(crate::TableSortingFlags::empty());
         let mut samples = SamplesInfo::new();
         for (i, n) in tables.nodes().iter().enumerate() {
             if n.time == num_generations {