refactor: unify vector/matrix dispatch with Tree enum

docs/development.md

@@ -238,5 +238,5 @@ devtools::load_all('./r-phylo2vec')
 
 # A small demo
 v = sample_vector(5, FALSE)
-to_newick_from_vector(v)
+to_newick(v)
 ```

phylo2vec/src/matrix/ops.rs

@@ -65,13 +65,13 @@ pub fn get_node_depths(m: &ArrayView2<f64>) -> Vec<f64> {
 ///     [0.0, 0.5, 0.2],
 /// ];
 /// // Root (node 4) has depth 0
-/// assert_eq!(get_node_depth(&m.view(), 4), 0.0);
+/// assert_eq!(get_node_depth(&m.view(), 4).unwrap(), 0.0);
 /// // Node 3 is 0.5 from root
-/// assert!((get_node_depth(&m.view(), 3) - 0.5).abs() < 1e-10);
+/// assert!((get_node_depth(&m.view(), 3).unwrap() - 0.5).abs() < 1e-10);
 /// // Leaf 0 is 0.5 + 0.3 = 0.8 from root
-/// assert!((get_node_depth(&m.view(), 0) - 0.8).abs() < 1e-10);
+/// assert!((get_node_depth(&m.view(), 0).unwrap() - 0.8).abs() < 1e-10);
 /// ```
-pub fn get_node_depth(m: &ArrayView2<f64>, node: usize) -> f64 {
+pub fn get_node_depth(m: &ArrayView2<f64>, node: usize) -> Result<f64, String> {
     let (v, bls) = parse_matrix(m);
     _get_node_depth(&v, Some(&bls), node)
 }
@@ -160,7 +160,7 @@ mod tests {
         #[case] node: usize,
         #[case] expected_depth: f64,
     ) {
-        let depth = get_node_depth(&m.view(), node);
+        let depth = get_node_depth(&m.view(), node).unwrap();
         assert!(
             (depth - expected_depth).abs() < 1e-10,
             "Expected depth {expected_depth}, got {depth}"
@@ -175,7 +175,7 @@ mod tests {
     fn test_get_node_depth_root_is_zero(#[case] n_leaves: usize) {
         let m = sample_matrix(n_leaves, false);
         let root = 2 * m.nrows(); // Root node index
-        let depth = get_node_depth(&m.view(), root);
+        let depth = get_node_depth(&m.view(), root).unwrap();
         assert_eq!(depth, 0.0, "Root should have depth 0");
     }
 

phylo2vec/src/vector/ops.rs

@@ -160,28 +160,30 @@ fn min_common_ancestor(path1: &[usize], path2: &[usize]) -> usize {
 /// use phylo2vec::vector::ops::get_common_ancestor;
 ///
 /// let v = vec![0, 1, 2, 3, 4];
-/// let mrca = get_common_ancestor(&v, 2, 3);
+/// let mrca = get_common_ancestor(&v, 2, 3).unwrap();
 /// // Newick string of v = (0,(1,(2,(3,(4,5)6)7)8)9)10;
 /// // So 2 and 3 share the MRCA 8.
 /// assert_eq!(mrca, 8);
 /// ```
-pub fn get_common_ancestor(v: &[usize], node1: usize, node2: usize) -> usize {
+pub fn get_common_ancestor(v: &[usize], node1: usize, node2: usize) -> Result<usize, String> {
     let node_max = 2 * v.len();
 
     if node1 > node_max || node2 > node_max {
-        panic!("Node indices must be within the range of the Phylo2Vec vector. Max = {node_max}, got node1 = {node1} and node2 = {node2}");
+        return Err(format!(
+              "Node indices must be within range. Max = {node_max}, got node1 = {node1} and node2 = {node2}"
+          ));
     }
 
     if node1 == node2 {
-        return node1;
+        return Ok(node1);
     }
 
     let pairs = to_pairs(v);
 
     let path1 = get_ancestry_path_of_node(&pairs, node1);
     let path2 = get_ancestry_path_of_node(&pairs, node2);
 
-    min_common_ancestor(&path1, &path2)
+    Ok(min_common_ancestor(&path1, &path2))
 }
 
 /// Generic function to calculate the depths of all nodes in a Phylo2Vec tree.
@@ -238,20 +240,24 @@ pub fn _get_node_depths(v: &[usize], bls: Option<&Vec<[f64; 2]>>) -> Vec<f64> {
 /// The root has depth 0, and depths increase as you move toward the leaves.
 ///
 /// When `bls` is `None`, all branch lengths are assumed to be 1.0 (topological depth).
-pub fn _get_node_depth(v: &[usize], bls: Option<&Vec<[f64; 2]>>, node: usize) -> f64 {
+pub fn _get_node_depth(
+    v: &[usize],
+    bls: Option<&Vec<[f64; 2]>>,
+    node: usize,
+) -> Result<f64, String> {
     let k = v.len();
     let n_leaves = k + 1;
     let n_nodes = 2 * n_leaves - 1;
 
     if node >= n_nodes {
-        panic!(
+        return Err(format!(
             "Node index out of bounds. Max node = {}, got node = {}",
             n_nodes - 1,
             node
-        );
+        ));
     }
 
-    _get_node_depths(v, bls)[node]
+    Ok(_get_node_depths(v, bls)[node])
 }
 
 /// Get the depths of all nodes in a Phylo2Vec vector (topological).
@@ -297,15 +303,15 @@ pub fn get_node_depths(v: &[usize]) -> Vec<f64> {
 /// // Tree: (0,(1,(2,3)4)5)6
 /// let v = vec![0, 1, 2];
 /// // Root (node 6) has depth 0
-/// assert_eq!(get_node_depth(&v, 6), 0.0);
+/// assert_eq!(get_node_depth(&v, 6).unwrap(), 0.0);
 /// // Node 5 is 1 edge from root
-/// assert_eq!(get_node_depth(&v, 5), 1.0);
+/// assert_eq!(get_node_depth(&v, 5).unwrap(), 1.0);
 /// // Node 4 is 2 edges from root
-/// assert_eq!(get_node_depth(&v, 4), 2.0);
+/// assert_eq!(get_node_depth(&v, 4).unwrap(), 2.0);
 /// // Leaf 0 is 1 edge from root
-/// assert_eq!(get_node_depth(&v, 0), 1.0);
+/// assert_eq!(get_node_depth(&v, 0).unwrap(), 1.0);
 /// ```
-pub fn get_node_depth(v: &[usize], node: usize) -> f64 {
+pub fn get_node_depth(v: &[usize], node: usize) -> Result<f64, String> {
     _get_node_depth(v, None, node)
 }
 
@@ -565,7 +571,7 @@ mod tests {
         #[case] node2: usize,
         #[case] expected_mrca: usize,
     ) {
-        let mrca = get_common_ancestor(&v, node1, node2);
+        let mrca = get_common_ancestor(&v, node1, node2).unwrap();
         assert_eq!(
             mrca, expected_mrca,
             "Expected mrca of nodes {node1} and {node2} for v = {v:?} to be {expected_mrca}, but got {mrca}"
@@ -599,7 +605,7 @@ mod tests {
         #[case] node: usize,
         #[case] expected_depth: f64,
     ) {
-        let depth = get_node_depth(&v, node);
+        let depth = get_node_depth(&v, node).unwrap();
         assert!(
             (depth - expected_depth).abs() < 1e-10,
             "Expected depth {expected_depth} for node {node} in v = {v:?}, got {depth}"
@@ -616,7 +622,7 @@ mod tests {
 
         let v = sample_vector(n_leaves, false);
         let root = 2 * v.len(); // Root node index
-        let depth = get_node_depth(&v, root);
+        let depth = get_node_depth(&v, root).unwrap();
         assert_eq!(depth, 0.0, "Root should have depth 0");
     }
 
@@ -656,9 +662,9 @@ mod tests {
     #[rstest]
     #[case(vec![0, 1, 2], 7)] // Max node is 6
     #[case(vec![0], 3)] // Max node is 2
-    #[should_panic]
+    #[should_panic(expected = "Node index out of bounds")]
     fn test_get_node_depth_out_of_bounds(#[case] v: Vec<usize>, #[case] node: usize) {
-        get_node_depth(&v, node);
+        get_node_depth(&v, node).unwrap();
     }
 
     #[rstest]

py-phylo2vec/benchmarks/test_bench.py

@@ -18,7 +18,7 @@ def test_vector_to_newick_ordered(benchmark, sample_size):
     sample_size : int
         Number of leaves in the vector
     """
-    benchmark(core.to_newick_from_vector, core.sample_vector(sample_size, True))
+    benchmark(core.to_newick, core.sample_vector(sample_size, True))
 
 
 @pytest.mark.parametrize("sample_size", BIG_RANGE)
@@ -32,7 +32,7 @@ def test_vector_to_newick(benchmark, sample_size):
     sample_size : int
         Number of leaves in the vector
     """
-    benchmark(core.to_newick_from_vector, core.sample_vector(sample_size, False))
+    benchmark(core.to_newick, core.sample_vector(sample_size, False))
 
 
 @pytest.mark.parametrize("sample_size", BIG_RANGE)
@@ -46,7 +46,7 @@ def test_matrix_to_newick(benchmark, sample_size):
     sample_size : int
         Number of leaves in the vector
     """
-    benchmark(core.to_newick_from_matrix, core.sample_matrix(sample_size, False))
+    benchmark(core.to_newick, core.sample_matrix(sample_size, False))
 
 
 @pytest.mark.parametrize("sample_size", BIG_RANGE)
@@ -62,7 +62,7 @@ def test_vector_from_newick(benchmark, sample_size):
     """
     benchmark(
         core.to_vector,
-        core.to_newick_from_vector(core.sample_vector(sample_size, False)),
+        core.to_newick(core.sample_vector(sample_size, False)),
     )
 
 
@@ -79,5 +79,5 @@ def test_matrix_from_newick(benchmark, sample_size):
     """
     benchmark(
         core.to_matrix,
-        core.to_newick_from_matrix(core.sample_matrix(sample_size, False)),
+        core.to_newick(core.sample_matrix(sample_size, False)),
     )

py-phylo2vec/phylo2vec/base/newick.py

@@ -39,13 +39,4 @@ def to_newick(vector_or_matrix: np.ndarray) -> str:
     newick : str
         Newick tree
     """
-    if vector_or_matrix.ndim == 2:
-        newick = core.to_newick_from_matrix(vector_or_matrix)
-    elif vector_or_matrix.ndim == 1:
-        newick = core.to_newick_from_vector(vector_or_matrix)
-    else:
-        raise ValueError(
-            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
-        )
-
-    return newick
+    return core.to_newick(vector_or_matrix)

py-phylo2vec/phylo2vec/stats/nodewise.py

@@ -7,36 +7,28 @@
 from phylo2vec.utils.vector import check_vector
 
 
-def cophenetic_distances(vector_or_matrix, unrooted=False):
+def cophenetic_distances(tree, unrooted=False):
     """
     Compute the cophenetic distance matrix of a Phylo2Vec
     vector (topological) or matrix (from branch lengths).
 
     Parameters
     ----------
-    vector_or_matrix : numpy.ndarray
+    tree : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
 
     Returns
     -------
     numpy.ndarray
         Cophenetic distance matrix
     """
-    if vector_or_matrix.ndim == 2:
-        coph = core.cophenetic_distances_with_bls(vector_or_matrix, unrooted=unrooted)
-    elif vector_or_matrix.ndim == 1:
-        coph = core.cophenetic_distances(vector_or_matrix, unrooted=unrooted)
-    else:
-        raise ValueError(
-            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
-        )
-    return coph
+    return core.cophenetic_distances(tree, unrooted=unrooted)
 
 
 NODEWISE_DISTANCES = {"cophenetic": cophenetic_distances}
 
 
-def pairwise_distances(vector_or_matrix, metric="cophenetic"):
+def pairwise_distances(tree, metric="cophenetic"):
     """
     Compute a pairwise distance matrix
     for tree nodes from a Phylo2Vec vector.
@@ -45,7 +37,7 @@ def pairwise_distances(vector_or_matrix, metric="cophenetic"):
 
     Parameters
     ----------
-    vector_or_matrix : numpy.ndarray
+    tree : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
     metric : str, optional
         Pairwise distance metric, by default "cophenetic"
@@ -55,74 +47,50 @@ def pairwise_distances(vector_or_matrix, metric="cophenetic"):
     numpy.ndarray
         Distance matrix
     """
-    if vector_or_matrix.ndim == 2:
-        check_matrix(vector_or_matrix)
-    elif vector_or_matrix.ndim == 1:
-        check_vector(vector_or_matrix)
-    else:
-        raise ValueError(
-            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
-        )
-
     func = NODEWISE_DISTANCES[metric]
 
-    return func(vector_or_matrix)
+    return func(tree)
 
 
-def cov(vector_or_matrix):
+def cov(tree):
     """
     Compute the covariance matrix of a Phylo2Vec vector or matrix.
 
     Adapted from `vcv.phylo` in <https://github.com/emmanuelparadis/ape>
 
     Parameters
     ----------
-    vector_or_matrix : numpy.ndarray
+    tree : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
 
     Returns
     -------
     vcv : numpy.ndarray
         Covariance matrix
     """
-    if vector_or_matrix.ndim == 2:
-        vcv = core.vcv_with_bls(vector_or_matrix)
-    elif vector_or_matrix.ndim == 1:
-        vcv = core.vcv(vector_or_matrix)
-    else:
-        raise ValueError(
-            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
-        )
-    return vcv
-
-
-def precision(vector_or_matrix):
+    return core.vcv(tree)
+
+
+def precision(tree):
     """
     Compute the precision matrix of a Phylo2Vec vector or matrix.
 
     Adapted from: `inverseA.R` in <https://github.com/cran/MCMCglmm>
 
     Parameters
     ----------
-    vector_or_matrix : numpy.ndarray
+    tree : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
 
     Returns
     -------
     precision: numpy.ndarray
         Precision matrix
     """
-    if vector_or_matrix.ndim == 2:
-        precursor = core.pre_precision_with_bls(vector_or_matrix)
-    elif vector_or_matrix.ndim == 1:
-        precursor = core.pre_precision(vector_or_matrix)
-    else:
-        raise ValueError(
-            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
-        )
+    precursor = core.pre_precision(tree)
 
     # Schur complement of the precursor matrix
-    n_leaves = vector_or_matrix.shape[0] + 1
+    n_leaves = tree.shape[0] + 1
     a = precursor[:n_leaves, :n_leaves]
     b = precursor[:n_leaves, n_leaves:]
     c = precursor[n_leaves:, n_leaves:]

py-phylo2vec/phylo2vec/stats/treewise.py

@@ -55,8 +55,4 @@ def robinson_foulds(
     ete3.Tree.robinson_foulds : Reference implementation in ete3
     ape::dist.topo : Reference implementation in R's ape package
     """
-    # Extract topology (column 0) if matrix input
-    v1 = tree1[:, 0].astype(int).tolist() if tree1.ndim == 2 else tree1.tolist()
-    v2 = tree2[:, 0].astype(int).tolist() if tree2.ndim == 2 else tree2.tolist()
-
-    return core.robinson_foulds(v1, v2, normalize)
+    return core.robinson_foulds(tree1, tree2, normalize)