Repository: scikit-bio/scikit-bio

Issue:
Expose _vectorize_counts_and_tree
As shown in the [PhyloRPCA paper](https://journals.asm.org/doi/10.1128/msystems.00050-22), leveraging internal node counts in the tree (in addition to tip abundances) can double the accuracy in downstream ML pipelines. It makes a lot of sense to directly expose this functionality, which is leveraged internally within fast unifrac

https://github.com/scikit-bio/scikit-bio/blob/main/skbio/diversity/_util.py#L122

Hints from issue comments:
@mataton: @mortonjt This has been completed in #2173.

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 846f994e..5d974d55 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@
 * Added `TreeNode.is_bifurcating` to check if a tree is bifurcating (i.e., binary) ([#2117](https://github.com/scikit-bio/scikit-bio/pull/2117)).
 * Added support for Python's `pathlib` module in the IO system ([#2119](https://github.com/scikit-bio/scikit-bio/pull/2119)).
 * Added `TreeNode.path` to return a list of nodes representing the path from one node to another ([#2131](https://github.com/scikit-bio/scikit-bio/pull/2131)).
+* Exposed `vectorize_counts_and_tree` function from the `diversity` module to allow use for improving ML accuracy in downstream pipelines ([#2173](https://github.com/scikit-bio/scikit-bio/pull/2173))
 
 ### Performance enhancements
 
diff --git a/skbio/diversity/__init__.py b/skbio/diversity/__init__.py
index a4789aa8..7a641a1a 100644
--- a/skbio/diversity/__init__.py
+++ b/skbio/diversity/__init__.py
@@ -53,6 +53,15 @@ Beta diversity
    block_beta_diversity
 
 
+Utility functions
+-----------------
+
+.. autosummary::
+   :toctree: generated/
+
+   vectorize_counts_and_tree
+
+
 Introduction
 ------------
 
@@ -444,6 +453,7 @@ from ._driver import (
     get_beta_diversity_metrics,
 )
 from ._block import block_beta_diversity
+from ._util import vectorize_counts_and_tree
 
 __all__ = [
     "alpha_diversity",
@@ -452,4 +462,5 @@ __all__ = [
     "get_beta_diversity_metrics",
     "partial_beta_diversity",
     "block_beta_diversity",
+    "vectorize_counts_and_tree",
 ]
diff --git a/skbio/diversity/_util.py b/skbio/diversity/_util.py
index 4b4aea15..e92a2879 100644
--- a/skbio/diversity/_util.py
+++ b/skbio/diversity/_util.py
@@ -183,7 +183,7 @@ def _validate_taxa_and_tree(counts, taxa, tree, rooted=True):
         )
 
 
-def _vectorize_counts_and_tree(counts, taxa, tree):
+def vectorize_counts_and_tree(counts, taxa, tree):
     """Index tree and convert counts to np.array in corresponding order.
 
     Parameters
@@ -209,6 +209,18 @@ def _vectorize_counts_and_tree(counts, taxa, tree):
     --------
     skbio.tree.TreeNode.to_array
 
+    Notes
+    -----
+    Leveraging internal node counts in the tree (in addition to tip abundances) can
+    double the accuracy in downstream machine learning pipelines [1]_.
+
+    References
+    ----------
+    .. [1] Martino C., McDonald D., Cantrell K., Dilmore AH., Vázquez-Baeza Y.,
+       Shenhav L., Shaffer J.P., Rahman G., Armstrong G., Allaband C., Song S.J.,
+       Knight R. Compositionally aware phylogenetic beta-diversity measures better
+       resolve microbiomes associated with phenotype. mSystems. 7(3) (2022).
+
     """
     tree_index = tree.to_array(nan_length_value=0.0)
     taxa = np.asarray(taxa)
diff --git a/skbio/diversity/alpha/_pd.py b/skbio/diversity/alpha/_pd.py
index 6597cd31..5bbe0640 100644
--- a/skbio/diversity/alpha/_pd.py
+++ b/skbio/diversity/alpha/_pd.py
@@ -11,7 +11,7 @@ import numpy as np
 from skbio.diversity._util import (
     _validate_counts_vector,
     _validate_taxa_and_tree,
-    _vectorize_counts_and_tree,
+    vectorize_counts_and_tree,
     _check_taxa_alias,
 )
 
@@ -26,7 +26,7 @@ def _setup_pd(counts, taxa, tree, validate, rooted, single_sample):
         else:
             _validate_taxa_and_tree(counts[0], taxa, tree, rooted)
 
-    counts_by_node, _, branch_lengths = _vectorize_counts_and_tree(counts, taxa, tree)
+    counts_by_node, _, branch_lengths = vectorize_counts_and_tree(counts, taxa, tree)
 
     return counts_by_node, branch_lengths
 
diff --git a/skbio/diversity/beta/_unifrac.py b/skbio/diversity/beta/_unifrac.py
index 4f8ff6d8..d44046f2 100644
--- a/skbio/diversity/beta/_unifrac.py
+++ b/skbio/diversity/beta/_unifrac.py
@@ -13,7 +13,7 @@ import numpy as np
 from skbio.diversity._util import (
     _validate_counts_matrix,
     _validate_taxa_and_tree,
-    _vectorize_counts_and_tree,
+    vectorize_counts_and_tree,
     _check_taxa_alias,
 )
 from skbio.diversity._phylogenetic import _tip_distances
@@ -341,11 +341,11 @@ def _setup_pairwise_unifrac(
         _validate(u_counts, v_counts, taxa, tree)
 
     # temporarily store u_counts and v_counts in a 2-D array as that's what
-    # _vectorize_counts_and_tree takes
+    # vectorize_counts_and_tree takes
     u_counts = np.asarray(u_counts)
     v_counts = np.asarray(v_counts)
     counts = np.vstack([u_counts, v_counts])
-    counts_by_node, tree_index, branch_lengths = _vectorize_counts_and_tree(
+    counts_by_node, tree_index, branch_lengths = vectorize_counts_and_tree(
         counts, taxa, tree
     )
     # unpack counts vectors for single pairwise UniFrac calculation
@@ -496,7 +496,7 @@ def _setup_multiple_unifrac(counts, taxa, tree, validate):
     if validate:
         _validate_taxa_and_tree(counts[0], taxa, tree)
 
-    counts_by_node, tree_index, branch_lengths = _vectorize_counts_and_tree(
+    counts_by_node, tree_index, branch_lengths = vectorize_counts_and_tree(
         counts, taxa, tree
     )
 


Test patch:
diff --git a/skbio/diversity/tests/test_util.py b/skbio/diversity/tests/test_util.py
index b1e3366b..8db50414 100644
--- a/skbio/diversity/tests/test_util.py
+++ b/skbio/diversity/tests/test_util.py
@@ -18,7 +18,7 @@ from skbio.table import example_table
 from skbio.diversity._util import (_validate_counts_vector,
                                    _validate_counts_matrix,
                                    _validate_taxa_and_tree,
-                                   _vectorize_counts_and_tree,
+                                   vectorize_counts_and_tree,
                                    _quantitative_to_qualitative_counts,
                                    _check_taxa_alias,
                                    _table_to_numpy,
@@ -270,7 +270,7 @@ class ValidationTests(TestCase):
         tree = TreeNode.read(io.StringIO("((a:1, b:2)c:3)root;"))
         counts = np.array([[0, 1], [1, 5], [10, 1]])
         count_array, indexed, branch_lengths = \
-            _vectorize_counts_and_tree(counts, np.array(['a', 'b']), tree)
+            vectorize_counts_and_tree(counts, np.array(['a', 'b']), tree)
         exp_counts = np.array([[0, 1, 10], [1, 5, 1], [1, 6, 11], [1, 6, 11]])
         npt.assert_equal(count_array, exp_counts.T)