pylocuszoom 1.2.0__py3-none-any.whl → 1.3.1__py3-none-any.whl

pylocuszoom/__init__.py

@@ -40,6 +40,9 @@ __version__ = "1.1.2"
 # Backend types
 from .backends import BackendType, get_backend
 
+# Colocalization visualization
+from .coloc_plotter import ColocPlotter
+
 # Colors and LD
 from .colors import (
     LEAD_SNP_COLOR,
@@ -101,10 +104,13 @@ from .forest import validate_forest_df
 from .gene_track import get_nearest_gene, plot_gene_track
 
 # Labels
-from .labels import add_snp_labels
+from .labels import add_snp_labels, adjust_snp_labels
 
 # LD calculation
-from .ld import calculate_ld
+from .ld import calculate_ld, calculate_pairwise_ld
+
+# LD heatmap visualization
+from .ld_heatmap_plotter import LDHeatmapPlotter
 
 # File format loaders
 from .loaders import (
@@ -137,6 +143,9 @@ from .logging import disable_logging, enable_logging
 # Manhattan and QQ plotting
 from .manhattan_plotter import ManhattanPlotter
 
+# Miami plot (mirrored Manhattan comparison)
+from .miami_plotter import MiamiPlotter
+
 # PheWAS support
 from .phewas import validate_phewas_df
 from .plotter import LocusZoomPlotter
@@ -161,7 +170,10 @@ __all__ = [
     "__version__",
     "LocusZoomPlotter",
     "ManhattanPlotter",
+    "MiamiPlotter",
     "StatsPlotter",
+    "LDHeatmapPlotter",
+    "ColocPlotter",
     # Backends
     "BackendType",
     "get_backend",
@@ -180,8 +192,10 @@ __all__ = [
     "plot_gene_track",
     # LD
     "calculate_ld",
+    "calculate_pairwise_ld",
     # Labels
     "add_snp_labels",
+    "adjust_snp_labels",
     # Recombination
     "add_recombination_overlay",
     "ensure_recomb_maps",

pylocuszoom/backends/base.py

@@ -3,10 +3,13 @@
 Defines the interface that matplotlib, plotly, and bokeh backends must implement.
 """
 
-from typing import Any, Callable, List, Optional, Protocol, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Protocol, Tuple, Union
 
 import pandas as pd
 
+if TYPE_CHECKING:
+    pass
+
 
 class PlotBackend(Protocol):
     """Protocol defining the backend interface for LocusZoom plots.
@@ -330,7 +333,8 @@ class PlotBackend(Protocol):
         label_top_n: int,
         genes_df: Optional[pd.DataFrame],
         chrom: int,
-    ) -> None:
+        adjust: bool = True,
+    ) -> List[Any]:
         """Add SNP labels to plot.
 
         No-op if supports_snp_labels=False. Matplotlib uses adjustText
@@ -345,6 +349,23 @@ class PlotBackend(Protocol):
             label_top_n: Number of top SNPs to label.
             genes_df: Gene annotations (unused, for signature compatibility).
             chrom: Chromosome number (unused, for signature compatibility).
+            adjust: If True, run adjustText immediately. If False, caller
+                must call adjust_snp_labels() after setting axis limits.
+
+        Returns:
+            List of text annotation objects (empty for non-matplotlib backends).
+        """
+        ...
+
+    def adjust_snp_labels(self, ax: Any, texts: List[Any]) -> None:
+        """Adjust SNP label positions after axis limits are set.
+
+        This should be called AFTER all axis limits have been finalized,
+        as adjustText needs to know the final plot bounds.
+
+        Args:
+            ax: Axes or panel.
+            texts: List of text annotation objects from add_snp_labels().
         """
         ...
 
@@ -686,6 +707,21 @@ class PlotBackend(Protocol):
         """
         ...
 
+    def add_effect_legend(
+        self,
+        ax: Any,
+        effect_bins: List[Tuple[float, str, str]],
+    ) -> None:
+        """Add effect direction legend for colocalization plots.
+
+        Shows effect direction categories (same direction, opposite, missing).
+
+        Args:
+            ax: Axes or panel.
+            effect_bins: List of (threshold, label, color) tuples.
+        """
+        ...
+
     def add_eqtl_legend(
         self,
         ax: Any,
@@ -831,3 +867,59 @@ class PlotBackend(Protocol):
             fig: Figure object.
         """
         ...
+
+    # =========================================================================
+    # Heatmap Rendering (for LD heatmaps)
+    # =========================================================================
+
+    def add_heatmap(
+        self,
+        ax: Any,
+        data: Any,
+        x_coords: List[float],
+        y_coords: List[float],
+        cmap_colors: Optional[List[str]] = None,
+        vmin: float = 0.0,
+        vmax: float = 1.0,
+        mask_upper: bool = True,
+    ) -> Any:
+        """Render heatmap with optional triangular masking.
+
+        Used for LD heatmap visualization where data is a symmetric matrix
+        and typically only the lower triangle is displayed.
+
+        Args:
+            ax: Axes or panel to plot on.
+            data: 2D numpy array of values (use NaN for missing data).
+            x_coords: X coordinates for cell positions.
+            y_coords: Y coordinates for cell positions.
+            cmap_colors: Color gradient endpoints [start_color, end_color].
+                Defaults to white-to-red ["#FFFFFF", "#FF0000"].
+            vmin: Minimum value for color scale.
+            vmax: Maximum value for color scale.
+            mask_upper: If True, mask upper triangle for lower-triangular display.
+
+        Returns:
+            Heatmap object (mappable for colorbar attachment).
+        """
+        ...
+
+    def add_colorbar(
+        self,
+        ax: Any,
+        mappable: Any,
+        label: str = "R²",
+        orientation: str = "vertical",
+    ) -> Any:
+        """Add colorbar legend for heatmap.
+
+        Args:
+            ax: Axes or panel (or figure for some backends).
+            mappable: Heatmap object returned by add_heatmap.
+            label: Colorbar label (e.g., "R²" or "D'").
+            orientation: "vertical" or "horizontal".
+
+        Returns:
+            Colorbar object.
+        """
+        ...

pylocuszoom/backends/bokeh_backend.py

@@ -623,7 +623,12 @@ class BokehBackend:
         label_top_n: int,
         genes_df: Optional[pd.DataFrame],
         chrom: int,
-    ) -> None:
+        adjust: bool = True,
+    ) -> List[Any]:
+        """No-op: Bokeh uses hover tooltips instead of text labels."""
+        return []
+
+    def adjust_snp_labels(self, ax: Any, texts: List[Any]) -> None:
         """No-op: Bokeh uses hover tooltips instead of text labels."""
         pass
 
