8 tree edit operations nnisprtbr

src/lib.rs

@@ -6,7 +6,7 @@ pub mod iter;
 
 #[cfg(test)]
 mod tests {
-    use crate::tree::{RootedPhyloTree, simple_rtree::SimpleRTree};
+    use crate::tree::{RootedPhyloTree, simple_rtree::*};
     #[test]
     fn read_small_tree() {
         let input_str = String::from("((A,B),(C,D));");
@@ -55,5 +55,21 @@ mod tests {
         tree.reroot_at_edge((&1, &0), (None, None));
         dbg!(tree.to_newick());
     }
+    #[test]
+    fn spr_small_tree() {
+        let input_str = String::from("((A,B),(C,D));");
+        let mut tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+        tree.spr((&0, &1), (&4, &5), (None, None));
+        dbg!(tree.to_newick());
+    }
+    #[test]
+    fn induce_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 mut tree = RootedPhyloTree::from_newick(input_str);
+        dbg!(tree.to_newick());
+        let x = tree.induce_tree(vec!["2".to_string(), "4".to_string(), "8".to_string(), "10".to_string()]);
+        dbg!(x.to_newick());
+    }
 
 }

src/tree.rs

@@ -127,7 +127,7 @@ impl RootedPhyloTree{
 impl SimpleRTree for RootedPhyloTree{
     fn add_node(&mut self)->NodeID{
         // New node id
-        let node_id = self.nodes.len();
+        let node_id = self.nodes.keys().max().unwrap_or(self.get_root())+&1;
         // add entry of node in parents and children fields
         self.nodes.insert(node_id, NodeType::Internal(None));
         self.parents.insert(node_id, None);
@@ -261,9 +261,14 @@ impl SimpleRTree for RootedPhyloTree{
 
     fn prune(&mut self, node_id: &NodeID)-> Box<dyn SimpleRTree>{
         let root= *node_id;
+        let root_parent = self.get_node_parent(node_id).expect("Node has no parent! Clean tree first...");
+        self.children.entry(*root_parent).or_default().retain(|(child_id, _w)| *child_id!=root);
+        self.parents.insert(root, None);
+
         let mut nodes: HashMap<NodeID, NodeType>= HashMap::new();
         let mut children: HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>> = HashMap::new();
         let mut parents: HashMap<NodeID, Option<NodeID>> = HashMap::new();
+
         for decsendant_node_id in self.iter_node_pre(node_id){
             nodes.insert(decsendant_node_id, self.nodes.remove(&decsendant_node_id).expect("Invalid NodeID!").clone());
             children.insert(decsendant_node_id, self.children.remove(&decsendant_node_id).expect("Invalid NodeID!").clone());
@@ -347,6 +352,8 @@ impl SimpleRTree for RootedPhyloTree{
         self.parents.insert(new_node_id, Some(edge.0.clone()));
         self.children.entry(new_node_id).or_default().push((edge.1.clone(), edge_weights.1));
         self.parents.insert(edge.1.clone(), Some(new_node_id));
+        self.children.entry(edge.0.clone()).or_default().retain(|(id, _w)| id!=edge.1);
+        self.children.entry(edge.0.clone()).or_default().push((new_node_id, edge_weights.0));
         new_node_id
     }
 
@@ -432,4 +439,35 @@ impl SimpleRTree for RootedPhyloTree{
         }).collect();
     }
 
+}
+
+impl RPhyTree for RootedPhyloTree{
+    fn induce_tree(&self, taxa: Vec<String>)->Box<dyn RPhyTree>{
+        let mut nodes: HashMap<NodeID, NodeType> = HashMap::new();
+        let leaf_ids = self.get_nodes()
+            .iter()
+            .filter(|(_id, n_type)| taxa.contains(&n_type.taxa()))
+            .map(|(id, _)| (id));
+        for id in leaf_ids{
+            nodes.insert(id.clone(), self.get_node(id).clone());
+            nodes.extend(self.get_ancestors_pre(id).iter().map(|node_id| (node_id.clone(), self.get_node(node_id).clone())).collect::<HashMap<NodeID, NodeType>>());
+        }
+        let root = self.get_mrca(nodes.keys().collect_vec());
+        let children: HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>> = nodes.keys()
+            .map(|id| (id.clone(), self.get_node_children(id).into_iter().filter(|(child_id, _)| nodes.contains_key(child_id)).map(|i| i.clone()).collect_vec()))
+            .collect();
+        let mut parents: HashMap<NodeID, Option<NodeID>> = nodes.keys()
+            .map(|id| (id.clone(), self.get_node_parent(id).cloned()))
+            .collect();
+        parents.insert(root.clone(), None);
+        Box::new(
+            RootedPhyloTree{
+                root,
+                nodes,
+                children,
+                parents,
+            }
+        )
+    }
+
 }

src/tree/simple_rtree.rs

@@ -62,6 +62,14 @@ pub trait SimpleRTree {
         false
     }
 
+    /// Remove all weights
+    fn unweight(&mut self){
+        let edge_iter = self.iter_edges_pre(self.get_root());
+        for edge in edge_iter{
+            self.set_edge_weight(&edge.0, &edge.1, None);
+        }
+    }
+
     /// Get all node-child relationships
     fn get_children(&self)->&HashMap<NodeID, Vec<(NodeID, Option<EdgeWeight>)>>;
 
@@ -191,13 +199,17 @@ pub trait SimpleRTree {
         let node  = self.get_node(node_id);
         let mut tmp = String::new();
         if !self.get_node_children(node_id).is_empty(){
-            tmp.push('(');
+            if self.get_node_children(node_id).len()>1{
+                tmp.push('(');
+            }
             for (child_id, w) in self.get_node_children(node_id){
                 let child_str = format!("{},", self.subtree_to_newick(child_id, *w));
                 tmp.push_str(&child_str);
             }
             tmp.pop();
-            tmp.push(')');
+            if self.get_node_children(node_id).len()>1{
+                tmp.push(')');
+            }
         }
         tmp.push_str(&print_node(node, edge_weight));
         tmp
@@ -211,3 +223,15 @@ pub trait SimpleRTree {
     /// Increment all node_ids
     fn incerement_ids(&mut self, value: &usize);
 }
+
+pub trait RPhyTree:SimpleRTree {
+    /// SPR function
+    fn spr(&mut self, edge1: (&NodeID, &NodeID), edge2: (&NodeID, &NodeID), edge2_weights: (Option<EdgeWeight>, Option<EdgeWeight>)){
+        let graft_edge_weight = self.get_edge_weight(edge1.0, edge1.1).cloned();
+        let pruned_tree = self.prune(edge1.1);
+        self.graft(pruned_tree, edge2, edge2_weights, graft_edge_weight);
+    }
+
+    /// Induce tree
+    fn induce_tree(&self, taxa: Vec<String>)->Box<dyn RPhyTree>;
+}