linkml-store 0.2.6__py3-none-any.whl → 0.2.10rc1__py3-none-any.whl

linkml_store/plotting/heatmap.py

@@ -0,0 +1,356 @@
+"""
+Heatmap visualization module for LinkML data.
+
+This module provides functions to generate heatmaps from pandas DataFrames or tabular data files.
+"""
+
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from matplotlib.colors import LinearSegmentedColormap
+from scipy.cluster import hierarchy
+from scipy.spatial import distance
+
+from linkml_store.utils.format_utils import Format, load_objects, write_output
+
+logger = logging.getLogger(__name__)
+
+
+def create_heatmap(
+    data: pd.DataFrame,
+    x_column: str,
+    y_column: str,
+    value_column: Optional[str] = None,
+    title: Optional[str] = None,
+    figsize: Tuple[int, int] = (10, 8),
+    cmap: Union[str, LinearSegmentedColormap] = "YlGnBu",
+    annot: bool = True,
+    fmt: Optional[str] = None,  # Dynamically determined based on data
+    linewidths: float = 0.5,
+    linecolor: str = "white",
+    square: bool = False,
+    output_file: Optional[str] = None,
+    dpi: int = 300,
+    missing_value: Any = np.nan,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+    robust: bool = False,
+    remove_duplicates: bool = True,
+    font_size: int = 10,
+    cluster: Union[bool, Literal["both", "x", "y"]] = False,
+    cluster_method: str = "complete",  # linkage method: complete, average, single, etc.
+    cluster_metric: str = "euclidean",  # distance metric: euclidean, cosine, etc.
+    **kwargs,
+) -> Tuple[plt.Figure, plt.Axes]:
+    """
+    Create a heatmap from a pandas DataFrame.
+
+    Args:
+        data: Input DataFrame containing the data to plot
+        x_column: Column to use for x-axis categories
+        y_column: Column to use for y-axis categories
+        value_column: Column containing values for the heatmap. If None, frequency counts will be used.
+        title: Title for the heatmap
+        figsize: Figure size as (width, height) in inches
+        cmap: Colormap for the heatmap
+        annot: Whether to annotate cells with values
+        fmt: String formatting code for annotations (auto-detected if None)
+        linewidths: Width of lines between cells
+        linecolor: Color of lines between cells
+        square: Whether to make cells square
+        output_file: File path to save the figure (optional)
+        dpi: Resolution for saved figure
+        missing_value: Value to use for missing data (defaults to NaN)
+        vmin: Minimum value for colormap scaling
+        vmax: Maximum value for colormap scaling
+        robust: If True, compute colormap limits using robust quantiles instead of min/max
+        remove_duplicates: If True, removes duplicate rows before creating the heatmap
+        font_size: Font size for annotations
+        cluster: Whether and which axes to cluster:
+                - False: No clustering (default)
+                - True or "both": Cluster both x and y axes
+                - "x": Cluster only x-axis
+                - "y": Cluster only y-axis
+        cluster_method: Linkage method for hierarchical clustering
+                      (e.g., "single", "complete", "average", "ward")
+        cluster_metric: Distance metric for clustering (e.g., "euclidean", "correlation", "cosine")
+        **kwargs: Additional keyword arguments to pass to seaborn's heatmap function
+
+    Returns:
+        Tuple containing the figure and axes objects
+    """
+    # Validate input
+    if x_column not in data.columns:
+        raise ValueError(f"x_column '{x_column}' not found in DataFrame columns: {list(data.columns)}")
+    if y_column not in data.columns:
+        raise ValueError(f"y_column '{y_column}' not found in DataFrame columns: {list(data.columns)}")
+    if value_column and value_column not in data.columns:
+        raise ValueError(f"value_column '{value_column}' not found in DataFrame columns: {list(data.columns)}")
+
+    # Remove duplicates by default (assume they're accidents unless user overrides)
+    if remove_duplicates:
+        data = data.drop_duplicates()
+    
+    # Prepare the data
+    if value_column:
+        # Use the provided value column
+        pivot_data = data.pivot_table(
+            index=y_column, 
+            columns=x_column, 
+            values=value_column, 
+            aggfunc='mean',
+            fill_value=missing_value
+        )
+    else:
+        # Use frequency counts
+        cross_tab = pd.crosstab(data[y_column], data[x_column])
+        pivot_data = cross_tab
+    
+    # Auto-detect format string if not provided
+    if fmt is None:
+        # Check if the pivot table contains integers only
+        if pivot_data.dtypes.apply(lambda x: pd.api.types.is_integer_dtype(x)).all():
+            fmt = 'd'  # Integer format
+        else:
+            fmt = '.1f'  # One decimal place for floats
+    
+    # Make sure all cells have a reasonable minimum size
+    min_height = max(4, 80 / len(pivot_data.index) if len(pivot_data.index) > 0 else 10)
+    min_width = max(4, 80 / len(pivot_data.columns) if len(pivot_data.columns) > 0 else 10)
+    
+    # Adjust figure size based on the number of rows and columns 
+    adjusted_height = max(figsize[1], min_height * len(pivot_data.index) / 10)
+    adjusted_width = max(figsize[0], min_width * len(pivot_data.columns) / 10)
+    adjusted_figsize = (adjusted_width, adjusted_height)
+    
+    # Create figure and axes
+    fig, ax = plt.subplots(figsize=adjusted_figsize)
+    
+    # Apply clustering if requested
+    row_linkage = None
+    col_linkage = None
+    
+    if cluster:
+        cluster_axes = cluster
+        if cluster_axes is True:
+            cluster_axes = "both"
+            
+        # Fill NAs for clustering
+        pivot_data_for_clustering = pivot_data.fillna(0)
+        
+        # Cluster rows (y-axis)
+        if cluster_axes in ["both", "y"]:
+            try:
+                # Calculate distance matrix and linkage for rows
+                row_distances = distance.pdist(pivot_data_for_clustering.values, metric=cluster_metric)
+                row_linkage = hierarchy.linkage(row_distances, method=cluster_method)
+                
+                # Reorder rows based on clustering
+                row_dendrogram = hierarchy.dendrogram(row_linkage, no_plot=True)
+                row_order = row_dendrogram['leaves']
+                pivot_data = pivot_data.iloc[row_order]
+                
+                logger.info(f"Applied clustering to rows using {cluster_method} linkage and {cluster_metric} metric")
+            except Exception as e:
+                logger.warning(f"Failed to cluster rows: {e}")
+        
+        # Cluster columns (x-axis)
+        if cluster_axes in ["both", "x"]:
+            try:
+                # Calculate distance matrix and linkage for columns
+                col_distances = distance.pdist(pivot_data_for_clustering.values.T, metric=cluster_metric)
+                col_linkage = hierarchy.linkage(col_distances, method=cluster_method)
+                
+                # Reorder columns based on clustering
+                col_dendrogram = hierarchy.dendrogram(col_linkage, no_plot=True)
+                col_order = col_dendrogram['leaves']
+                pivot_data = pivot_data.iloc[:, col_order]
+                
+                logger.info(f"Applied clustering to columns using {cluster_method} linkage and {cluster_metric} metric")
+            except Exception as e:
+                logger.warning(f"Failed to cluster columns: {e}")
+
+    # Create the heatmap
+    sns.heatmap(
+        pivot_data,
+        cmap=cmap,
+        annot=annot,
+        fmt=fmt,
+        linewidths=linewidths,
+        linecolor=linecolor,
+        square=square,
+        vmin=vmin,
+        vmax=vmax,
+        robust=robust,
+        ax=ax,
+        annot_kws={'fontsize': font_size},
+        **kwargs
+    )
+    
+    # Set title if provided
+    if title:
+        ax.set_title(title, fontsize=font_size + 4)
+    
+    # Improve display of tick labels
+    plt.xticks(rotation=45, ha="right", fontsize=font_size)
+    plt.yticks(rotation=0, fontsize=font_size)
+    
+    # Add grid lines to make the table more readable
+    ax.grid(False)
+    
+    # Improve contrast for better readability
+    for _, spine in ax.spines.items():
+        spine.set_visible(True)
+        spine.set_color('black')
+        spine.set_linewidth(1)
+    
+    # Adjust layout
+    plt.tight_layout()
+    
+    # Save the figure if output file is specified
+    if output_file:
+        output_path = Path(output_file)
+        output_dir = output_path.parent
+        if not output_dir.exists():
+            output_dir.mkdir(parents=True, exist_ok=True)
+        plt.savefig(output_file, dpi=dpi, bbox_inches="tight")
+        logger.info(f"Heatmap saved to {output_file}")
+    
+    return fig, ax
+
+
+def heatmap_from_file(
+    file_path: str,
+    x_column: str,
+    y_column: str,
+    value_column: Optional[str] = None,
+    format: Optional[Union[Format, str]] = None,
+    compression: Optional[str] = None,
+    output_file: Optional[str] = None,
+    remove_duplicates: bool = True,
+    **kwargs,
+) -> Tuple[plt.Figure, plt.Axes]:
+    """
+    Create a heatmap from a file (CSV, TSV, etc.).
+
+    Args:
+        file_path: Path to the input file or "-" for stdin
+        x_column: Column to use for x-axis categories
+        y_column: Column to use for y-axis categories
+        value_column: Column containing values for the heatmap. If None, frequency counts will be used.
+        format: Format of the input file (auto-detected if None)
+        compression: Compression format ('gz' or 'tgz')
+        output_file: File path to save the figure (optional)
+        remove_duplicates: If True, removes duplicate rows before creating the heatmap
+        **kwargs: Additional arguments to pass to create_heatmap
+
+    Returns:
+        Tuple containing the figure and axes objects
+    """
+    # Handle stdin input safely
+    import sys
+    import io
+    import pandas as pd
+    import click
+    
+    # Load the data
+    if file_path == "-":
+        # Read directly from stdin since format_utils will use sys.stdin which may already be consumed
+        if not format or str(format).lower() in ['csv', 'tsv']:
+            # Default to CSV if no format specified
+            delimiter = ',' if not format or str(format).lower() == 'csv' else '\t'
+            df = pd.read_csv(sys.stdin, delimiter=delimiter)
+        else:
+            # Try to use format_utils but with a backup plan
+            try:
+                objs = load_objects(file_path, format=format, compression=compression)
+                df = pd.DataFrame(objs)
+            except ValueError as e:
+                if "I/O operation on closed file" in str(e):
+                    logger.warning("Could not read from stdin. It may have been consumed already.")
+                    raise click.UsageError("Error reading from stdin. Please provide a file path or ensure stdin has data.")
+                else:
+                    raise
+    else:
+        # For regular files, use format_utils as normal
+        if (not format or format in ["csv", "tsv"]) and not compression:
+            df = pd.read_csv(file_path)
+        else:
+            objs = load_objects(file_path, format=format, compression=compression)
+            df = pd.DataFrame(objs)
+
+    # Create the heatmap
+    return create_heatmap(
+        data=df,
+        x_column=x_column,
+        y_column=y_column,
+        value_column=value_column,
+        output_file=output_file,
+        remove_duplicates=remove_duplicates,
+        **kwargs
+    )
+
+
+def export_heatmap_data(
+    data: pd.DataFrame,
+    x_column: str,
+    y_column: str,
+    value_column: Optional[str] = None,
+    output_file: Optional[str] = None,
+    format: Union[Format, str] = Format.CSV,
+    missing_value: Any = np.nan,
+    remove_duplicates: bool = True,
+) -> pd.DataFrame:
+    """
+    Export heatmap data to a file or return it as a DataFrame.
+
+    Args:
+        data: Input DataFrame containing the data
+        x_column: Column to use for x-axis categories
+        y_column: Column to use for y-axis categories
+        value_column: Column containing values for the heatmap. If None, frequency counts will be used.
+        output_file: File path to save the data (optional)
+        format: Output format for the file
+        missing_value: Value to use for missing data
+        remove_duplicates: If True, removes duplicate rows before creating the pivot table
+
+    Returns:
+        DataFrame containing the pivot table data
+    """
+    # Remove duplicates by default (assume they're accidents unless user overrides)
+    if remove_duplicates:
+        # Keep the first occurrence of each x_column, y_column combination
+        data = data.drop_duplicates(subset=[x_column, y_column])
+
+    # Prepare the data
+    if value_column:
+        # Use the provided value column
+        pivot_data = data.pivot_table(
+            index=y_column, 
+            columns=x_column, 
+            values=value_column, 
+            aggfunc='mean',
+            fill_value=missing_value
+        )
+    else:
+        # Use frequency counts
+        cross_tab = pd.crosstab(data[y_column], data[x_column])
+        pivot_data = cross_tab
+    
+    # Reset index to make the y_column a regular column
+    result_df = pivot_data.reset_index()
+    
+    # Write to file if output_file is provided
+    if output_file:
+        # Convert to records format for writing
+        records = result_df.to_dict(orient='records')
+        write_output(records, format=format, target=output_file)
+        logger.info(f"Heatmap data saved to {output_file}")
+    
+    return result_df

