2 create rooted phylo tree struct and impl generic rooted tree traits

src/iter/edge_iter.rs

@@ -1,49 +1,99 @@
 use crate::node::*;
 use crate::tree::simple_rtree::*;
-use std::collections::{HashMap, HashSet};
-use itertools::Itertools;
-
+use std::collections::HashMap;
+use crate::iter::node_iter::{PostOrdNodes, PreOrdNodes};
+use crate::tree::RootedPhyloTree;
 
 pub struct PreOrdEdges
 {
-    stack: Vec<(NodeID, NodeID)>,
-    nodes: HashMap<NodeID, HashSet<NodeID>>
+    node_iter: PreOrdNodes,
+    parents: HashMap<NodeID, (Option<NodeID>, Option<EdgeWeight>)>,
 }
 
 impl PreOrdEdges
 {
-    pub fn new(_tree: &HashMap<NodeID, (EdgeWeight, NodeID)>)->Self{
-        Self { stack:vec![], nodes: HashMap::new()}
+    pub fn new(tree: &RootedPhyloTree, start_node: &NodeID)->Self{
+        Self { 
+            node_iter: PreOrdNodes::new(
+                start_node,
+                tree.get_children(),
+            ),
+            parents: tree.get_parents().into_iter()
+            .filter(|(_child_id, parent_id)| parent_id!=&&None)
+            .map(|(child_id, parent_id)| (child_id.clone(), (parent_id.clone(), tree.get_edge_weight(parent_id.as_ref().unwrap(), child_id).cloned()))).collect(),
+        }
     }
 }
 
 impl Iterator for PreOrdEdges
 {
-    type Item = (NodeID, NodeID);
+    type Item = (NodeID, NodeID, Option<EdgeWeight>);
 
     fn next(&mut self)->Option<Self::Item>{
-        todo!();
+        while let Some(next_node) = self.node_iter.next() {
+            match next_node {
+                0 => {
+                    continue;
+                }
+                _ => {
+                    let parents = self.parents.get(&next_node).unwrap();
+                    return Some((parents.0.unwrap(), next_node, parents.1));
+                }
+            }
+        }
+        None
     }
 }
 
 pub struct PostOrdEdges
 {
-    stack: Vec<(NodeID, NodeID)>,
-    nodes: HashMap<NodeID, HashSet<NodeID>>
+    stack: Vec<(NodeID, NodeID, Option<EdgeWeight>)>,
+    node_iter: PostOrdNodes,
+    children: HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>>,
+    parents: HashMap<NodeID, Option<NodeID>>,
 }
 
 impl PostOrdEdges
 {
-    pub fn new(_tree: &HashMap<NodeID, (EdgeWeight, NodeID)>)->Self{
-        Self { stack:vec![], nodes: HashMap::new()}
+    pub fn new(tree: &RootedPhyloTree, start_node: &NodeID)->Self{
+        Self { 
+            stack:vec![], 
+            node_iter: PostOrdNodes::new(
+                start_node,
+                tree.get_children(),
+            ),
+            children: tree.get_children().clone(),
+            parents: tree.get_parents().clone(),
+        }
     }
 }
 
 impl Iterator for PostOrdEdges
 {
-    type Item = (NodeID, NodeID);
+    type Item = (NodeID, NodeID, Option<EdgeWeight>);
 
     fn next(&mut self)->Option<Self::Item>{
-        todo!();
-    }
+        match self.stack.pop(){
+            Some((n1, n2, w)) => Some((n1, n2, w)),
+            None => {
+                match self.node_iter.next(){
+                    Some(node_id) => {
+                        let node_id_parent = self.parents.get(&node_id).unwrap();
+                        match node_id_parent {
+                            Some(parent_id) => {
+                                let mut w: Option<EdgeWeight> = None;
+                                for (child_node_id, weight) in self.children.get(parent_id).unwrap(){
+                                    if child_node_id==&node_id{
+                                        w = *weight;
+                                    }
+                                }
+                                Some((*parent_id, node_id, w))
+                            },
+                            None => None,
+                        }
+                    },
+                    None => None
+                }
+            }
+        }    }
 }

src/iter/node_iter.rs

@@ -3,7 +3,6 @@ use crate::tree::simple_rtree::*;
 use std::collections::{HashMap, HashSet};
 use itertools::Itertools;
 
-
 pub struct PreOrdNodes
 {
     stack: Vec<NodeID>,
@@ -12,8 +11,11 @@ pub struct PreOrdNodes
 
 impl PreOrdNodes
 {
-    pub fn new(start_node_id: &NodeID, _tree: &HashMap<NodeID, (EdgeWeight, NodeID)>)->Self{
-        Self { stack:vec![*start_node_id], nodes: HashMap::new()}
+    pub fn new(start_node_id: &NodeID, children: &HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>>)->Self{
+        Self { stack:vec![*start_node_id], nodes: children.iter()
+            .map(|(k, v)| (*k, v.iter()
+                .map(|ni| ni.0).collect::<HashSet<NodeID>>()))
+            .collect()}
     }
 }
 
@@ -22,7 +24,16 @@ impl Iterator for PreOrdNodes
     type Item = NodeID;
 
     fn next(&mut self)->Option<Self::Item>{
-        todo!()
+        match self.stack.pop() {
+            Some(node_id) => {
+                let children_ids:HashSet<NodeID> = self.nodes.get(&node_id).cloned().expect("Invalid Node ID!");
+                for child_node_id in children_ids.into_iter().sorted(){
+                    self.stack.push(child_node_id)
+            }
+            Some(node_id)
+            }
+            None => None,
+        }
     }
 }
 
@@ -34,8 +45,11 @@ pub struct PostOrdNodes
 
 impl PostOrdNodes
 {
-    pub fn new(start_node_id: &NodeID, _tree: &HashMap<NodeID, (EdgeWeight, NodeID)>)->Self{
-        Self { stack:vec![*start_node_id], nodes: HashMap::new()}
+    pub fn new(start_node_id: &NodeID, children: &HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>>)->Self{
+        Self { stack:vec![*start_node_id], nodes: children.iter()
+                .map(|(k, v)| (*k, v.iter()
+                    .map(|ni| ni.0).collect::<HashSet<NodeID>>()))
+                .collect()}
     }
 }
 
@@ -44,6 +58,18 @@ impl Iterator for PostOrdNodes
     type Item = NodeID;
 
     fn next(&mut self)->Option<Self::Item>{
-        todo!()
+        while let Some(node_id) = self.stack.pop()  {
+            if self.nodes.contains_key(&node_id){
+                self.stack.push(node_id);
+                let children = self.nodes.remove(&node_id).unwrap();
+                for child_id in children.into_iter().sorted(){
+                    self.stack.push(child_id)
+                }
+            }
+            else{
+                return Some(node_id)
+            }
+        }
+        None
     }
-}
+}

src/lib.rs

@@ -1,4 +1,59 @@
 pub mod node;
 pub mod tree;
 pub mod taxa;
-pub mod iter;
+pub mod iter;
+
+
+#[cfg(test)]
+mod tests {
+    use crate::tree::{RootedPhyloTree, simple_rtree::SimpleRTree};
+    #[test]
+    fn read_small_tree() {
+        let input_str = String::from("((A,B),(C,D));");
+        let tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn read_big_tree() {
+        let input_str = String::from("(0,(1,(2,(3,(4,(5,(6,(7,(8,(9,(10,(11,(12,(13,(14,(15,(16,(17,(18,(19,(20,(21,(22,(23,(24,(25,(26,(27,(28,(29,(30,(31,(32,(33,(34,(35,(36,(37,(38,(39,(40,(41,(42,(43,(44,(45,(46,(47,(48,(49,(50,(51,(52,(53,(54,(55,(56,(57,(58,(59,(60,(61,(62,(63,(64,(65,(66,(67,(68,(69,(70,(71,(72,(73,(74,(75,(76,(77,(78,(79,(80,(81,(82,(83,(84,(85,(86,(87,(88,(89,(90,(91,(92,(93,(94,(95,(96, (97,98))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))");
+        let tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn read_smalllw_tree() {
+        let input_str = String::from("((A:0.12,B:12),(C:10,D:0.001));");
+        let tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn read_smallfw_tree() {
+        let input_str = String::from("((A:0.12,B:12):10,(C:15,D:0.001):20);");
+        let tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn read_smallfwfl_tree() {
+        let input_str = String::from("((A:0.12,B:12)E:10,(C:15,D:0.001)F:20)G;");
+        let mut tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+        tree.reroot_at_node(&1);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn reroot_node_smallfwfl_tree() {
+        let input_str = String::from("((A:0.12,B:12)E:10,(C:15,D:0.001)F:20)G;");
+        let mut tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+        tree.reroot_at_node(&1);
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn reroot_edge_smallfwfl_tree() {
+        let input_str = String::from("((A:0.12,B:12)E:10,(C:15,D:0.001)F:20)G;");
+        let mut tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+        tree.reroot_at_edge((&1, &0), (None, None));
+        dbg!(tree.to_newick());
+    }
+
+}

src/node.rs

@@ -1,12 +1,59 @@
+use std::fmt::{Debug, Display};
+
 pub type NodeID = usize;
-pub type NodeType = bool; // True for leaves, false for internal nodes
+// pub type NodeType = bool; // True for leaves, false for internal nodes
+
+#[derive(Clone, PartialEq, Eq)]
+pub enum NodeType{
+    Internal(Option<String>),
+    Leaf(Option<String>),
+}
+
+impl NodeType{
+    pub fn new(is_leaf: bool, taxa: Option<String>)->Self{
+        match is_leaf {
+            true => Self::Leaf(taxa),
+            false => Self::Internal(taxa),
+        }
+    }
+
+    pub fn is_leaf(&self)->bool{
+        match self {
+            NodeType::Internal(_taxa) => false,
+            NodeType::Leaf(_taxa) => true,
+        }
+    }
 
-pub trait Node {
-    fn is_leaf(&self)->bool;
+    pub fn flip(&mut self){
+        match self {
+            NodeType::Internal(taxa) => {*self = NodeType::Leaf(taxa.clone())},
+            NodeType::Leaf(taxa) => {*self = NodeType::Internal(taxa.clone())},
+        }
+    }
+
+    pub fn taxa(&self)->String{
+        match self {
+            NodeType::Internal(taxa) => taxa.clone().unwrap_or("".to_string()),
+            NodeType::Leaf(taxa) => taxa.clone().unwrap_or("".to_string()),
+        }
+    }
+
+    pub fn node_type(&self)->String{
+        match self.is_leaf() {
+            false => "Internal".to_string(),
+            true => "Leaf".to_string(),
+        }
+    }
+}
+
+impl Debug for NodeType{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.node_type(), self.taxa())
+    }
 }
 
-impl Node for NodeType{
-    fn is_leaf(&self)->bool {
-        *self
+impl Display for NodeType{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}", self.node_type(), self.taxa())
     }
 }