gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/report/visualize.py

@@ -0,0 +1,1409 @@
+import gc
+import logging
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Literal
+
+import distinctipy
+import matplotlib
+import matplotlib.axes
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+import scipy.stats as stats
+from rich import print
+from scipy.cluster.hierarchy import leaves_list, linkage
+from scipy.spatial import KDTree
+from tqdm import tqdm
+
+warnings.filterwarnings("ignore")
+
+
+def remove_outliers_MAD(data, threshold=3.5):
+    """
+    Remove outliers based on Median Absolute Deviation (MAD).
+    """
+    if isinstance(data, pd.Series | pd.DataFrame):
+        data_values = data.values.flatten()
+    else:
+        data_values = np.asarray(data).flatten()
+
+    if len(data_values) == 0:
+        return data, np.ones(len(data), dtype=bool)
+
+    median = np.nanmedian(data_values)
+    mad = np.nanmedian(np.abs(data_values - median))
+    if mad == 0:
+        return data, np.ones(len(data), dtype=bool)
+
+    modified_z_score = 0.6745 * (data_values - median) / mad
+    mask = np.abs(modified_z_score) <= threshold
+
+    if isinstance(data, pd.Series):
+        return data[mask], mask
+    elif isinstance(data, np.ndarray):
+        if len(data.shape) == 1:
+            return data[mask], mask
+        else:
+            return data.flatten()[mask], mask
+    return data[mask], mask
+
+
+def load_ldsc(ldsc_input_file):
+    ldsc = pd.read_csv(
+        ldsc_input_file,
+        compression="gzip",
+        dtype={"spot": str, "p": float},
+        index_col="spot",
+        usecols=["spot", "p"],
+    )
+    ldsc["logp"] = -np.log10(ldsc.p)
+    return ldsc
+
+
+# %%
+def load_st_coord(adata, feature_series: pd.Series, annotation):
+    spot_name = adata.obs_names.to_list()
+    assert "spatial" in adata.obsm.keys(), "spatial coordinates are not found in adata.obsm"
+
+    # to DataFrame
+    space_coord = adata.obsm["spatial"]
+    if isinstance(space_coord, np.ndarray):
+        space_coord = pd.DataFrame(space_coord, columns=["sx", "sy"], index=spot_name)
+    else:
+        space_coord = pd.DataFrame(space_coord.values, columns=["sx", "sy"], index=spot_name)
+
+    space_coord = space_coord[space_coord.index.isin(feature_series.index)]
+    space_coord_concat = pd.concat([space_coord.loc[feature_series.index], feature_series], axis=1)
+    space_coord_concat.head()
+    if annotation is not None:
+        annotation = pd.Series(
+            adata.obs[annotation].values, index=adata.obs_names, name="annotation"
+        )
+        space_coord_concat = pd.concat([space_coord_concat, annotation], axis=1)
+    return space_coord_concat
+
+
+def estimate_plotly_point_size(coordinates, DEFAULT_PIXEL_WIDTH=1000):
+    # Convert to numpy array if it's a DataFrame or other array-like object
+    if hasattr(coordinates, 'values'):
+        coordinates = coordinates.values
+    coordinates = np.asarray(coordinates)
+
+    tree = KDTree(coordinates)
+    distances, _ = tree.query(coordinates, k=2)
+    avg_min_distance = np.median(distances[:, 1])
+    # get the width and height of the plot
+    width = np.max(coordinates[:, 0]) - np.min(coordinates[:, 0])
+    height = np.max(coordinates[:, 1]) - np.min(coordinates[:, 1])
+
+    scale_factor = DEFAULT_PIXEL_WIDTH / max(width, height)
+    pixel_width = width * scale_factor
+    pixel_height = height * scale_factor
+
+    point_size = avg_min_distance * scale_factor
+
+    return (pixel_width, pixel_height), point_size
+
+
+def estimate_matplotlib_scatter_marker_size(ax: matplotlib.axes.Axes, coordinates: np.ndarray,
+                                            x_limits: tuple | None = None,
+                                            y_limits: tuple | None = None) -> float:
+    """
+    Estimates the appropriate marker size to make adjacent markers touch.
+
+    This function calculates the size 's' for a square marker (in points^2)
+    such that its diameter in the plot corresponds to the average distance
+    to the nearest neighbor for each point in the dataset. It accounts for
+    the plot's aspect ratio and final rendered dimensions.
+
+    Args:
+        ax (matplotlib.axes.Axes): The subplot object. The function will
+            temporarily set its limits and aspect ratio to ensure the
+            transformation from data units to display units is accurate.
+        coordinates (np.ndarray): A NumPy array of shape (n, 2)
+            containing the (x, y) coordinates of the points.
+        x_limits (Optional[tuple]): Optional (min, max) tuple to override automatic x-axis limits.
+        y_limits (Optional[tuple]): Optional (min, max) tuple to override automatic y-axis limits.
+
+    Returns:
+        float: The estimated marker size 's' (in points^2) for use with
+               ax.scatter().
+    """
+    # 1. Set up the axes' properties to ensure accurate transformations.
+    # The aspect ratio and data limits must be set to correctly
+    # calculate the relationship between data units and display units (inches/points).
+    ax.set_aspect('equal')
+
+    # Use provided limits if available, otherwise calculate from data
+    if x_limits is not None:
+        x_data_min, x_data_max = x_limits
+    else:
+        x_data_min, x_data_max = np.min(coordinates[:, 0]), np.max(coordinates[:, 0])
+
+    if y_limits is not None:
+        y_data_min, y_data_max = np.min(coordinates[:, 1]), np.max(coordinates[:, 1])
+    else:
+        y_data_min, y_data_max = np.min(coordinates[:, 1]), np.max(coordinates[:, 1])
+
+    ax.set_xlim(x_data_min, x_data_max)
+    ax.set_ylim(y_data_min, y_data_max)
+
+    # Force a draw of the canvas to finalize the transformations.
+    ax.figure.canvas.draw()
+
+    # 2. Calculate the required marker radius in data units.
+    # We find the average distance to the nearest neighbor for all points.
+    # The desired radius is half of this distance.
+    tree = KDTree(coordinates)
+    distances, _ = tree.query(coordinates, k=2)
+    radius_data = np.mean(distances[:, 1]) / 2
+
+    # 3. Convert the data radius to display units (points).
+    # This requires transforming the axes' bounding box from data coordinates
+    # to display coordinates (pixels), then to physical units (inches).
+
+    # Get the bounding box in display (pixel) coordinates
+    x_display_min, _ = ax.transData.transform((x_data_min, y_data_min))
+    x_display_max, _ = ax.transData.transform((x_data_max, y_data_max))
+
+    # Convert the display coordinates to inches
+    x_inch_min, _ = ax.figure.dpi_scale_trans.inverted().transform((x_display_min, 0))
+    x_inch_max, _ = ax.figure.dpi_scale_trans.inverted().transform((x_display_max, 0))
+
+    width_inch = x_inch_max - x_inch_min
+    width_data = x_data_max - x_data_min
+
+    # Calculate the radius in inches. This scales the data radius by the
+    # ratio of the plot's physical width to its data width.
+    # This works because the aspect ratio is 'equal'.
+    radius_inch = (radius_data / width_data) * width_inch
+
+    # Convert inches to points (1 inch = 72 points).
+    radius_points = radius_inch * 72
+
+    # 4. Calculate the marker size 's'.
+    # For ax.scatter, 's' is the marker area in points^2.
+    # For a square marker, the area is (side)^2, where side = 2 * radius.
+    square_marker_size = (2 * radius_points) ** 2
+
+    return square_marker_size * 1.2
+
+
+def draw_scatter(
+    space_coord_concat,
+    title=None,
+    fig_style: Literal["dark", "light"] = "light",
+    point_size: int = None,
+    width=800,
+    height=600,
+    annotation=None,
+    color_by="logp",
+    color_continuous_scale=None,
+    plot_origin="upper",
+):
+    # Set theme based on fig_style
+    if fig_style == "dark":
+        px.defaults.template = "plotly_dark"
+    else:
+        px.defaults.template = "plotly_white"
+
+    if color_continuous_scale is None:
+        custom_color_scale = [
+            (1, "#d73027"),  # Red
+            (7 / 8, "#f46d43"),  # Red-Orange
+            (6 / 8, "#fdae61"),  # Orange
+            (5 / 8, "#fee090"),  # Light Orange
+            (4 / 8, "#e0f3f8"),  # Light Blue
+            (3 / 8, "#abd9e9"),  # Sky Blue
+            (2 / 8, "#74add1"),  # Medium Blue
+            (1 / 8, "#4575b4"),  # Dark Blue
+            (0, "#313695"),  # Deep Blue
+        ]
+        custom_color_scale.reverse()
+        color_continuous_scale = custom_color_scale
+
+    # Create the scatter plot
+    fig = px.scatter(
+        space_coord_concat,
+        x="sx",
+        y="sy",
+        color=color_by,
+        symbol="annotation" if annotation is not None else None,
+        title=title,
+        color_continuous_scale=color_continuous_scale,
+        range_color=[0, max(space_coord_concat[color_by])],
+    )
+
+    # Update marker size if specified
+    if point_size is not None:
+        fig.update_traces(marker=dict(size=point_size, symbol="circle"))
+
+    # Update layout for figure size
+    fig.update_layout(
+        autosize=False,
+        width=width,
+        height=height,
+    )
+
+    # Adjusting the legend
+    fig.update_layout(
+        legend=dict(
+            yanchor="top",
+            y=0.95,
+            xanchor="left",
+            x=1.0,
+            font=dict(
+                size=10,
+            ),
+        )
+    )
+
+    # Update colorbar to be at the bottom and horizontal
+    fig.update_layout(
+        coloraxis_colorbar=dict(
+            orientation="h",  # Make the colorbar horizontal
+            x=0.5,  # Center the colorbar horizontally
+            y=-0.0,  # Position below the plot
+            xanchor="center",  # Anchor the colorbar at the center
+            yanchor="top",  # Anchor the colorbar at the top to keep it just below the plot
+            len=0.75,  # Length of the colorbar relative to the plot width
+            title=dict(
+                text="-log10(p)" if color_by == "logp" else color_by,  # Colorbar title
+                side="top",  # Place the title at the top of the colorbar
+            ),
+        )
+    )
+    # Remove gridlines, axis labels, and ticks
+    fig.update_xaxes(
+        showgrid=False,  # Hide x-axis gridlines
+        zeroline=False,  # Hide x-axis zero line
+        showticklabels=False,  # Hide x-axis tick labels
+        title=None,  # Remove x-axis title
+        scaleanchor="y",  # Link the x-axis scale to the y-axis scale
+    )
+
+    fig.update_yaxes(
+        showgrid=False,  # Hide y-axis gridlines
+        zeroline=False,  # Hide y-axis zero line
+        showticklabels=False,  # Hide y-axis tick labels
+        title=None,  # Remove y-axis title
+        autorange="reversed" if plot_origin == "upper" else True,
+    )
+
+    # Adjust margins to ensure no clipping and equal axis ratio
+    fig.update_layout(
+        margin=dict(l=0, r=0, t=20, b=10),  # Adjust margins to prevent clipping
+    )
+
+    # Adjust the title location and font size
+    fig.update_layout(
+        title=dict(
+            y=0.98,
+            x=0.5,  # Center the title horizontally
+            xanchor="center",  # Anchor the title at the center
+            yanchor="top",  # Anchor the title at the top
+            font=dict(
+                size=20  # Increase the title font size
+            ),
+        )
+    )
+
+    return fig
+
+
+def _create_color_map(category_list: list, hex=False, rng=42) -> dict[str, tuple]:
+    unique_categories = sorted(set(category_list), key=str)
+
+    # Check for 'NaN' or nan and handle separately
+    nan_values = [v for v in unique_categories if str(v).lower() in ['nan', 'none', 'null']]
+    other_categories = [v for v in unique_categories if v not in nan_values]
+
+    n_colors = len(other_categories)
+
+    # Generate N visually distinct colors for non-NaN categories
+    if n_colors > 0:
+        colors = distinctipy.get_colors(n_colors, rng=rng)
+        color_map = dict(zip(other_categories, colors, strict=False))
+    else:
+        color_map = {}
+
+    # Assign grey color to NaN values
+    grey_rgb = (0.827, 0.827, 0.827) # lightgrey
+    for v in nan_values:
+        color_map[v] = grey_rgb
+
+    if hex:
+        # Convert RGB tuples to hex format
+        color_map = {category: distinctipy.get_hex(color_map[category]) for category in color_map}
+        print("Generated color map in hex format")
+    return color_map
+
+
+class VisualizeRunner:
+    def __init__(self, config):
+        self.config = config
+
+    custom_colors_list = [
+        '#d73027', '#f46d43', '#fdae61', '#fee090', '#e0f3f8',
+        '#abd9e9', '#74add1', '#4575b4', '#313695'
+    ]
+
+    def _generate_visualizations(self, obs_ldsc_merged: pd.DataFrame):
+        """Generate all visualizations"""
+
+        # Create visualization directories
+        single_sample_folder = self.config.visualization_result_dir / 'single_sample_multi_trait_plot'
+        annotation_folder = self.config.visualization_result_dir / 'annotation_distribution'
+        annotation_folder.mkdir(exist_ok=True, parents=True)
+
+        sample_names_list = sorted(obs_ldsc_merged['sample_name'].unique())
+
+        for sample_name in tqdm(sample_names_list, desc='Generating visualizations'):
+            # Multi-trait plot
+            traits_png = single_sample_folder / 'static_png' / f'{sample_name}_gwas_traits_pvalues.jpg'
+            traits_pdf = single_sample_folder / 'static_pdf' / f'{sample_name}_gwas_traits_pvalues.pdf'  # Added PDF output path
+
+            # Create parent directories for the output files
+            traits_png.parent.mkdir(exist_ok=True, parents=True)
+            traits_pdf.parent.mkdir(exist_ok=True, parents=True)
+
+            # Call the modified matplotlib-based plotting function.
+            # This function saves files directly and does not return a figure object.
+            self._create_single_sample_multi_trait_plots(
+                obs_ldsc_merged=obs_ldsc_merged,
+                trait_names=self.config.trait_name_list,
+                sample_name=sample_name,
+                output_png_path=traits_png,
+                output_pdf_path=traits_pdf,
+                max_cols=self.config.single_sample_multi_trait_max_cols,
+                subsample_n_points=self.config.subsample_n_points,
+                # Use new parameters from the updated VisualizationConfig
+                subplot_width_inches=self.config.single_sample_multi_trait_subplot_width_inches,
+                dpi=self.config.single_sample_multi_trait_dpi,
+                enable_pdf_output=self.config.enable_pdf_output
+            )
+
+            # Annotation distribution plots
+            sample_data = obs_ldsc_merged.query(f'sample_name == "{sample_name}"')
+            (pixel_width, pixel_height), point_size = estimate_plotly_point_size(sample_data[['sx', 'sy']].values)
+
+            for annotation in self.config.cauchy_annotations:
+                annotation_dir = annotation_folder / annotation
+                annotation_dir.mkdir(exist_ok=True)
+
+                annotation_color_map = _create_color_map(obs_ldsc_merged[annotation].unique(), hex=True)
+                fig = self._draw_scatter(sample_data, title=f'{annotation}_{sample_name}',
+                                         point_size=point_size, width=pixel_width, height=pixel_height,
+                                         hover_text_list=self.config.hover_text_list,
+                                         color_by=annotation, color_map=annotation_color_map)
+
+                annotation_png = annotation_dir / f'{sample_name}_{annotation}.png'
+                annotation_html = annotation_dir / f'{sample_name}_{annotation}.html'
+                fig.write_image(annotation_png)
+                fig.write_html(annotation_html)
+
+        # Generate multi-sample annotation plots
+        print("Generating multi-sample annotation plots...")
+        sample_count = len(sample_names_list)
+        n_rows, n_cols = self._calculate_optimal_grid_layout(
+            item_count=sample_count,
+            max_cols=self.config.single_sample_multi_trait_max_cols
+        )
+
+        for annotation in tqdm(self.config.cauchy_annotations,
+                               desc='Generating multi-sample annotation plots'):
+            annotation_dir = annotation_folder / annotation
+            annotation_dir.mkdir(exist_ok=True)
+
+            self._create_multi_sample_annotation_plot(
+                obs_ldsc_merged=obs_ldsc_merged,
+                annotation=annotation,
+                sample_names_list=sample_names_list,
+                output_dir=annotation_dir,
+                n_rows=n_rows,
+                n_cols=n_cols
+            )
+
+    def _create_single_trait_multi_sample_plots(self, obs_ldsc_merged: pd.DataFrame):
+        """Generate single trait multi-sample visualizations using matplotlib"""
+
+        trait_names = self.config.trait_name_list
+
+        # Create output directory
+        single_trait_folder = self.config.visualization_result_dir / 'single_trait_multi_sample_plot'
+        single_trait_folder.mkdir(exist_ok=True, parents=True)
+
+        # Prepare coordinate columns (assuming sx, sy are the spatial coordinates)
+        obs_ldsc_merged = obs_ldsc_merged.copy()
+
+        # Get sample count to determine grid dimensions
+        sample_count = obs_ldsc_merged['sample_name'].nunique()
+        n_rows, n_cols = self._calculate_optimal_grid_layout(
+            item_count=sample_count,
+            max_cols=self.config.single_trait_multi_sample_max_cols
+        )
+
+        print(f"Generating plots for {len(trait_names)} traits with {sample_count} samples in {n_rows}x{n_cols} grid")
+
+        # Generate visualization for each trait
+        for trait in tqdm(trait_names, desc="Generating single trait multi-sample plots"):
+            if trait not in obs_ldsc_merged.columns:
+                print(f"Warning: Trait {trait} not found in data. Skipping.")
+                continue
+
+            self._create_single_trait_multi_sample_matplotlib_plot(
+                obs_ldsc_merged=obs_ldsc_merged,
+                trait_abbreviation=trait,
+                output_png_path=single_trait_folder / f'{trait}_multi_sample_plot.jpg',
+                output_pdf_path=single_trait_folder / f'{trait}_multi_sample_plot.pdf',
+                n_rows=n_rows,
+                n_cols=n_cols,
+                subplot_width_inches=self.config.single_trait_multi_sample_subplot_width_inches,
+                scaling_factor=self.config.single_trait_multi_sample_scaling_factor,
+                dpi=self.config.single_trait_multi_sample_dpi,
+                enable_pdf_output=self.config.enable_pdf_output,
+                share_coords=self.config.share_coords
+            )
+
+    def _calculate_optimal_grid_layout(self, item_count: int, max_cols: int = 8) -> tuple[int, int]:
+        """
+        Calculate optimal grid dimensions (rows, cols) for displaying items in a grid.
+
+        Args:
+            item_count: Number of items to display
+            max_cols: Maximum number of columns allowed
+
+        Returns:
+            tuple: (n_rows, n_cols) for optimal grid layout
+        """
+        import math
+
+        if item_count <= 0:
+            return 1, 1
+
+        # For small counts, use simple layouts favoring horizontal arrangement
+        if item_count <= 3:
+            return 1, item_count
+        elif item_count <= 6:
+            return 2, math.ceil(item_count / 2)
+        elif item_count <= 12:
+            return 3, math.ceil(item_count / 3)
+        else:
+            # For larger counts, try to create a roughly square grid
+            # but respect the max_cols constraint
+            optimal_cols = min(math.ceil(math.sqrt(item_count)), max_cols)
+            optimal_rows = math.ceil(item_count / optimal_cols)
+
+            # If we hit the max_cols limit, recalculate rows
+            if optimal_cols >= max_cols:
+                n_cols = max_cols
+                n_rows = math.ceil(item_count / max_cols)
+            else:
+                n_rows = optimal_rows
+                n_cols = optimal_cols
+
+        print(f"Calculated grid layout: {n_rows} rows × {n_cols} cols for {item_count} items")
+        return n_rows, n_cols
+
+    def _create_single_trait_multi_sample_matplotlib_plot(self, obs_ldsc_merged: pd.DataFrame, trait_abbreviation: str,
+                                                          sample_name_list: list[str] | None = None,
+                                                          output_png_path: Path | None = None,
+                                                          output_pdf_path: Path | None = None,
+                                                          n_rows: int = 6, n_cols: int = 8,
+                                                          subplot_width_inches: float = 4.0,
+                                                          scaling_factor: float = 1.0, dpi: int = 300,
+                                                          enable_pdf_output: bool = True,
+                                                          show=False,
+                                                          share_coords: bool = False
+                                                          ):
+        """
+        Create and save a visualization for a specific trait showing all samples
+        """
+
+        matplotlib.rcParams['figure.dpi'] = dpi
+        print(f"Creating visualization for {trait_abbreviation}")
+
+        # Check if trait exists in the dataframe
+        if trait_abbreviation not in obs_ldsc_merged.columns:
+            print(f"Warning: Trait {trait_abbreviation} not found in the data. Skipping.")
+            return
+
+        # Set font to Arial with fallbacks to avoid warnings
+        plt.rcParams['font.family'] = 'sans-serif'
+        plt.rcParams['font.sans-serif'] = ['Arial', 'DejaVu Sans', 'Liberation Sans', 'Bitstream Vera Sans', 'sans-serif']
+
+        custom_cmap = mcolors.LinearSegmentedColormap.from_list('custom_cmap', self.custom_colors_list)
+        custom_cmap = custom_cmap.reversed()
+
+        # Calculate figure size based on subplot dimensions
+        fig_width = n_cols * subplot_width_inches
+        fig_height = n_rows * subplot_width_inches
+
+        # Create figure with title
+        fig = plt.figure(figsize=(fig_width, fig_height))
+
+        # Add main title
+        fig.suptitle(trait_abbreviation, fontsize=24, fontweight='bold', y=0.98)
+
+        # Create grid of subplots
+        grid_specs = fig.add_gridspec(nrows=n_rows, ncols=n_cols, wspace=0.1, hspace=0.1)
+
+        _, pass_filter_mask = remove_outliers_MAD(obs_ldsc_merged[trait_abbreviation])
+        obs_ldsc_merged_filtered = obs_ldsc_merged[pass_filter_mask]
+
+        pd_min = 0
+        pd_max = obs_ldsc_merged_filtered[trait_abbreviation].quantile(0.999)
+
+        print(f"Color scale min: {pd_min}, max: {pd_max}")
+        # Get list of sample names - use provided list or fallback to sorted unique
+        if sample_name_list is None:
+            sample_name_list = sorted(obs_ldsc_merged_filtered['sample_name'].unique())
+
+        # get the x and y limit if share coordinates
+        if share_coords:
+            x_limits = (obs_ldsc_merged_filtered['sx'].min(), obs_ldsc_merged_filtered['sx'].max())
+            y_limits = (obs_ldsc_merged_filtered['sy'].min(), obs_ldsc_merged_filtered['sy'].max())
+        else:
+            x_limits = None
+            y_limits = None
+
+        # Create a scatter plot for each sample
+        for position_num, select_sample_name in enumerate(sample_name_list[:n_rows * n_cols], 1):
+            # Calculate row and column in the grid
+            row = (position_num - 1) // n_cols
+            col = (position_num - 1) % n_cols
+
+            # Create subplot
+            ax = fig.add_subplot(grid_specs[row, col])
+
+            # Get data for this sample
+            sample_data = obs_ldsc_merged_filtered[obs_ldsc_merged_filtered['sample_name'] == select_sample_name]
+
+            point_size = self.estimate_matplitlib_scatter_marker_size(ax, sample_data[['sx', 'sy']].values,
+                                                                     x_limits=x_limits, y_limits=y_limits)
+            point_size *= scaling_factor  # Apply scaling factor
+            # Create scatter plot
+            scatter = ax.scatter(
+                sample_data['sx'],
+                sample_data['sy'],
+                c=sample_data[trait_abbreviation],
+                cmap=custom_cmap,
+                s=point_size,
+                vmin=pd_min,
+                vmax=pd_max,
+                marker='o',
+                edgecolors='none',
+                rasterized=True if output_pdf_path is not None and enable_pdf_output else False
+            )
+
+            if self.config.plot_origin == 'upper':
+                ax.invert_yaxis()
+
+            ax.axis('off')
+            # Add sample label as title
+            ax.set_title(select_sample_name, fontsize=12, pad=None)
+
+        # Add colorbar to the right side
+        cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])  # [left, bottom, width, height]
+        cbar = fig.colorbar(scatter, cax=cbar_ax)
+        cbar.set_label('$-\\log_{10}p$', fontsize=12, fontweight='bold')
+
+        if output_png_path is not None:
+            output_png_path.parent.mkdir(parents=True, exist_ok=True)
+            plt.savefig(output_png_path, dpi=dpi, bbox_inches='tight', )
+
+        if output_pdf_path is not None and enable_pdf_output:
+            output_pdf_path.parent.mkdir(parents=True, exist_ok=True)
+            plt.savefig(output_pdf_path, bbox_inches='tight', )
+
+        # Close the figure to free memory only if not returning it
+        if show:
+            plt.show()
+
+        gc.collect()
+        return fig
+
+    def _create_single_sample_multi_trait_plots(self,
+                                                obs_ldsc_merged: pd.DataFrame,
+                                                trait_names: list[str],
+                                                sample_name: str,
+                                                # New arguments for output paths, as matplotlib saves directly
+                                                output_png_path: Path | None,
+                                                output_pdf_path: Path | None,
+                                                # Arguments from original function signature, adapted for matplotlib
+                                                max_cols: int = 5,
+                                                subsample_n_points: int | None = None,
+                                                # subplot_width is now interpreted as inches for figsize
+                                                subplot_width_inches: float = 4.0,
+                                                dpi: int = 300,
+                                                enable_pdf_output: bool = True
+                                                ):
+
+        print(f"Creating Matplotlib-based multi-trait visualization for sample: {sample_name}")
+
+        # 1. Filter data for the specific sample and subsample if requested
+        sample_plot_data = obs_ldsc_merged[obs_ldsc_merged['sample_name'] == sample_name].copy()
+        if subsample_n_points and len(sample_plot_data) > subsample_n_points:
+            print(f"Subsampling to {subsample_n_points} points for plotting.")
+            sample_plot_data = sample_plot_data.sample(n=subsample_n_points, random_state=42)
+
+        if sample_plot_data.empty:
+            print(f"Warning: No data found for sample '{sample_name}'. Skipping plot generation.")
+            return
+
+        # 2. Calculate optimal grid layout for subplots
+        n_traits = len(trait_names)
+        n_rows, n_cols = self._calculate_optimal_grid_layout(item_count=n_traits, max_cols=max_cols)
+        print(f"Plotting {n_traits} traits in a {n_rows}x{n_cols} grid.")
+
+        # 3. Determine figure size and create figure and axes
+        # Estimate subplot height based on data's aspect ratio to avoid distortion
+        x_range = sample_plot_data['sx'].max() - sample_plot_data['sx'].min()
+        y_range = sample_plot_data['sy'].max() - sample_plot_data['sy'].min()
+        aspect_ratio = y_range / x_range if x_range > 0 else 1.0
+        subplot_height_inches = subplot_width_inches * aspect_ratio
+
+        # Calculate total figure size, adding padding for titles and colorbars
+        fig_width = subplot_width_inches * n_cols
+        fig_height = (subplot_height_inches * n_rows) * 1.2  # Add 20% vertical space for titles/colorbars
+
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=(fig_width, fig_height), squeeze=False)
+        fig.suptitle(f"Sample: {sample_name}", fontsize=16, fontweight='bold')
+
+        # 4. Define custom colormap and font
+        # Set font to Arial with fallbacks to avoid warnings
+        plt.rcParams['font.family'] = 'sans-serif'
+        plt.rcParams['font.sans-serif'] = ['Arial', 'DejaVu Sans', 'Liberation Sans', 'Bitstream Vera Sans', 'sans-serif']
+
+        custom_cmap = matplotlib.colors.LinearSegmentedColormap.from_list('custom_cmap',
+                                                                          self.custom_colors_list).reversed()
+
+        # 5. Iterate through traits and create each subplot
+        axes_flat = axes.flatten()
+        for i, trait in enumerate(trait_names):
+            if i >= len(axes_flat):
+                break  # Should not happen with correct grid calculation, but is a safe guard
+
+            ax = axes_flat[i]
+
+            # Estimate marker size to fill space without overlap
+            point_size = estimate_matplotlib_scatter_marker_size(ax, sample_plot_data[['sx', 'sy']].values)
+
+            # Determine color scale, capping at the 99.9th percentile to handle outliers
+            sample_trait_data = sample_plot_data[['sx', 'sy', trait]].dropna()
+            trait_values, mask = remove_outliers_MAD(sample_trait_data[trait])
+            sample_trait_data = sample_trait_data[mask]  # filter out outliers
+
+            vmin = 0
+            vmax = trait_values.quantile(0.999)
+            if pd.isna(vmax) or vmax == 0:
+                vmax = trait_values.max() if trait_values.max() > 0 else 1.0
+
+            # Create the scatter plot
+            scatter = ax.scatter(
+                sample_trait_data['sx'],
+                sample_trait_data['sy'],
+                c=trait_values,
+                cmap=custom_cmap,
+                s=point_size,
+                vmin=vmin,
+                vmax=vmax,
+                marker='o',
+                edgecolors='none',
+                rasterized=True if output_pdf_path is not None and enable_pdf_output else False
+            )
+
+            ax.set_title(trait, fontsize=16, pad=10, fontweight='bold')
+            ax.set_aspect('equal', adjustable='box')
+
+            if self.config.plot_origin == 'upper':
+                ax.invert_yaxis()
+
+            ax.axis('off')
+
+            # Add a colorbar to each subplot
+            cbar = fig.colorbar(scatter, ax=ax, orientation='horizontal', pad=0.1, fraction=0.05)
+            cbar.set_label('$-\\log_{10}p$', fontsize=8)
+            cbar.ax.tick_params(labelsize=7)
+
+        # Hide any unused axes in the grid
+        for j in range(len(trait_names), len(axes_flat)):
+            axes_flat[j].axis('off')
+        #
+        # # 6. Adjust layout and save the figure
+        # fig.tight_layout(rect=[0, 0.03, 1, 0.95])  # Adjust for suptitle and bottom elements
+
+        # Only proceed with saving if paths are provided
+        if output_png_path is not None:
+            output_png_path.parent.mkdir(parents=True, exist_ok=True)
+            plt.savefig(output_png_path, dpi=dpi, bbox_inches='tight', facecolor='white')
+            print(f"Saved multi-trait plot for '{sample_name}' to:\n  - {output_png_path}")
+
+        if output_pdf_path is not None and enable_pdf_output:
+            output_pdf_path.parent.mkdir(parents=True, exist_ok=True)
+            plt.savefig(output_pdf_path, bbox_inches='tight', facecolor='white')
+            print(f"Saved multi-trait plot for '{sample_name}' to:\n  - {output_pdf_path}")
+
+        # Clean up to free memory
+        plt.close(fig)
+
+    def _draw_scatter(self, space_coord_concat: pd.DataFrame, title: str | None = None,
+                      fig_style: str = 'light', point_size: int | None = None,
+                      hover_text_list: list[str] | None = None,
+                      width: int = 800, height: int = 600, annotation: str | None = None,
+                      color_by: str = 'logp', color_map: dict | None = None):
+        """Create scatter plot (adapted from original draw_scatter function)"""
+        # Set theme based on fig_style
+        if fig_style == 'dark':
+            px.defaults.template = "plotly_dark"
+        else:
+            px.defaults.template = "plotly_white"
+
+        custom_color_scale = [
+            (1, '#d73027'),  # Red
+            (7 / 8, '#f46d43'),  # Red-Orange
+            (6 / 8, '#fdae61'),  # Orange
+            (5 / 8, '#fee090'),  # Light Orange
+            (4 / 8, '#e0f3f8'),  # Light Blue
+            (3 / 8, '#abd9e9'),  # Sky Blue
+            (2 / 8, '#74add1'),  # Medium Blue
+            (1 / 8, '#4575b4'),  # Dark Blue
+            (0, '#313695')  # Deep Blue
+        ]
+        custom_color_scale.reverse()
+
+        # if category data
+        if not pd.api.types.is_numeric_dtype(space_coord_concat[color_by]):
+            # Create the scatter plot
+            fig = px.scatter(
+                space_coord_concat,
+                x='sx',
+                y='sy',
+                color=color_by,
+                # symbol=annotation,
+                title=title,
+                color_discrete_map=color_map,
+                hover_name=color_by,
+                hover_data=hover_text_list,
+                # color_continuous_scale=custom_color_scale,
+                # range_color=[0, max(space_coord_concat[color_by])],
+            )
+        else:
+            fig = px.scatter(
+                space_coord_concat,
+                x='sx',
+                y='sy',
+                color=color_by,
+                symbol=annotation,
+                title=title,
+                hover_name=color_by,
+                hover_data=hover_text_list,
+                color_continuous_scale=custom_color_scale,
+                range_color=[0, space_coord_concat[color_by].max()],
+            )
+
+        # Update marker size if specified
+        if point_size is not None:
+            fig.update_traces(marker=dict(size=point_size, symbol='circle'))
+
+        # Update layout for figure size
+        fig.update_layout(
+            autosize=False,
+            width=width,
+            height=height,
+        )
+
+        # Adjusting the legend - Updated position and marker size
+        fig.update_layout(
+            legend=dict(
+                yanchor="middle",  # Anchor point for y
+                y=0.5,  # Center vertically
+                xanchor="left",  # Anchor point for x
+                x=1.02,  # Position just outside the plot
+                font=dict(
+                    size=10,
+                ),
+                itemsizing='constant',  # Makes legend markers a constant size
+                itemwidth=30,  # Adjust width of legend items
+            )
+        )
+
+        # Update colorbar to be at the bottom and horizontal
+        fig.update_layout(
+            coloraxis_colorbar=dict(
+                orientation='h',
+                x=0.5,
+                y=-0.0,
+                xanchor='center',
+                yanchor='top',
+                len=0.75,
+                title=dict(
+                    text='-log10(p)' if color_by == 'logp' else color_by,
+                    side='top'
+                )
+            )
+        )
+
+        # Remove gridlines, axis labels, and ticks
+        fig.update_xaxes(
+            showgrid=False,
+            zeroline=False,
+            showticklabels=False,
+            title=None,
+            scaleanchor='y',
+        )
+
+        fig.update_yaxes(
+            showgrid=False,
+            zeroline=False,
+            showticklabels=False,
+            title=None,
+            autorange='reversed' if self.config.plot_origin == 'upper' else True
+        )
+
+        # Adjust margins to ensure no clipping and equal axis ratio
+        fig.update_layout(
+            margin=dict(l=0, r=100, t=20, b=10),  # Increased right margin to accommodate legend
+            height=width
+        )
+
+        # Adjust the title location and font size
+        fig.update_layout(
+            title=dict(
+                y=0.98,
+                x=0.5,
+                xanchor='center',
+                yanchor='top',
+                font=dict(
+                    size=20
+                )
+            ))
+
+        return fig
+
+    @classmethod
+    def estimate_matplitlib_scatter_marker_size(cls, ax: matplotlib.axes.Axes, coordinates: np.ndarray,
+                                                x_limits: tuple | None = None,
+                                                y_limits: tuple | None = None) -> float:
+        """Alias for estimate_matplotlib_scatter_marker_size (with typo) for backward compatibility."""
+        return estimate_matplotlib_scatter_marker_size(ax, coordinates, x_limits, y_limits)
+
+    @classmethod
+    def estimate_matplotlib_scatter_marker_size(cls, ax: matplotlib.axes.Axes, coordinates: np.ndarray,
+                                                x_limits: tuple | None = None,
+                                                y_limits: tuple | None = None) -> float:
+        """Alias for estimate_matplotlib_scatter_marker_size for backward compatibility."""
+        return estimate_matplotlib_scatter_marker_size(ax, coordinates, x_limits, y_limits)
+
+    def _create_multi_sample_annotation_plot(self, obs_ldsc_merged: pd.DataFrame, annotation: str,
+                                             sample_names_list: list, output_dir: Path,
+                                             n_rows: int, n_cols: int,
+                                             fig_width: float = 20, fig_height: float = 15,
+                                             scaling_factor: float = 1.0, dpi: int = 300):
+        """Create multi-sample annotation plot using matplotlib with subplots for each sample"""
+
+        print(f"Creating multi-sample plot for annotation: {annotation}")
+
+        # Create figure
+        fig = plt.figure(figsize=(fig_width, fig_height))
+        fig.suptitle(f'{annotation} - All Samples', fontsize=24, fontweight='bold', y=0.98)
+
+        # Create grid of subplots
+        grid_specs = fig.add_gridspec(nrows=n_rows, ncols=n_cols, wspace=0.1, hspace=0.1)
+
+        # Get unique annotation values and create color map
+        unique_annotations = obs_ldsc_merged[annotation].unique()
+        if pd.api.types.is_numeric_dtype(obs_ldsc_merged[annotation]):
+
+            custom_cmap = matplotlib.colors.LinearSegmentedColormap.from_list('custom_cmap', self.custom_colors_list)
+            cmap = custom_cmap.reversed()
+            norm = plt.Normalize(vmin=obs_ldsc_merged[annotation].min(),
+                                 vmax=obs_ldsc_merged[annotation].max())
+        else:
+            # For categorical annotations, use discrete colors
+            color_map = _create_color_map(unique_annotations, hex=False)
+
+        # Create scatter plot for each sample
+        for position_num, sample_name in enumerate(sample_names_list[:n_rows * n_cols], 1):
+            # Calculate row and column in the grid
+            row = (position_num - 1) // n_cols
+            col = (position_num - 1) % n_cols
+
+            # Create subplot
+            ax = fig.add_subplot(grid_specs[row, col])
+
+            # Get data for this sample
+            sample_data = obs_ldsc_merged[obs_ldsc_merged['sample_name'] == sample_name]
+
+            # Estimate point size based on data density
+            point_size = estimate_matplotlib_scatter_marker_size(ax, sample_data[['sx', 'sy']].values)
+            point_size *= scaling_factor  # Apply scaling factor
+
+            # Create scatter plot
+            if pd.api.types.is_numeric_dtype(obs_ldsc_merged[annotation]):
+                ax.scatter(sample_data['sx'], sample_data['sy'],
+                                     c=sample_data[annotation], cmap=cmap, norm=norm,
+                                     s=point_size, alpha=1.0, edgecolors='none')
+            else:
+                # For categorical data, plot each category separately
+                for cat in unique_annotations:
+                    cat_data = sample_data[sample_data[annotation] == cat]
+                    if len(cat_data) > 0:
+                        ax.scatter(cat_data['sx'], cat_data['sy'],
+                                   c=[color_map[cat]], s=point_size, alpha=1.0,
+                                   edgecolors='none', label=cat)
+
+            # Set subplot title
+            ax.set_title(sample_name, fontsize=10)
+            ax.set_aspect('equal')
+            if self.config.plot_origin == 'upper':
+                ax.invert_yaxis()
+            ax.axis('off')
+
+        # Add colorbar for numeric annotations or legend for categorical
+        if pd.api.types.is_numeric_dtype(obs_ldsc_merged[annotation]):
+            # Create a colorbar on the right side of the figure
+            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+            sm.set_array([])
+            fig.subplots_adjust(right=0.85)
+            cbar_ax = fig.add_axes([0.88, 0.15, 0.03, 0.7])
+            cbar = fig.colorbar(sm, cax=cbar_ax, orientation='vertical')
+            cbar.set_label(annotation, fontsize=14)
+        else:
+            # Create a legend on the right side of the figure
+            handles = [plt.Line2D([0], [0], marker='o', color='w', label=label,
+                                  markerfacecolor=color, markersize=getattr(self.config, 'legend_marker_size', 10))
+                       for label, color in color_map.items()]
+            fig.subplots_adjust(right=0.8)
+            fig.legend(handles=handles, title=annotation, loc='center left', bbox_to_anchor=(0.85, 0.5))
+
+        # Save the plot if output_dir is provided
+        if output_dir:
+            output_path = output_dir / f'multi_sample_{annotation}.png'
+            plt.savefig(output_path, dpi=dpi, bbox_inches='tight', facecolor='white')
+            print(f"Saved multi-sample annotation plot: {output_path}")
+
+        return fig
+
+    def _run_cauchy_analysis(self, obs_ldsc_merged: pd.DataFrame):
+        """Run Cauchy combination analysis"""
+
+        trait_names = self.config.trait_name_list
+
+        for annotation_col in self.config.cauchy_annotations:
+            print(f"Running Cauchy analysis for {annotation_col}...")
+
+            cauchy_results = self._run_cauchy_combination_per_annotation(
+                obs_ldsc_merged, annotation_col=annotation_col, trait_cols=trait_names)
+
+            # Save results
+            output_file = self.config.visualization_result_dir / f'cauchy_results_{annotation_col}.csv'
+            self._save_cauchy_results_to_csv(cauchy_results, output_file)
+
+            # Generate heatmaps
+            self._generate_cauchy_heatmaps(cauchy_results, annotation_col)
+
+    def _run_cauchy_combination_per_annotation(self, df: pd.DataFrame, annotation_col: str,
+                                               trait_cols: list[str], max_workers=None):
+        """
+        Runs the Cauchy combination on each annotation category for each given trait in parallel.
+        Also calculates odds ratios with confidence intervals for significant spots in each annotation.
+        """
+        from functools import partial
+
+        import statsmodels.api as sm
+        from scipy.stats import fisher_exact
+
+        results_dict = {}
+        annotations = df[annotation_col].unique()
+
+        # Helper function to process a single trait for a given annotation
+        def process_trait(trait, anno_data, all_data, annotation):
+            # Calculate significance threshold (Bonferroni correction)
+            sig_threshold = 0.05 / len(all_data)
+
+            # Get p-values for this annotation and trait
+            log10p = anno_data[trait].values
+            log10p, mask = remove_outliers_MAD(log10p, )
+            p_values = 10 ** (-log10p)  # convert from log10(p) to p
+
+            # Calculate Cauchy combination and median
+            p_cauchy_val = self._acat_test(p_values)
+            p_median_val = np.median(p_values)
+
+            # Calculate significance statistics
+            sig_spots_in_anno = np.sum(p_values < sig_threshold)
+            total_spots_in_anno = len(p_values)
+
+            # Get p-values for other annotations
+            other_annotations_mask = all_data[annotation_col] != annotation
+            other_p_values = 10 ** (-all_data.loc[other_annotations_mask, trait].values)
+            sig_spots_elsewhere = np.sum(other_p_values < sig_threshold)
+            total_spots_elsewhere = len(other_p_values)
+
+            # Odds ratio calculation using Fisher's exact test
+            try:
+                # Create contingency table
+                contingency_table = np.array([
+                    [sig_spots_in_anno, total_spots_in_anno - sig_spots_in_anno],
+                    [sig_spots_elsewhere, total_spots_elsewhere - sig_spots_elsewhere]
+                ])
+
+                # Calculate odds ratio and p-value using Fisher's exact test
+                odds_ratio, p_value = fisher_exact(contingency_table)
+
+                # if odds_ratio is infinite, set it to a large number
+                if odds_ratio == np.inf:
+                    odds_ratio = 1e4  # Set to a large number to avoid overflow
+
+                # Calculate confidence intervals
+                table = sm.stats.Table2x2(contingency_table)
+                conf_int = table.oddsratio_confint()
+                ci_low, ci_high = conf_int
+            except Exception as e:
+                # Handle calculation errors
+                odds_ratio = 0
+                p_value = 1
+                ci_low, ci_high = 0, 0
+                print(f"Fisher's exact test failed for {trait} in {annotation}: {e}")
+
+            return {
+                'trait': trait,
+                'p_cauchy': p_cauchy_val,
+                'p_median': p_median_val,
+                'odds_ratio': odds_ratio,
+                'ci_low': ci_low,
+                'ci_high': ci_high,
+                'p_odds_ratio': p_value,
+                'sig_spots': sig_spots_in_anno,
+                'total_spots': total_spots_in_anno,
+                'sig_ratio': sig_spots_in_anno / total_spots_in_anno if total_spots_in_anno > 0 else 0,
+                'overall_sig_spots': sig_spots_in_anno + sig_spots_elsewhere,
+                'overall_spots': total_spots_in_anno + total_spots_elsewhere
+            }
+
+        # Process each annotation (sequential)
+        for anno in tqdm(annotations, desc="Processing annotations"):
+            df_anno = df[df[annotation_col] == anno]
+
+            # Create a partial function with fixed parameters
+            process_trait_for_anno = partial(process_trait,
+                                             anno_data=df_anno,
+                                             all_data=df,
+                                             annotation=anno)
+
+            # Process traits in parallel with progress bar
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                # Create list for results and submit all tasks
+                futures = list(tqdm(
+                    executor.map(process_trait_for_anno, trait_cols),
+                    total=len(trait_cols),
+                    desc=f"Processing traits for {anno}",
+                    leave=False
+                ))
+                trait_results = list(futures)
+
+            # Create a DataFrame for this annotation
+            anno_results_df = pd.DataFrame(trait_results).sort_values(by='p_cauchy')
+            results_dict[anno] = anno_results_df
+
+        return results_dict
+
+    def _acat_test(self, pvalues: np.ndarray, weights=None):
+        logger = logging.getLogger('gsMap.post_analysis.cauchy')
+        if np.any(np.isnan(pvalues)):
+            raise ValueError("Cannot have NAs in the p-values.")
+        if np.any((pvalues > 1) | (pvalues < 0)):
+            raise ValueError("P-values must be between 0 and 1.")
+        if np.any(pvalues == 0) and np.any(pvalues == 1):
+            raise ValueError("Cannot have both 0 and 1 p-values.")
+        if np.any(pvalues == 0):
+            logger.info("Warn: p-values are exactly 0.")
+            return 0
+        if np.any(pvalues == 1):
+            logger.info("Warn: p-values are exactly 1.")
+            return 1
+
+        if weights is None:
+            weights = np.full(len(pvalues), 1 / len(pvalues))
+        else:
+            if len(weights) != len(pvalues):
+                raise Exception("Length of weights and p-values differs.")
+            if any(weights < 0):
+                raise Exception("All weights must be positive.")
+            weights = np.array(weights) / np.sum(weights)
+
+        is_small = pvalues < 1e-16
+        is_large = ~is_small
+
+        if not np.any(is_small):
+            cct_stat = np.sum(weights * np.tan((0.5 - pvalues) * np.pi))
+        else:
+            cct_stat = np.sum((weights[is_small] / pvalues[is_small]) / np.pi) + \
+                       np.sum(weights[is_large] * np.tan((0.5 - pvalues[is_large]) * np.pi))
+
+        if cct_stat > 1e15:
+            pval = (1 / cct_stat) / np.pi
+        else:
+            pval = 1 - stats.cauchy.cdf(cct_stat)
+
+        return pval
+
+    def _save_cauchy_results_to_csv(self, cauchy_results: dict, output_path: Path):
+        """Save Cauchy results to CSV"""
+        all_results = []
+        for annotation, df in cauchy_results.items():
+            df_copy = df.copy()
+            df_copy['annotation'] = annotation
+            all_results.append(df_copy)
+
+        combined_results = pd.concat(all_results, ignore_index=True)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        combined_results.to_csv(output_path, index=False)
+
+        return combined_results
+
+    def _generate_cauchy_heatmaps(self, cauchy_results: dict, annotation_col: str):
+        """Generate multiple types of Cauchy combination heatmaps"""
+        # Convert results to pivot table format for different metrics
+        table_cauchy = self._results_dict_to_log10_table(cauchy_results, value_col='p_cauchy', log10_transform=True)
+        table_median = self._results_dict_to_log10_table(cauchy_results, value_col='p_median', log10_transform=True)
+        table_odds_ratio = self._results_dict_to_log10_table(cauchy_results, value_col='odds_ratio',
+                                                             log10_transform=False)
+
+        # Create heatmap directories
+        cauchy_heatmap_base = self.config.visualization_result_dir / 'cauchy_heatmap'
+        static_folder = cauchy_heatmap_base / 'static_png'
+        interactive_folder = cauchy_heatmap_base / 'interactive_html'
+
+        for folder in [static_folder, interactive_folder]:
+            folder.mkdir(exist_ok=True, parents=True)
+
+        # Calculate dimensions
+        num_annotations, num_traits = table_cauchy.shape
+        width = 50 * num_traits
+        height = 50 * num_annotations
+
+        # 1. Cauchy combination heatmap (non-normalized)
+        fig = self._plot_p_cauchy_heatmap(
+            df=table_cauchy,
+            title=f"Cauchy Combination Heatmap -- By {annotation_col}",
+            cluster_rows=True, cluster_cols=True,
+            width=width, height=height,
+            text_format=".2f", font_size=10, margin_pad=150
+        )
+        fig.write_image(static_folder / f'cauchy_combination_by_{annotation_col}.png', scale=2)
+        fig.write_html(interactive_folder / f'cauchy_combination_by_{annotation_col}.html')
+
+        # 2. Cauchy combination heatmap (normalized)
+        fig = self._plot_p_cauchy_heatmap(
+            df=table_cauchy,
+            title=f"Cauchy Combination Heatmap -- By {annotation_col}",
+            normalize_axis='column',
+            cluster_rows=True, cluster_cols=True,
+            width=width, height=height,
+            text_format=".2f", font_size=10, margin_pad=150
+        )
+        fig.write_image(static_folder / f'cauchy_combination_by_{annotation_col}_normalized.png', scale=2)
+        fig.write_html(interactive_folder / f'cauchy_combination_by_{annotation_col}_normalized.html')
+
+        # 3. Median p-value heatmap (non-normalized)
+        fig = self._plot_p_cauchy_heatmap(
+            df=table_median,
+            title=f"Median log 10 pvalue Heatmap -- By {annotation_col}",
+            cluster_rows=True, cluster_cols=True,
+            width=width, height=height,
+            text_format=".2f", font_size=10, margin_pad=150
+        )
+        fig.write_image(static_folder / f'median_pvalue_{annotation_col}.png', scale=2)
+        fig.write_html(interactive_folder / f'median_pvalue_{annotation_col}.html')
+
+        # 4. Median p-value heatmap (normalized)
+        fig = self._plot_p_cauchy_heatmap(
+            df=table_median,
+            title=f"Median log 10 pvalue Heatmap -- By {annotation_col}",
+            normalize_axis='column',
+            cluster_rows=True, cluster_cols=True,
+            width=width, height=height,
+            text_format=".2f", font_size=10, margin_pad=150
+        )
+        fig.write_image(static_folder / f'median_pvalue_{annotation_col}_normalized.png', scale=2)
+        fig.write_html(interactive_folder / f'median_pvalue_{annotation_col}_normalized.html')
+
+        # 5. Odds ratio heatmap
+        fig = self._plot_p_cauchy_heatmap(
+            df=table_odds_ratio,
+            title=f"Odds Ratio Heatmap -- By {annotation_col}",
+            cluster_rows=True, cluster_cols=True,
+            width=width, height=height,
+            text_format=".2f", font_size=10, margin_pad=150
+        )
+        fig.write_image(static_folder / f'odds_ratio_{annotation_col}.png', scale=2)
+        fig.write_html(interactive_folder / f'odds_ratio_{annotation_col}.html')
+
+    def _results_dict_to_log10_table(self, results_dict: dict, value_col: str = 'p_cauchy',
+                                     log10_transform: bool = True, epsilon: float = 1e-300) -> pd.DataFrame:
+        """Convert results dict to pivot table"""
+        all_data = []
+        for anno, df in results_dict.items():
+            temp = df.copy()
+            temp['annotation'] = anno
+            all_data.append(temp)
+
+        combined_df = pd.concat(all_data, ignore_index=True)
+
+        if log10_transform:
+            combined_df.loc[combined_df[value_col] == 0, value_col] = epsilon
+            combined_df['transformed'] = -np.log10(combined_df[value_col])
+        else:
+            combined_df['transformed'] = combined_df[value_col]
+
+        pivot_df = combined_df.pivot(index='annotation', columns='trait', values='transformed')
+        return pivot_df
+
+    def _plot_p_cauchy_heatmap(self, df: pd.DataFrame, title: str = "Cauchy Combination Heatmap",
+                               normalize_axis: Literal["row", "column"] | None = None,
+                               cluster_rows: bool = False, cluster_cols: bool = False,
+                               color_continuous_scale: str | list = "RdBu_r",
+                               width: int | None = None, height: int | None = None,
+                               text_format: str = ".2f",
+                               show_text: bool = True, font_size: int = 10, margin_pad: int = 150) -> go.Figure:
+        """
+        Create an enhanced heatmap visualization for trait-annotation relationships.
+        """
+        data = df.copy()
+
+        # Input validation
+        if not isinstance(data, pd.DataFrame):
+            raise TypeError("Input must be a pandas DataFrame")
+        if data.empty:
+            raise ValueError("Input DataFrame is empty")
+        if not np.issubdtype(data.values.dtype, np.number):
+            raise ValueError("DataFrame must contain numeric values")
+
+        n_rows, n_cols = data.shape
+        # Set dynamic width/height if not provided to ensure good aspect ratio
+        # Previously we used 50 and 30, which led to vertical stretching.
+        # Let's use more balanced units.
+        if width is None:
+            width = max(600, n_cols * 150 + margin_pad * 2)
+        if height is None:
+            height = max(500, n_rows * 60 + margin_pad * 2)
+
+        # Normalization with error handling
+        if normalize_axis in ['row', 'column']:
+            axis = 1 if normalize_axis == 'row' else 0
+            try:
+                # Store original data for text annotations
+                original_data = data.copy()
+
+                # Calculate min and max along specified axis
+                min_vals = data.min(axis=axis)
+                max_vals = data.max(axis=axis)
+                range_vals = max_vals - min_vals
+
+                # Replace zero range with 1 to avoid division by zero
+                range_vals = range_vals.replace(0, 1)
+
+                # Normalize using broadcasting
+                if normalize_axis == 'row':
+                    data = data.sub(min_vals, axis=0).div(range_vals, axis=0)
+                else:  # column
+                    data = data.sub(min_vals, axis=1).div(range_vals, axis=1)
+
+                data = data.fillna(0)
+            except Exception as e:
+                raise ValueError(f"Normalization failed: {str(e)}")
+        else:
+            # No normalization, use original data for both color and text
+            original_data = data
+
+        # Clustering with error handling
+        try:
+            if cluster_rows:
+                row_linkage = linkage(data.fillna(0).values, method='average', metric='euclidean')
+                row_order = leaves_list(row_linkage)
+                data = data.iloc[row_order, :]
+                original_data = original_data.iloc[row_order, :]  # Apply the same order to original data
+
+            if cluster_cols:
+                col_linkage = linkage(data.fillna(0).values.T, method='average', metric='euclidean')
+                col_order = leaves_list(col_linkage)
+                data = data.iloc[:, col_order]
+                original_data = original_data.iloc[:, col_order]  # Apply the same order to original data
+        except Exception as e:
+            raise ValueError(f"Clustering failed: {str(e)}")
+
+        # Create heatmap with enhanced formatting
+        if normalize_axis is None:
+            # Use original settings for speed when no normalization is applied
+            fig = px.imshow(
+                data,
+                color_continuous_scale=color_continuous_scale,
+                aspect='auto',
+                width=width,
+                height=height,
+                text_auto=text_format if show_text else False  # Automatic text generation
+            )
+        else:
+            # Use custom logic for normalization (manual text annotations)
+            fig = px.imshow(
+                data,
+                color_continuous_scale=color_continuous_scale,
+                aspect='auto',
+                width=width,
+                height=height,
+                text_auto=False  # Disable automatic text generation
+            )
+
+            # Add manual text annotations using original data
+            if show_text:
+                for i, row in enumerate(original_data.values):
+                    for j, value in enumerate(row):
+                        fig.add_annotation(
+                            x=j,
+                            y=i,
+                            text=f"{value:{text_format}}",
+                            showarrow=False,
+                            font=dict(size=font_size, color='black')
+                        )
+
+        # Enhanced layout configuration
+        fig.update_layout(
+            title={
+                'text': title,
+                'y': 0.98,
+                'x': 0.5,
+                'xanchor': 'center',
+                'yanchor': 'bottom',
+                'font': {'size': font_size + 4}
+            },
+            xaxis={
+                'title': "Trait",
+                'tickangle': 45,
+                'side': 'bottom',
+                'tickfont': {'size': font_size},
+                'title_font': {'size': font_size + 2}
+            },
+            yaxis={
+                'title': "Annotation",
+                'tickfont': {'size': font_size},
+                'title_font': {'size': font_size + 2}
+            },
+            width=width,
+            height=height,
+            template='plotly_white',
+            margin=dict(l=margin_pad, r=margin_pad, t=margin_pad, b=margin_pad),
+            coloraxis_colorbar={
+                'title': '-log10(p)' if not normalize_axis else 'Normalized Value',
+                'title_side': 'right',
+                'title_font': {'size': font_size + 2},
+                'tickfont': {'size': font_size}
+            }
+        )
+        return fig