cstag.to_mutation_percentages

  1from __future__ import annotations
  2
  3import re
  4from collections import defaultdict
  5from collections.abc import Mapping, Sequence
  6from pathlib import Path
  7from typing import Any
  8
  9import matplotlib.pyplot as plt
 10from matplotlib.colors import is_color_like, to_rgb
 11from matplotlib.ticker import MaxNLocator
 12
 13CS_PREFIX = "cs:Z:"
 14CS_TOKEN_PATTERN = re.compile(
 15    r"=[ACGTN]+|:[0-9]+|\*[acgtn][acgtn]|\+[acgtn]+|-[acgtn]+"
 16)
 17PROFILE_KEYS = (
 18    "position",
 19    "coverage",
 20    "total_pct",
 21    "insertion_pct",
 22    "deletion_pct",
 23    "substitution_pct",
 24)
 25REGION_KEYS = ("name", "start", "end", "color")
 26
 27
 28def _percentage(reads: set[int], coverage_count: int) -> float:
 29    return 100.0 * len(reads) / coverage_count
 30
 31
 32def _tokenize_cs(cs_tag: str, tag_index: int) -> list[str]:
 33    """Validate and tokenize one minimap2 short- or long-form cs tag."""
 34    if not isinstance(cs_tag, str):
 35        raise TypeError(f"cs_tags[{tag_index}] must be a string")
 36
 37    if cs_tag.startswith(CS_PREFIX):
 38        normalized = cs_tag[len(CS_PREFIX) :]
 39        if CS_PREFIX in normalized:
 40            raise ValueError(
 41                f"Invalid cs tag at index {tag_index}: repeated cs:Z: prefix"
 42            )
 43    else:
 44        if CS_PREFIX in cs_tag:
 45            raise ValueError(
 46                f"Invalid cs tag at index {tag_index}: cs:Z: must be at the start"
 47            )
 48        normalized = cs_tag
 49
 50    if not normalized:
 51        raise ValueError(f"Invalid cs tag at index {tag_index}: tag is empty")
 52    if "~" in normalized:
 53        raise ValueError(
 54            f"cs_tags[{tag_index}] contains unsupported splice operation '~'"
 55        )
 56
 57    tokens = CS_TOKEN_PATTERN.findall(normalized)
 58    if "".join(tokens) != normalized:
 59        raise ValueError(f"Invalid cs tag at index {tag_index}: {cs_tag}")
 60    if any(token.startswith(":") for token in tokens) and any(
 61        token.startswith("=") for token in tokens
 62    ):
 63        raise ValueError(f"cs_tags[{tag_index}] mixes short and long match operations")
 64    return tokens
 65
 66
 67def _events_for_tag(
 68    cs_tag: str, tag_index: int
 69) -> tuple[set[int], dict[str, set[int]]]:
 70    """Return covered positions and per-category event positions for one tag."""
 71    tokens = _tokenize_cs(cs_tag, tag_index)
 72    reference_position = 1
 73    coverage: set[int] = set()
 74    events: dict[str, set[int]] = {
 75        "insertion": set(),
 76        "deletion": set(),
 77        "substitution": set(),
 78    }
 79
 80    for token in tokens:
 81        operation = token[0]
 82        if operation in {":", "="}:
 83            length = int(token[1:]) if operation == ":" else len(token) - 1
 84            covered_positions = range(reference_position, reference_position + length)
 85            coverage.update(covered_positions)
 86            reference_position += length
 87        elif operation == "*":
 88            coverage.add(reference_position)
 89            events["substitution"].add(reference_position)
 90            reference_position += 1
 91        elif operation == "+":
 92            anchor_position = reference_position - 1
 93            if anchor_position < 1 or anchor_position not in coverage:
 94                raise ValueError(
 95                    f"cs_tags[{tag_index}] starts with an insertion that has "
 96                    "no left reference anchor"
 97                )
 98            events["insertion"].add(anchor_position)
 99        elif operation == "-":
