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

pylocuszoom/plotter.py

@@ -23,7 +23,9 @@ from .colors import (
     EQTL_NEGATIVE_BINS,
     EQTL_POSITIVE_BINS,
     LD_BINS,
+    LD_HEATMAP_COLORS,
     LEAD_SNP_COLOR,
+    LEAD_SNP_HIGHLIGHT_COLOR,
     get_credible_set_color,
     get_eqtl_color,
     get_ld_bin,
@@ -223,6 +225,10 @@ class LocusZoomPlotter:
         genes_df: Optional[pd.DataFrame] = None,
         exons_df: Optional[pd.DataFrame] = None,
         recomb_df: Optional[pd.DataFrame] = None,
+        ld_heatmap_df: Optional[pd.DataFrame] = None,
+        ld_heatmap_snp_ids: Optional[List[str]] = None,
+        ld_heatmap_height: float = 0.25,
+        ld_heatmap_metric: str = "r2",
     ) -> Any:
         """Create a regional association plot.
 
@@ -246,12 +252,21 @@ class LocusZoomPlotter:
             exons_df: Exon annotations with chr, start, end, gene_name.
             recomb_df: Pre-loaded recombination rate data.
                 If None and show_recombination=True, loads from species default.
+            ld_heatmap_df: Pairwise LD matrix (square DataFrame) from
+                calculate_pairwise_ld. If provided with ld_heatmap_snp_ids,
+                renders heatmap panel below association plot.
+            ld_heatmap_snp_ids: List of SNP IDs in matrix order. Required if
+                ld_heatmap_df is provided.
+            ld_heatmap_height: Height ratio of heatmap panel relative to
+                association panel. Default 0.25.
+            ld_heatmap_metric: LD metric label for colorbar ("r2" or "dprime").
 
         Returns:
             Figure object (type depends on backend).
 
         Raises:
             ValidationError: If parameters or DataFrame columns are invalid.
+            ValueError: If ld_heatmap_df provided without ld_heatmap_snp_ids.
 
         Example:
             >>> fig = plotter.plot(
@@ -280,6 +295,12 @@ class LocusZoomPlotter:
         # Validate inputs
         validate_gwas_df(gwas_df, pos_col=pos_col, p_col=p_col)
 
+        # Validate LD heatmap parameters
+        if ld_heatmap_df is not None and ld_heatmap_snp_ids is None:
+            raise ValueError(
+                "ld_heatmap_snp_ids is required when ld_heatmap_df is provided"
+            )
+
         # Auto-fetch genes if enabled and not provided
         if genes_df is None and self._auto_genes:
             logger.debug(
@@ -352,8 +373,33 @@ class LocusZoomPlotter:
         if show_recombination and recomb_df is None:
             recomb_df = self._get_recomb_for_region(chrom, start, end)
 
+        # Transform heatmap to genomic coordinates if provided
+        heatmap_data = None
+        if ld_heatmap_df is not None and ld_heatmap_snp_ids is not None:
+            heatmap_data = self._transform_heatmap_to_genomic_coords(
+                ld_matrix=ld_heatmap_df,
+                snp_ids=ld_heatmap_snp_ids,
+                gwas_df=df,
+                start=start,
+                end=end,
+                rs_col=rs_col,
+                pos_col=pos_col,
+            )
+            if heatmap_data is None:
+                logger.warning(
+                    "No SNPs from LD heatmap overlap with region - heatmap not rendered"
+                )
+
         # Create figure layout
-        fig, ax, gene_ax = self._create_figure(genes_df, chrom, start, end, figsize)
+        fig, ax, gene_ax, heatmap_ax = self._create_figure_with_heatmap(
+            genes_df=genes_df,
+            chrom=chrom,
+            start=start,
+            end=end,
+            figsize=figsize,
+            heatmap_data=heatmap_data,
+            heatmap_height=ld_heatmap_height,
+        )
 
         # Plot association data
         self._plot_association(ax, df, pos_col, ld_col, lead_pos, rs_col, p_col)
@@ -370,9 +416,11 @@ class LocusZoomPlotter:
         )
 
         # Add SNP labels (capability check - interactive backends use hover tooltips)
+        # Create labels without adjusting - we'll adjust after axis limits are set
+        snp_label_texts: list = []
         if snp_labels and rs_col in df.columns and label_top_n > 0 and not df.empty:
             if self._backend.supports_snp_labels:
-                self._backend.add_snp_labels(
+                snp_label_texts = self._backend.add_snp_labels(
                     ax,
                     df,
                     pos_col=pos_col,
@@ -381,6 +429,7 @@ class LocusZoomPlotter:
                     label_top_n=label_top_n,
                     genes_df=genes_df,
                     chrom=chrom,
+                    adjust=False,  # Defer adjustment until after axis limits set
                 )
 
         # Add recombination overlay (all backends with secondary axis support)
@@ -407,11 +456,38 @@ class LocusZoomPlotter:
             plot_gene_track_generic(
                 gene_ax, self._backend, genes_df, chrom, start, end, exons_df
             )
-            self._backend.set_xlabel(gene_ax, f"Chromosome {chrom} (Mb)")
             self._backend.hide_spines(gene_ax, ["top", "right", "left"])
             # Format both axes for interactive backends (they don't share x-axis)
             self._backend.format_xaxis_mb(gene_ax)
-        else:
+            # Only set x-label on gene track if no heatmap below
+            if heatmap_ax is None:
+                self._backend.set_xlabel(gene_ax, f"Chromosome {chrom} (Mb)")
+
+        # Render LD heatmap panel if data available
+        if heatmap_ax is not None and heatmap_data is not None:
+            filtered_matrix, x_positions, filtered_snp_ids = heatmap_data
+            # Find lead SNP ID if lead_pos is set
+            lead_snp_id = None
+            if lead_pos is not None and rs_col in df.columns:
+                lead_row = df[df[pos_col] == lead_pos]
+                if not lead_row.empty:
+                    lead_snp_id = lead_row[rs_col].iloc[0]
+            self._render_heatmap_panel(
+                ax=heatmap_ax,
+                fig=fig,
+                ld_matrix=filtered_matrix,
+                x_positions=x_positions,
+                snp_ids=filtered_snp_ids,
+                metric=ld_heatmap_metric,
+                lead_snp_id=lead_snp_id,
+                start=start,
+                end=end,
+            )
+            # Heatmap is at bottom - set x-label on it
+            self._backend.set_xlabel(heatmap_ax, f"Chromosome {chrom} (Mb)")
+            self._backend.format_xaxis_mb(heatmap_ax)
+        elif gene_ax is None and heatmap_ax is None:
+            # No gene track and no heatmap - set x-label on association plot
             self._backend.set_xlabel(ax, f"Chromosome {chrom} (Mb)")
 
         # Format x-axis with Mb labels (association axis always needs formatting)
@@ -420,6 +496,11 @@ class LocusZoomPlotter:
         # Adjust layout
         self._backend.finalize_layout(fig, hspace=0.1)
 
+        # Adjust SNP labels AFTER all axis limits and layout are finalized
+        # adjustText needs final plot bounds to position labels correctly
+        if snp_label_texts:
+            self._backend.adjust_snp_labels(ax, snp_label_texts)
+
         return fig
 
     def _create_figure(
@@ -471,6 +552,304 @@ class LocusZoomPlotter:
             )
             return fig, axes[0], None
 
+    def _create_figure_with_heatmap(
+        self,
+        genes_df: Optional[pd.DataFrame],
+        chrom: int,
+        start: int,
+        end: int,
+        figsize: Tuple[float, float],
+        heatmap_data: Optional[Tuple[pd.DataFrame, List[int], List[str]]],
+        heatmap_height: float = 0.25,
+    ) -> Tuple[Any, Any, Optional[Any], Optional[Any]]:
+        """Create figure with optional gene track and heatmap panel.
+
+        Args:
+            genes_df: Gene annotations DataFrame.
+            chrom: Chromosome number.
+            start: Region start position.
+            end: Region end position.
+            figsize: Base figure size as (width, height).
+            heatmap_data: Tuple of (filtered_matrix, x_positions, snp_ids) from
+                _transform_heatmap_to_genomic_coords, or None.
+            heatmap_height: Height ratio of heatmap relative to association panel.
+
+        Returns:
+            Tuple of (fig, assoc_ax, gene_ax, heatmap_ax). gene_ax and heatmap_ax
+            are None if not included.
+        """
+        # Calculate base heights
+        assoc_height = figsize[1] * 0.6
+
+        # Calculate gene track height if needed
+        gene_track_height = 0.0
+        if genes_df is not None:
+            chrom_str = normalize_chrom(chrom)
+            region_genes = genes_df[
+                (
+                    genes_df["chr"].astype(str).str.replace("chr", "", regex=False)
+                    == chrom_str
+                )
+                & (genes_df["end"] >= start)
+                & (genes_df["start"] <= end)
+            ]
+            if not region_genes.empty:
+                temp_positions = assign_gene_positions(
+                    region_genes.sort_values("start"), start, end
+                )
+                n_gene_rows = max(temp_positions) + 1 if temp_positions else 1
+            else:
+                n_gene_rows = 1
+
+            base_gene_height = 1.0
+            per_row_height = 0.5
+            gene_track_height = base_gene_height + (n_gene_rows - 1) * per_row_height
+
+        # Calculate heatmap height if needed
+        actual_heatmap_height = 0.0
+        if heatmap_data is not None:
+            actual_heatmap_height = assoc_height * heatmap_height
+
+        # Build panel list (top to bottom): assoc, gene track, heatmap
+        n_panels = 1  # Association panel always present
+        height_ratios = [assoc_height]
+
+        if genes_df is not None:
+            n_panels += 1
+            height_ratios.append(gene_track_height)
+
+        if heatmap_data is not None:
+            n_panels += 1
+            height_ratios.append(actual_heatmap_height)
+
+        total_height = sum(height_ratios)
+
+        # Create figure
+        fig, axes = self._backend.create_figure(
+            n_panels=n_panels,
+            height_ratios=height_ratios,
+            figsize=(figsize[0], total_height),
+            sharex=True,
+        )
+
+        # Assign axes
+        assoc_ax = axes[0]
+        gene_ax = None
+        heatmap_ax = None
+
+        panel_idx = 1
+        if genes_df is not None:
+            gene_ax = axes[panel_idx]
+            panel_idx += 1
+        if heatmap_data is not None:
+            heatmap_ax = axes[panel_idx]
+
+        return fig, assoc_ax, gene_ax, heatmap_ax
+
+    def _transform_heatmap_to_genomic_coords(
+        self,
+        ld_matrix: pd.DataFrame,
+        snp_ids: List[str],
+        gwas_df: pd.DataFrame,
+        start: int,
+        end: int,
+        rs_col: str,
+        pos_col: str,
+    ) -> Optional[Tuple[pd.DataFrame, List[int], List[str]]]:
+        """Transform heatmap matrix to genomic coordinates.
+
+        Args:
+            ld_matrix: Square LD matrix from calculate_pairwise_ld.
+            snp_ids: SNP IDs in matrix order.
+            gwas_df: GWAS DataFrame with position column.
+            start: Region start position.
+            end: Region end position.
+            rs_col: SNP ID column name.
+            pos_col: Position column name.
+
+        Returns:
+            Tuple of (filtered_matrix, x_positions, filtered_snp_ids), or None
+            if no SNPs overlap with the region.
+        """
+        # Build SNP-to-position mapping from GWAS data
+        if rs_col not in gwas_df.columns:
+            logger.warning(
+                f"Cannot map heatmap to genomic coords: column '{rs_col}' not in GWAS data"
+            )
+            return None
+
+        snp_to_pos = dict(zip(gwas_df[rs_col], gwas_df[pos_col]))
+
+        # Filter to SNPs present in GWAS and within region
+        filtered_indices = []
+        filtered_snp_ids = []
+        x_positions = []
+
+        for i, snp_id in enumerate(snp_ids):
+            if snp_id in snp_to_pos:
+                pos = snp_to_pos[snp_id]
+                if start <= pos <= end:
+                    filtered_indices.append(i)
+                    filtered_snp_ids.append(snp_id)
+                    x_positions.append(int(pos))
+
+        if not filtered_indices:
+            return None
+
+        # Filter matrix to matching rows/columns
+        filtered_matrix = ld_matrix.iloc[filtered_indices, filtered_indices].copy()
+
+        return filtered_matrix, x_positions, filtered_snp_ids
+
+    def _render_heatmap_panel(
+        self,
+        ax: Any,
+        fig: Any,
+        ld_matrix: pd.DataFrame,
+        x_positions: List[int],
+        snp_ids: List[str],
+        metric: str,
+        lead_snp_id: Optional[str],
+        start: int,
+        end: int,
+    ) -> None:
+        """Render LD heatmap panel with genomic x-coordinates.
+
+        Args:
+            ax: Axes object for heatmap panel.
+            fig: Figure object.
+            ld_matrix: Filtered LD matrix.
+            x_positions: Genomic positions for each SNP (x-axis).
+            snp_ids: SNP IDs in filtered order.
+            metric: LD metric label ("r2" or "dprime").
+            lead_snp_id: Lead SNP ID to highlight (if present in snp_ids).
+            start: Region start for x-axis limits.
+            end: Region end for x-axis limits.
+        """
+        data = ld_matrix.values
+        n_snps = len(snp_ids)
+
+        # Skip rendering if only one SNP (can't show pairwise LD)
+        if n_snps < 2:
+            logger.debug("Skipping heatmap: fewer than 2 SNPs after filtering")
+            return
+
+        # Render triangular heatmap at genomic positions
+        mappable = self._backend.add_heatmap(
+            ax,
+            data=data,
+            x_coords=x_positions,
+            y_coords=list(range(n_snps)),  # Keep y as indices (0, 1, 2, ...)
+            cmap_colors=LD_HEATMAP_COLORS,
+            vmin=0.0,
+            vmax=1.0,
+            mask_upper=True,
+        )
+
+        # Add colorbar
+        label = "R²" if metric == "r2" else "D'"
+        self._backend.add_colorbar(ax, mappable, label=label)
+
+        # Highlight lead SNP if present
+        if lead_snp_id is not None and lead_snp_id in snp_ids:
+            lead_idx = snp_ids.index(lead_snp_id)
+            self._highlight_heatmap_snp(ax, fig, lead_idx, n_snps)
+
+        # Set x-axis limits to match regional plot
+        self._backend.set_xlim(ax, start, end)
+
+        # Hide y-axis (SNP indices are not meaningful for viewer)
+        self._backend.set_yticks(ax, [], [])
+        self._backend.hide_spines(ax, ["top", "right", "left"])
+
+    def _highlight_heatmap_snp(
+        self, ax: Any, fig: Any, snp_idx: int, n_snps: int
+    ) -> None:
+        """Highlight a SNP's row/column in the heatmap.
+
+        Args:
+            ax: Axes object.
+            fig: Figure object.
+            snp_idx: Index of SNP to highlight.
+            n_snps: Total number of SNPs in matrix.
+        """
+        if self._backend_name == "matplotlib":
+            from matplotlib.patches import Rectangle
+
+            # Highlight row (cells in row snp_idx, columns 0 to snp_idx)
+            for j in range(snp_idx + 1):
+                rect = Rectangle(
+                    (j - 0.5, snp_idx - 0.5),
+                    1.0,
+                    1.0,
+                    fill=False,
+                    edgecolor=LEAD_SNP_HIGHLIGHT_COLOR,
+                    linewidth=2,
+                    zorder=10,
+                )
+                ax.add_patch(rect)
+
+            # Highlight column (cells in column snp_idx, rows snp_idx to n_snps-1)
+            for i in range(snp_idx + 1, n_snps):
+                rect = Rectangle(
+                    (snp_idx - 0.5, i - 0.5),
+                    1.0,
+                    1.0,
+                    fill=False,
+                    edgecolor=LEAD_SNP_HIGHLIGHT_COLOR,
+                    linewidth=2,
+                    zorder=10,
+                )
+                ax.add_patch(rect)
+
+        elif self._backend_name == "plotly":
+            # For plotly, add shapes for row and column highlights
+            for j in range(snp_idx + 1):
+                fig.add_shape(
+                    type="rect",
+                    x0=j - 0.5,
+                    x1=j + 0.5,
+                    y0=snp_idx - 0.5,
+                    y1=snp_idx + 0.5,
+                    line=dict(color=LEAD_SNP_HIGHLIGHT_COLOR, width=2),
+                    fillcolor="rgba(0,0,0,0)",
+                )
+
+            for i in range(snp_idx + 1, n_snps):
+                fig.add_shape(
+                    type="rect",
+                    x0=snp_idx - 0.5,
+                    x1=snp_idx + 0.5,
+                    y0=i - 0.5,
+                    y1=i + 0.5,
+                    line=dict(color=LEAD_SNP_HIGHLIGHT_COLOR, width=2),
+                    fillcolor="rgba(0,0,0,0)",
+                )
+
+        elif self._backend_name == "bokeh":
+            # For bokeh, add rect glyphs for highlights
+            for j in range(snp_idx + 1):
+                ax.rect(
+                    x=j,
+                    y=snp_idx,
+                    width=1,
+                    height=1,
+                    fill_alpha=0,
+                    line_color=LEAD_SNP_HIGHLIGHT_COLOR,
+                    line_width=2,
+                )
+
+            for i in range(snp_idx + 1, n_snps):
+                ax.rect(
+                    x=snp_idx,
+                    y=i,
+                    width=1,
+                    height=1,
+                    fill_alpha=0,
+                    line_color=LEAD_SNP_HIGHLIGHT_COLOR,
+                    line_width=2,
+                )
+
     def _plot_association(
         self,
         ax: Any,
@@ -648,6 +1027,10 @@ class LocusZoomPlotter:
         finemapping_df: Optional[pd.DataFrame] = None,
         finemapping_cs_col: Optional[str] = "cs",
         recomb_df: Optional[pd.DataFrame] = None,
+        ld_heatmap_df: Optional[pd.DataFrame] = None,
+        ld_heatmap_snp_ids: Optional[List[str]] = None,
+        ld_heatmap_height: float = 0.25,
+        ld_heatmap_metric: str = "r2",
     ) -> Any:
         """Create stacked regional association plots for multiple GWAS.
 