linkml_store/utils/dat_parser.py

@@ -0,0 +1,95 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+ENTRY = Dict[str, Any]
+
+
+def parse_sib_format(text) -> Tuple[Optional[ENTRY], List[ENTRY]]:
+    """
+    Parse SIB/Swiss-Prot format data into a structured dictionary.
+
+    Args:
+        text (str): The text in SIB/Swiss-Prot format
+
+    Returns:
+        dict: A dictionary with entry IDs as keys and parsed data as values
+    """
+    # Split the text into entries (separated by //)
+    entries = text.split("//\n")
+    header = None
+
+    # Initialize results dictionary
+    results = []
+
+    # Parse each entry
+    for entry in entries:
+        if not entry.strip():
+            continue
+
+        # Initialize dictionary for current entry
+        current_entry = {}
+        current_code = None
+
+        # Process each line
+        for line in entry.strip().split("\n"):
+            if not line.strip():
+                continue
+
+            # Check if this is a new field (starts with a 2-letter code followed by space)
+            if len(line) > 2 and line[2] == " ":
+                current_code = line[0:2]
+                # Remove the code and the following space(s)
+                value = line[3:].strip()
+
+                # Initialize as list if needed for multi-line fields
+                if current_code not in current_entry:
+                    current_entry[current_code] = []
+
+                current_entry[current_code].append(value)
+
+            # Continuation of previous field
+            elif current_code is not None:
+                # Handle continuation lines (typically indented)
+                if current_code == "CC":
+                    # For comments, preserve the indentation
+                    current_entry[current_code].append(line)
+                else:
+                    # For other fields, strip and append
+                    current_entry[current_code].append(line.strip())
+
+        # Combine multiline comments; e.g
+        # -!- ...
+        #     ...
+        # -!- ...
+        ccs = current_entry.get("CC", [])
+        new_ccs = []
+        for cc in ccs:
+            if not cc.startswith("-!-") and new_ccs:
+                new_ccs[-1] += " " + cc
+            else:
+                new_ccs.append(cc)
+        current_entry["CC"] = new_ccs
+        for k, vs in current_entry.items():
+            if k != "CC":
+                combined = "".join(vs)
+                combined = combined.strip()
+                if combined.endswith("."):
+                    combined = combined.split(".")
+                    combined = [c.strip() for c in combined if c.strip()]
+                    if k == "DE":
+                        combined = combined[0]
+                current_entry[k] = combined
+
+        if "ID" in current_entry:
+            results.append(current_entry)
+        else:
+            header = current_entry
+
+    return header, results
+
+
+# Example usage:
+# data = parse_sib_format(text)
+# for entry_id, entry_data in data.items():
+#     print(f"Entry: {entry_id}")
+#     for code, values in entry_data.items():
+#         print(f"  {code}: {values}")

