masster 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

masster/study/plot.py

@@ -7,14 +7,6 @@ import holoviews as hv
 import numpy as np
 import panel
 import polars as pl
-
-from bokeh.io.export import export_png
-from bokeh.models import ColumnDataSource
-from bokeh.models import HoverTool
-from bokeh.palettes import Turbo256
-from bokeh.plotting import figure
-from bokeh.plotting import output_file
-from bokeh.plotting import show
 from tqdm import tqdm
 
 hv.extension("bokeh")
@@ -163,11 +155,11 @@ def plot_consensus_2d(
     width=900,
     height=900,
     mz_range=None,
-    rt_range=None
+    rt_range=None,
 ):
     """
     Plot consensus features in a 2D scatter plot with retention time vs m/z.
-    
+
     Parameters:
         filename (str, optional): Path to save the plot
         colorby (str): Column name to use for color mapping (default: "number_samples")
@@ -187,13 +179,13 @@ def plot_consensus_2d(
         self.logger.error("No consensus map found.")
         return
     data = self.consensus_df.clone()
-    
+
     # Filter by mz_range and rt_range if provided
     if mz_range is not None:
         data = data.filter((pl.col("mz") >= mz_range[0]) & (pl.col("mz") <= mz_range[1]))
     if rt_range is not None:
         data = data.filter((pl.col("rt") >= rt_range[0]) & (pl.col("rt") <= rt_range[1]))
-    
+
     if colorby not in data.columns:
         self.logger.error(f"Column {colorby} not found in consensus_df.")
         return
@@ -342,13 +334,13 @@ def plot_samples_2d(
     width=900,
     height=900,
     mz_range=None,
-    rt_range=None
+    rt_range=None,
 ):
     """
     Plot all feature maps for sample_uid in parameter uids in an overlaid scatter plot.
     Each sample is a different color. Alpha scales with intensity.
     OPTIMIZED VERSION: Uses vectorized operations and batch processing.
-    
+
     Parameters:
         samples: Sample UIDs to plot
         filename (str, optional): Path to save the plot