@@ -645,11 +650,16 @@ class BokehBackend:
             if hasattr(x_range, "start") and x_range.start is not None
             else 0
         )
-        y = (
-            y_range.start + y_frac * (y_range.end - y_range.start)
-            if hasattr(y_range, "start") and y_range.start is not None
-            else 0
-        )
+
+        # Handle both normal and inverted y-axis
+        # For inverted axis (start > end), y_frac=0.95 should be near start (top visually)
+        if hasattr(y_range, "start") and y_range.start is not None:
+            y_min = min(y_range.start, y_range.end)
+            y_max = max(y_range.start, y_range.end)
+            # Position label at top of visible range regardless of axis direction
+            y = y_min + y_frac * (y_max - y_min)
+        else:
+            y = 0
 
         label_obj = Label(
             x=x,
@@ -734,6 +744,18 @@ class BokehBackend:
             items.append(self._add_legend_item(ax, source, label, color, "square"))
         self._create_legend(ax, items, "r²")
 
+    def add_effect_legend(
+        self,
+        ax: figure,
+        effect_bins: List[Tuple[float, str, str]],
+    ) -> None:
+        """Add effect direction legend for colocalization plots."""
+        source = self._ensure_legend_range(ax)
+        items = []
+        for _, label, color in effect_bins:
+            items.append(self._add_legend_item(ax, source, label, color, "circle"))
+        self._create_legend(ax, items, "Effect")
+
     def add_legend(
         self,
         ax: figure,
@@ -968,3 +990,135 @@ class BokehBackend:
         """
         # Bokeh handles layout differently - column spacing is fixed
         pass
+
+    def add_heatmap(
+        self,
+        ax: figure,
+        data: Any,
+        x_coords: List[float],
+        y_coords: List[float],
+        cmap_colors: Optional[List[str]] = None,
+        vmin: float = 0.0,
+        vmax: float = 1.0,
+        mask_upper: bool = True,
+    ) -> Any:
+        """Render heatmap with optional triangular masking.
+
+        Args:
+            ax: Bokeh figure.
+            data: 2D numpy array of values (NaN for missing).
+            x_coords: X coordinates for cells.
+            y_coords: Y coordinates for cells.
+            cmap_colors: Color gradient endpoints [start, end].
+            vmin: Minimum value for color scale.
+            vmax: Maximum value for color scale.
+            mask_upper: If True, mask upper triangle.
+
+        Returns:
+            LinearColorMapper for colorbar attachment.
+        """
+        import numpy as np
+        from bokeh.models import LinearColorMapper
+
+        if cmap_colors is None:
+            cmap_colors = ["#FFFFFF", "#FF0000"]
+
+        # Create custom palette from start to end color
+        # For a simple 2-color gradient, create a palette of intermediate colors
+        palette = _create_color_palette(cmap_colors[0], cmap_colors[1], 256)
+
+        mapper = LinearColorMapper(
+            palette=palette,
+            low=vmin,
+            high=vmax,
+            nan_color="#808080",  # Grey for missing
+        )
+
+        # Build rect data for lower triangle
+        xs, ys, values = [], [], []
+        n = data.shape[0]
+        for i in range(n):
+            for j in range(n):
+                # Lower triangle including diagonal
+                if mask_upper and j > i:
+                    continue
+                val = data[i, j]
+                xs.append(j)
+                ys.append(i)
+                values.append(val if not np.isnan(val) else float("nan"))
+
+        source = ColumnDataSource({"x": xs, "y": ys, "value": values})
+
+        ax.rect(
+            x="x",
+            y="y",
+            width=1,
+            height=1,
+            fill_color={"field": "value", "transform": mapper},
+            line_color=None,
+            source=source,
+        )
+        return mapper
+
+    def add_colorbar(
+        self,
+        ax: figure,
+        mappable: Any,
+        label: str = "R²",
+        orientation: str = "vertical",
+    ) -> Any:
+        """Add colorbar legend for heatmap.
+
+        Args:
+            ax: Bokeh figure.
+            mappable: LinearColorMapper from add_heatmap.
+            label: Colorbar label.
+            orientation: "vertical" or "horizontal".
+
+        Returns:
+            ColorBar object.
+        """
+        from bokeh.models import BasicTicker, ColorBar
+
+        orientation_map = {"vertical": "vertical", "horizontal": "horizontal"}
+        color_bar = ColorBar(
+            color_mapper=mappable,
+            ticker=BasicTicker(),
+            label_standoff=6,
+            title=label,
+            orientation=orientation_map.get(orientation, "vertical"),
+        )
+        ax.add_layout(color_bar, "right")
+        return color_bar
+
+
+def _create_color_palette(start_color: str, end_color: str, n_colors: int) -> List[str]:
+    """Create a linear color palette between two colors.
+
+    Args:
+        start_color: Starting hex color.
+        end_color: Ending hex color.
+        n_colors: Number of colors in palette.
+
+    Returns:
+        List of hex color strings.
+    """
+
+    def hex_to_rgb(hex_color: str) -> Tuple[int, int, int]:
+        hex_color = hex_color.lstrip("#")
+        return tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4))
+
+    def rgb_to_hex(rgb: Tuple[int, int, int]) -> str:
+        return "#{:02x}{:02x}{:02x}".format(*rgb)
+
+    start_rgb = hex_to_rgb(start_color)
+    end_rgb = hex_to_rgb(end_color)
+
+    palette = []
+    for i in range(n_colors):
+        t = i / (n_colors - 1) if n_colors > 1 else 0
+        r = int(start_rgb[0] + t * (end_rgb[0] - start_rgb[0]))
+        g = int(start_rgb[1] + t * (end_rgb[1] - start_rgb[1]))
+        b = int(start_rgb[2] + t * (end_rgb[2] - start_rgb[2]))
+        palette.append(rgb_to_hex((r, g, b)))
+    return palette

pylocuszoom/backends/matplotlib_backend.py

@@ -264,7 +264,8 @@ class MatplotlibBackend:
         label_top_n: int,
         genes_df: Optional[pd.DataFrame],
         chrom: int,
-    ) -> None:
+        adjust: bool = True,
+    ) -> List[Any]:
         """Add SNP labels using adjustText.
 
         Args:
@@ -276,10 +277,15 @@ class MatplotlibBackend:
             label_top_n: Number of top SNPs to label.
             genes_df: Gene annotations (unused, for signature compatibility).
             chrom: Chromosome number (unused, for signature compatibility).
+            adjust: If True, run adjustText immediately. If False, caller
+                must call adjust_snp_labels() after setting axis limits.
+
+        Returns:
+            List of text annotation objects.
         """
         from ..labels import add_snp_labels as _add_snp_labels
 
-        _add_snp_labels(
+        return _add_snp_labels(
             ax,
             df,
             pos_col=pos_col,
@@ -288,8 +294,20 @@ class MatplotlibBackend:
             label_top_n=label_top_n,
             genes_df=genes_df,
             chrom=chrom,
+            adjust=adjust,
         )
 
+    def adjust_snp_labels(self, ax: Axes, texts: List[Any]) -> None:
+        """Adjust SNP label positions after axis limits are set.
+
+        Args:
+            ax: Matplotlib axes.
+            texts: List of text annotation objects from add_snp_labels().
+        """
+        from ..labels import adjust_snp_labels as _adjust_snp_labels
+
+        _adjust_snp_labels(ax, texts)
+
     def add_rectangle(
         self,
         ax: Axes,
@@ -714,6 +732,41 @@ class MatplotlibBackend:
             labelspacing=0.4,
         )
 
+    def add_effect_legend(
+        self,
+        ax: Axes,
+        effect_bins: List[Tuple[float, str, str]],
+    ) -> None:
+        """Add effect direction legend for colocalization plots.
+
+        Args:
+            ax: Matplotlib axes.
+            effect_bins: List of (threshold, label, color) tuples.
+        """
+        from matplotlib.lines import Line2D
+
+        legend_elements = []
+        for _threshold, label, color in effect_bins:
+            legend_elements.append(
+                Line2D(
+                    [0],
+                    [0],
+                    marker="o",
+                    color="w",
+                    markerfacecolor=color,
+                    markeredgecolor="black",
+                    markersize=8,
+                    label=label,
+                )
+            )
+
+        ax.legend(
+            handles=legend_elements,
+            loc="upper right",
+            fontsize=8,
+            framealpha=0.9,
+        )
+
     def axvline(
         self,
         ax: Axes,
@@ -806,3 +859,90 @@ class MatplotlibBackend:
             left=left, right=right, top=top, bottom=bottom, hspace=hspace
         )
         plt.ion()
+
+    def add_heatmap(
+        self,
+        ax: Axes,
+        data: Any,
+        x_coords: List[float],
+        y_coords: List[float],
+        cmap_colors: Optional[List[str]] = None,
+        vmin: float = 0.0,
+        vmax: float = 1.0,
+        mask_upper: bool = True,
+    ) -> Any:
+        """Render heatmap with optional triangular masking.
+
+        Args:
+            ax: Matplotlib axes.
+            data: 2D numpy array of values (NaN for missing).
+            x_coords: X coordinates for cells.
+            y_coords: Y coordinates for cells.
+            cmap_colors: Color gradient endpoints [start, end].
+            vmin: Minimum value for color scale.
+            vmax: Maximum value for color scale.
+            mask_upper: If True, mask upper triangle.
+
+        Returns:
+            AxesImage object for colorbar attachment.
+        """
+        import numpy as np
+        from matplotlib.colors import LinearSegmentedColormap
+
+        # Default white-to-red gradient
+        if cmap_colors is None:
+            cmap_colors = ["#FFFFFF", "#FF0000"]
+
+        cmap = LinearSegmentedColormap.from_list("ld_heatmap", cmap_colors, N=256)
+
+        # Mask upper triangle if requested
+        plot_data = data.copy()
+        if mask_upper:
+            mask = np.triu(np.ones_like(data, dtype=bool), k=1)
+            plot_data = np.ma.array(data, mask=mask)
+
+        # Compute extent from coordinates if non-trivial (genomic coordinates)
+        # This allows heatmap to align with regional plot x-axis
+        extent = None
+        if x_coords and len(x_coords) > 1:
+            x_min, x_max = min(x_coords), max(x_coords)
+            # Only use extent if coordinates are not simple indices
+            if x_max - x_min > len(x_coords):
+                # Genomic coordinates - use extent for alignment
+                # Add half-cell padding for proper cell centering
+                y_min, y_max = min(y_coords), max(y_coords)
+                extent = [x_min, x_max, y_min - 0.5, y_max + 0.5]
+
+        # Use imshow for heatmap rendering
+        im = ax.imshow(
+            plot_data,
+            cmap=cmap,
+            vmin=vmin,
+            vmax=vmax,
+            aspect="auto",
+            origin="lower",
+            extent=extent,
+        )
+        return im
+
+    def add_colorbar(
+        self,
+        ax: Axes,
+        mappable: Any,
+        label: str = "R²",
+        orientation: str = "vertical",
+    ) -> Any:
+        """Add colorbar legend for heatmap.
+
+        Args:
+            ax: Matplotlib axes.
+            mappable: AxesImage from add_heatmap.
+            label: Colorbar label.
+            orientation: "vertical" or "horizontal".
+
+        Returns:
+            Colorbar object.
+        """
+        cbar = plt.colorbar(mappable, ax=ax, orientation=orientation)
+        cbar.set_label(label)
+        return cbar

pylocuszoom/backends/plotly_backend.py

@@ -841,7 +841,12 @@ class PlotlyBackend:
         label_top_n: int,
         genes_df: Optional[pd.DataFrame],
         chrom: int,
-    ) -> None:
+        adjust: bool = True,
+    ) -> List[Any]:
+        """No-op: Plotly uses hover tooltips instead of text labels."""
+        return []
+
+    def adjust_snp_labels(self, ax: Any, texts: List[Any]) -> None:
         """No-op: Plotly uses hover tooltips instead of text labels."""
         pass
 
@@ -887,6 +892,19 @@ class PlotlyBackend:
 
         self._configure_legend(fig, row, "legend", "r²")
 
+    def add_effect_legend(
+        self,
+        ax: Tuple[go.Figure, int],
+        effect_bins: List[Tuple[float, str, str]],
+    ) -> None:
+        """Add effect direction legend for colocalization plots."""
+        fig, row, col, _ = self._extract_row_col(ax)
+
+        for _, label, color in effect_bins:
+            self._add_legend_item(fig, row, label, color, "circle", 10, "legend")
+
+        self._configure_legend(fig, row, "legend", "Effect")
+
     def add_legend(
         self,
         ax: Tuple[go.Figure, int],
@@ -1219,3 +1237,85 @@ class PlotlyBackend:
                             )
                         }
                     )
+
+    def add_heatmap(
+        self,
+        ax: Tuple[go.Figure, int],
+        data: Any,
+        x_coords: List[float],
+        y_coords: List[float],
+        cmap_colors: Optional[List[str]] = None,
+        vmin: float = 0.0,
+        vmax: float = 1.0,
+        mask_upper: bool = True,
+    ) -> Any:
+        """Render heatmap with optional triangular masking.
+
+        Args:
+            ax: Tuple of (figure, row_number).
+            data: 2D numpy array of values (NaN for missing).
+            x_coords: X coordinates for cells.
+            y_coords: Y coordinates for cells.
+            cmap_colors: Color gradient endpoints [start, end].
+            vmin: Minimum value for color scale.
+            vmax: Maximum value for color scale.
+            mask_upper: If True, mask upper triangle.
+
+        Returns:
+            Heatmap trace.
+        """
+        import numpy as np
+
+        fig, row, col, _ = self._extract_row_col(ax)
+
+        if cmap_colors is None:
+            cmap_colors = ["#FFFFFF", "#FF0000"]
+
+        # Mask upper triangle by setting to NaN
+        plot_data = data.copy()
+        if mask_upper:
+            for i in range(data.shape[0]):
+                for j in range(i + 1, data.shape[1]):
+                    plot_data[i, j] = np.nan
+
+        # Replace NaN with None for Plotly
+        z = np.where(np.isnan(plot_data), None, plot_data)
+
+        colorscale = [[0, cmap_colors[0]], [1, cmap_colors[1]]]
+
+        trace = go.Heatmap(
+            z=z.tolist(),
+            x=list(range(len(x_coords))),
+            y=list(range(len(y_coords))),
+            colorscale=colorscale,
+            zmin=vmin,
+            zmax=vmax,
+            showscale=True,
+            colorbar=dict(title="R²"),
+        )
+        fig.add_trace(trace, row=row, col=col)
+        return trace
+
+    def add_colorbar(
+        self,
+        ax: Tuple[go.Figure, int],
+        mappable: Any,
+        label: str = "R²",
+        orientation: str = "vertical",
+    ) -> Any:
+        """Add colorbar legend for heatmap.
+
+        For Plotly, colorbar is added via showscale=True in add_heatmap.
+        This method is a no-op but kept for API consistency.
+
+        Args:
+            ax: Tuple of (figure, row_number).
+            mappable: Heatmap trace (unused).
+            label: Colorbar label (unused, set in add_heatmap).
+            orientation: Orientation (unused).
+
+        Returns:
+            None.
+        """
+        # Colorbar is configured in add_heatmap via showscale=True
+        pass