pylocuszoom 0.1.0__py3-none-any.whl

pylocuszoom/eqtl.py

@@ -0,0 +1,218 @@
+"""eQTL data handling and validation for pyLocusZoom.
+
+Provides utilities for loading, validating, and preparing expression
+quantitative trait loci (eQTL) data for overlay on regional plots.
+"""
+
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+
+from .logging import logger
+
+
+REQUIRED_EQTL_COLS = ["pos", "p_value"]
+OPTIONAL_EQTL_COLS = ["gene", "effect_size", "rs", "se"]
+
+
+class EQTLValidationError(ValueError):
+    """Raised when eQTL DataFrame validation fails."""
+
+    pass
+
+
+def validate_eqtl_df(
+    df: pd.DataFrame,
+    pos_col: str = "pos",
+    p_col: str = "p_value",
+) -> None:
+    """Validate eQTL DataFrame has required columns.
+
+    Args:
+        df: eQTL DataFrame to validate.
+        pos_col: Column name for genomic position.
+        p_col: Column name for p-value.
+
+    Raises:
+        EQTLValidationError: If required columns are missing.
+    """
+    missing = []
+    if pos_col not in df.columns:
+        missing.append(pos_col)
+    if p_col not in df.columns:
+        missing.append(p_col)
+
+    if missing:
+        raise EQTLValidationError(
+            f"eQTL DataFrame missing required columns: {missing}. "
+            f"Required: {pos_col} (position), {p_col} (p-value)"
+        )
+
+
+def filter_eqtl_by_gene(
+    df: pd.DataFrame,
+    gene: str,
+    gene_col: str = "gene",
+) -> pd.DataFrame:
+    """Filter eQTL data to a specific target gene.
+
+    Args:
+        df: eQTL DataFrame.
+        gene: Target gene name to filter for.
+        gene_col: Column containing gene names.
+
+    Returns:
+        Filtered DataFrame containing only eQTLs for the target gene.
+
+    Raises:
+        EQTLValidationError: If gene column doesn't exist.
+    """
+    if gene_col not in df.columns:
+        raise EQTLValidationError(
+            f"Cannot filter by gene: column '{gene_col}' not found. "
+            f"Available columns: {list(df.columns)}"
+        )
+
+    filtered = df[df[gene_col] == gene].copy()
+    logger.debug(f"Filtered eQTL data to {len(filtered)} variants for gene {gene}")
+    return filtered
+
+
+def filter_eqtl_by_region(
+    df: pd.DataFrame,
+    chrom: int,
+    start: int,
+    end: int,
+    pos_col: str = "pos",
+    chrom_col: Optional[str] = "chr",
+) -> pd.DataFrame:
+    """Filter eQTL data to a genomic region.
+
+    Args:
+        df: eQTL DataFrame.
+        chrom: Chromosome number.
+        start: Start position.
+        end: End position.
+        pos_col: Column name for position.
+        chrom_col: Column name for chromosome (if present).
+
+    Returns:
+        Filtered DataFrame containing only eQTLs in the region.
+    """
+    mask = (df[pos_col] >= start) & (df[pos_col] <= end)
+
+    # Filter by chromosome if column exists
+    if chrom_col and chrom_col in df.columns:
+        chrom_str = str(chrom).replace("chr", "")
+        df_chrom = df[chrom_col].astype(str).str.replace("chr", "", regex=False)
+        mask = mask & (df_chrom == chrom_str)
+
+    filtered = df[mask].copy()
+    logger.debug(f"Filtered eQTL data to {len(filtered)} variants in region chr{chrom}:{start}-{end}")
+    return filtered
+
+
+def prepare_eqtl_for_plotting(
+    df: pd.DataFrame,
+    pos_col: str = "pos",
+    p_col: str = "p_value",
+    gene: Optional[str] = None,
+    chrom: Optional[int] = None,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+) -> pd.DataFrame:
+    """Prepare eQTL data for plotting.
+
+    Validates, filters, and adds computed columns needed for plotting.
+
+    Args:
+        df: Raw eQTL DataFrame.
+        pos_col: Column name for position.
+        p_col: Column name for p-value.
+        gene: Optional gene to filter for.
+        chrom: Optional chromosome for region filtering.
+        start: Optional start position for region filtering.
+        end: Optional end position for region filtering.
+
+    Returns:
+        Prepared DataFrame with neglog10p column added.
+    """
+    validate_eqtl_df(df, pos_col=pos_col, p_col=p_col)
+
+    result = df.copy()
+
+    # Filter by gene if specified
+    if gene:
+        result = filter_eqtl_by_gene(result, gene)
+
+    # Filter by region if specified
+    if chrom is not None and start is not None and end is not None:
+        result = filter_eqtl_by_region(result, chrom, start, end, pos_col=pos_col)
+
+    # Add -log10(p) column
+    result["neglog10p"] = -np.log10(result[p_col].clip(lower=1e-300))
+
+    return result
+
+
+def get_eqtl_genes(df: pd.DataFrame, gene_col: str = "gene") -> List[str]:
+    """Get list of unique genes in eQTL data.
+
+    Args:
+        df: eQTL DataFrame.
+        gene_col: Column containing gene names.
+
+    Returns:
+        Sorted list of unique gene names.
+    """
+    if gene_col not in df.columns:
+        return []
+    return sorted(df[gene_col].dropna().unique().tolist())
+
+
+def calculate_colocalization_overlap(
+    gwas_df: pd.DataFrame,
+    eqtl_df: pd.DataFrame,
+    gwas_pos_col: str = "ps",
+    eqtl_pos_col: str = "pos",
+    gwas_p_col: str = "p_wald",
+    eqtl_p_col: str = "p_value",
+    p_threshold: float = 1e-5,
+) -> pd.DataFrame:
+    """Find SNPs significant in both GWAS and eQTL.
+
+    Simple overlap analysis - for formal colocalization,
+    use dedicated tools like coloc or eCAVIAR.
+
+    Args:
+        gwas_df: GWAS results DataFrame.
+        eqtl_df: eQTL results DataFrame.
+        gwas_pos_col: Position column in GWAS data.
+        eqtl_pos_col: Position column in eQTL data.
+        gwas_p_col: P-value column in GWAS data.
+        eqtl_p_col: P-value column in eQTL data.
+        p_threshold: P-value threshold for significance.
+
+    Returns:
+        DataFrame with overlapping significant SNPs from both datasets.
+    """
+    # Filter to significant SNPs
+    sig_gwas = gwas_df[gwas_df[gwas_p_col] < p_threshold][[gwas_pos_col, gwas_p_col]]
+    sig_eqtl = eqtl_df[eqtl_df[eqtl_p_col] < p_threshold][[eqtl_pos_col, eqtl_p_col]]
+
+    # Merge on position
+    overlap = sig_gwas.merge(
+        sig_eqtl,
+        left_on=gwas_pos_col,
+        right_on=eqtl_pos_col,
+        how="inner",
+        suffixes=("_gwas", "_eqtl"),
+    )
+
+    logger.info(
+        f"Found {len(overlap)} SNPs significant in both GWAS and eQTL "
+        f"(p < {p_threshold})"
+    )
+
+    return overlap

pylocuszoom/gene_track.py

@@ -0,0 +1,311 @@
+"""Gene track visualization for regional association plots.
+
+Provides LocusZoom-style gene track plotting with:
+- Thin horizontal lines for introns
+- Thick rectangles for exons
+- Arrows indicating strand direction
+- Gene name labels
+"""
+
+from typing import List, Optional, Union
+
+import pandas as pd
+from matplotlib.axes import Axes
+from matplotlib.patches import Polygon, Rectangle
+
+from .utils import normalize_chrom
+
+# Strand-specific colors (bold, distinct)
+STRAND_COLORS: dict[Optional[str], str] = {
+    "+": "#6A3D9A",  # Bold purple for forward strand
+    "-": "#1F78B4",  # Bold teal/blue for reverse strand
+    None: "#666666",  # Grey if no strand info
+}
+
+# Layout constants
+ROW_HEIGHT = 0.40  # Total height per row
+GENE_AREA = 0.28  # Bottom portion for gene drawing
+EXON_HEIGHT = 0.22  # Exon rectangle height
+INTRON_HEIGHT = 0.02  # Thin intron line
+
+
+def assign_gene_positions(genes_df: pd.DataFrame, start: int, end: int) -> List[int]:
+    """Assign row indices to genes to minimize overlap.
+
+    Uses a greedy algorithm to stack genes vertically, placing each gene
+    in the lowest row where it doesn't overlap with existing genes.
+
+    Args:
+        genes_df: Gene annotations DataFrame sorted by start position.
+        start: Region start position.
+        end: Region end position.
+
+    Returns:
+        List of integer row indices (0, 1, 2, ...) for each gene.
+    """
+    positions = []
+    occupied = []  # List of (end_pos, row)
+    region_width = end - start
+
+    for _, gene in genes_df.iterrows():
+        gene_start = max(gene["start"], start)
+        gene_end = min(gene["end"], end)
+
+        # Find first available row with buffer for label spacing
+        row = 0
+        label_buffer = region_width * 0.08  # Extra space for labels
+        for occ_end, occ_row in occupied:
+            if occ_row == row and occ_end > gene_start - label_buffer:
+                row = occ_row + 1
+
+        positions.append(row)
+        occupied.append((gene_end, row))
+
+    return positions
+
+
+def get_nearest_gene(
+    genes_df: pd.DataFrame,
+    chrom: Union[int, str],
+    pos: int,
+    window: int = 50000,
+) -> Optional[str]:
+    """Get the nearest gene name for a genomic position.
+
+    Searches for genes that overlap or are within the specified window
+    of the given position, returning the closest by midpoint distance.
+
+    Args:
+        genes_df: Gene annotations DataFrame with chr, start, end, gene_name.
+        chrom: Chromosome number or string.
+        pos: Position in base pairs.
+        window: Window size in bp for searching nearby genes.
+
+    Returns:
+        Gene name string or None if no gene found within window.
+
+    Example:
+        >>> gene = get_nearest_gene(genes_df, chrom=1, pos=1500000)
+        >>> gene
+        'BRCA1'
+    """
+    chrom_str = normalize_chrom(chrom)
+    chrom_genes = genes_df[
+        genes_df["chr"].astype(str).str.replace("chr", "", regex=False) == chrom_str
+    ]
+
+    if chrom_genes.empty:
+        return None
+
+    # Find genes that overlap or are within window
+    nearby = chrom_genes[
+        (chrom_genes["start"] - window <= pos) & (chrom_genes["end"] + window >= pos)
+    ]
+
+    if nearby.empty:
+        return None
+
+    # Return the closest gene (by midpoint distance)
+    nearby = nearby.copy()
+    nearby["dist"] = abs((nearby["start"] + nearby["end"]) / 2 - pos)
+    return nearby.loc[nearby["dist"].idxmin(), "gene_name"]
+
+
+def plot_gene_track(
+    ax: Axes,
+    genes_df: pd.DataFrame,
+    chrom: Union[int, str],
+    start: int,
+    end: int,
+    exons_df: Optional[pd.DataFrame] = None,
+) -> None:
+    """Plot gene annotations as a LocusZoom-style track.
+
+    Creates a gene track with:
+    - Thin horizontal lines for introns (gene body)
+    - Thick rectangles for exons
+    - Arrows indicating strand direction
+    - Gene name labels
+
+    Args:
+        ax: Matplotlib axes for gene track.
+        genes_df: Gene annotations with chr, start, end, gene_name,
+            and optionally strand (+/-) column.
+        chrom: Chromosome number or string.
+        start: Region start position.
+        end: Region end position.
+        exons_df: Exon annotations with chr, start, end, gene_name
+            columns for drawing exon structure. Optional.
+    """
+    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)
+    ].copy()
+
+    ax.set_xlim(start, end)
+    ax.set_ylabel("Genes", fontsize=10)
+    ax.set_yticks([])
+
+    # theme_classic: only bottom spine
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.spines["left"].set_visible(False)
+    ax.spines["bottom"].set_linewidth(0.5)
+
+    if region_genes.empty:
+        ax.set_ylim(0, 1)
+        ax.text(
+            (start + end) / 2,
+            0.5,
+            "No genes",
+            ha="center",
+            va="center",
+            fontsize=9,
+            color="grey",
+            style="italic",
+        )
+        return
+
+    # Assign vertical positions to avoid overlap
+    region_genes = region_genes.sort_values("start")
+    positions = assign_gene_positions(region_genes, start, end)
+
+    # Set y-axis limits - small bottom margin for gene body, tight top
+    max_row = max(positions) if positions else 0
+    bottom_margin = EXON_HEIGHT / 2 + 0.02  # Room for bottom gene
+    top_margin = 0.15  # Small space above top label
+    ax.set_ylim(
+        -bottom_margin,
+        (max_row + 1) * ROW_HEIGHT - ROW_HEIGHT + GENE_AREA + top_margin,
+    )
+
+    # Filter exons for this region if available
+    region_exons = None
+    if exons_df is not None and not exons_df.empty:
+        region_exons = exons_df[
+            (
+                exons_df["chr"].astype(str).str.replace("chr", "", regex=False)
+                == chrom_str
+            )
+            & (exons_df["end"] >= start)
+            & (exons_df["start"] <= end)
+        ].copy()
+
+    for idx, (_, gene) in enumerate(region_genes.iterrows()):
+        gene_start = max(int(gene["start"]), start)
+        gene_end = min(int(gene["end"]), end)
+        row = positions[idx]
+        gene_name = gene.get("gene_name", "")
+
+        # Get strand-specific color
+        strand = gene.get("strand") if "strand" in gene.index else None
+        gene_col = STRAND_COLORS.get(strand, STRAND_COLORS[None])
+
+        # Y position: bottom of row + offset for gene area
+        y_gene = row * ROW_HEIGHT + 0.05
+        y_label = y_gene + EXON_HEIGHT / 2 + 0.01  # Just above gene top
+
+        # Check if we have exon data for this gene
+        gene_exons = None
+        if region_exons is not None and not region_exons.empty and gene_name:
+            gene_exons = region_exons[region_exons["gene_name"] == gene_name].copy()
+
+        if gene_exons is not None and not gene_exons.empty:
+            # Draw intron line (thin horizontal line spanning gene)
+            ax.add_patch(
+                Rectangle(
+                    (gene_start, y_gene - INTRON_HEIGHT / 2),
+                    gene_end - gene_start,
+                    INTRON_HEIGHT,
+                    facecolor=gene_col,
+                    edgecolor=gene_col,
+                    linewidth=0.5,
+                    zorder=1,
+                )
+            )
+
+            # Draw exons (thick rectangles)
+            for _, exon in gene_exons.iterrows():
+                exon_start = max(int(exon["start"]), start)
+                exon_end = min(int(exon["end"]), end)
+                ax.add_patch(
+                    Rectangle(
+                        (exon_start, y_gene - EXON_HEIGHT / 2),
+                        exon_end - exon_start,
+                        EXON_HEIGHT,
+                        facecolor=gene_col,
+                        edgecolor=gene_col,
+                        linewidth=0.5,
+                        zorder=2,
+                    )
+                )
+        else:
+            # No exon data - draw full gene body as rectangle (fallback)
+            ax.add_patch(
+                Rectangle(
+                    (gene_start, y_gene - EXON_HEIGHT / 2),
+                    gene_end - gene_start,
+                    EXON_HEIGHT,
+                    facecolor=gene_col,
+                    edgecolor=gene_col,
+                    linewidth=0.5,
+                    zorder=2,
+                )
+            )
+
+        # Add strand direction triangle at gene tip
+        if "strand" in gene.index:
+            strand = gene["strand"]
+            region_width = end - start
+            arrow_dir = 1 if strand == "+" else -1
+
+            # Triangle dimensions - whole arrow past gene end
+            tri_height = EXON_HEIGHT * 0.35
+            tri_width = region_width * 0.006
+
+            # Triangle entirely past gene tip
+            if arrow_dir == 1:  # Forward strand: arrow starts at gene end
+                base_x = gene_end
+                tip_x = base_x + tri_width
+                tri_points = [
+                    [tip_x, y_gene],  # Tip pointing right
+                    [base_x, y_gene + tri_height],
+                    [base_x, y_gene - tri_height],
+                ]
+            else:  # Reverse strand: arrow starts at gene start
+                base_x = gene_start
+                tip_x = base_x - tri_width
+                tri_points = [
+                    [tip_x, y_gene],  # Tip pointing left
+                    [base_x, y_gene + tri_height],
+                    [base_x, y_gene - tri_height],
+                ]
+
+            triangle = Polygon(
+                tri_points,
+                closed=True,
+                facecolor="black",
+                edgecolor="black",
+                linewidth=0.5,
+                zorder=5,
+            )
+            ax.add_patch(triangle)
+
+        # Add gene name label in the gap above gene
+        if gene_name:
+            label_pos = (gene_start + gene_end) / 2
+            ax.text(
+                label_pos,
+                y_label,
+                gene_name,
+                ha="center",
+                va="bottom",
+                fontsize=5.5,
+                color="#000000",
+                fontweight="medium",
+                style="italic",
+                zorder=4,
+                clip_on=True,
+            )

pylocuszoom/labels.py

@@ -0,0 +1,118 @@
+"""SNP label placement for regional association plots.
+
+Provides automatic labeling of top significant SNPs with:
+- SNP ID (rs number)
+- Nearest gene name (if gene annotations provided)
+- Automatic overlap avoidance (if adjustText installed)
+"""
+
+from typing import List, Optional, Union
+
+import pandas as pd
+from matplotlib.axes import Axes
+from matplotlib.text import Annotation
+
+from .gene_track import get_nearest_gene
+
+
+def add_snp_labels(
+    ax: Axes,
+    df: pd.DataFrame,
+    pos_col: str = "ps",
+    neglog10p_col: str = "neglog10p",
+    rs_col: str = "rs",
+    label_top_n: int = 5,
+    genes_df: Optional[pd.DataFrame] = None,
+    chrom: Optional[Union[int, str]] = None,
+    max_label_length: int = 15,
+) -> List[Annotation]:
+    """Add text labels to top SNPs in the regional plot.
+
+    Labels the most significant SNPs with either their SNP ID
+    or the nearest gene name (if genes_df provided).
+
+    Args:
+        ax: Matplotlib axes object.
+        df: DataFrame with SNP data. Must have the specified position,
+            neglog10p, and rs columns.
+        pos_col: Column name for position.
+        neglog10p_col: Column name for -log10(p-value).
+        rs_col: Column name for SNP ID.
+        label_top_n: Number of top SNPs to label.
+        genes_df: Optional gene annotations for gene-based labels.
+            If provided with chrom, labels will show nearest gene name
+            instead of SNP ID.
+        chrom: Chromosome number. Required if genes_df is provided.
+        max_label_length: Maximum label length before truncation.
+
+    Returns:
+        List of matplotlib text annotation objects.
+
+    Example:
+        >>> fig, ax = plt.subplots()
+        >>> # ... plot your data ...
+        >>> texts = add_snp_labels(ax, df, label_top_n=5)
+    """
+    if neglog10p_col not in df.columns:
+        raise ValueError(
+            f"Column '{neglog10p_col}' not found in DataFrame. "
+            "Ensure -log10(p) values are calculated before calling add_snp_labels."
+        )
+
+    # Get top N SNPs by -log10(p)
+    top_snps = df.nlargest(label_top_n, neglog10p_col)
+
+    texts = []
+    for _, snp in top_snps.iterrows():
+        x = snp[pos_col]
+        y = snp[neglog10p_col]
+
+        # Determine label text
+        label = str(snp[rs_col])
+
+        # Try to get gene name if genes_df provided
+        if genes_df is not None and chrom is not None:
+            nearest_gene = get_nearest_gene(genes_df, chrom, int(x))
+            if nearest_gene:
+                label = nearest_gene
+
+        # Truncate long labels
+        if len(label) > max_label_length:
+            label = label[: max_label_length - 3] + "..."
+
+        # Add text annotation with offset
+        text = ax.annotate(
+            label,
+            xy=(x, y),
+            xytext=(5, 5),
+            textcoords="offset points",
+            fontsize=8,
+            fontweight="bold",
+            color="#333333",
+            ha="left",
+            va="bottom",
+            zorder=15,
+            bbox=dict(
+                boxstyle="round,pad=0.2",
+                facecolor="white",
+                edgecolor="none",
+                alpha=0.8,
+            ),
+        )
+        texts.append(text)
+
+    # Try to adjust text positions to avoid overlap
+    try:
+        from adjustText import adjust_text
+
+        adjust_text(
+            texts,
+            ax=ax,
+            arrowprops=dict(arrowstyle="-", color="gray", lw=0.5),
+            expand_points=(1.5, 1.5),
+        )
+    except ImportError:
+        # adjustText not installed, labels may overlap
+        pass
+
+    return texts