linkml_store/utils/enrichment_analyzer.py

@@ -0,0 +1,217 @@
+from collections import Counter
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+from pydantic import BaseModel
+from scipy import stats
+
+from linkml_store.api import Collection
+
+
+class EnrichedCategory(BaseModel):
+    """
+    Information about a category enriched in a sample
+    """
+
+    category: str
+    fold_change: float
+    original_p_value: float
+    adjusted_p_value: float
+
+
+class EnrichmentAnalyzer:
+    def __init__(self, df: pd.DataFrame, sample_key: str, classification_key: str):
+        """
+        Initialize the analyzer with a DataFrame and key column names.
+        Precomputes category frequencies for the entire dataset.
+
+        Args:
+            df: DataFrame containing the data
+            sample_key: Column name for sample IDs
+            classification_key: Column name for category lists
+        """
+        self.df = df
+        self.sample_key = sample_key
+        self.classification_key = classification_key
+
+        # Precompute global category statistics
+        self.global_stats = self._compute_global_stats()
+
+        # Cache for sample-specific category counts
+        self.sample_cache: Dict[str, Counter] = {}
+
+    @classmethod
+    def from_collection(cls, collection: Collection, sample_key: str, classification_key: str) -> "EnrichmentAnalyzer":
+        """
+        Initialize the analyzer with a Collection and key column names.
+        Precomputes category frequencies for the entire dataset.
+
+        Args:
+            collection: Collection containing the data
+            sample_key: Column name for sample IDs
+            classification_key: Column name for category lists
+        """
+        column_atts = [sample_key, classification_key]
+        results = collection.find(select_cols=column_atts, limit=-1)
+        df = results.rows_dataframe
+        ea = cls(df, sample_key=sample_key, classification_key=classification_key)
+        return ea
+
+    def _compute_global_stats(self) -> Dict[str, int]:
+        """
+        Compute global category frequencies across all samples.
+        Returns a dictionary of category -> count
+        """
+        global_counter = Counter()
+
+        # Flatten all categories and count
+        for categories in self.df[self.classification_key]:
+            if isinstance(categories, list):
+                global_counter.update(categories)
+            else:
+                # Handle case where categories might be a string
+                global_counter.update([categories])
+
+        return global_counter
+
+    @property
+    def sample_ids(self) -> List[str]:
+        df = self.df
+        return df[self.sample_key].unique().tolist()
+
+    def _get_sample_stats(self, sample_id: str) -> Counter:
+        """
+        Get category frequencies for a specific sample.
+        Uses caching to avoid recomputation.
+        """
+        if sample_id in self.sample_cache:
+            return self.sample_cache[sample_id]
+
+        sample_data = self.df[self.df[self.sample_key] == sample_id]
+        if sample_data.empty:
+            raise KeyError(f"Sample ID '{sample_id}' not found")
+        sample_data = sample_data.dropna()
+        # if sample_data.empty:
+        #    raise ValueError(f"Sample ID '{sample_id}' has missing values after dropping NA")
+        counter = Counter()
+
+        for categories in sample_data[self.classification_key]:
+            if isinstance(categories, list):
+                counter.update(categories)
+            else:
+                counter.update([categories])
+
+        self.sample_cache[sample_id] = counter
+        return counter
+
+    def find_enriched_categories(
+        self,
+        sample_id: str,
+        min_occurrences: int = 5,
+        p_value_threshold: float = 0.05,
+        multiple_testing_correction: str = "bh",
+    ) -> List[EnrichedCategory]:
+        """
+        Find categories that are enriched in the given sample.
+
+        Args:
+            sample_id: ID of the sample to analyze
+            min_occurrences: Minimum number of occurrences required for a category
+            p_value_threshold: P-value threshold for significance
+
+        Returns:
+            List of tuples (category, fold_change, p_value) sorted by significance
+        """
+        sample_stats = self._get_sample_stats(sample_id)
+        total_sample_annotations = sum(sample_stats.values())
+        total_global_annotations = sum(self.global_stats.values())
+
+        results = []
+
+        for category, sample_count in sample_stats.items():
+            global_count = self.global_stats[category]
+
+            # Skip rare categories
+            if global_count < min_occurrences:
+                continue
+
+            # Calculate fold change
+            sample_freq = sample_count / total_sample_annotations
+            global_freq = global_count / total_global_annotations
+            fold_change = sample_freq / global_freq if global_freq > 0 else float("inf")
+
+            # Perform Fisher's exact test
+            contingency_table = np.array(
+                [
+                    [sample_count, global_count - sample_count],
+                    [
+                        total_sample_annotations - sample_count,
+                        total_global_annotations - total_sample_annotations - (global_count - sample_count),
+                    ],
+                ]
+            )
+
+            _, p_value = stats.fisher_exact(contingency_table)
+
+            if p_value < p_value_threshold:
+                results.append((category, fold_change, p_value))
+
+        if not results:
+            return results
+
+        # Sort by p-value
+        results.sort(key=lambda x: x[2])
+
+        # Apply multiple testing correction
+        categories, fold_changes, p_values = zip(*results)
+
+        if multiple_testing_correction.lower() == "bonf":
+            # Bonferroni correction
+            n_tests = len(self.global_stats)  # Total number of categories tested
+            adjusted_p_values = [min(1.0, p * n_tests) for p in p_values]
+
+        elif multiple_testing_correction.lower() == "bh":
+            # Benjamini-Hochberg correction
+            n = len(p_values)
+            sorted_indices = np.argsort(p_values)
+            sorted_p_values = np.array(p_values)[sorted_indices]
+
+            # Calculate BH adjusted p-values
+            adjusted_p_values = np.zeros(n)
+            for i, p in enumerate(sorted_p_values):
+                adjusted_p_values[i] = p * n / (i + 1)
+
+            # Ensure monotonicity
+            for i in range(n - 2, -1, -1):
+                adjusted_p_values[i] = min(adjusted_p_values[i], adjusted_p_values[i + 1])
+
+            # Restore original order
+            inverse_indices = np.argsort(sorted_indices)
+            adjusted_p_values = adjusted_p_values[inverse_indices]
+
+            # Ensure we don't exceed 1.0
+            adjusted_p_values = np.minimum(adjusted_p_values, 1.0)
+
+        else:
+            # No correction
+            adjusted_p_values = p_values
+
+        # Filter by adjusted p-value threshold and create final results
+        # Create EnrichedCategory objects
+        final_results = [
+            EnrichedCategory(category=cat, fold_change=fc, original_p_value=p, adjusted_p_value=adj_p)
+            for cat, fc, p, adj_p in zip(categories, fold_changes, p_values, adjusted_p_values)
+            if adj_p < p_value_threshold
+        ]
+
+        # Sort by adjusted p-value
+        final_results.sort(key=lambda x: x.adjusted_p_value)
+        return final_results
+
+
+# Example usage:
+# analyzer = EnrichmentAnalyzer(df, 'sample_id', 'categories')
+# enriched = analyzer.find_enriched_categories('sample1')
+# for category, fold_change, p_value in enriched:
+#     print(f"{category}: {fold_change:.2f}x enrichment (p={p_value:.2e})")