From a0043d748555b8fce1a62fe267b88803b58f7399 Mon Sep 17 00:00:00 2001
From: Joel Hellewell <joel@ebi.ac.uk>
Date: Wed, 1 May 2024 15:49:35 +0100
Subject: [PATCH] More tidying and renaming

---
 build.rs                       |  4 ++--
 src/{dspsa.rs => hillclimb.rs} | 18 +++++++++++++-----
 src/lib.rs                     |  4 ++--
 src/likelihoods.rs             |  3 ++-
 src/tests.rs                   |  2 +-
 5 files changed, 20 insertions(+), 11 deletions(-)
 rename src/{dspsa.rs => hillclimb.rs} (91%)

diff --git a/build.rs b/build.rs
index e6177f2..9f6fe83 100644
--- a/build.rs
+++ b/build.rs
@@ -1,10 +1,10 @@
 fn main() {
-    cxx_build::bridge("src/newick.rs")
+    cxx_build::bridge("src/build_tree.rs")
         .file("src/phylo2vec.cpp")
         .std("c++17")
         .compile("phylo2vec");
 
-        println!("cargo:rerun-if-changed=src/newick.rs");
+        println!("cargo:rerun-if-changed=src/build_tree.rs");
         println!("cargo:rerun-if-changed=src/phylo2vec.cpp");
         println!("cargo:rerun-if-changed=src/phylo2vec.hpp");
 }
diff --git a/src/dspsa.rs b/src/hillclimb.rs
similarity index 91%
rename from src/dspsa.rs
rename to src/hillclimb.rs
index 6a575f9..a573860 100644
--- a/src/dspsa.rs
+++ b/src/hillclimb.rs
@@ -39,24 +39,32 @@ impl Tree {
         };
 
         let mut candidate_vec: Vec<usize>;
+        let mut best_vec: Option<Vec<usize>> = None;
+        let mut best_likelihood: f64 = working_tree.get_tree_likelihood();
         for k in 0..=iterations {
             
             println!("Optimisation step {} out of {}", k, iterations);
             // println!("Old vector {:?}", self.tree_vec);
             println!("Tree log likelihood: {}", self.get_tree_likelihood());
 
-            candidate_vec = hill_peturb(&self.tree_vec, 28);
+            candidate_vec = hill_peturb(&self.tree_vec, self.tree_vec.len());
             working_tree.update(&candidate_vec);
             working_tree.update_likelihood(q);
+            let new_likelihood = working_tree.get_tree_likelihood();
 
             // println!("New vector {:?}", candidate_vec);
-            println!("New likelihood {}", working_tree.get_tree_likelihood());
+            println!("New likelihood {}", new_likelihood);
 
-            if working_tree.get_tree_likelihood() > self.get_tree_likelihood() {
+            if new_likelihood > best_likelihood {
                 println!("Climbing hill!");
-                self.update(&working_tree.tree_vec);
-                self.update_likelihood(q);
+                best_vec = Some(working_tree.tree_vec.clone());
+                best_likelihood = new_likelihood;
             }
+        };
+
+        if best_vec.is_some() {
+            self.update(&best_vec.unwrap());
+            self.update_likelihood(q);
         }
     }
 
diff --git a/src/lib.rs b/src/lib.rs
index 1d42903..224e93c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,7 +5,7 @@ mod node;
 mod build_tree;
 mod tests;
 mod tree;
-mod dspsa;
+mod hillclimb;
 mod tree_iterators;
 mod tree_to_newick;
 
@@ -39,7 +39,7 @@ pub fn main() {
     // let end = Instant::now();
     tr.add_genetic_data(&args.alignment);
 
-    tr.update_likelihood_postorder(&q);
+    tr.initialise_likelihood(&q);
 
     println!("{}", tr.get_tree_likelihood());
     println!("{:?}", tr.newick());
diff --git a/src/likelihoods.rs b/src/likelihoods.rs
index ea4433b..acdae11 100644
--- a/src/likelihoods.rs
+++ b/src/likelihoods.rs
@@ -49,7 +49,7 @@ impl Tree {
 
     // Traverses tree below given node (except leaves), updating likelihood
     // Used after initial tree constructions to fill in likelihood at all internal nodes
-    pub fn update_likelihood_postorder(&mut self, rate_matrix: &na::Matrix4<f64>) {
+    pub fn initialise_likelihood(&mut self, rate_matrix: &na::Matrix4<f64>) {
         let nodes: Vec<usize> = self
             .postorder_notips(self.get_root())
             .map(|n| n.index)
@@ -72,6 +72,7 @@ impl Tree {
         }
     }
 
+    // Fetches likelihood value for a tree
     pub fn get_tree_likelihood(&self) -> f64 {
         self.mutation_lists
             .get(self.get_root().unwrap().index)
diff --git a/src/tests.rs b/src/tests.rs
index 9daafcc..dfe9b8c 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -196,7 +196,7 @@ mod tests {
 
     tr.mutation_lists = genetic_data;
 
-    tr.update_likelihood_postorder(&q);
+    tr.initialise_likelihood(&q);
     
     let old_likelihood = tr.get_tree_likelihood();