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

pylocuszoom/finemapping.py

@@ -4,18 +4,20 @@ Provides utilities for loading, validating, and preparing statistical
 fine-mapping results (SuSiE, FINEMAP, etc.) for visualization.
 """
 
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import pandas as pd
 
+from .backends.base import PlotBackend
+from .backends.hover import HoverConfig, HoverDataBuilder
+from .colors import PIP_LINE_COLOR, get_credible_set_color
 from .exceptions import FinemappingValidationError, ValidationError
 from .logging import logger
 from .utils import filter_by_region
 from .validation import DataFrameValidator
 
-# Required columns for fine-mapping data
+# Required columns for fine-mapping data (default column names)
 REQUIRED_FINEMAPPING_COLS = ["pos", "pip"]
-OPTIONAL_FINEMAPPING_COLS = ["rs", "cs", "cs_id", "effect", "se"]
 
 
 def validate_finemapping_df(
@@ -207,3 +209,109 @@ def calculate_credible_set_coverage(
         coverage[cs_id] = cs_data[pip_col].sum()
 
     return coverage
+
+
+def plot_finemapping(
+    backend: PlotBackend,
+    ax: Any,
+    df: pd.DataFrame,
+    pos_col: str = "pos",
+    pip_col: str = "pip",
+    cs_col: Optional[str] = "cs",
+    show_credible_sets: bool = True,
+    pip_threshold: float = 0.0,
+) -> None:
+    """Plot fine-mapping results (PIP line with credible set coloring).
+
+    Renders posterior inclusion probabilities as a line plot, with optional
+    scatter points colored by credible set membership.
+
+    Args:
+        backend: Plotting backend implementing PlotBackend protocol.
+        ax: Axes or panel to plot on.
+        df: Fine-mapping DataFrame with pos and pip columns.
+        pos_col: Column name for position.
+        pip_col: Column name for posterior inclusion probability.
+        cs_col: Column name for credible set assignment (optional).
+        show_credible_sets: Whether to color points by credible set.
+        pip_threshold: Minimum PIP to display as scatter point.
+    """
+    # Build hover data using HoverDataBuilder
+    extra_cols = {pip_col: "PIP"}
+    if cs_col and cs_col in df.columns:
+        extra_cols[cs_col] = "Credible Set"
+    hover_config = HoverConfig(
+        pos_col=pos_col if pos_col in df.columns else None,
+        extra_cols=extra_cols,
+    )
+    hover_builder = HoverDataBuilder(hover_config)
+
+    # Sort by position for line plotting
+    df = df.sort_values(pos_col)
+
+    # Plot PIP as line
+    backend.line(
+        ax,
+        df[pos_col],
+        df[pip_col],
+        color=PIP_LINE_COLOR,
+        linewidth=1.5,
+        alpha=0.8,
+        zorder=1,
+    )
+
+    # Check if credible sets are available
+    has_cs = cs_col is not None and cs_col in df.columns and show_credible_sets
+    credible_sets = get_credible_sets(df, cs_col) if has_cs else []
+
+    if credible_sets:
+        # Plot points colored by credible set
+        for cs_id in credible_sets:
+            cs_data = df[df[cs_col] == cs_id]
+            color = get_credible_set_color(cs_id)
+            backend.scatter(
+                ax,
+                cs_data[pos_col],
+                cs_data[pip_col],
+                colors=color,
+                sizes=50,
+                marker="o",
+                edgecolor="black",
+                linewidth=0.5,
+                zorder=3,
+                hover_data=hover_builder.build_dataframe(cs_data),
+            )
+        # Plot variants not in any credible set (only if threshold is set)
+        if pip_threshold > 0:
+            non_cs_data = df[(df[cs_col].isna()) | (df[cs_col] == 0)]
+            non_cs_data = non_cs_data[non_cs_data[pip_col] >= pip_threshold]
+            if not non_cs_data.empty:
+                backend.scatter(
+                    ax,
+                    non_cs_data[pos_col],
+                    non_cs_data[pip_col],
+                    colors="#BEBEBE",
+                    sizes=30,
+                    marker="o",
+                    edgecolor="black",
+                    linewidth=0.3,
+                    zorder=2,
+                    hover_data=hover_builder.build_dataframe(non_cs_data),
+                )
+    else:
+        # No credible sets - show all points above threshold
+        if pip_threshold > 0:
+            high_pip = df[df[pip_col] >= pip_threshold]
+            if not high_pip.empty:
+                backend.scatter(
+                    ax,
+                    high_pip[pos_col],
+                    high_pip[pip_col],
+                    colors=PIP_LINE_COLOR,
+                    sizes=50,
+                    marker="o",
+                    edgecolor="black",
+                    linewidth=0.5,
+                    zorder=3,
+                    hover_data=hover_builder.build_dataframe(high_pip),
+                )

pylocuszoom/labels.py

@@ -24,6 +24,7 @@ def add_snp_labels(
     genes_df: Optional[pd.DataFrame] = None,
     chrom: Optional[Union[int, str]] = None,
     max_label_length: int = 15,
+    adjust: bool = True,
     **kwargs: Any,
 ) -> List[Annotation]:
     """Add text labels to top SNPs in the regional plot.