100            length = len(token) - 1
101            deleted_positions = set(
102                range(reference_position, reference_position + length)
103            )
104            coverage.update(deleted_positions)
105            events["deletion"].update(deleted_positions)
106            reference_position += length
107
108    if not coverage:
109        raise ValueError(f"cs_tags[{tag_index}] has a reference span of zero")
110    return coverage, events
111
112
113def summarize_cs(cs_tags: Sequence[str]) -> list[dict[str, int | float]]:
114    """Calculate mutation percentages at each relative reference position.
115
116    Short- and long-form cs tags are accepted, with or without the ``cs:Z:``
117    prefix. All tags are assumed to begin at relative reference position 1.
118    Each percentage uses the number of tags covering that position as its
119    denominator. ``total_pct`` is the union of reads with any mutation at the
120    position, so it cannot exceed 100 percent.
121
122    Args:
123        cs_tags: Non-empty sequence of minimap2 cs tag strings.
124
125    Returns:
126        One dictionary per 1-based reference position. Each dictionary contains
127        ``position``, ``coverage``, ``total_pct``, ``insertion_pct``,
128        ``deletion_pct``, and ``substitution_pct``.
129    """
130    if isinstance(cs_tags, (str, bytes)) or not isinstance(cs_tags, Sequence):
131        raise TypeError("cs_tags must be a non-empty sequence of strings")
132    if not cs_tags:
133        raise ValueError("cs_tags must not be empty")
134
135    coverage_reads: dict[int, set[int]] = defaultdict(set)
136    event_reads: dict[str, dict[int, set[int]]] = {
137        "insertion": defaultdict(set),
138        "deletion": defaultdict(set),
139        "substitution": defaultdict(set),
140    }
141
142    for tag_index, cs_tag in enumerate(cs_tags):
143        coverage, events = _events_for_tag(cs_tag, tag_index)
144        for position in coverage:
145            coverage_reads[position].add(tag_index)
146        for category, positions in events.items():
147            for position in positions:
148                event_reads[category][position].add(tag_index)
149
150    max_position = max(coverage_reads)
151    profile: list[dict[str, int | float]] = []
152    for position in range(1, max_position + 1):
153        covered = coverage_reads[position]
154        coverage_count = len(covered)
155        insertion_reads = event_reads["insertion"][position]
156        deletion_reads = event_reads["deletion"][position]
157        substitution_reads = event_reads["substitution"][position]
158        total_reads = insertion_reads | deletion_reads | substitution_reads
159
160        profile.append(
161            {
162                "position": position,
163                "coverage": coverage_count,
164                "total_pct": _percentage(total_reads, coverage_count),
165                "insertion_pct": _percentage(insertion_reads, coverage_count),
166                "deletion_pct": _percentage(deletion_reads, coverage_count),
167                "substitution_pct": _percentage(substitution_reads, coverage_count),
168            }
169        )
170    return profile
171
172
173def _validate_regions(
174    regions: Sequence[Mapping[str, object]] | None, max_position: int
175) -> list[dict[str, str | int]]:
176    """Validate optional 1-based inclusive region annotations."""
177    if regions is None:
178        return []
179    if isinstance(regions, (str, bytes)) or not isinstance(regions, Sequence):
180        raise TypeError("regions must be a sequence of dictionaries or None")
181
182    validated: list[dict[str, str | int]] = []
183    for region_index, region in enumerate(regions):
184        if not isinstance(region, Mapping):
185            raise TypeError(f"regions[{region_index}] must be a dictionary")
186        missing = set(REGION_KEYS) - set(region)
187        if missing:
188            missing_list = ", ".join(sorted(missing))
189            raise ValueError(f"regions[{region_index}] is missing keys: {missing_list}")
190
191        name = region["name"]
192        start = region["start"]
193        end = region["end"]
194        color = region["color"]
195        if not isinstance(name, str) or not name:
196            raise ValueError(
197                f"regions[{region_index}]['name'] must be a non-empty string"
198            )
199        if isinstance(start, bool) or not isinstance(start, int):
200            raise TypeError(f"regions[{region_index}]['start'] must be an integer")
201        if isinstance(end, bool) or not isinstance(end, int):
202            raise TypeError(f"regions[{region_index}]['end'] must be an integer")
203        if not 1 <= start <= end <= max_position:
204            raise ValueError(
205                f"regions[{region_index}] must satisfy "
206                f"1 <= start <= end <= {max_position}"
207            )
208        if not isinstance(color, str) or not is_color_like(color):
209            raise ValueError(f"regions[{region_index}]['color'] is not a valid color")
210
211        validated.append({"name": name, "start": start, "end": end, "color": color})
212    return validated
213
214
215def plot_mutation_percentages(
216    records: Sequence[dict[str, int | float]],
217    output_path: Path | None = None,
218    regions: Sequence[Mapping[str, object]] | None = None,
219) -> tuple[Any, Any]:
220    """Create four mutation plots with optional region highlights."""
221    if isinstance(records, (str, bytes)) or not isinstance(records, Sequence):
222        raise TypeError("records must be a non-empty sequence of dictionaries")
223    if not records:
224        raise ValueError("records must not be empty")
225    if output_path is not None and not isinstance(output_path, Path):
226        raise TypeError("output_path must be a pathlib.Path or None")
227
228    for record_index, record in enumerate(records):
229        if not isinstance(record, dict):
230            raise TypeError(f"records[{record_index}] must be a dictionary")
231        missing = set(PROFILE_KEYS) - set(record)
232        if missing:
233            missing_list = ", ".join(sorted(missing))
234            raise ValueError(f"records[{record_index}] is missing keys: {missing_list}")
235
236    positions = [record["position"] for record in records]
237    series = (
238        ("total_pct", "Total mutations", "#1f77b4"),
239        ("insertion_pct", "Insertions", "#ff7f0e"),
240        ("deletion_pct", "Deletions", "#d62728"),
241        ("substitution_pct", "Substitutions", "#2ca02c"),
242    )
243    validated_regions = _validate_regions(regions, int(positions[-1]))
244    figure, axes = plt.subplots(
245        4,
246        1,
247        figsize=(12, 10),
248        sharex=True,
249        constrained_layout=True,
250    )
251
252    for axis, (key, title, color) in zip(axes, series, strict=True):
253        values = [record[key] for record in records]
254        axis.bar(
255            positions,
256            values,
257            width=0.8,
258            align="center",
259            color=color,
260            edgecolor="none",
261            alpha=0.85,
262            zorder=3,
263        )
264        for region_index, region in enumerate(validated_regions):
265            start = int(region["start"])
266            end = int(region["end"])
267            region_color = str(region["color"])
268            highlight = axis.axvspan(
269                start - 0.5,
270                end + 0.5,
271                color=region_color,
272                alpha=0.18,
273                linewidth=0,
274                zorder=0.5,
275            )
276            highlight.set_gid(f"cstag-region-{region_index}")
277            red, green, blue = to_rgb(region_color)
278            axis.text(
279                (start + end) / 2,
280                0.97,
281                str(region["name"]),
282                transform=axis.get_xaxis_transform(),
283                ha="center",
284                va="top",
285                color=(red * 0.65, green * 0.65, blue * 0.65),
286                fontsize="small",
287                clip_on=True,
288                zorder=3,
289            )
290        axis.set_title(title)
291        axis.set_ylabel("Mutation (%)")
292        axis.set_ylim(0, 100)
293        axis.grid(axis="both", alpha=0.25)
294
295    axes[-1].set_xlabel("Reference position (1-based)")
296    axes[-1].set_xlim(positions[0] - 0.5, positions[-1] + 0.5)
297    if len(positions) <= 20:
298        axes[-1].set_xticks(positions)
299    else:
300        axes[-1].xaxis.set_major_locator(MaxNLocator(integer=True))
301
302    if output_path is not None:
303        output_path.parent.mkdir(parents=True, exist_ok=True)
304        if output_path.suffix.lower() == ".pdf":
305            # Keep plot elements vector-based and embed editable TrueType text.
306            with plt.rc_context({"pdf.fonttype": 42}):
307                figure.savefig(output_path, bbox_inches="tight")
308        else:
309            figure.savefig(output_path, dpi=150, bbox_inches="tight")
310
311    return figure, axes
312
313
314def to_mutation_percentages(
315    cs_tags: Sequence[str],
316    output_path: Path,
317    regions: Sequence[Mapping[str, object]] | None = None,
318) -> list[dict[str, int | float]]:
319    """Report and plot mutation percentages at each reference position.
320
321    Args:
322        cs_tags: Python sequence containing short- or long-form cs tag strings.
323        output_path: Path of the plot file to create. The format is inferred
324            from its extension. Use ``.pdf`` for an editable vector PDF with
325            embedded TrueType fonts, or ``.png`` for a raster image.
326        regions: Optional sequence of dictionaries describing regions to highlight
327            in the mutation plots. Each dictionary must contain ``name``
328            (the displayed label), ``start`` and ``end`` (1-based, inclusive
329            reference positions), and ``color`` (a Matplotlib-compatible color).
330            For example::
331
332                [
333                    {"name": "crRNA", "start": 80, "end": 99,
334                     "color": "lightblue"},
335                    {"name": "index", "start": 207, "end": 214,
336                     "color": "lightgreen"},
337                ]
338
339            Each region is drawn as a translucent vertical band behind the data
340            in all four mutation panels, with its name shown at the top of the
341            band. Overlapping regions remain visible through blended colors.
342
343    Returns:
344        One dictionary per 1-based relative reference position containing
345        coverage and total, insertion, deletion, and substitution percentages.
346    """
347    if not isinstance(output_path, Path):
348        raise TypeError("output_path must be a pathlib.Path")
349
350    profile = summarize_cs(cs_tags)
351    figure, _ = plot_mutation_percentages(profile, output_path, regions=regions)
352    plt.close(figure)
353    return profile