Merge pull request #17 from FinnOD/load-branch-lengths

Load branch lengths

.gitignore

@@ -145,3 +145,7 @@ wheelhouse/
 
 # Node modules
 node_modules/
+
+# venv
+bin/*
+pyvenv.cfg

phylodm/__init__.py

@@ -80,6 +80,9 @@ def add_edge(self, parent_id: int, child_id: int, length: float):
         """
         return self._rs.add_edge(parent_id=parent_id, child_id=child_id, length=length)
 
+    def update_edge_lengths(self, new_edge_lengths: np.ndarray):
+        return self._rs.update_edge_lengths(lengths=new_edge_lengths)
+
     def get_nodes(self) -> List[int]:
         """Return all node indexes in the tree."""
         return self._rs.get_nodes()

src/pdm.rs

@@ -37,7 +37,7 @@ pub struct PDM {
 impl PDM {
     /// Return the number of nodes (leaf + internal) in the tree.
     #[must_use]
-    fn n_nodes(&self) -> usize {
+    pub fn n_nodes(&self) -> usize {
         self.nodes.len()
     }
 
@@ -160,6 +160,20 @@ impl PDM {
         self.get_node_mut(child).set_parent(parent, length);
     }
 
+    /// Update the edge lengths of a tree
+    ///
+    /// # Arguments
+    ///
+    /// * `lengths`: - Slice of `Edge`s
+    ///
+    pub fn update_edge_lengths(&mut self, lengths: &[Edge]) -> Result<(), PhyloErr>{
+        for i in 0..self.n_nodes() {
+            self.get_node_mut((&self.nodes[i]).id).set_parent_distance(lengths[i]);
+        }
+        self.compute_row_vec()?; // This appears to be unnecessary.
+        Ok(())
+    }
+
     /// Set the depth of each node in the tree.
     fn assign_node_depth(&mut self) -> Result<(), PhyloErr> {
         // Iterate over each node to make sure there is only one root node.

src/python.rs

@@ -1,4 +1,4 @@
-use numpy::{PyArray2, ToPyArray};
+use numpy::{PyArray1, PyArray2, ToPyArray};
 use pyo3::{Py, pyclass, pymethods, pymodule, PyResult, Python, types::PyModule};
 use pyo3::exceptions::PyValueError;
 
@@ -46,6 +46,18 @@ impl PhyloDM {
         );
     }
 
+    pub fn update_edge_lengths(&mut self, lengths: &PyArray1<f64>) -> PyResult<()> {
+        let lengths_len = lengths.len();
+        let n_nodes = self.tree.n_nodes();
+        if lengths_len != n_nodes {
+            return Err(PyValueError::new_err(format!("Incorrect length of branch lengths. Tree has {n_nodes}. Received {lengths_len}.")));
+        }
+        let binding = lengths.to_vec().unwrap();
+        let new_lengths_vec: Vec<Edge> = binding.iter().map(|x| Edge(*x)).collect();
+        let _ = self.tree.update_edge_lengths(&new_lengths_vec);
+        Ok(())
+    }
+
     pub fn get_nodes(&self) -> Vec<usize> {
         let mut out = Vec::new();
         for node in &self.tree.nodes {

src/tree/node.rs

@@ -54,6 +54,12 @@ impl Node {
         self.parent_distance = Some(length);
     }
 
+    // Set the distance to parent.
+    // Can the root have a distance?
+    pub fn set_parent_distance(&mut self, length: Edge) {
+        self.parent_distance = Some(length);
+    }
+
     /// Set the depth of this node.
     pub fn set_depth(&mut self, depth: NodeDepth) {
         self.depth = Some(depth);

test/test_phylodm.py

@@ -146,6 +146,32 @@ def test_load_from_dendropy_with_trifurication(self):
         self.assertTrue(test_tree['taxa'] == tuple(pdm.taxa()))
         return
 
+    def test_tree_set_branch_lengths(self):
+        test_tree = get_test_tree(10, trifurication=True)
+        pdm = PhyloDM.load_from_dendropy(test_tree['tree'])
+
+        print('\n\n')
+        print('default')
+        dm = pdm.dm(norm=False)
+        print(dm)
+        print('pdm length ', pdm.length())
+        print('test length', test_tree['length'])
+
+        print('\nset 1')
+        num_nodes = len(pdm.get_nodes())
+        pdm.update_edge_lengths(np.ones((num_nodes)))
+        dm2 = pdm.dm(norm=False)
+        print('pdm length ', pdm.length())
+        print(dm2)
+
+        print('\nset 0')
+        num_nodes = len(pdm.get_nodes())
+        pdm.update_edge_lengths(np.zeros((num_nodes)))
+        dm3 = pdm.dm(norm=False)
+        print('pdm length ', pdm.length())
+        print(dm3)
+
+
     # def test_tree_with_bootstraps_from_newick(self):
     #     with tempfile.TemporaryDirectory() as tmp_dir:
     #         tmp_dir = Path(tmp_dir)

tests/test_bl_7.tree

@@ -0,0 +1 @@
+((T8:7,T2:7):7,((T9:7,(T10:7,T4:7):7):7,((T6:7,((T3:7,T1:7):7,T7:7):7):7,T5:7):7):7):7;

tests/tree.rs

@@ -1,7 +1,7 @@
 #[cfg(test)]
 mod tests {
     use phylodm::PDM;
-    use phylodm::tree::Taxon;
+    use phylodm::tree::{Taxon, Edge};
 
     #[test]
     fn test_tree_dm_twice() {
@@ -169,4 +169,49 @@ mod tests {
             ]
         );
     }
+
+    #[test]
+    fn test_set_lengths() {
+
+        // Load in the default tree from test.tree
+        // Then calculate the distance matrix as arr_normal
+        let mut tree_normal = PDM::default();
+        let _ = tree_normal.load_from_newick_path("tests/test.tree");
+        let (_taxon, arr_normal) = tree_normal.matrix(false).unwrap();
+        let tree_normal_length = tree_normal.length();
+
+        // Load in the modified tree with all branch lengths = 7
+        // Then calculate the distance matrix as arr_bl7
+        let mut tree_bl7 = PDM::default();
+        let _ = tree_bl7.load_from_newick_path("tests/test_bl_7.tree");
+        let (_taxon, arr_bl7) = tree_bl7.matrix(false).unwrap();
+        let tree_bl7_length = tree_bl7.length();
+
+        // Use the default tree and modify all branch lengths to be 7
+        // Then calculate the distance matrix as arr_modified_7
+        let new_lengths: Vec<Edge> = vec![Edge(7.0); tree_normal.n_nodes()];
+        let _ = tree_normal.update_edge_lengths(&new_lengths);
+        let (_taxon, arr_modified_7) = tree_normal.matrix(false).unwrap();
+        let tree_modified_7_length: Edge = tree_normal.length();
+
+
+        println!("\n\nNormal");
+        println!("{:?}", arr_normal);
+        println!("{:?}", tree_normal_length);
+
+        println!("Branch lengths 7 (from file)");
+        println!("{:?}", arr_bl7);
+        println!("{:?}", tree_bl7_length);
+
+        println!("Branch lengths 7 (modified)");
+        println!("{:?}", arr_modified_7);
+        println!("{:?}", tree_modified_7_length);
+
+        // Somehow the modified tree is 7 units longer that the one read from file.
+        assert_eq!(arr_bl7, arr_modified_7);
+        assert_eq!(tree_bl7_length.0, tree_modified_7_length.0);
+
+
+    }
+
 }