diff --git a/snapatac2/plotting/__init__.py b/snapatac2/plotting/__init__.py index c6adec189..6a0ace922 100644 --- a/snapatac2/plotting/__init__.py +++ b/snapatac2/plotting/__init__.py @@ -1,7 +1,9 @@ from __future__ import annotations -import numpy as np import logging +from pathlib import Path + +import numpy as np import snapatac2 from snapatac2._snapatac2 import AnnData, AnnDataSet @@ -553,11 +555,401 @@ def motif_enrichment( **kwargs, ) +def _parse_gtf_attributes(attr_str: str) -> dict[str, str]: + """Parse a GTF ``key "value";`` attribute string.""" + attrs = {} + for part in attr_str.rstrip(";").split(";"): + part = part.strip() + if not part: + continue + if '"' in part: + key, val = part.split('"', 1) + attrs[key.strip()] = val.strip('"').strip() + return attrs + + +def _open_annotation(f: str | Path): + import gzip + path = Path(f) + return gzip.open(path, "rt") if path.suffix == ".gz" else open(path, "r") + + +def _lookup_gene_region(annotation: str | Path, gene_name: str) -> tuple[str, int, int] | None: + """Look up *(chrom, start, end)* of *gene_name* in a GTF file.""" + for seqid, feat_type, feat_start, feat_end, _strand, attrs in _iter_annotation(annotation): + if feat_type != "gene": + continue + if attrs.get("gene_name") == gene_name: + return (seqid, feat_start, feat_end) + return None + + +def _gen_to_pixel(pos: int, region_start: int, region_end: int, n_points: int) -> float: + """Map a genomic coordinate to the x-pixel range [0, n_points-1].""" + return (pos - region_start) / (region_end - region_start) * (n_points - 1) + + +def _style_annotation_axis(ax, label: str): + """Strip spines/ticks from *ax* and set its ylabel.""" + ax.set_yticks([]) + ax.spines[["top", "right", "left"]].set_visible(False) + ax.tick_params(left=False, labelleft=False) + ax.set_xticks([]) + ax.set_ylabel(label, fontsize=7, labelpad=2) + + +def _greedy_row_layout(intervals: list[tuple[float, float]]) -> list[int]: + """Assign non-overlapping rows to a list of *(x1, x2)* pixel intervals. + Returns row indices (same order as *intervals*). + """ + row_ends: list[float] = [] + placements: list[int] = [] + for x1, x2 in intervals: + row = 0 + while row < len(row_ends) and x1 < row_ends[row]: + row += 1 + if row == len(row_ends): + row_ends.append(x2) + else: + row_ends[row] = max(row_ends[row], x2) + placements.append(row) + return placements + + +def _iter_annotation( + annotation: str | Path, + chrom: str | None = None, + start: int | None = None, + end: int | None = None, +): + """Yield ``(seqid, feat_type, feat_start, feat_end, strand, attrs)`` + from a coordinate-sorted GTF file, optionally filtered by region.""" + found_chrom = False + with _open_annotation(annotation) as fh: + for line in fh: + if line.startswith("#") or line.strip() == "": + continue + cols = line.strip().split("\t") + if len(cols) < 9: + continue + seqid, _source, feat_type, feat_start, feat_end, _score, strand, _phase, attrs_str = cols + if chrom is not None and seqid != chrom: + if found_chrom: + break + continue + found_chrom = True + feat_start = int(feat_start) + feat_end = int(feat_end) + if start is not None and feat_end < start: + continue + if end is not None and feat_start > end: + break + yield (seqid, feat_type, feat_start, feat_end, strand, + _parse_gtf_attributes(attrs_str)) + + +def _normalize_gene_types( + gene_type: str | list[str] | None, +) -> set[str] | None: + """Normalize gene type selection for annotation track filtering.""" + if gene_type is None: + return None + if isinstance(gene_type, str): + return {gene_type} + return set(gene_type) + + +def _get_gene_models( + annotation: str | Path, + chrom: str, + start: int, + end: int, + gene_types: set[str] | None = None, +) -> list[dict]: + """Read gene/exon models overlapping *chrom:start-end* from a GTF file.""" + genes: dict[str, dict] = {} + for seqid, feat_type, feat_start, feat_end, strand, attrs in _iter_annotation( + annotation, chrom=chrom, start=start, end=end, + ): + gene_id = attrs.get("gene_id") + if not gene_id: + continue + if gene_id not in genes: + genes[gene_id] = { + "name": attrs.get("gene_name", gene_id), + "gene_type": attrs.get("gene_type", attrs.get("gene_biotype")), + "strand": strand, + "exons": [], + "tx_start": feat_start, + "tx_end": feat_end, + "chrom": seqid, + } + gene = genes[gene_id] + gene["tx_start"] = min(gene["tx_start"], feat_start) + gene["tx_end"] = max(gene["tx_end"], feat_end) + if feat_type == "exon": + gene["exons"].append((feat_start, feat_end)) + + result = [] + for gene in genes.values(): + if gene_types is not None and gene["gene_type"] not in gene_types: + continue + exons = _merge_intervals(gene["exons"]) + result.append({ + "name": gene["name"], + "strand": gene["strand"], + "tx_start": max(gene["tx_start"], start), + "tx_end": min(gene["tx_end"], end), + "exons": [(max(s, start), min(e, end)) for s, e in exons + if max(s, start) < min(e, end)], + }) + return result + + +def _merge_intervals( + intervals: list[tuple[int, int]], +) -> list[tuple[int, int]]: + """Merge duplicate and overlapping genomic intervals.""" + merged: list[tuple[int, int]] = [] + for start, end in sorted(set(intervals)): + if merged and start <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + return merged + + +def _draw_gene_track( + ax, + models: list[dict], + region_start: int, + region_end: int, + n_points: int, +): + """Draw gene models on *ax* spanning pixel range [0, n_points-1].""" + import matplotlib.patches as mpatches + + if not models: + _style_annotation_axis(ax, "Genes") + return + + y_center = 0.5 + exon_height = 0.35 + + intervals = [(_gen_to_pixel(g["tx_start"], region_start, region_end, n_points), + _gen_to_pixel(g["tx_end"], region_start, region_end, n_points)) + for g in models] + placements = _greedy_row_layout(intervals) + + n_rows = max(placements) + 1 if placements else 1 + ax.set_ylim(-0.3, n_rows + 0.3) + + for g, row, (gx1, gx2) in zip(models, placements, intervals): + y_base = row + y_center + ax.hlines(y_base, gx1, gx2, colors="0.2", linewidth=1.5, zorder=2) + + for ex_start, ex_end in g["exons"]: + ex1 = _gen_to_pixel(ex_start, region_start, region_end, n_points) + ex2 = _gen_to_pixel(ex_end, region_start, region_end, n_points) + if ex2 - ex1 < 1: + continue + rect = mpatches.Rectangle( + (ex1, y_base - exon_height / 2), ex2 - ex1, exon_height, + facecolor="0.2", edgecolor="none", linewidth=0, zorder=3, + ) + ax.add_patch(rect) + + if g["strand"] == "+": + ax.annotate("▶", xy=(gx2, y_base), fontsize=5, ha="center", + va="center", color="0.2", zorder=4) + elif g["strand"] == "-": + ax.annotate("◀", xy=(gx1, y_base), fontsize=5, ha="center", + va="center", color="0.2", zorder=4) + + ax.text(gx1, y_base + exon_height / 2 + 0.06, g["name"], + fontsize=5, ha="left", va="bottom", color="0.2", zorder=5) + + ax.set_xlim(0, n_points - 1) + _style_annotation_axis(ax, "Genes") + + +def _add_highlights( + axes: list, + highlights: list[tuple[int, int, object]], + region_start: int, + region_end: int, + n_points: int, + alpha: float, +): + """Overlay normalized highlighted regions across all *axes*.""" + if not 0 <= alpha <= 1: + raise ValueError("highlight_alpha must be between 0 and 1.") + for start, end, color in highlights: + x1 = _gen_to_pixel(start, region_start, region_end, n_points) + x2 = _gen_to_pixel(end, region_start, region_end, n_points) + for ax in axes: + ax.axvspan(x1, x2, color=color, alpha=alpha, zorder=0) + + +def _normalize_highlights( + highlights: list, + region_start: int, + region_end: int, +) -> list[tuple[int, int, object]]: + """Normalize and clip tuple- or dict-based highlight definitions.""" + result = [] + default_color = (1, 0.85, 0.6) + for index, item in enumerate(highlights): + if isinstance(item, dict): + if "start" not in item or "end" not in item: + raise ValueError( + f"highlight[{index}] must contain 'start' and 'end'." + ) + start, end = item["start"], item["end"] + color = item.get("color", default_color) + elif isinstance(item, (list, tuple)) and len(item) in (2, 3): + start, end = item[:2] + color = item[2] if len(item) > 2 else default_color + else: + raise ValueError( + f"highlight[{index}] must be a 2- or 3-item sequence or a dict." + ) + + if not isinstance(start, (int, float)) or not isinstance(end, (int, float)): + raise ValueError(f"highlight[{index}] coordinates must be numeric.") + if start >= end: + raise ValueError(f"highlight[{index}] start must be smaller than end.") + clipped_start = max(start, region_start) + clipped_end = min(end, region_end) + if clipped_start < clipped_end: + result.append((clipped_start, clipped_end, color)) + return result + + +def _get_peaks_in_region( + adata: AnnData, + chrom: str, + start: int, + end: int, + max_peaks: int = 2000, +) -> list[tuple[int, int, str]]: + """Return peaks (var_names ``"chrN:start-end"``) overlapping *chrom:start-end*. + + Chromosome and coordinate order in ``adata.var_names`` is used to stop the + scan as soon as it passes the requested region. + """ + peaks: list[tuple[int, int, str]] = [] + prefix = f"{chrom}:" + found_chrom = False + for name in adata.var_names: + if not name.startswith(prefix): + if found_chrom: + break + continue + found_chrom = True + coord_part = name[len(prefix):] + if "-" not in coord_part: + continue + try: + ps, pe = coord_part.split("-", 1) + ps = int(ps) + pe = int(pe) + except (ValueError, IndexError): + continue + if ps >= end: + break + if ps < end and pe > start: + peaks.append((ps, pe, name)) + if len(peaks) >= max_peaks: + break + return peaks + + +def _draw_peak_track( + ax, + peaks: list[tuple[int, int, str]], + region_start: int, + region_end: int, + n_points: int, +): + """Draw overlapping peaks as horizontal bars in a compact track.""" + import matplotlib.patches as mpatches + from matplotlib.colors import to_rgba + + if not peaks: + _style_annotation_axis(ax, "Peaks") + return + + bar_height = 0.25 + + intervals = [(_gen_to_pixel(max(ps, region_start), region_start, region_end, n_points), + _gen_to_pixel(min(pe, region_end), region_start, region_end, n_points)) + for ps, pe, _name in peaks] + placements = _greedy_row_layout(intervals) + + n_rows = max(placements) + 1 if placements else 1 + ax.set_ylim(-0.3, n_rows + 0.3) + ax.set_xlim(0, n_points - 1) + + for (ps, pe, _name), row, (px1, px2) in zip(peaks, placements, intervals): + y_base = row + 0.5 + c = to_rgba("tab:red", alpha=0.7) + rect = mpatches.Rectangle( + (px1, y_base - bar_height / 2), max(px2 - px1, 1.0), bar_height, + facecolor=c, edgecolor="none", linewidth=0, zorder=3, + ) + ax.add_patch(rect) + + _style_annotation_axis(ax, "Peaks") + + +def _resolve_annotation_path(gene_annotation) -> Path | None: + """Resolve a *gene_annotation* parameter to a Path, or ``None``.""" + if gene_annotation is None: + return None + if not isinstance(gene_annotation, (str, Path)): + raise TypeError("gene_annotation must be a path to a GTF file.") + return Path(gene_annotation) + + +def _parse_region( + region: str, + gene_annotation_path: Path | None, +) -> tuple[str, int, int, str]: + """Parse *region* (``"chrom:start-end"`` or gene name) into + ``(chrom, start, end, original_region_string)``. + + If *region* is a gene name and *gene_annotation_path* is provided, + the gene coordinates are looked up from the annotation file. + """ + if ":" in region and "-" in region: + chrom, coord_part = region.split(":", 1) + start_s, end_s = coord_part.split("-", 1) + start, end = int(start_s), int(end_s) + if start >= end: + raise ValueError("Region start must be smaller than end.") + return (chrom, start, end, region) + if gene_annotation_path is not None: + result = _lookup_gene_region(gene_annotation_path, region.strip()) + if result is not None: + chrom, gs, ge = result + return (chrom, gs, ge, region) + raise ValueError( + f"Could not parse region {region!r}. Format as 'chr:start-end' or " + "provide gene_annotation for gene name lookup." + ) + + def coverage( adata: AnnData, region: str, groupby: str | list[str], out_file: str | None = None, + gene_annotation: str | Path | None = None, + gene_type: str | list[str] | None = "protein_coding", + highlight: list | None = None, + highlight_alpha: float = 0.25, + peak_track: bool = True, ): """Plot coverage tracks for grouped cells across one genomic region. @@ -581,13 +973,39 @@ def coverage( Annotated data matrix with fragments available for coverage retrieval. region : str Genomic interval to plot, formatted as ``"chrom:start-end"``; for - example, ``"chr1:100000-200000"``. + example, ``"chr1:100000-200000"``. If ``gene_annotation`` is provided + and *region* does not match that format, it is treated as a gene name + and the gene's coordinates are looked up from the annotation file. groupby : str or list of str Cell grouping definition. If a string, groups are read from ``adata.obs[groupby]``. out_file : str or None Output path for saving the Matplotlib figure. If ``None``, display the plot with ``matplotlib.pyplot.show``. + gene_annotation : str or pathlib.Path, optional + Path to a coordinate-sorted GTF annotation file. When provided, a gene + annotation track is drawn below the coverage tracks showing gene + introns, exons, and strand direction. Additionally, *region* can be + given as a gene name (e.g. ``"GENE_A"``) to automatically resolve the + coordinates. + gene_type : str, list of str, or None + Gene types shown in the annotation track. Defaults to + ``"protein_coding"``. Pass multiple types as a list, or ``None`` to + show all annotated gene types. This does not affect gene-name lookup. + highlight : list, optional + Regions to highlight across coverage tracks. Each element can be: + + * ``(start, end)`` — highlighted in pale orange. + * ``(start, end, color)`` — highlighted in a custom color (any + matplotlib-compatible color or RGBA tuple). + * ``{"start": s, "end": e, "color": c}`` — dict form. + + Coordinates are genomic positions within the plotted interval. + highlight_alpha : float + Opacity of highlighted regions, between 0 (transparent) and 1 (opaque). + peak_track : bool + If ``True``, draw a peak annotation track at the top showing the + overlapping peak regions from ``adata.var_names``. Returns ------- @@ -612,46 +1030,74 @@ def coverage( from matplotlib import pyplot as plt + annotation_path = _resolve_annotation_path(gene_annotation) + gene_types = _normalize_gene_types(gene_type) + chrom, region_start, region_end, resolved_region = _parse_region(region, annotation_path) + coverage_region_str = f"{chrom}:{region_start}-{region_end}" + if not 0 <= highlight_alpha <= 1: + raise ValueError("highlight_alpha must be between 0 and 1.") + groupby = adata.obs[groupby] if isinstance(groupby, str) else groupby groupby = [x for x in groupby] signal_values = [] track_names = [] - for k, v in sorted(list(internal.get_coverage(adata, region, groupby).items())): + for k, v in sorted(list(internal.get_coverage(adata, coverage_region_str, groupby).items())): track_names.append(k) signal_values.append(v) signal_values = np.array(signal_values) - - start, end = region.split(":")[1].split("-") - start = int(start) - end = int(end) height_per_track = 1.2 width = 6 n_tracks, n_points = signal_values.shape + + n_rows = n_tracks + int(peak_track) + int(annotation_path is not None) + fig, axes = plt.subplots( - n_tracks, 1, - figsize=[width, n_tracks * height_per_track], - sharex=True, - constrained_layout=True, + n_rows, 1, + figsize=[width, n_rows * height_per_track], + sharex=True, constrained_layout=True, ) - - if n_tracks == 1: + if n_rows == 1: axes = [axes] + elif not isinstance(axes, (list, tuple)): + axes = list(axes) - # Compute global max for y-axis scaling global_max = signal_values.max() + coverage_start = int(peak_track) + coverage_axes = axes[coverage_start:coverage_start + n_tracks] + if peak_track: + peaks = _get_peaks_in_region(adata, chrom, region_start, region_end) + _draw_peak_track(axes[0], peaks, region_start, region_end, n_points) + cmap = plt.get_cmap("tab10") - for i, (ax, signal) in enumerate(zip(axes, signal_values)): - color = cmap(i % 10) # cycle through colors if >10 tracks + for i, (ax, signal) in enumerate(zip(coverage_axes, signal_values)): + color = cmap(i % 10) ax.fill_between(range(n_points), 0, signal, color=color) ax.set_title(track_names[i], fontsize=7) ax.spines[['top', 'right']].set_visible(False) ax.set_ylim(0, global_max) - axes[-1].set_xticks([0, n_points - 1]) - axes[-1].set_xticklabels([str(start), str(end)]) - axes[-1].set_xlabel(region) + if annotation_path is not None: + models = _get_gene_models( + annotation_path, chrom, region_start, region_end, gene_types, + ) + _draw_gene_track(axes[-1], models, region_start, region_end, n_points) + + for ax in axes[:-1]: + ax.set_xticks([]) + ax.set_xlabel("") + bottom_ax = axes[-1] + bottom_ax.set_xticks([0, n_points - 1]) + bottom_ax.set_xticklabels([str(region_start), str(region_end)]) + bottom_ax.set_xlabel(resolved_region) + + if highlight is not None: + highlights = _normalize_highlights(highlight, region_start, region_end) + _add_highlights( + coverage_axes, highlights, region_start, region_end, n_points, + alpha=highlight_alpha, + ) fig.supylabel("RPM")