diff --git a/src/branchlength.rs b/src/branchlength.rs index 9bbb152..d51ee45 100644 --- a/src/branchlength.rs +++ b/src/branchlength.rs @@ -1,9 +1,40 @@ use crate::{rate_matrix::RateMatrix, treestate::TreeMove}; - -pub struct BranchMove{} +use rand::prelude::Distribution; +// use rand::distributions::Normal; +use statrs::distribution::Normal; +use crate::topology::Topology; +use crate::TreeState; +pub struct BranchMove{ + indices: Vec, +} impl TreeMove for BranchMove { fn generate(&self, ts: &crate::treestate::TreeState) -> crate::treestate::TreeState { - todo!() + let normal = Normal::new(0.0, 1.0).unwrap(); + let mut changes: Vec = Vec::new(); + + // This is not ideal + let mut nodes = ts.top.nodes.clone(); + + for i in self.indices.iter() { + let ind = *i; + let mut bl = nodes[ind].get_branchlen(); + bl = bl.ln() + normal.sample(&mut rand::thread_rng()); + nodes[ind].set_branchlen(bl.exp()); + changes.push(ind); + } + + let new_top = Topology{ + nodes: nodes, + tree_vec: ts.top.tree_vec.clone(), + likelihood: ts.top.likelihood, + }; + + TreeState{ + top: new_top, + mat: ts.mat, + ll: ts.ll, + changed_nodes: Some(changes), + } } } \ No newline at end of file