@@ -680,10 +1063,21 @@ class LocusZoomPlotter:
                 Displayed as PIP line with optional credible set coloring.
             finemapping_cs_col: Column name for credible set assignment.
             recomb_df: Pre-loaded recombination rate data.
+            ld_heatmap_df: Pairwise LD matrix (square DataFrame) from
+                calculate_pairwise_ld. If provided with ld_heatmap_snp_ids,
+                renders heatmap panel at the very bottom of the stack.
+            ld_heatmap_snp_ids: List of SNP IDs in matrix order. Required if
+                ld_heatmap_df is provided.
+            ld_heatmap_height: Height ratio of heatmap panel relative to
+                association panel. Default 0.25.
+            ld_heatmap_metric: LD metric label for colorbar ("r2" or "dprime").
 
         Returns:
             Figure object (type depends on backend).
 
+        Raises:
+            ValueError: If ld_heatmap_df provided without ld_heatmap_snp_ids.
+
         Example:
             >>> fig = plotter.plot_stacked(
             ...     [gwas_height, gwas_bmi, gwas_whr],
@@ -739,6 +1133,12 @@ class LocusZoomPlotter:
         if eqtl_df is not None:
             validate_eqtl_df(eqtl_df)
 
+        # Validate LD heatmap parameters
+        if ld_heatmap_df is not None and ld_heatmap_snp_ids is None:
+            raise ValueError(
+                "ld_heatmap_snp_ids is required when ld_heatmap_df is provided"
+            )
+
         # Handle lead positions
         if lead_positions is None:
             lead_positions = []
@@ -762,10 +1162,31 @@ class LocusZoomPlotter:
         if ld_reference_files is None and ld_reference_file is not None:
             ld_reference_files = [ld_reference_file] * n_gwas
 
+        # Transform heatmap to genomic coordinates if provided (use first GWAS for mapping)
+        heatmap_data = None
+        if ld_heatmap_df is not None and ld_heatmap_snp_ids is not None:
+            # Use first GWAS DataFrame for SNP-to-position mapping
+            first_gwas = gwas_dfs[0].copy()
+            first_gwas = self._transform_pvalues(first_gwas, p_col)
+            heatmap_data = self._transform_heatmap_to_genomic_coords(
+                ld_matrix=ld_heatmap_df,
+                snp_ids=ld_heatmap_snp_ids,
+                gwas_df=first_gwas,
+                start=start,
+                end=end,
+                rs_col=rs_col,
+                pos_col=pos_col,
+            )
+            if heatmap_data is None:
+                logger.warning(
+                    "No SNPs from LD heatmap overlap with region - heatmap not rendered"
+                )
+
         # Calculate panel layout
         panel_height = 2.5  # inches per GWAS panel
         eqtl_height = 2.0 if eqtl_df is not None else 0
         finemapping_height = 1.5 if finemapping_df is not None else 0
+        heatmap_height_inches = panel_height * ld_heatmap_height if heatmap_data else 0
 
         # Gene track height
         if genes_df is not None:
@@ -790,11 +1211,13 @@ class LocusZoomPlotter:
             gene_track_height = 0
 
         # Calculate total panels and heights
+        # Order from top to bottom: GWAS, finemapping, eQTL, gene track, heatmap
         n_panels = (
             n_gwas
             + (1 if finemapping_df is not None else 0)
             + (1 if eqtl_df is not None else 0)
             + (1 if genes_df is not None else 0)
+            + (1 if heatmap_data is not None else 0)
         )
         height_ratios = [panel_height] * n_gwas
         if finemapping_df is not None:
@@ -803,6 +1226,8 @@ class LocusZoomPlotter:
             height_ratios.append(eqtl_height)
         if genes_df is not None:
             height_ratios.append(gene_track_height)
+        if heatmap_data is not None:
+            height_ratios.append(heatmap_height_inches)
 
         # Calculate figure height
         total_height = figsize[1] if figsize[1] else sum(height_ratios)
@@ -824,6 +1249,9 @@ class LocusZoomPlotter:
             sharex=True,
         )
 
