diff --git a/book/src/SUMMARY.md b/book/src/SUMMARY.md index 033f77d73..ca59c79c5 100644 --- a/book/src/SUMMARY.md +++ b/book/src/SUMMARY.md @@ -12,3 +12,4 @@ - [Working with tree sequences](./working_with_tree_sequences.md) - [Initialization from a table collection](./tree_sequence_from_table_collection.md) - [Iterating over trees](./tree_sequence_iterate_trees.md) + - [Working with trees](./tree_sequence_tree.md) diff --git a/book/src/tree_sequence_tree.md b/book/src/tree_sequence_tree.md new file mode 100644 index 000000000..9a92e0472 --- /dev/null +++ b/book/src/tree_sequence_tree.md @@ -0,0 +1,38 @@ +## Working with trees + +Iterating over a tree sequence returns instances of `tskit::Tree`. +This type is immutable. +No access is provided to the underlying pointers. + +The API provides a set of iterators over the data in a tree. + +For example, to collect the siblings of each node into a vector: + +```rust, noplaygound, ignore +{{#include ../../tests/book_trees.rs:iterate_node_siblings}} +``` + +We may do the same calculation using API elements giving `&[NodeId]` +(slices of node ids). +This method more closely matches the `tskit-c` API. + +```rust, noplaygound, ignore +{{#include ../../tests/book_trees.rs:iterate_node_siblings_via_arrays}} +``` + +This approach is more complex: + +* Slice element access is via `usize` (`size_t` in C/C++). +* Row ids are `i32` (`tsk_id_t` in `tskit-c`) behind the newtypes. +* `tskit` implements `TryFrom` to help you out, forcing + you to reckon with the fact that conversion from `i32` + to `usize` is fallible. + +These conversion semantics require us to manually handle all possible error +paths at each step. + +We can have an intermediate level of complexity using getters from the tree arrays: + +```rust, noplaygound, ignore +{{#include ../../tests/book_trees.rs:iterate_node_siblings_via_array_getters}} +``` diff --git a/tests/book_trees.rs b/tests/book_trees.rs index 894c4bb69..d9a4b9995 100644 --- a/tests/book_trees.rs +++ b/tests/book_trees.rs @@ -73,4 +73,77 @@ fn initialize_from_table_collection() { // _tree is a tskit::Tree } // ANCHOR_END: iterate_trees + + let mut tree_iterator = treeseq.tree_iterator(TreeFlags::default()).unwrap(); + // ANCHOR: iterate_node_siblings + // This is an enum defining supported + // traversal orders through a Tree. + use tskit::NodeTraversalOrder; + while let Some(tree) = tree_iterator.next() { + for node in tree.traverse_nodes(NodeTraversalOrder::Preorder) { + if let Some(parent) = tree.parent(node) { + // Collect the siblings of node into a Vec + // The children function returns another iterator + let _siblings = if let Some(child_iterator) = tree.children(parent) { + child_iterator + .filter(|child| child != &node) + .collect::>() + } else { + // assign empty vector + vec![] + }; + } + } + } + // ANCHOR_END: iterate_node_siblings + + let mut tree_iterator = treeseq.tree_iterator(TreeFlags::default()).unwrap(); + // ANCHOR: iterate_node_siblings_via_arrays + while let Some(tree) = tree_iterator.next() { + let parents = tree.parent_array(); + let rsibs = tree.right_sib_array(); + let lchildren = tree.left_child_array(); + for node in tree.traverse_nodes(NodeTraversalOrder::Preorder) { + let mut siblings = vec![]; + assert!(!node.is_null()); + if let Some(parent) = parents.get(usize::try_from(node).unwrap()) { + if !parent.is_null() { + if let Some(child) = lchildren.get(usize::try_from(*parent).unwrap()) { + let mut u = *child; + while !u.is_null() { + if u != node { + siblings.push(u); + } + if let Some(sib) = rsibs.get(usize::try_from(u).unwrap()) { + u = *sib; + } + } + } + } + } + } + } + // ANCHOR_END: iterate_node_siblings_via_arrays + + let mut tree_iterator = treeseq.tree_iterator(TreeFlags::default()).unwrap(); + // ANCHOR: iterate_node_siblings_via_array_getters + while let Some(tree) = tree_iterator.next() { + for node in tree.traverse_nodes(NodeTraversalOrder::Preorder) { + let mut siblings = vec![]; + if let Some(parent) = tree.parent(node) { + if let Some(child) = tree.left_child(parent) { + let mut u = child; + while !u.is_null() { + if u != node { + siblings.push(u); + } + if let Some(sib) = tree.right_sib(u) { + u = sib; + } + } + } + } + } + } + // ANCHOR_END: iterate_node_siblings_via_array_getters }