Add validation step to simplification.

Make existing error types provide PartialEq.

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/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/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,
             },
         )
     }

src/wright_fisher.rs

@@ -7,7 +7,7 @@
 //! to others.  Feel free to copy them!
 use crate::simplify_from_edge_buffer::simplify_from_edge_buffer;
 use crate::simplify_tables::*;
-use crate::tables::{validate_edge_table, TableCollection};
+use crate::tables::TableCollection;
 use crate::tsdef::*;
 use crate::EdgeBuffer;
 use crate::ForrusttsError;
@@ -300,47 +300,31 @@ fn fill_samples(parents: &[Parent], samples: &mut SamplesInfo) {
 
 fn sort_and_simplify(
     flags: SimulationFlags,
+    simplification_flags: SimplificationFlags,
     samples: &SamplesInfo,
     state: &mut SimplificationBuffers,
     pop: &mut PopulationState,
     output: &mut SimplificationOutput,
 ) {
     if !flags.contains(SimulationFlags::BUFFER_EDGES) {
         pop.tables.sort_tables(crate::TableSortingFlags::empty());
-        debug_assert!(validate_edge_table(
-            pop.tables.genome_length(),
-            pop.tables.edges(),
-            pop.tables.nodes()
-        )
-        .unwrap());
         if flags.contains(SimulationFlags::USE_STATE) {
             simplify_tables(
                 samples,
-                SimplificationFlags::empty(),
+                simplification_flags,
                 state,
                 &mut pop.tables,
                 output,
             )
             .unwrap();
         } else {
-            simplify_tables_without_state(
-                samples,
-                SimplificationFlags::empty(),
-                &mut pop.tables,
-                output,
-            )
-            .unwrap();
+            simplify_tables_without_state(samples, simplification_flags, &mut pop.tables, output)
+                .unwrap();
         }
-        debug_assert!(validate_edge_table(
-            pop.tables.genome_length(),
-            pop.tables.edges(),
-            pop.tables.nodes()
-        )
-        .unwrap());
     } else {
         simplify_from_edge_buffer(
             samples,
-            SimplificationFlags::empty(),
+            simplification_flags,
             state,
             &mut pop.edge_buffer,
             &mut pop.tables,
@@ -352,13 +336,14 @@ fn sort_and_simplify(
 
 fn simplify_and_remap_nodes(
     flags: SimulationFlags,
+    simplification_flags: SimplificationFlags,
     samples: &mut SamplesInfo,
     state: &mut SimplificationBuffers,
     pop: &mut PopulationState,
     output: &mut SimplificationOutput,
 ) {
     fill_samples(&pop.parents, samples);
-    sort_and_simplify(flags, samples, state, pop, output);
+    sort_and_simplify(flags, simplification_flags, samples, state, pop, output);
 
     for p in &mut pop.parents {
         p.node0 = output.idmap[p.node0 as usize];
@@ -451,10 +436,14 @@ pub struct SimulationParams {
     /// Bitwise flag tweaking the behavior of the
     /// simplification algorithm.
     pub flags: SimulationFlags,
+    /// Flags to affect simplification behavior
+    pub simplification_flags: SimplificationFlags,
 }
 
 impl SimulationParams {
     /// Create a new instance
+    /// [``SimulationParams::simplification_flags``]
+    /// is initialized to empty.
     pub fn new(
         simplification_interval: Option<Time>,
         seed: usize,
@@ -466,6 +455,7 @@ impl SimulationParams {
             seed,
             nsteps,
             flags,
+            simplification_flags: SimplificationFlags::empty(),
         }
     }
 }
@@ -561,6 +551,7 @@ pub fn neutral_wf(
         {
             simplify_and_remap_nodes(
                 params.flags,
+                params.simplification_flags,
                 &mut samples,
                 &mut state,
                 &mut pop,
@@ -575,6 +566,7 @@ pub fn neutral_wf(
     if !simplified && actual_simplification_interval != -1 {
         simplify_and_remap_nodes(
             params.flags,
+            params.simplification_flags,
             &mut samples,
             &mut state,
             &mut pop,