+        # Collect label texts for deferred adjustment
+        all_snp_label_texts: list[tuple] = []
+
         # Plot each GWAS panel
         for i, (gwas_df, lead_pos) in enumerate(zip(gwas_dfs, lead_positions)):
             ax = axes[i]
@@ -874,9 +1302,10 @@ class LocusZoomPlotter:
             )
 
             # Add SNP labels (capability check - interactive backends use hover tooltips)
+            # Create labels without adjusting - we'll adjust after axis limits are set
             if snp_labels and rs_col in df.columns and label_top_n > 0 and not df.empty:
                 if self._backend.supports_snp_labels:
-                    self._backend.add_snp_labels(
+                    texts = self._backend.add_snp_labels(
                         ax,
                         df,
                         pos_col=pos_col,
@@ -885,7 +1314,10 @@ class LocusZoomPlotter:
                         label_top_n=label_top_n,
                         genes_df=genes_df,
                         chrom=chrom,
+                        adjust=False,  # Defer adjustment until after axis limits set
                     )
+                    if texts:
+                        all_snp_label_texts.append((ax, texts))
 
             # Add recombination overlay (only on first panel, all backends)
             if i == 0 and recomb_df is not None and not recomb_df.empty:
@@ -1070,8 +1502,37 @@ class LocusZoomPlotter:
             plot_gene_track_generic(
                 gene_ax, self._backend, genes_df, chrom, start, end, exons_df
             )
-            self._backend.set_xlabel(gene_ax, f"Chromosome {chrom} (Mb)")
             self._backend.hide_spines(gene_ax, ["top", "right", "left"])
+            panel_idx += 1
+
+        # Plot LD heatmap panel if provided (at very bottom)
+        if heatmap_data is not None:
+            heatmap_ax = axes[panel_idx]
+            filtered_matrix, x_positions, filtered_snp_ids = heatmap_data
+            # Find lead SNP ID from first GWAS panel if lead_positions set
+            lead_snp_id = None
+            if lead_positions and lead_positions[0] is not None:
+                first_gwas = gwas_dfs[0]
+                if rs_col in first_gwas.columns:
+                    lead_row = first_gwas[first_gwas[pos_col] == lead_positions[0]]
+                    if not lead_row.empty:
+                        lead_snp_id = lead_row[rs_col].iloc[0]
+            self._render_heatmap_panel(
+                ax=heatmap_ax,
+                fig=fig,
+                ld_matrix=filtered_matrix,
+                x_positions=x_positions,
+                snp_ids=filtered_snp_ids,
+                metric=ld_heatmap_metric,
+                lead_snp_id=lead_snp_id,
+                start=start,
+                end=end,
+            )
+            # Heatmap is at very bottom - set x-label here
+            self._backend.set_xlabel(heatmap_ax, f"Chromosome {chrom} (Mb)")
+        elif genes_df is not None:
+            # Gene track is at bottom (no heatmap) - set x-label on gene track
+            self._backend.set_xlabel(gene_ax, f"Chromosome {chrom} (Mb)")
         else:
             # Set x-label on bottom panel
             self._backend.set_xlabel(axes[-1], f"Chromosome {chrom} (Mb)")
@@ -1083,4 +1544,9 @@ class LocusZoomPlotter:
         # Adjust layout
         self._backend.finalize_layout(fig, hspace=0.1)
 
+        # Adjust SNP labels AFTER all axis limits and layout are finalized
+        # adjustText needs final plot bounds to position labels correctly
+        for ax, texts in all_snp_label_texts:
+            self._backend.adjust_snp_labels(ax, texts)
+
         return fig