tax_finder methods

Author: palec87

README.md

@@ -12,7 +12,7 @@ Package of utilities for FAIR-Ease demo workflows.
 
 These methods are currently use only for FAIR-EASE [pilot demos](https://github.com/palec87/momics-demos), but eventually they can serve for general purpose manipulation of metagenomic data, locally and in VREs.
 
-The idea is to provide testable methods to allow as much flexibility and remixing of the functionalities provided
+The idea is to provide testable methods to allow as much flexibility and remixing of the functionalities provided.
 
 ## Installation
 

momics/diversity.py

@@ -1,7 +1,7 @@
 import logging
 import pandas as pd
 import numpy as np
-from typing import List, Dict
+from typing import Union, List, Dict
 
 import skbio
 from skbio.diversity import beta_diversity
@@ -11,7 +11,7 @@
 from .utils import (
     check_index_names,
 )
-
+from momics.constants import TAXONOMY_RANKS
 
 # logger setup
 FORMAT = "%(levelname)s | %(name)s | %(message)s"
@@ -131,9 +131,65 @@ def calculate_shannon_index(df: pd.DataFrame) -> pd.Series:
     return df.apply(shannon_index, axis=1)
 
 
+####################
+# Search functions #
+####################
+def find_taxa_in_table(
+        table: pd.DataFrame,
+        tax_level: str,
+        search_term: Union[str, int],
+        ncbi_tax_id: bool=False,
+        exact_match:bool=False,
+    ) -> pd.DataFrame:
+    """
+    Find taxa in the given table at the specified taxonomic level matching the search term.
+
+    args:
+        table (pd.DataFrame): DataFrame containing taxonomic data.
+        tax_level (str): Taxonomic level to search ('all' for all levels).
+        search_term (str|int): Term to search for.
+        ncbi_tax_id (bool): If True, search by NCBI taxonomic ID.
+        exact_match (bool): If True, perform exact match; otherwise, use substring match.
+
+    returns:
+        pd.DataFrame: DataFrame containing matching taxa.
+    """
+    # ncbi_tax_id search
+    index_names = getattr(table.index, "names", [])
+    if ncbi_tax_id and ('ncbi_tax_id' not in table.columns and 'ncbi_tax_id' not in index_names):
+        raise ValueError("The table does not contain 'ncbi_tax_id' column or index level.")
+
+    # if ncbi_tax_id is an index level, bring it into a column for uniform handling
+    if ncbi_tax_id and ('ncbi_tax_id' in index_names):
+        table = table.reset_index()
+
+    if ncbi_tax_id:
+        # Search by NCBI taxonomic ID
+        matching_taxa = table[table['ncbi_tax_id'].astype(str) == str(search_term)]
+        return matching_taxa.set_index(index_names) if index_names else matching_taxa
+
+    # search by taxonomic level, all ranks
+    if tax_level == 'all':
+        found = []
+        for tax_level in TAXONOMY_RANKS:
+            if exact_match:
+                found.append(table[table[tax_level].str.lower().fillna('') == search_term.lower()])
+            else:
+                found.append(table[table[tax_level].str.contains(search_term, case=False, na=False)])
+        matching_taxa = pd.concat(found)
+    # specific taxonomic level
+    else:
+        if exact_match:
+            matching_taxa = table[table[tax_level].str.lower().fillna('') == search_term.lower()]
+        else:
+            matching_taxa = table[table[tax_level].str.contains(search_term, case=False, na=False)]
+
+    return matching_taxa
+
 #######################
 # diversity functions #
 #######################
+
 def calculate_alpha_diversity(df: pd.DataFrame, factors: pd.DataFrame) -> pd.DataFrame:
     """
     Calculates the alpha diversity (Shannon index) for a DataFrame.

momics/panel_utils.py

@@ -64,6 +64,63 @@ def is_port_in_use(port: int) -> bool:
         return s.connect_ex(("localhost", port)) == 0
 
 
+def tax_finder_selector() -> Tuple[
+    pn.widgets.Select,
+    pn.widgets.Select,
+    pn.widgets.TextInput,
+    pn.widgets.Checkbox,
+    pn.widgets.Checkbox,
+]:
+    select_table_tax = pn.widgets.Select(
+        name="Taxonomic table",
+        value="ssu",
+        options=["ssu", "lsu"],
+        description="Select a table for taxonomic search",
+    )
+
+    tax_level = pn.widgets.Select(
+        name="Taxonomic level",
+        value="all",
+        options=[
+            "all",
+            "ncbi_tax_id",
+            "superkingdom",
+            "kingdom",
+            "phylum",
+            "class",
+            "order",
+            "family",
+            "genus",
+            "species",
+        ],
+        description="Select a taxonomic search level",
+    )
+
+    search_term = pn.widgets.TextInput(
+        name="Search term",
+        value="",
+        description="Enter a search term (string or NCBI tax ID)",
+    )
+
+    checkbox_exact_match = pn.widgets.Checkbox(
+        name="Exact match of the search term",
+        value=False,
+    )
+
+    log_scale_checkbox = pn.widgets.Checkbox(
+        name="Log scale for abundance coloring",
+        value=True,
+    )
+
+    return (
+        select_table_tax,
+        tax_level,
+        search_term,
+        checkbox_exact_match,
+        log_scale_checkbox,
+    )
+
+
 def diversity_select_widgets(cat_columns: List[str], num_columns: List[str]) -> Tuple[
     pn.widgets.Select,
     pn.widgets.Select,