Merge pull request #27 from gibsramen/level-limits

Restructure level limits in DiversityHandler

Author: gibsramen

README.md

@@ -84,6 +84,10 @@ The main data structure in evident is the 'DiversityHandler'.
 This is the way that evident stores the diversity data and metadata for power calculations.
 For our alpha diversity example, we'll load the `AlphaDiversityHandler` class from evident.
 `AlphaDiversityHandler` takes as input the pandas Series with the diversity values and the pandas DataFrame containing the sample metadata.
+By default, evident will only consider metadata columns with, at max, 5 levels.
+To modify this behavior, provide a value for the `max_levels_per_category` argument.
+Additionally, evident will not consider any category levels represented by fewer than 3 samples.
+To modify this behavior, use the `min_count_per_level` argument.
 
 ```python
 adh = evident.AlphaDiversityHandler(faith_pd, metadata)
@@ -158,10 +162,6 @@ bokeh serve --show app
 ```
 
 This should open up a browser window where you can modify the chosen column, significance, level, and observations.
-By default, this interactive view will only consider metadata columns with, at max, 5 levels.
-To modify this behavior, use the `max_levels_per_category` argument in `create_bokeh_app`.
-Additionally, this interactive view will not consider any category levels represented by fewer than 3 samples.
-To modify this behavior, use the `min_count_per_level` argument.
 We also provide a command line script to generate an interactive app using some test data.
 You can access this script at `evident/tests/make_interactive.py`.
 

evident/diversity_handler.py

@@ -19,9 +19,62 @@
 
 class _BaseDiversityHandler(ABC):
     """Abstract class for handling diversity data and metadata."""
-    def __init__(self, data=None, metadata: pd.DataFrame = None):
+    def __init__(
+        self,
+        data=None,
+        metadata: pd.DataFrame = None,
+        max_levels_per_category: int = 5,
+        min_count_per_level: int = 3
+    ):
         self.data = data
-        self.metadata = metadata
+        metadata = metadata.copy()
+
+        cols_to_drop = []
+        levels_to_drop = dict()
+
+        warn_msg_num_levels = False
+        warn_msg_level_count = False
+        for col in metadata.columns:
+            # Drop non-categorical columns
+            if metadata[col].dtype != np.dtype("object"):
+                cols_to_drop.append(col)
+                continue
+
+            # Drop columns with only one level or more than max
+            num_uniq_cols = len(metadata[col].dropna().unique())
+            if not (1 < num_uniq_cols <= max_levels_per_category):
+                cols_to_drop.append(col)
+                warn_msg_num_levels = True
+                continue
+
+            # Drop levels that have fewer than min_count_per_level samples
+            level_count = metadata[col].value_counts()
+            under_thresh = level_count[level_count < min_count_per_level]
+            if not under_thresh.empty:
+                levels_under_thresh = list(under_thresh.index)
+                metadata[col].replace(
+                    {x: np.nan for x in levels_under_thresh},
+                    inplace=True
+                )
+                levels_to_drop[col] = levels_under_thresh
+                warn_msg_level_count = True
+
+        if warn_msg_num_levels:
+            warn(
+                "Some categories have been dropped because they had either "
+                "only one level or too many. Use the max_levels_per_category "
+                "argument to modify this threshold.\n"
+                f"Dropped columns: {cols_to_drop}"
+            )
+        if warn_msg_level_count:
+            warn(
+                "Some categorical levels have been dropped because they "
+                "did not have enough samples. Use the min_count_per_level "
+                "argument to modify this threshold.\n"
+                f"Dropped levels: {levels_to_drop}"
+            )
+
+        self.metadata = metadata.drop(columns=cols_to_drop)
 
     @property
     def samples(self):