@@ -41,6 +42,8 @@ def add_snp_labels(
         genes_df: Unused, kept for backward compatibility.
         chrom: Unused, kept for backward compatibility.
         max_label_length: Maximum label length before truncation.
+        adjust: If True, run adjustText immediately. If False, caller must
+            call adjust_snp_labels() after setting axis limits.
 
     Returns:
         List of matplotlib text annotation objects.
@@ -101,21 +104,43 @@ def add_snp_labels(
         )
         texts.append(text)
 
-    # Only use adjustText when there are multiple labels to avoid overlap
-    if len(texts) > 1:
-        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:
-            logger.warning(
-                "adjustText not installed - SNP labels may overlap. "
-                "Install with: pip install adjustText"
-            )
+    if adjust:
+        adjust_snp_labels(ax, texts)
 
     return texts
+
+
+def adjust_snp_labels(ax: Axes, texts: List[Annotation]) -> None:
+    """Adjust SNP label positions to avoid overlaps.
+
+    This function should be called AFTER all axis limits have been set,
+    as adjustText needs to know the final plot bounds to position labels
+    correctly within the visible area.
+
+    Args:
+        ax: Matplotlib axes object.
+        texts: List of text annotation objects from add_snp_labels().
+
+    Example:
+        >>> texts = add_snp_labels(ax, df, adjust=False)
+        >>> ax.set_xlim(start, end)
+        >>> ax.set_ylim(0, max_y)
+        >>> adjust_snp_labels(ax, texts)
+    """
+    if len(texts) <= 1:
+        return
+
+    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:
+        logger.warning(
+            "adjustText not installed - SNP labels may overlap. "
+            "Install with: pip install adjustText"
+        )

pylocuszoom/ld.py

@@ -16,6 +16,72 @@ from .logging import logger
 from .utils import validate_plink_files
 
 
+def build_pairwise_ld_command(
+    plink_path: str,
+    bfile_path: str,
+    output_path: str,
+    snp_list_file: Optional[str] = None,
+    chrom: Optional[int] = None,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    species: Optional[str] = "canine",
+    metric: str = "r2",
+) -> list:
+    """Build PLINK command for pairwise LD matrix computation.
+
+    Generates command for computing an N x N LD matrix using PLINK's
+    --r2 square (or --r dprime square) command.
+
+    Args:
+        plink_path: Path to PLINK executable.
+        bfile_path: Input binary fileset prefix (.bed/.bim/.fam).
+        output_path: Output prefix (creates .ld and .snplist files).
+        snp_list_file: Path to file with SNP IDs to extract (one per line).
+        chrom: Chromosome number for region-based extraction.
+        start: Start position (bp) for region-based extraction.
+        end: End position (bp) for region-based extraction.
+        species: Species flag ('canine', 'feline', or None for human).
+        metric: LD metric ('r2' or 'dprime').
+
+    Returns:
+        List of command arguments for subprocess.
+    """
+    cmd = [plink_path]
+
+    # Species flag
+    if species == "canine":
+        cmd.append("--dog")
+    elif species == "feline":
+        cmd.extend(["--chr-set", "18"])
+
+    # Input and output
+    cmd.extend(["--bfile", bfile_path])
+    cmd.extend(["--out", output_path])
+
+    # LD metric and square matrix flag
+    if metric == "dprime":
+        cmd.extend(["--r", "dprime", "square"])
+    else:
+        cmd.extend(["--r2", "square"])
+
+    # Track SNP order in output
+    cmd.append("--write-snplist")
+
+    # SNP extraction mode
+    if snp_list_file:
+        cmd.extend(["--extract", snp_list_file])
+
+    # Region-based extraction
+    if chrom is not None:
+        cmd.extend(["--chr", str(chrom)])
+    if start is not None:
+        cmd.extend(["--from-bp", str(start)])
+    if end is not None:
+        cmd.extend(["--to-bp", str(end)])
+
+    return cmd
+
+
 def find_plink() -> Optional[str]:
     """Find PLINK executable on PATH.
 
@@ -84,6 +150,51 @@ def build_ld_command(
     return cmd
 
 
+def parse_pairwise_ld_output(
+    ld_file: str, snplist_file: str
+) -> tuple[pd.DataFrame, list[str]]:
+    """Parse PLINK pairwise LD matrix output files.
+
+    PLINK --r2 square outputs:
+    - .ld file: N x N matrix of R2/D' values (whitespace-separated, no headers)
+    - .snplist file: SNP IDs in order (one per line)
+
+    Args:
+        ld_file: Path to .ld output file (square matrix).
+        snplist_file: Path to .snplist output file (SNP IDs).
+
+    Returns:
+        Tuple of (DataFrame with R2/D' values, list of SNP IDs).
+        DataFrame has SNP IDs as both index and columns.
+        Returns (empty DataFrame, empty list) if files not found.
+    """
+    # Check if files exist
+    if not os.path.exists(ld_file) or not os.path.exists(snplist_file):
+        return pd.DataFrame(), []
+
+    # Read SNP list
+    with open(snplist_file) as f:
+        snp_ids = [line.strip() for line in f if line.strip()]
+
+    if not snp_ids:
+        return pd.DataFrame(), []
+
+    # Read LD matrix (whitespace-separated, no headers)
+    # Values can be numbers or 'nan'
+    matrix = pd.read_csv(
+        ld_file,
+        sep=r"\s+",
+        header=None,
+        names=snp_ids,
+        index_col=False,
+    )
+
+    # Set SNP IDs as row index
+    matrix.index = snp_ids
+
+    return matrix, snp_ids
+
+
 def parse_ld_output(ld_file: str, lead_snp: str) -> pd.DataFrame:
     """Parse PLINK .ld output file.
 
@@ -208,3 +319,131 @@ def calculate_ld(
         # Clean up temp directory
         if cleanup_working_dir and os.path.exists(working_dir):
             shutil.rmtree(working_dir, ignore_errors=True)
+
+
+def calculate_pairwise_ld(
+    bfile_path: str,
+    snp_list: list[str] | None = None,
+    chrom: int | None = None,
+    start: int | None = None,
+    end: int | None = None,
+    plink_path: str | None = None,
+    working_dir: str | None = None,
+    species: str = "canine",
+    metric: str = "r2",
+) -> tuple[pd.DataFrame, list[str]]:
+    """Calculate pairwise LD matrix for a set of variants.
+
+    Runs PLINK --r2 square to compute an N x N LD matrix, suitable for
+    LD heatmap visualization.
+
+    Args:
+        bfile_path: Path to PLINK binary fileset (.bed/.bim/.fam prefix).
+        snp_list: List of SNP IDs to compute pairwise LD between.
+        chrom: Chromosome number for region-based extraction.
+        start: Start position (bp) for region-based extraction.
+        end: End position (bp) for region-based extraction.
+        plink_path: Path to PLINK executable. Auto-detects if None.
+        working_dir: Directory for PLINK output files. Uses temp dir if None.
+        species: Species flag ('canine', 'feline', or None for human).
+        metric: LD metric ('r2' or 'dprime').
+
+    Returns:
+        Tuple of (LD matrix DataFrame, list of SNP IDs).
+        DataFrame has SNP IDs as both index and columns.
+        Returns (empty DataFrame, empty list) if PLINK fails.
+
+    Raises:
+        FileNotFoundError: If PLINK executable not found.
+        ValidationError: If PLINK binary files (.bed/.bim/.fam) are missing.
+        ValidationError: If requested SNPs are not found in reference panel.
+
+    Example:
+        >>> matrix, snp_ids = calculate_pairwise_ld(
+        ...     bfile_path="/path/to/genotypes",
+        ...     snp_list=["rs1", "rs2", "rs3"],
+        ... )
+        >>> # matrix is 3x3 DataFrame with LD values
+        >>> matrix.loc["rs1", "rs2"]  # LD between rs1 and rs2
+    """
+    from .utils import ValidationError
+
+    # Find PLINK
+    if plink_path is None:
+        plink_path = find_plink()
+    if plink_path is None:
+        raise FileNotFoundError(
+            "PLINK not found. Install PLINK 1.9 or specify plink_path."
+        )
+
+    logger.debug(f"Using PLINK at {plink_path}")
+
+    # Validate PLINK files exist
+    validate_plink_files(bfile_path)
+
+    # Use temp directory if working_dir not specified
+    cleanup_working_dir = False
+    if working_dir is None:
+        working_dir = tempfile.mkdtemp(prefix="snp_scope_pairwise_ld_")
+        cleanup_working_dir = True
+
+    try:
+        os.makedirs(working_dir, exist_ok=True)
+        output_prefix = os.path.join(working_dir, "pairwise_ld")
+
+        # Write SNP list to file if provided
+        snp_list_file = None
+        if snp_list:
+            snp_list_file = os.path.join(working_dir, "snp_list.txt")
+            with open(snp_list_file, "w") as f:
+                for snp in snp_list:
+                    f.write(f"{snp}\n")
+
+        # Build and run PLINK command
+        cmd = build_pairwise_ld_command(
+            plink_path=plink_path,
+            bfile_path=bfile_path,
+            output_path=output_prefix,
+            snp_list_file=snp_list_file,
+            chrom=chrom,
+            start=start,
+            end=end,
+            species=species,
+            metric=metric,
+        )
+
+        logger.debug(f"Running PLINK command: {' '.join(cmd)}")
+
+        result = subprocess.run(
+            cmd,
+            cwd=working_dir,
+            capture_output=True,
+            text=True,
+        )
+
+        if result.returncode != 0:
+            logger.warning(
+                f"PLINK pairwise LD calculation failed: {result.stderr[:200]}"
+            )
+            return pd.DataFrame(), []
+
+        # Parse output
+        ld_file = f"{output_prefix}.ld"
+        snplist_file = f"{output_prefix}.snplist"
+
+        matrix, found_snps = parse_pairwise_ld_output(ld_file, snplist_file)
+
+        # Validate all requested SNPs were found
+        if snp_list:
+            missing_snps = set(snp_list) - set(found_snps)
+            if missing_snps:
+                raise ValidationError(
+                    f"SNPs not found in reference panel: {', '.join(sorted(missing_snps))}"
+                )
+
+        return matrix, found_snps
+
+    finally:
+        # Clean up temp directory
+        if cleanup_working_dir and os.path.exists(working_dir):
+            shutil.rmtree(working_dir, ignore_errors=True)