From fbbe614878400a2f3da79b1428fb2e5dd149f140 Mon Sep 17 00:00:00 2001
From: jhellewell14 <joel@ebi.ac.uk>
Date: Wed, 9 Oct 2024 10:09:23 +0100
Subject: [PATCH] Added Tree method to do NNI. At the moment it just does a
 NNI. What I really need is for it to suggest a NNI, calculate the likelihood,
 and then accept or reject it

---
 src/branch_optimise.rs | 15 +++++++++
 src/lib.rs             |  5 +++
 src/tree.rs            |  3 +-
 src/tree_moves.rs      | 76 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 97 insertions(+), 2 deletions(-)
 create mode 100644 src/branch_optimise.rs
 create mode 100644 src/tree_moves.rs

diff --git a/src/branch_optimise.rs b/src/branch_optimise.rs
new file mode 100644
index 0000000..e5a3bb6
--- /dev/null
+++ b/src/branch_optimise.rs
@@ -0,0 +1,15 @@
+use argmin::core::{Error, CostFunction, Gradient, Hessian};
+use crate::tree::Tree;
+use crate::rate_matrix::RateMatrix;
+
+// impl<T: RateMatrix> CostFunction for Tree<T> {
+//     type Param = Vec<f64>;
+//     type Output = f64;
+
+//     fn cost(&self, param: &Self::Param) -> Result<Self::Output, Error> {
+        
+//         self.rate_matrix.update_params(*param);
+//         self.initialise_likelihood();
+//         return Ok(self.get_tree_likelihood())
+//     }
+// }
\ No newline at end of file
diff --git a/src/lib.rs b/src/lib.rs
index af86f2a..cd4ddc6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,10 @@ mod tree;
 mod tree_iterators;
 mod tree_to_newick;
 mod rate_matrix;
+mod branch_optimise;
+mod tree_moves;
+
+use rand::Rng;
 
 use crate::build_tree::*;
 use crate::tree::Tree;
@@ -41,4 +45,5 @@ pub fn main() {
         eprintln!("Done in {}s", end.duration_since(start).as_secs());
         eprintln!("Done in {}ms", end.duration_since(start).as_millis());
     }
+
 }
diff --git a/src/tree.rs b/src/tree.rs
index 5c8d412..2ca7c32 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -1,7 +1,6 @@
 use crate::mutation::{create_list, Mutation};
 use crate::node::Node;
-use crate::rate_matrix::{self, RateMatrix};
-use crate::rate_matrix::GTR;
+use crate::rate_matrix::{RateMatrix};
 use crate::vector_to_tree;
 use needletail::*;
 use std::collections::HashMap;
diff --git a/src/tree_moves.rs b/src/tree_moves.rs
new file mode 100644
index 0000000..8846b68
--- /dev/null
+++ b/src/tree_moves.rs
@@ -0,0 +1,76 @@
+use crate::newick_to_vector;
+use crate::tree::Tree;
+use crate::rate_matrix::RateMatrix;
+use rand::Rng;
+
+// Rooted Nearest Neighbour Interchange
+// Select random internal node - v
+// Assign new parent as parent-of-parent
+// Assign parent as having this node as parent
+
+impl<T: RateMatrix> Tree<T> {
+
+    pub fn nni(&mut self) -> () {
+        // Get vector of all internal nodes that can be swapped
+        let mut int_nodes: Vec<usize> = self.postorder_notips(self.get_root()).map(|node| node.index).collect();
+        int_nodes.pop(); // Last element is root, which we can't swap with parent so pop it off
+        let ind = rand::thread_rng().gen_range(0..int_nodes.len());
+        let child_index = int_nodes[ind]; // was sample
+
+        // Children and parent of selected node
+        let child_node_children = self.nodes[child_index].children;
+        let child_node_parent = self.nodes[child_index].parent.unwrap();
+
+        // Children and parent of selected node's parent
+        let (lp, rp) = self.nodes[child_node_parent].children;
+        let parent_node_parent = self.nodes[child_node_parent].parent;
+
+        // Assign new parent and children to selected node's parent
+        self.nodes[child_node_parent].children = child_node_children;
+        let (lc, rc) = child_node_children;
+
+        // Update parents parent's children
+        let parent_node_parent = self.nodes[child_node_parent].parent.unwrap();
+        // if pp.is_none() {
+            // Needed for when parent's parent is root
+        // }
+
+        let (lpc, rpc) = self.nodes[parent_node_parent].children;
+        if lpc == Some(child_node_parent) {
+            self.nodes[parent_node_parent].children = (Some(child_index), rpc);
+        } else {
+            self.nodes[parent_node_parent].children = (lpc, Some(child_index));
+        }
+
+        self.nodes[child_node_parent].parent = Some(child_index);
+        // Update parent of moved children
+        self.nodes[lc.unwrap()].parent = Some(child_node_parent);
+        self.nodes[rc.unwrap()].parent = Some(child_node_parent);
+
+        // Assign new parent and children to selected node
+        self.nodes[child_index].parent = Some(parent_node_parent);
+
+        if lp == Some(child_index) {
+            self.nodes[child_index].children = (Some(child_node_parent), rp);
+            self.nodes[rp.unwrap()].parent = Some(child_index);
+        } else {
+            self.nodes[child_index].children = (lp, Some(child_node_parent));
+            self.nodes[lp.unwrap()].parent = Some(child_index);
+        }
+
+        // The new parent needs adding to changes for likelihood recalculation
+        let d: usize = self.nodes[child_node_parent].depth;
+        match self.changes.get(&d) {
+            None => {
+                self.changes.insert(d, vec![child_node_parent]);
+            }
+            Some(_) => {
+                self.changes.get_mut(&d).unwrap().push(child_node_parent);
+            }
+        }
+
+        // Update tree vector
+        self.tree_vec = newick_to_vector(&self.newick(), self.count_leaves()); 
+    }
+
+}