@@ -168,7 +221,7 @@ def _single_power_analysis(
         :type power: float
 
         :returns: Collection of values from power analysis
-        :rtype: evident.power.PowerAnalysisResult
+        :rtype: evident.results.PowerAnalysisResult
         """
         power_func = self._create_partial_power_func(
             column=column,
@@ -235,7 +288,7 @@ def _bulk_power_analysis(
         :type power: sequence of floats
 
         :returns: Collection of values from power analyses
-        :rtype: evident.power.PowerAnalysisResults
+        :rtype: evident.results.PowerAnalysisResults
         """
         # Convert all to list so we can use Cartesian product
         difference = _listify(difference)
@@ -311,12 +364,31 @@ def _create_partial_power_func(
 
 
 class AlphaDiversityHandler(_BaseDiversityHandler):
-    """Handler for alpha diversity data."""
     def __init__(
         self,
         data: pd.Series,
-        metadata: pd.DataFrame
+        metadata: pd.DataFrame,
+        max_levels_per_category: int = 5,
+        min_count_per_level: int = 3
     ):
+        """Handler for alpha diversity data.
+
+        :param data: Alpha diversity vector
+        :type data: pd.Series
+
+        :param metadata: Sample metadata
+        :type metadata: pd.DataFrame
+
+        :param max_levels_per_category: Max number of levels in a category to
+            keep. Any categorical columns that have more than this number of
+            unique levels will not be saved, defaults to 5.
+        :type max_levels_per_category: int
+
+        :param min_count_per_level: Min number of samples in a given category
+            level to keep. Any levels that have fewer than this many samples
+            will not be saved, defaults to 3.
+        :type min_count_per_level: int
+        """
         if not isinstance(data, pd.Series):
             raise ValueError("data must be of type pandas.Series")
         if data.isna().any():
@@ -329,7 +401,9 @@ def __init__(
 
         super().__init__(
             data=data.loc[samps_in_common],
-            metadata=metadata.loc[samps_in_common]
+            metadata=metadata.loc[samps_in_common],
+            max_levels_per_category=max_levels_per_category,
+            min_count_per_level=min_count_per_level
         )
 
     def subset_values(self, ids: list) -> np.array:
@@ -338,12 +412,31 @@ def subset_values(self, ids: list) -> np.array:
 
 
 class BetaDiversityHandler(_BaseDiversityHandler):
-    """Handler for beta diversity data."""
     def __init__(
         self,
         data: DistanceMatrix,
-        metadata: pd.DataFrame
+        metadata: pd.DataFrame,
+        max_levels_per_category: int = 5,
+        min_count_per_level: int = 3
     ):
+        """Handler for beta diversity data.
+
+        :param data: Beta diversity distance matrix
+        :type data: skbio.DistanceMatrix
+
+        :param metadata: Sample metadata
+        :type metadata: pd.DataFrame
+
+        :param max_levels_per_category: Max number of levels in a category to
+            keep. Any categorical columns that have more than this number of
+            unique levels will not be saved, defaults to 5.
+        :type max_levels_per_category: int
+
+        :param min_count_per_level: Min number of samples in a given category
+            level to keep. Any levels that have fewer than this many samples
+            will not be saved, defaults to 3.
+        :type min_count_per_level: int
+        """
         if not isinstance(data, DistanceMatrix):
             raise ValueError("data must be of type skbio.DistanceMatrix")
 
@@ -353,7 +446,9 @@ def __init__(
 
         super().__init__(
             data=data.filter(samps_in_common),
-            metadata=metadata.loc[samps_in_common]
+            metadata=metadata.loc[samps_in_common],
+            max_levels_per_category=max_levels_per_category,
+            min_count_per_level=min_count_per_level
         )
 
     def subset_values(self, ids: list) -> np.array:

evident/interactive.py

@@ -1,8 +1,5 @@
 import os
 import shutil
-from warnings import warn
-
-import numpy as np
 
 from evident.diversity_handler import (_BaseDiversityHandler,
                                        AlphaDiversityHandler,
@@ -12,8 +9,6 @@
 def create_bokeh_app(
     diversity_handler: _BaseDiversityHandler,
     output: os.PathLike,
-    max_levels_per_category: int = 5,
-    min_count_per_level: int = 3
 ) -> None:
     """Creates interactive power analysis using Bokeh.
 
@@ -22,15 +17,6 @@ def create_bokeh_app(
 
     :param output: Location to create Bokeh app
     :type output: os.PathLike
-
-    :param max_levels_per_category: Max number of levels in a category to
-        keep. Any categorical columns that have more than this number of
-        unique levels will not be saved, defaults to 5.
-    :type max_levels_per_category: int
-
-    :param min_count_per_level: Min number of samples in a given category
-        level to keep. Any levels that have fewer than this many samples
-        will not be saved, defaults = 3.
     """
     curr_path = os.path.dirname(__file__)
     support_files = os.path.join(curr_path, "support_files")
@@ -40,49 +26,9 @@ def create_bokeh_app(
     data_dir = os.path.join(output, "data")
     os.mkdir(data_dir)
 
-    # Process metadata
     md = diversity_handler.metadata.copy()
-    cols_to_drop = []
-    warn_msg_num_levels = False
-    warn_msg_level_count = False
-    for col in md.columns:
-        # Drop non-categorical columns
-        if md[col].dtype != np.dtype("object"):
-            cols_to_drop.append(col)
-            continue
-
-        # Drop columns with only one level or more than max
-        if not (1 < len(md[col].dropna().unique()) <= max_levels_per_category):
-            cols_to_drop.append(col)
-            warn_msg_num_levels = True
-            continue
-
-        # Drop levels that have fewer than min_count_per_level samples
-        level_count = md[col].value_counts()
-        under_thresh = level_count[level_count < min_count_per_level]
-        if not under_thresh.empty:
-            levels_under_thresh = list(under_thresh.index)
-            md[col].replace(
-                {x: np.nan for x in levels_under_thresh},
-                inplace=True
-            )
-            warn_msg_level_count = True
-
-    if warn_msg_num_levels:
-        warn(
-            "Some categories have been dropped because they had either only "
-            "one level or too many. Use the max_levels_per_category "
-            "argument to modify this threshold."
-        )
-    if warn_msg_level_count:
-        warn(
-            "Some categorical levels have been dropped because they "
-            "did not have enough samples. Use the min_count_per_level "
-            "argument to modify this threshold."
-        )
-
     md_loc = os.path.join(data_dir, "metadata.tsv")
-    md.drop(columns=cols_to_drop).to_csv(md_loc, sep="\t", index=True)
+    md.to_csv(md_loc, sep="\t", index=True)
 
     data = diversity_handler.data
     if isinstance(diversity_handler, AlphaDiversityHandler):

evident/q2/_methods.py

@@ -12,13 +12,17 @@
 def alpha_power_analysis(
     alpha_diversity: pd.Series,
     sample_metadata: CategoricalMetadataColumn,
+    max_levels_per_category: int = 5,
+    min_count_per_level: int = 3,
     alpha: list = None,
     power: list = None,
     total_observations: list = None,
-    difference: list = None
+    difference: list = None,
 ) -> pd.DataFrame:
     res = _power_analysis(alpha_diversity, sample_metadata,
-                          AlphaDiversityHandler, alpha=alpha, power=power,
+                          AlphaDiversityHandler,
+                          max_levels_per_category, min_count_per_level,
+                          alpha=alpha, power=power,
                           total_observations=total_observations,
                           difference=difference)
     return res
@@ -27,22 +31,28 @@ def alpha_power_analysis(
 def beta_power_analysis(
     beta_diversity: DistanceMatrix,
     sample_metadata: CategoricalMetadataColumn,
+    max_levels_per_category: int = 5,
+    min_count_per_level: int = 3,
     alpha: list = None,
     power: list = None,
     total_observations: list = None,
-    difference: list = None
+    difference: list = None,
 ) -> pd.DataFrame:
     res = _power_analysis(beta_diversity, sample_metadata,
-                          BetaDiversityHandler, alpha=alpha, power=power,
+                          BetaDiversityHandler,
+                          max_levels_per_category, min_count_per_level,
+                          alpha=alpha, power=power,
                           total_observations=total_observations,
                           difference=difference)
     return res
 
 
-def _power_analysis(data, metadata, handler, **kwargs):
+def _power_analysis(data, metadata, handler, max_levels_per_category,
+                    min_count_per_level, **kwargs):
     md = metadata.to_series()
     column = md.name
-    dh = handler(data, md.to_frame())
+    dh = handler(data, md.to_frame(), max_levels_per_category,
+                 min_count_per_level)
     res = dh.power_analysis(column, **kwargs)
     return res.to_dataframe()
 
@@ -52,11 +62,14 @@ def alpha_effect_size_by_category(
     sample_metadata: Metadata,
     columns: List[str],
     pairwise: bool = False,
-    n_jobs: int = None
+    n_jobs: int = None,
+    max_levels_per_category: int = 5,
+    min_count_per_level: int = 3
 ) -> pd.DataFrame:
     res = _effect_size_by_category(alpha_diversity, sample_metadata,
                                    AlphaDiversityHandler, columns, pairwise,
-                                   n_jobs)
+                                   n_jobs, max_levels_per_category,
+                                   min_count_per_level)
     return res
 
 
@@ -65,17 +78,22 @@ def beta_effect_size_by_category(
     sample_metadata: Metadata,
     columns: List[str],
     pairwise: bool = False,
-    n_jobs: int = None
+    n_jobs: int = None,
+    max_levels_per_category: int = 5,
+    min_count_per_level: int = 3
 ) -> pd.DataFrame:
     res = _effect_size_by_category(beta_diversity, sample_metadata,
                                    BetaDiversityHandler, columns, pairwise,
-                                   n_jobs)
+                                   n_jobs, max_levels_per_category,
+                                   min_count_per_level)
     return res
 
 
 def _effect_size_by_category(data, metadata, handler, columns, pairwise,
-                             n_jobs):
-    dh = handler(data, metadata.to_dataframe())
+                             n_jobs, max_levels_per_category,
+                             min_count_per_level):
+    dh = handler(data, metadata.to_dataframe(), max_levels_per_category,
+                 min_count_per_level)
     if pairwise:
         res = pairwise_effect_size_by_category(dh, columns, n_jobs=n_jobs)
     else: