changes to add functionality for filtering

qiime2 · May 14, 2024 · 3a76163 · 3a76163
1 parent 76e3380
commit 3a76163
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 5 deletions.
diff --git a/q2_diversity/__init__.py b/q2_diversity/__init__.py
@@ -14,7 +14,7 @@
 from ._ordination import pcoa, pcoa_biplot, tsne, umap
 from ._procrustes import procrustes_analysis, partial_procrustes
 from ._core_metrics import core_metrics_phylogenetic, core_metrics
-from ._filter import filter_distance_matrix
+from ._filter import filter_distance_matrix, filter_alpha_diversity_artifact
 from ._version import get_versions
 
 __version__ = get_versions()['version']
@@ -25,7 +25,7 @@
            'pcoa', 'tsne', 'umap', 'pcoa_biplot', 'alpha_group_significance',
            'bioenv', 'beta_group_significance', 'alpha_correlation',
            'core_metrics_phylogenetic', 'core_metrics',
-           'filter_distance_matrix', 'mantel', 'alpha_rarefaction',
-           'beta_rarefaction', 'procrustes_analysis', 'beta_correlation',
-           'adonis', 'partial_procrustes'
+           'filter_alpha_diversity_artifact', 'filter_distance_matrix',
+           'alpha_rarefaction', 'beta_rarefaction', 'procrustes_analysis',
+           'beta_correlation', 'adonis', 'partial_procrustes', 'mantel'
            ]
diff --git a/q2_diversity/_filter.py b/q2_diversity/_filter.py
@@ -9,6 +9,8 @@
 import skbio
 import qiime2
 
+import pandas as pd
+
 
 def filter_distance_matrix(distance_matrix: skbio.DistanceMatrix,
                            metadata: qiime2.Metadata,
@@ -25,3 +27,18 @@ def filter_distance_matrix(distance_matrix: skbio.DistanceMatrix,
     except skbio.stats.distance.DissimilarityMatrixError:
         raise ValueError(
             "All samples were filtered out of the distance matrix.")
+
+
+# This function filters the SampleData[AlphaDiversity] table by the metadata
+# only sample present in the metadata will remain
+def filter_alpha_diversity_artifact(alpha_diversity: pd.Series,
+                                    metadata: qiime2.Metadata,
+                                    exclude_ids: bool = False) -> pd.Series:
+    ids_to_keep = metadata.get_ids()
+    if exclude_ids:
+        ids_to_keep = set(alpha_diversity.ids) - set(ids_to_keep)
+        filtered_table = alpha_diversity.filter(ids_to_keep, inplace=False)
+    if filtered_table.is_empty():
+        raise ValueError(
+            "All samples were filtered out of the alpha diversity artifact.")
+    return filtered_table
diff --git a/q2_diversity/plugin_setup.py b/q2_diversity/plugin_setup.py
@@ -164,7 +164,6 @@
                  'metric for all samples in a feature table.'),
 )
 
-
 plugin.pipelines.register_function(
     function=q2_diversity.alpha,
     inputs={'table':
@@ -645,6 +644,42 @@
     }
 )
 
+plugin.methods.register_function(
+    function=q2_diversity.filter_alpha_diversity_artifact,
+    inputs={
+        'alpha_diversity': SampleData[AlphaDiversity]
+    },
+    parameters={
+        'metadata': Metadata,
+        'exclude_ids': Bool
+    },
+    outputs=[
+        ('filtered_alpha_diversity_artifact', SampleData[AlphaDiversity])
+    ],
+    name="Filter samples from an alpha diversity metric.",
+    description="Filter samples from an alpha diversity metric, retaining "
+                "the samples with corresponding `metadata` (or retaining only "
+                "only samples without metadata, if `exclude_ids` is "
+                "True). See the filtering tutorial on "
+                "https://docs.qiime2.org for additional details.",
+    input_descriptions={
+        'alpha_diversity': 'Alpha diversity sample data to filter by sample'
+    },
+    parameter_descriptions={
+       'metadata': 'Sample metadata is used to select samples to retain from '
+                   'the sample data (default) or select samples to exclude '
+                   'using the `exclude_ids` parameter.',
+       'exclude_ids': 'If `True`, the samples selected by `metadata` '
+                      'will be excluded from the filtered '
+                      'sample data instead of being retained.'
+    }
+    # output_descriptions={
+    #     'filtered_alpha_diversity_artifact': 'Alpha diversity sample data '
+    #                                          'filtered to include samples '
+    #                                          'matching search criteria'
+    # }
+)
+
 plugin.visualizers.register_function(
     function=q2_diversity.alpha_group_significance,
     inputs={'alpha_diversity': SampleData[AlphaDiversity]},