@@ -366,6 +358,12 @@ def plot_samples_2d(
         rt_range (tuple, optional): Retention time range for filtering features (min_rt, max_rt)
     """
 
+    # Local bokeh imports to avoid heavy top-level dependency
+    from bokeh.plotting import figure, show, output_file
+    from bokeh.io.export import export_png
+    from bokeh.models import ColumnDataSource, HoverTool
+    from bokeh.palettes import Turbo256
+
     sample_uids = self._get_sample_uids(samples)
 
     if not sample_uids:
@@ -385,7 +383,7 @@ def plot_samples_2d(
 
     # OPTIMIZATION 1: Batch filter all features for selected samples at once
     features_batch = self.features_df.filter(pl.col("sample_uid").is_in(sample_uids))
-    
+
     # Filter by mz_range and rt_range if provided
     if mz_range is not None:
         features_batch = features_batch.filter((pl.col("mz") >= mz_range[0]) & (pl.col("mz") <= mz_range[1]))
@@ -560,6 +558,9 @@ def plot_chrom(
         self.logger.error("No chromatogram data found.")
         return
 
+    # Local import for color palette
+    from bokeh.palettes import Turbo256
+
     # Assign a fixed color to each sample/column
     sample_names = [col for col in chroms.columns if col not in ["consensus_uid"]]
     if not sample_names:
@@ -569,12 +570,12 @@ def plot_chrom(
 
     plots = []
     self.logger.info(f"Plotting {chroms.shape[0]} chromatograms...")
-    tdqm_disable = self.log_level not in ["TRACE", "DEBUG", "INFO"]
+    tqdm_disable = self.log_level not in ["TRACE", "DEBUG", "INFO"]
     for row in tqdm(
         chroms.iter_rows(named=True),
         total=chroms.shape[0],
         desc=f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]} | INFO     | {self.log_label}Plot chromatograms",
-        disable=tdqm_disable,
+        disable=tqdm_disable,
     ):
         consensus_uid = row["consensus_uid"]  # Get consensus_uid from the row
         consensus_id = consensus_uid  # Use the same value for consensus_id
@@ -698,3 +699,438 @@ def plot_chrom(
         # In a server context, return the panel object instead of showing or saving directly
         # return panel.panel(layout)
         panel.panel(layout).show()
+
+
+def plot_consensus_stats(
+    self,
+    filename=None,
+    width=1200,
+    height=1200,
+    alpha=0.6,
+    markersize=3,
+):
+    """
+    Plot a scatter plot matrix (SPLOM) of consensus statistics using Bokeh.
+
+    Parameters:
+        filename (str, optional): Output filename for saving the plot
+        width (int): Overall width of the plot (default: 1200)
+        height (int): Overall height of the plot (default: 1200)
+        alpha (float): Point transparency (default: 0.6)
+        markersize (int): Size of points (default: 5)
+    """
+    from bokeh.layouts import gridplot
+    from bokeh.models import ColumnDataSource, HoverTool
+    from bokeh.plotting import figure, show, output_file
+
+    # Check if consensus_df exists and has data
+    if self.consensus_df is None or self.consensus_df.is_empty():
+        self.logger.error("No consensus data available. Run merge/find_consensus first.")
+        return
+
+    # Define the columns to plot
+    columns = [
+        "rt",
+        "mz",
+        "number_samples",
+        "log10_quality",
+        "mz_delta_mean",
+        "rt_delta_mean",
+        "chrom_coherence_mean",
+        "chrom_prominence_scaled_mean",
+        "inty_mean",
+        "number_ms2",
+    ]
+
+    # Check which columns exist in the dataframe and compute missing ones
+    available_columns = self.consensus_df.columns
+    data_df = self.consensus_df.clone()
+
+    # Add log10_quality if quality exists
+    if "quality" in available_columns and "log10_quality" not in available_columns:
+        data_df = data_df.with_columns(
+            pl.col("quality").log10().alias("log10_quality"),
+        )
+
+    # Filter columns that actually exist
+    final_columns = [col for col in columns if col in data_df.columns]
+
+    if len(final_columns) < 2:
+        self.logger.error(f"Need at least 2 columns for SPLOM. Available: {final_columns}")
+        return
+
+    self.logger.debug(f"Creating SPLOM with columns: {final_columns}")
+
+    # Add important ID columns for tooltips even if not plotting them
+    tooltip_columns = []
+    for id_col in ["consensus_uid", "consensus_id"]:
+        if id_col in data_df.columns and id_col not in final_columns:
+            tooltip_columns.append(id_col)
+
+    # Select plotting columns plus tooltip columns
+    all_columns = final_columns + tooltip_columns
+    data_pd = data_df.select(all_columns).to_pandas()
+
+    # Remove any infinite or NaN values
+    data_pd = data_pd.replace([np.inf, -np.inf], np.nan).dropna()
+
+    if data_pd.empty:
+        self.logger.error("No valid data after removing NaN/infinite values.")
+        return
+
+    source = ColumnDataSource(data_pd)
+
+    n_vars = len(final_columns)
+
+    # Fixed dimensions - override user input to ensure consistent layout
+    total_width = 1200
+    total_height = 1200
+
+    # Calculate plot sizes to ensure uniform inner plot areas
+    # First column needs extra width for y-axis labels
+    plot_width_first = 180  # Wider to account for y-axis labels
+    plot_width_others = 120  # Standard width for other columns
+    plot_height_normal = 120  # Standard height
+    plot_height_last = 155  # Taller last row to accommodate x-axis labels while keeping inner plot area same size
+
+    # Create grid of plots with variable outer sizes but equal inner areas
+    plots = []
+
+    for i, y_var in enumerate(final_columns):
+        row = []
+        for j, x_var in enumerate(final_columns):
+            # Determine if this plot needs axis labels
+            has_x_label = i == n_vars - 1  # bottom row
+            has_y_label = j == 0  # left column
+
+            # First column wider to accommodate y-axis labels, ensuring equal inner plot areas
+            current_width = plot_width_first if has_y_label else plot_width_others
+            current_height = plot_height_last if has_x_label else plot_height_normal
+
+            p = figure(
+                width=current_width,
+                height=current_height,
+                title=None,  # No title on any plot
+                toolbar_location=None,
+                # Adjusted borders - first column has more space, others minimal
+                min_border_left=70 if has_y_label else 15,
+                min_border_bottom=50 if has_x_label else 15,
+                min_border_right=15,
+                min_border_top=15,
+            )
+
+            # Ensure subplot background and border are explicitly white so the plot looks
+            # correct in dark and light themes.
+            p.outline_line_color = None
+            p.border_fill_color = "white"
+            p.border_fill_alpha = 1.0
+            p.background_fill_color = "white"
+
+            # Remove axis lines to eliminate black lines between plots
+            p.xaxis.axis_line_color = None
+            p.yaxis.axis_line_color = None
+
+            # Keep subtle grid lines for data reference
+            p.grid.visible = True
+            p.grid.grid_line_color = "#E0E0E0"  # Light gray grid lines
+
+            # Set axis labels and formatting
+            if has_x_label:  # bottom row
+                p.xaxis.axis_label = x_var
+                p.xaxis.axis_label_text_font_size = "12pt"
+                p.xaxis.major_label_text_font_size = "9pt"
+                p.xaxis.axis_label_standoff = 15
+            else:
+                p.xaxis.major_label_text_font_size = "0pt"
+                p.xaxis.minor_tick_line_color = None
+                p.xaxis.major_tick_line_color = None
+
+            if has_y_label:  # left column
+                p.yaxis.axis_label = y_var
+                p.yaxis.axis_label_text_font_size = "10pt"  # Smaller y-axis title
+                p.yaxis.major_label_text_font_size = "8pt"
+                p.yaxis.axis_label_standoff = 12
+            else:
+                p.yaxis.major_label_text_font_size = "0pt"
+                p.yaxis.minor_tick_line_color = None
+                p.yaxis.major_tick_line_color = None
+
+            if i == j:
+                # Diagonal: histogram
+                hist, edges = np.histogram(data_pd[x_var], bins=30)
+                p.quad(
+                    top=hist,
+                    bottom=0,
+                    left=edges[:-1],
+                    right=edges[1:],
+                    fill_color="green",
+                    line_color="white",
+                    alpha=alpha,
+                )
+            else:
+                # Off-diagonal: scatter plot
+                scatter = p.scatter(
+                    x=x_var,
+                    y=y_var,
+                    size=markersize,
+                    alpha=alpha,
+                    color="blue",
+                    source=source,
+                )
+
+                # Add hover tool
+                hover = HoverTool(
+                    tooltips=[
+                        (x_var, f"@{x_var}{{0.0000}}"),
+                        (y_var, f"@{y_var}{{0.0000}}"),
+                        (
+                            "consensus_uid",
+                            "@consensus_uid"
+                            if "consensus_uid" in data_pd.columns
+                            else "@consensus_id"
+                            if "consensus_id" in data_pd.columns
+                            else "N/A",
+                        ),
+                        ("rt", "@rt{0.00}" if "rt" in data_pd.columns else "N/A"),
+                        ("mz", "@mz{0.0000}" if "mz" in data_pd.columns else "N/A"),
+                    ],
+                    renderers=[scatter],
+                )
+                p.add_tools(hover)
+
+            row.append(p)
+        plots.append(row)
+
+    # Link axes for same variables
+    for i in range(n_vars):
+        for j in range(n_vars):
+            if i != j:  # Don't link diagonal plots
+                # Link x-axis to other plots in same column
+                for k in range(n_vars):
+                    if k != i and k != j:
+                        plots[i][j].x_range = plots[k][j].x_range
+
+                # Link y-axis to other plots in same row
+                for k in range(n_vars):
+                    if k != j and k != i:
+                        plots[i][j].y_range = plots[i][k].y_range
+
+    # Create grid layout and force overall background/border to white so the outer
+    # container doesn't show dark UI colors in night mode.
+    grid = gridplot(plots)
+
+    # Set overall background and border to white when supported
+    if hasattr(grid, "background_fill_color"):
+        grid.background_fill_color = "white"
+    if hasattr(grid, "border_fill_color"):
+        grid.border_fill_color = "white"
+
+    # Output and show
+    if filename:
+        output_file(filename)
+
+    show(grid)
+    return grid
+
+
+def plot_pca(
+    self,
+    filename=None,
+    width=600,
+    height=600,
+    alpha=0.8,
+    markersize=8,
+    n_components=2,
+    color_by=None,
+    title="PCA of Consensus Matrix",
+):
+    """
+    Plot PCA (Principal Component Analysis) of the consensus matrix using Bokeh.
+
+    Parameters:
+        filename (str, optional): Output filename for saving the plot
+        width (int): Plot width (default: 800)
+        height (int): Plot height (default: 600)
+        alpha (float): Point transparency (default: 0.8)
+        markersize (int): Size of points (default: 8)
+        n_components (int): Number of PCA components to compute (default: 2)
+        color_by (str, optional): Column from samples_df to color points by
+        title (str): Plot title (default: "PCA of Consensus Matrix")
+    """
+    from bokeh.models import ColumnDataSource, HoverTool, ColorBar, LinearColorMapper
+    from bokeh.plotting import figure, show, output_file
+    from bokeh.palettes import Category20, viridis
+    from bokeh.transform import factor_cmap
+    from sklearn.decomposition import PCA
+    from sklearn.preprocessing import StandardScaler
+    import pandas as pd
+    import numpy as np
+
+    # Check if consensus matrix and samples_df exist
+    try:
+        consensus_matrix = self.get_consensus_matrix()
+        samples_df = self.samples_df
+    except Exception as e:
+        self.logger.error(f"Error getting consensus matrix or samples_df: {e}")
+        return
+
+    if consensus_matrix is None or consensus_matrix.shape[0] == 0:
+        self.logger.error("No consensus matrix available. Run merge/find_consensus first.")
+        return
+
+    if samples_df is None or samples_df.is_empty():
+        self.logger.error("No samples dataframe available.")
+        return
+
+    self.logger.info(f"Performing PCA on consensus matrix with shape: {consensus_matrix.shape}")
+
+    # Convert consensus matrix to numpy if it's not already
+    if hasattr(consensus_matrix, "values"):
+        matrix_data = consensus_matrix.values
+    elif hasattr(consensus_matrix, "to_numpy"):
+        matrix_data = consensus_matrix.to_numpy()
+    else:
+        matrix_data = np.array(consensus_matrix)
+
+    # Transpose matrix so samples are rows and features are columns
+    matrix_data = matrix_data.T
+
+    # Handle missing values by replacing with 0
+    matrix_data = np.nan_to_num(matrix_data, nan=0.0, posinf=0.0, neginf=0.0)
+
+    # Standardize the data
+    scaler = StandardScaler()
+    matrix_scaled = scaler.fit_transform(matrix_data)
+
+    # Perform PCA
+    pca = PCA(n_components=n_components)
+    pca_result = pca.fit_transform(matrix_scaled)
+
+    # Get explained variance ratios
+    explained_var = pca.explained_variance_ratio_
+
+    self.logger.info(f"PCA explained variance ratios: {explained_var}")
+
+    # Convert samples_df to pandas for easier manipulation
+    samples_pd = samples_df.to_pandas()
+
+    # Create dataframe with PCA results and sample information
+    pca_df = pd.DataFrame({
+        "PC1": pca_result[:, 0],
+        "PC2": pca_result[:, 1] if n_components > 1 else np.zeros(len(pca_result)),
+    })
+
+    # Add sample information to PCA dataframe
+    if len(samples_pd) == len(pca_df):
+        for col in samples_pd.columns:
+            pca_df[col] = samples_pd[col].values
+    else:
+        self.logger.warning(
+            f"Sample count mismatch: samples_df has {len(samples_pd)} rows, "
+            f"but consensus matrix has {len(pca_df)} samples"
+        )
+
+    # Prepare color mapping
+    color_column = None
+    color_mapper = None
+
+    if color_by and color_by in pca_df.columns:
+        color_column = color_by
+        unique_values = pca_df[color_by].unique()
+
+        # Handle categorical vs numeric coloring
+        if pca_df[color_by].dtype in ["object", "string", "category"]:
+            # Categorical coloring
+            if len(unique_values) <= 20:
+                palette = Category20[min(20, max(3, len(unique_values)))]
+            else:
+                palette = viridis(min(256, len(unique_values)))
+            color_mapper = factor_cmap(color_by, palette, unique_values)
+        else:
+            # Numeric coloring
+            palette = viridis(256)
+            color_mapper = LinearColorMapper(
+                palette=palette,
+                low=pca_df[color_by].min(),
+                high=pca_df[color_by].max(),
+            )
+
+    # Create Bokeh plot
+    p = figure(
+        width=width,
+        height=height,
+        title=f"{title} (PC1: {explained_var[0]:.1%}, PC2: {explained_var[1]:.1%})",
+        tools="pan,wheel_zoom,box_zoom,reset,save",
+    )
+
+    p.xaxis.axis_label = f"PC1 ({explained_var[0]:.1%} variance)"
+    p.yaxis.axis_label = f"PC2 ({explained_var[1]:.1%} variance)"
+
+    # Create data source
+    source = ColumnDataSource(pca_df)
+
+    # Create scatter plot
+    if color_mapper:
+        if isinstance(color_mapper, LinearColorMapper):
+            scatter = p.scatter(
+                "PC1",
+                "PC2",
+                size=markersize,
+                alpha=alpha,
+                color={"field": color_by, "transform": color_mapper},
+                source=source,
+            )
+            # Add colorbar for numeric coloring
+            color_bar = ColorBar(color_mapper=color_mapper, width=8, location=(0, 0))
+            p.add_layout(color_bar, "right")
+        else:
+            scatter = p.scatter(
+                "PC1",
+                "PC2",
+                size=markersize,
+                alpha=alpha,
+                color=color_mapper,
+                source=source,
+                legend_field=color_by,
+            )
+    else:
+        scatter = p.scatter(
+            "PC1",
+            "PC2",
+            size=markersize,
+            alpha=alpha,
+            color="blue",
+            source=source,
+        )
+
+    # Create comprehensive hover tooltips with all sample information
+    tooltip_list = [
+        ("PC1", "@PC1{0.00}"),
+        ("PC2", "@PC2{0.00}"),
+    ]
+
+    # Add all sample dataframe columns to tooltips
+    for col in samples_pd.columns:
+        if col in pca_df.columns:
+            if pca_df[col].dtype in ["float64", "float32"]:
+                tooltip_list.append((col, f"@{col}{{0.00}}"))
+            else:
+                tooltip_list.append((col, f"@{col}"))
+
+    hover = HoverTool(
+        tooltips=tooltip_list,
+        renderers=[scatter],
+    )
+    p.add_tools(hover)
+
+    # Add legend if using categorical coloring
+    if color_mapper and not isinstance(color_mapper, LinearColorMapper) and color_by:
+        p.legend.location = "top_left"
+        p.legend.click_policy = "hide"
+
+    # Output and show
+    if filename:
+        output_file(filename)
+
+    show(p)
+    return p