Skip to content

Commit

Permalink
Update to latest tskit (rust). Closes #82 (#84)
Browse files Browse the repository at this point in the history 
* Pin to precise tskit version. Closes #79.

* Update to new tskit_rust API.

* Remove use of unsafe code in tests.
  • Loading branch information
@molpopgen
molpopgen committed Apr 22, 2021
1 parent 7702c39 commit c82a7f0
Show file tree
Hide file tree
Showing 5 changed files with 71 additions and 106 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
2 changes: 1 addition & 1 deletion forrustts_examples/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ edition = "2018"
path = "../"

[dependencies]
tskit = "0.1.1"
tskit = "=0.2.0"
clap = "~2.27.0"

[[bin]]
Expand Down
17 changes: 10 additions & 7 deletions forrustts_examples/neutral_wf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
149 changes: 54 additions & 95 deletions src/test_simplify_tables.rs
View file
Original file line number Diff line number Diff line change
@@ -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::*;
Expand All @@ -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<tskr::tsk_treeseq_t> {
let mut tsk_ts: MaybeUninit<tskr::tsk_treeseq_t> = 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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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());
}
}
7 changes: 5 additions & 2 deletions src/tskit_tools.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ pub fn simple_time_reverser(x: Time) -> Box<dyn Fn(Time) -> 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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -138,6 +139,7 @@ fn swap_with_empty<T>(v: &mut Vec<T>) {
/// # 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
Expand Down Expand Up @@ -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
Expand All @@ -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() {
Expand Down

0 comments on commit c82a7f0

Please sign in to comment.