Skip to content

Commit

Permalink
Lots of cleaning up old code and fixing clippy suggestions
Browse files Browse the repository at this point in the history
  • Loading branch information
@jhellewell14
jhellewell14 committed Nov 29, 2024
1 parent 22ead15 commit cadebdc
Show file tree
Hide file tree
Showing 9 changed files with 81 additions and 409 deletions.
111 changes: 8 additions & 103 deletions src/genetic_data.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ use needletail::parse_fastx_file;
use crate::topology::Topology;
use ndarray::s;
use logaddexp::LogAddExp;
use crate::moves::MoveFn;
use rand::{thread_rng, Rng};

const NEGINF: f64 = -f64::INFINITY;
Expand Down Expand Up @@ -76,13 +75,11 @@ pub fn create_genetic_data(filename: &str, topology: &Topology, rate_matrix: &na
let mut seq_i = 0;
while let Some(record) = reader2.next() {
let seqrec = record.expect("Invalid record");
let mut loc_i = 0;
for e in seqrec.seq().iter() {
for (loc_i, e) in seqrec.seq().iter().enumerate() {
let cur = char_to_likelihood(&(*e as char));
for j in 0..4 {
gen_data[[seq_i, loc_i, j]] = *cur.get(j).unwrap();
}
loc_i += 1;
}
seq_i += 1;
}
Expand All @@ -95,9 +92,9 @@ pub fn create_internal_data(mut data: ndarray::ArrayBase<ndarray::OwnedRepr<f64>
topology: &Topology, rate_matrix: &na::Matrix4<f64>) ->
ndarray::ArrayBase<ndarray::OwnedRepr<f64>, ndarray::Dim<[usize; 3]>> {
// Iterate over internal nodes postorder
let mut nodes = topology.postorder_notips(topology.get_root());
let nodes = topology.postorder_notips(topology.get_root());

while let Some(node) = nodes.next() {
for node in nodes {
let i = node.get_id();
// Calculate node likelihood
let lchild = node.get_lchild().unwrap();
Expand Down Expand Up @@ -160,33 +157,20 @@ pub fn node_likelihood(seql: ndarray::ArrayBase<ndarray::ViewRepr<&f64>, ndarray
) -> ndarray::ArrayBase<ndarray::OwnedRepr<f64>, ndarray::Dim<[usize; 2]>> {

let out = ndarray::Array2::from_shape_fn((seql.dim().0, 4), |(i, j)|
child_likelihood_i(j, seql.slice(s![i, ..]), &matrixl) +
child_likelihood_i(j, seqr.slice(s![i, ..]), &matrixr));
child_likelihood_i(j, seql.slice(s![i, ..]), matrixl) +
child_likelihood_i(j, seqr.slice(s![i, ..]), matrixr));

out
}

pub const BF_DEFAULT: [f64; 4] = [0.25, 0.25, 0.25, 0.25];

pub fn likelihood(top: &Topology, gen_data: &ndarray::ArrayBase<ndarray::OwnedRepr<f64>, ndarray::Dim<[usize; 3]>>) -> f64 {
let root_likelihood = gen_data.slice(s![top.get_root().get_id(), .., .. ]);

root_likelihood
.rows()
.into_iter()
.fold(0.0, |acc, base | acc + base_freq_logse(base, BF_DEFAULT))
}

pub fn base_freq_logse(muta: ndarray::ArrayBase<ndarray::ViewRepr<&f64>, ndarray::Dim<[usize; 1]>>, bf: [f64; 4]) -> f64 {
muta.iter()
.zip(bf.iter())
.fold(0.0, |tot, (muta, bf)| tot + muta.exp() * bf)
.ln()
}
pub struct CandidateTopology{
pub new_topology: Topology,
pub changes: Option<Vec<usize>>,
}

impl Topology {
pub fn find_changes(&self, other: &Topology) -> Option<Vec<usize>> {
Expand All @@ -196,91 +180,12 @@ impl Topology {
.filter(|(a, b)| a.get_parent().ne(&b.get_parent()))
.map(|(a, b)| a.get_id())
.collect();
if out.len() > 0 {
return Some(out);
if out.is_empty() {
None
} else {
return None;
Some(out)
}
}

pub fn apply_move<T: MoveFn>(&mut self,
move_fn: T,
accept_fn: fn(&f64, &f64) -> bool,
gen_data: &mut ndarray::ArrayBase<ndarray::OwnedRepr<f64>, ndarray::Dim<[usize; 3]>>,
rate_matrix: &na::Matrix4<f64>) -> () {

// Get current likelihood, calculating if needed
if self.likelihood.is_none() {
self.likelihood = Some(likelihood(&self, gen_data));
}
let old_ll = self.likelihood.unwrap();


// Generate new candidate Topology using move function
let new_top = move_fn.generate_move(&self);

println!("Changes: {:?}", new_top.changes);

// Move did nothing, no changes needed
if new_top.changes.is_none() {
return ()
}

// Do minimal likelihood updates (and push new values into HashMap temporarily)
let nodes = new_top.new_topology.changes_iter_notips(new_top.changes.unwrap());
let mut temp_likelihoods: HashMap<usize, ndarray::ArrayBase<ndarray::OwnedRepr<f64>, ndarray::Dim<[usize; 2]>>> = HashMap::new();
for node in nodes {
// check if in HM
let lchild = node.get_lchild().unwrap();
let rchild = node.get_rchild().unwrap();
let seql: ndarray::ArrayBase<ndarray::ViewRepr<&f64>, ndarray::Dim<[usize; 2]>>;
let seqr: ndarray::ArrayBase<ndarray::ViewRepr<&f64>, ndarray::Dim<[usize; 2]>>;

match (temp_likelihoods.contains_key(&lchild), temp_likelihoods.contains_key(&rchild)) {
(true, true) => {
seql = temp_likelihoods.get(&lchild).unwrap().slice(s![.., ..]);
seqr = temp_likelihoods.get(&rchild).unwrap().slice(s![.., ..]);
},
(true, false) => {
seql = temp_likelihoods.get(&lchild).unwrap().slice(s![.., ..]);
seqr = slice_data(rchild, &gen_data);
},
(false, true) => {
seql = slice_data(lchild, &gen_data);
seqr = temp_likelihoods.get(&rchild).unwrap().slice(s![.., ..]);
},
(false, false) => {
seql = slice_data(lchild, &gen_data);
seqr = slice_data(rchild, &gen_data);
},
};

let node_ll = node_likelihood(seql, seqr,
&matrix_exp(rate_matrix, new_top.new_topology.nodes[lchild].get_branchlen()),
&matrix_exp(rate_matrix, new_top.new_topology.nodes[rchild].get_branchlen()));

temp_likelihoods.insert(node.get_id(), node_ll);
}

// Calculate whole new topology likelihood at root
let new_ll = temp_likelihoods
.get(&new_top.new_topology.get_root().get_id())
.unwrap()
.rows()
.into_iter()
.fold(0.0, |acc, base | acc + base_freq_logse(base, BF_DEFAULT));

// Likelihood decision rule
if accept_fn(&old_ll, &new_ll) {
// Drain hashmap into gen_data
for (i, ll_data) in temp_likelihoods.drain() {
gen_data.slice_mut(s![i, .., ..]).assign(&ll_data);
}
// Update Topology
self.nodes = new_top.new_topology.nodes;
self.tree_vec = new_top.new_topology.tree_vec;
self.likelihood = Some(new_ll);
}
}
}

14 changes: 6 additions & 8 deletions src/iterators.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,18 +82,17 @@ impl<'a> Topology {
};

if node.get_id().eq(&parent_left_child.unwrap()) {
return Handedness::Left;
Handedness::Left
} else {
return Handedness::Right;
Handedness::Right
}
}

pub fn swap_to_right_child(&self, node: &NodeTuple) -> Option<&NodeTuple> {
let parent = self.get_parent(node);
if parent.is_some() {
return self.get_rchild(parent.unwrap());
if let Some(par) = self.get_parent(node) {
self.get_rchild(par)
} else {
return self.get_rchild(self.get_root());
self.get_rchild(self.get_root())
}
}
}
Expand Down Expand Up @@ -139,8 +138,7 @@ impl<'a> Iterator for ChangeIter<'a> {
};
}


return Some(current_node);
Some(current_node)
}
}

Expand Down
31 changes: 11 additions & 20 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,12 @@ mod treestate;
use rate_matrix::RateMatrix;
use state_data::create_dummy_statedata;
use topology::Topology;
use treestate::TreeState;
use treestate::{TreeState, hillclimb_accept};

use crate::newick_to_vec::*;
extern crate nalgebra as na;
pub mod cli;
use crate::cli::*;
use std::env::args;
use std::time::Instant;
use crate::genetic_data::*;
use crate::moves::*;
Expand All @@ -33,11 +32,14 @@ pub fn main() {

let v = random_vector(28);

let mut t: Topology = Topology::from_vec(&v);
let t: Topology = Topology::from_vec(&v);

let p = rate_matrix::GTR::default();
let p = rate_matrix::Gtr::default();
let mut gen_data = create_genetic_data(&args.alignment, &t, &p.get_matrix());

let mge_mat = na::Matrix2::new(0.4, 0.6, 0.6, 0.4);
let mut st = create_dummy_statedata(1, &t, &mge_mat);

let mut ts = TreeState{
top: t,
mat: p,
Expand All @@ -49,26 +51,15 @@ pub fn main() {
println!("{:?}", ts.top.get_newick());
println!("{:?}", ts.top.tree_vec);

let mge_mat = na::Matrix2::new(0.4, 0.6, 0.6, 0.4);
// let mut st = create_dummy_statedata(1, &t, &mge_mat);

// let mut pp = rate_matrix::GTR::default();
// println!("{:?}", pp.get_matrix());
// update_matrix(&mut t, always_accept, &mut gen_data, &mut pp);
// println!("{:?}", pp.get_matrix());
// let mv = ChildSwap{};
// t.apply_move(mv2, hillclimb_accept, &mut gen_data, &mut p.get_matrix());

if !args.no_optimise {
let start = Instant::now();
for i in 0..0 {
for i in 0..5 {
println!{"Step {}", i};
let new_v = random_vector(27);
let mv = ExactMove{target_vector: new_v};
// let mv = ChildSwap{};
// let new_v = random_vector(27);
// let mv = ExactMove{target_vector: new_v};
let mv = ChildSwap{};
// let mv = PeturbVec{n: 1};
ts.apply_move(mv, always_accept, &mut gen_data);
// t.apply_move(mv, hillclimb_accept, &mut gen_data, &mut p.get_matrix());
ts.apply_move(mv, hillclimb_accept, &mut gen_data);

}
let end = Instant::now();
Expand Down
Loading

0 comments on commit cadebdc

Please sign in to comment.