Skip to content

Commit

Permalink
Moving optimisation code into Tree impl
Browse files Browse the repository at this point in the history
  • Loading branch information
@jhellewell14
jhellewell14 committed Feb 26, 2024
1 parent fc450d4 commit be110b7
Show file tree
Hide file tree
Showing 4 changed files with 140 additions and 59 deletions.
7 changes: 0 additions & 7 deletions listeria_simple.fasta
View file

This file was deleted.

2 changes: 1 addition & 1 deletion src/cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use clap::Parser;
#[command(version, about, long_about = None)]
pub struct Args {
/// Alignment file in FASTA format
#[arg(short, long, default_value = "listeria0.aln")]
#[arg(short, long, default_value = "tests/test_files_in/listeria0.aln")]
pub alignment: String,

/// Write the likelihood of the tree and alignment, do not optimise
Expand Down
86 changes: 86 additions & 0 deletions src/dspsa.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use rand::Rng;
use crate::Tree;

pub fn phi(v: &[f64]) -> Vec<f64> {
v.iter().enumerate().map(|(i, value)| {
Expand Down Expand Up @@ -27,4 +28,89 @@ pub fn peturbation_vec(n: usize) -> Vec<f64> {
}).collect();
delta[0] = 0.0;
delta
}

pub fn theta_change(pivec: &Vec<f64>, delta: &Vec<f64>, plus: bool) -> Vec<usize> {

let zip = pivec.iter().zip(delta.iter());

match plus {
true => {
zip
.map(|(x, y)| (x + (y / 2.0)).round() as usize)
.collect()
},
false => {
zip
.map(|(x, y)| (x - (y / 2.0)).round() as usize)
.collect()
}
}
}

impl Tree {
pub fn optimise(&mut self, q: &na::Matrix4<f64>, iterations: usize) {
// Convert tree vector to Vec<f64>
let mut theta: Vec<f64> = self.tree_vec.iter().map(|x| *x as f64).collect();
println!("Current tree vector is: {:?}", theta);
println!("Current likelihood is: {}", self.get_tree_likelihood());
let n: usize = theta.len();

// Tuning parameters for optimisation, will
// eventually have defaults or be passed in
let a: f64 = 2.0;
let cap_a: f64 = 2.0;
let alpha: f64 = 0.75;

// Pre-allocate vectors
let mut delta: Vec<f64> = Vec::with_capacity(n);
let mut pivec: Vec<f64> = Vec::with_capacity(n);
let mut thetaplus: Vec<usize> = Vec::with_capacity(n);
let mut thetaminus: Vec<usize> = Vec::with_capacity(n);

// Optimisation loop
for k in 0..=iterations {
println!("Optimisation step {} out of {}", k, iterations);
// Generate peturbation vector
delta = peturbation_vec(n);

// Generate pi vector
pivec = piv(&theta);

// Calculate theta+ and theta-,
// New tree vectors based on peturbation
thetaplus = theta_change(&pivec, &delta, true);
thetaminus = theta_change(&pivec, &delta, false);

// Update tree and calculate likelihoods
self.update_tree(Some(thetaplus), false);
self.update_likelihood(&q);
let lplus: f64 = self.get_tree_likelihood();

self.update_tree(Some(thetaminus), false);
self.update_likelihood(&q);
let lminus: f64 = self.get_tree_likelihood();

// Update theta based on likelihoods of theta+/-
let ldiff = lplus - lminus;

let ghat: Vec<f64> = delta.iter()
.map(|el| if !el.eq(&0.0) {el * ldiff} else {0.0}).collect();

let ak: f64 = a / (1.0 + cap_a + k as f64).powf(alpha);

// Set new theta
theta = theta.iter().zip(ghat.iter())
.map(|(theta, g)| *theta - ak * g).collect();
println!("New tree vector is: {:?}", theta);
}

// Update final tree after finishing optimisation
println!("New tree vector is: {:?}", theta);
let new_tree_vec: Vec<usize> = theta.iter().map(|x| *x as usize).collect();
println!("New tree vector is: {:?}", new_tree_vec);
self.update_tree(Some(new_tree_vec), false);
self.update_likelihood(&q);
println!("New tree likelihood is {}", self.get_tree_likelihood());
}
}
104 changes: 53 additions & 51 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,74 +51,76 @@ pub fn main() {
println!("{:?}", tr.newick());
println!("{:?}", tr.tree_vec);

if !args.no_optimise {
let mut theta: Vec<f64> = tr.tree_vec.iter().map(|x| *x as f64).collect();
let n = theta.len();
tr.optimise(&q, 5);

let a = 2.0;
let A = 2.0;
let alpha = 0.75;
// let k = 0;
// if !args.no_optimise {
// let mut theta: Vec<f64> = tr.tree_vec.iter().map(|x| *x as f64).collect();
// let n = theta.len();

let mut llvec: Vec<f64> = Vec::new();
// let a = 2.0;
// let A = 2.0;
// let alpha = 0.75;
// // let k = 0;

let start = Instant::now();
for k in 0..=200 {
println!("k: {:?}", k);
// println!("theta: {:?}", theta);
// let mut llvec: Vec<f64> = Vec::new();

// // Peturbation vector
let delta = peturbation_vec(n);
// println!("delta: {:?}", delta);
// let start = Instant::now();
// for k in 0..=200 {
// println!("k: {:?}", k);
// // println!("theta: {:?}", theta);

// // Pi vector
let pivec: Vec<f64> = piv(&theta);
// // println!("pivec: {:?}", pivec);
// // // Peturbation vector
// let delta = peturbation_vec(n);
// // println!("delta: {:?}", delta);

// // theta+/-
let thetaplus: Vec<usize> = pivec.iter().zip(delta.iter()).map(|(x, y)| (x + (y / 2.0)).round() as usize).collect();
let thetaminus: Vec<usize> = pivec.iter().zip(delta.iter()).map(|(x, y)| (x - (y / 2.0)).round() as usize).collect();
// // // Pi vector
// let pivec: Vec<f64> = piv(&theta);
// // // println!("pivec: {:?}", pivec);

// // println!("thetaplus: {:?}", thetaplus);
// // println!("thetaminus: {:?}", thetaminus);
// // // theta+/-
// let thetaplus: Vec<usize> = pivec.iter().zip(delta.iter()).map(|(x, y)| (x + (y / 2.0)).round() as usize).collect();
// let thetaminus: Vec<usize> = pivec.iter().zip(delta.iter()).map(|(x, y)| (x - (y / 2.0)).round() as usize).collect();

// // Calculate likelihood at theta trees
tr.update_tree(Some(thetaplus), false);
// // println!("tree changes: {:?}", tr.changes);
tr.update_likelihood(&q);
let x1 = tr.get_tree_likelihood();
// // println!("thetaplus ll: {:?}", x1);
// // // println!("thetaplus: {:?}", thetaplus);
// // // println!("thetaminus: {:?}", thetaminus);

tr.update_tree(Some(thetaminus), false);
// // println!("tree changes: {:?}", tr.changes);
tr.update_likelihood(&q);
let x2 = tr.get_tree_likelihood();
// // println!("thetaminus ll: {:?}", x2);
// // // Calculate likelihood at theta trees
// tr.update_tree(Some(thetaplus), false);
// // // println!("tree changes: {:?}", tr.changes);
// tr.update_likelihood(&q);
// let x1 = tr.get_tree_likelihood();
// // // println!("thetaplus ll: {:?}", x1);

// // Calculations to work out new theta
let ldiff = x1 - x2;
let ghat: Vec<f64> = delta.iter().map(|el| if !el.eq(&0.0) {el * ldiff} else {0.0}).collect();
// tr.update_tree(Some(thetaminus), false);
// // // println!("tree changes: {:?}", tr.changes);
// tr.update_likelihood(&q);
// let x2 = tr.get_tree_likelihood();
// // // println!("thetaminus ll: {:?}", x2);

let ak = a / (1.0 + A + k as f64).powf(alpha);
// // // Calculations to work out new theta
// let ldiff = x1 - x2;
// let ghat: Vec<f64> = delta.iter().map(|el| if !el.eq(&0.0) {el * ldiff} else {0.0}).collect();

theta = theta.iter().zip(ghat.iter()).map(|(theta, g)| *theta - ak * g).collect();
// let ak = a / (1.0 + A + k as f64).powf(alpha);

llvec.push(x1);
// theta = theta.iter().zip(ghat.iter()).map(|(theta, g)| *theta - ak * g).collect();

// // println!("ghat: {:?}", ghat);
// llvec.push(x1);

}
// // // println!("ghat: {:?}", ghat);

let out: Vec<f64> = phi(&theta).iter().map(|x| x.round()).collect();
println!("final theta: {:?}", out);
// }

println!("{:?}", &llvec);
// println!("{:?}", &llvec[95..100]);
}
let end = Instant::now();
// let out: Vec<f64> = phi(&theta).iter().map(|x| x.round()).collect();
// println!("final theta: {:?}", out);

eprintln!("Done in {}s", end.duration_since(start).as_secs());
eprintln!("Done in {}ms", end.duration_since(start).as_millis());
eprintln!("Done in {}ns", end.duration_since(start).as_nanos());
// println!("{:?}", &llvec);
// // println!("{:?}", &llvec[95..100]);
// }
// let end = Instant::now();

// eprintln!("Done in {}s", end.duration_since(start).as_secs());
// eprintln!("Done in {}ms", end.duration_since(start).as_millis());
// eprintln!("Done in {}ns", end.duration_since(start).as_nanos());

}

0 comments on commit be110b7

Please sign in to comment.