diff --git a/Cargo.toml b/Cargo.toml index 091bda9c..0c95db53 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,4 +15,4 @@ keywords = ["simulation", "tree_sequences", "tskit", "population_genetics"] bitflags = "1.2.1" thiserror = "1.0" GSL = "4.0.0" -tskit = "0.1.1" +tskit = "=0.2.0" diff --git a/forrustts_examples/Cargo.toml b/forrustts_examples/Cargo.toml index 8f368e6d..e919d81a 100644 --- a/forrustts_examples/Cargo.toml +++ b/forrustts_examples/Cargo.toml @@ -10,7 +10,7 @@ edition = "2018" path = "../" [dependencies] -tskit = "0.1.1" +tskit = "=0.2.0" clap = "~2.27.0" [[bin]] diff --git a/forrustts_examples/neutral_wf.rs b/forrustts_examples/neutral_wf.rs index 96af6202..bb1e28b1 100644 --- a/forrustts_examples/neutral_wf.rs +++ b/forrustts_examples/neutral_wf.rs @@ -125,11 +125,14 @@ fn main() { &mut tables, ); - if simplify.is_some() { - tskit_tables.dump(&outfile, 0).unwrap(); - } else { - tskit_tables - .dump(&outfile, tskit::TSK_NO_BUILD_INDEXES) - .unwrap(); - } + match simplify.is_some() { + true => (), + false => { + let _ = tskit_tables.build_index().unwrap(); + } + }; + + tskit_tables + .dump(&outfile, tskit::TableOutputOptions::default()) + .unwrap(); } diff --git a/src/test_simplify_tables.rs b/src/test_simplify_tables.rs index bdabbbf6..406365a9 100644 --- a/src/test_simplify_tables.rs +++ b/src/test_simplify_tables.rs @@ -1,9 +1,5 @@ #[cfg(test)] mod test { - // NOTE: Currently, these tests are both testing - // stuff from tskit and forrusts, which isn't great. - // We'll clean this up later when we get better abstractions - // into tskit. use crate::simplify_tables_without_state; use crate::tsdef::{IdType, Position, Time}; use crate::wright_fisher::*; @@ -12,25 +8,7 @@ mod test { use crate::SimplificationFlags; use crate::SimplificationOutput; use crate::TableCollection; - use std::mem::MaybeUninit; - use tskit::bindings as tskr; - use tskit::ffi::TskitType; - - fn tables_to_treeseq(tables: &mut tskit::TableCollection) -> MaybeUninit { - let mut tsk_ts: MaybeUninit = MaybeUninit::uninit(); - unsafe { - let rv = tskr::tsk_table_collection_build_index(tables.as_mut_ptr(), 0); - assert_eq!(rv, 0); - let rv = tskr::tsk_treeseq_init( - tsk_ts.as_mut_ptr(), - tables.as_ptr(), - tskr::TSK_SAMPLE_LISTS, - ); - assert_eq!(rv, 0); - } - - tsk_ts - } + use tskit::TableAccess; fn simulate_data( num_generations: Time, @@ -111,52 +89,44 @@ mod test { *i = 1; } - let mut simplified_rust_tables = crate::tskit_tools::convert_to_tskit_minimal( + let simplified_rust_tables = crate::tskit_tools::convert_to_tskit_minimal( &tables, &is_sample, crate::tskit_tools::simple_time_reverser(num_generations), true, ); - unsafe { - let rv = tskr::tsk_table_collection_sort(tsk_tables.as_mut_ptr(), std::ptr::null(), 0); - assert!(rv == 0); - let rv = tskr::tsk_table_collection_simplify( - tsk_tables.as_mut_ptr(), - samples.samples.as_ptr(), - samples.samples.len() as u32, - 0, - std::ptr::null_mut(), - ); - assert!(rv == 0); - } + tsk_tables + .full_sort(tskit::TableSortOptions::default()) + .unwrap(); + tsk_tables + .simplify( + &samples.samples, + tskit::SimplificationOptions::default(), + false, + ) + .unwrap(); // Get tree sequences now - let mut tsk_ts = tables_to_treeseq(&mut tsk_tables); - let mut rust_ts = tables_to_treeseq(&mut simplified_rust_tables); - - unsafe { - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(tsk_ts.as_ptr())); - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(rust_ts.as_ptr())); - let ne = tskr::tsk_treeseq_get_num_edges(tsk_ts.as_ptr()); - let ne2 = tskr::tsk_treeseq_get_num_edges(rust_ts.as_ptr()); - assert_eq!(ne, ne2); - let nn = tskr::tsk_treeseq_get_num_nodes(tsk_ts.as_ptr()); - let nn2 = tskr::tsk_treeseq_get_num_nodes(rust_ts.as_ptr()); - assert_eq!(nn, nn2); - - // We expect the rv to be 0, - // so let's init it to something else - let mut kc: f64 = -1.; - let kcp: *mut f64 = &mut kc; - let rv = - tskr::tsk_treeseq_kc_distance(tsk_ts.as_mut_ptr(), rust_ts.as_mut_ptr(), 0., kcp); - assert_eq!(rv, 0); - assert!((kc - 0.).abs() < f64::EPSILON); - - tskr::tsk_treeseq_free(tsk_ts.as_mut_ptr()); - tskr::tsk_treeseq_free(rust_ts.as_mut_ptr()); - } + let tsk_ts = tsk_tables + .tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES) + .unwrap(); + + let rust_ts = simplified_rust_tables + .tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES) + .unwrap(); + + assert_eq!(500, tsk_ts.num_samples()); + assert_eq!(500, rust_ts.num_samples()); + let ne = tsk_ts.edges().num_rows(); + let ne2 = rust_ts.edges().num_rows(); + assert_eq!(ne, ne2); + let nn = tsk_ts.nodes().num_rows(); + let nn2 = rust_ts.nodes().num_rows(); + assert_eq!(nn, nn2); + + let kc = tsk_ts.kc_distance(&&rust_ts, 0.).unwrap(); + assert!((kc - 0.).abs() < f64::EPSILON); } #[test] @@ -192,38 +162,30 @@ mod test { let sum_times_buffered: Time = tables_buffered.nodes_.iter().map(|x| x.time).sum(); assert_eq!(sum_times_sorted, sum_times_buffered); - let mut tables_sorted_tskit = crate::tskit_tools::convert_to_tskit_minimal( + let tables_sorted_tskit = crate::tskit_tools::convert_to_tskit_minimal( &tables_sorted, &is_sample_sorted, crate::tskit_tools::simple_time_reverser(num_generations), true, ); - let mut tables_buffered_tskit = crate::tskit_tools::convert_to_tskit_minimal( + let tables_buffered_tskit = crate::tskit_tools::convert_to_tskit_minimal( &tables_buffered, &is_sample_buffered, crate::tskit_tools::simple_time_reverser(num_generations), true, ); - let mut sorted_ts = tables_to_treeseq(&mut tables_sorted_tskit); - let mut buffered_ts = tables_to_treeseq(&mut tables_buffered_tskit); - unsafe { - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(sorted_ts.as_ptr())); - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(buffered_ts.as_ptr())); - let mut kc: f64 = -1.; - let kcp: *mut f64 = &mut kc; - let rv = tskr::tsk_treeseq_kc_distance( - sorted_ts.as_mut_ptr(), - buffered_ts.as_mut_ptr(), - 0., - kcp, - ); - assert_eq!(rv, 0); - assert!((kc - 0.).abs() < f64::EPSILON); - tskr::tsk_treeseq_free(sorted_ts.as_mut_ptr()); - tskr::tsk_treeseq_free(buffered_ts.as_mut_ptr()); - } + let sorted_ts = tables_sorted_tskit + .tree_sequence(tskit::TreeSequenceFlags::default()) + .unwrap(); + let buffered_ts = tables_buffered_tskit + .tree_sequence(tskit::TreeSequenceFlags::default()) + .unwrap(); + assert_eq!(500, sorted_ts.num_samples()); + assert_eq!(500, buffered_ts.num_samples()); + let kc = sorted_ts.kc_distance(&&buffered_ts, 0.).unwrap(); + assert!((kc - 0.).abs() < f64::EPSILON); } // The KC distance code will barf on trees where samples @@ -264,31 +226,28 @@ mod test { let sum_times_buffered: Time = tables_buffered.nodes_.iter().map(|x| x.time).sum(); assert_eq!(sum_times_sorted, sum_times_buffered); - let mut tables_sorted_tskit = crate::tskit_tools::convert_to_tskit_minimal( + let tables_sorted_tskit = crate::tskit_tools::convert_to_tskit_minimal( &tables_sorted, &is_sample_sorted, crate::tskit_tools::simple_time_reverser(num_generations), true, ); - let mut tables_buffered_tskit = crate::tskit_tools::convert_to_tskit_minimal( + let tables_buffered_tskit = crate::tskit_tools::convert_to_tskit_minimal( &tables_buffered, &is_sample_buffered, crate::tskit_tools::simple_time_reverser(num_generations), true, ); - let mut sorted_ts = tables_to_treeseq(&mut tables_sorted_tskit); - let mut buffered_ts = tables_to_treeseq(&mut tables_buffered_tskit); - unsafe { - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(sorted_ts.as_ptr())); - assert_eq!(500, tskr::tsk_treeseq_get_num_samples(buffered_ts.as_ptr())); - assert_eq!( - tskr::tsk_treeseq_get_num_trees(sorted_ts.as_ptr()), - tskr::tsk_treeseq_get_num_trees(buffered_ts.as_ptr()) - ); - tskr::tsk_treeseq_free(sorted_ts.as_mut_ptr()); - tskr::tsk_treeseq_free(buffered_ts.as_mut_ptr()); - } + let sorted_ts = tables_sorted_tskit + .tree_sequence(tskit::TreeSequenceFlags::default()) + .unwrap(); + let buffered_ts = tables_buffered_tskit + .tree_sequence(tskit::TreeSequenceFlags::default()) + .unwrap(); + assert_eq!(500, sorted_ts.num_samples()); + assert_eq!(500, buffered_ts.num_samples()); + assert_eq!(sorted_ts.num_trees(), buffered_ts.num_trees()); } } diff --git a/src/tskit_tools.rs b/src/tskit_tools.rs index df075290..86cb2cb4 100644 --- a/src/tskit_tools.rs +++ b/src/tskit_tools.rs @@ -49,6 +49,7 @@ pub fn simple_time_reverser(x: Time) -> Box f64> { /// # Example /// /// ``` +/// use tskit::TableAccess; /// let mut tables = forrustts::TableCollection::new(100).unwrap(); /// tables.add_node(0, 0).unwrap(); // Add a parent node at time 0 /// tables.add_node(1, 0).unwrap(); // Add a child node at time 1 @@ -96,7 +97,7 @@ pub fn convert_to_tskit_minimal( } if build_indexes { - tsk_tables.build_index(0).unwrap(); + tsk_tables.build_index().unwrap(); } tsk_tables @@ -138,6 +139,7 @@ fn swap_with_empty(v: &mut Vec) { /// # Example /// /// ``` +/// use tskit::TableAccess; /// let mut tables = forrustts::TableCollection::new(100).unwrap(); /// tables.add_node(0, 0).unwrap(); // Add a parent node at time 0 /// tables.add_node(1, 0).unwrap(); // Add a child node at time 1 @@ -197,7 +199,7 @@ pub fn convert_to_tskit_and_drain_minimal( } if build_indexes { - tsk_tables.build_index(0).unwrap(); + tsk_tables.build_index().unwrap(); } tsk_tables @@ -207,6 +209,7 @@ pub fn convert_to_tskit_and_drain_minimal( mod tests { use super::*; + use tskit::TableAccess; #[test] fn test_convert_to_tskit() {