Blame for src/genome_annotation/analysis.py

"""Error analyses of a predicted annotation, as in the NTv3 annotation paper.

Everything works on ``evaluation.ScoringSets``, i.e. after the benchmark filters.
A locus is one reference gene; a prediction overlaps it when the two share coding
sequence on the same strand (a prediction inside an intron does not).
"""

import bisect
from collections import defaultdict
from collections.abc import Iterable
from dataclasses import dataclass

from .evaluation import CdsKey, Gene

LOCUS_CLASSES = ("correct", "boundary", "missed", "fragmented")
EXON_CLASSES = ("exact", "shifted_one", "shifted_both", "missed")
POSITIONS = ("first", "internal", "last", "single")

# Inclusive bins unless noted.
EXON_COUNT_BINS = ((1, 1), (2, 3), (4, 6), (7, 10), (11, 20), (21, 10**9))
SPAN_BINS = (  # half-open, bp of CDS envelope
    (0, 5_000),
    (5_000, 20_000),
    (20_000, 50_000),
    (50_000, 131_072),
    (131_072, 500_000),
    (500_000, 10**12),
)
EXON_LENGTH_BINS = (
    (1, 50),
    (51, 100),
    (101, 200),
    (201, 500),
    (501, 1_000),
    (1_001, 10**9),
)
WRONG_EXON_BUCKETS = ((0, 0), (1, 1), (2, 2), (3, 3), (4, 5), (6, 10**9))
OFFSET_BINS = (
    (1, 3),
    (4, 6),
    (7, 12),
    (13, 25),
    (26, 100),
    (101, 1_000),
    (1_001, 10**9),
)
MERGE_SIZE_BINS = ((2, 2), (3, 3), (4, 5), (6, 10**9))


def bin_label(lo: int, hi: int) -> str:
    """``"4-6"``, ``"1"``, or ``"21+"`` for the open-ended bin."""
    if lo == hi:
        return str(lo)
    return f"{lo}+" if hi >= 10**9 else f"{lo}-{hi}"


def _bin_index(value: int, bins: tuple[tuple[int, int], ...]) -> int:
    return next(
        (i for i, (lo, hi) in enumerate(bins) if lo <= value <= hi), len(bins) - 1
    )


def _half_open_bin_index(value: int, bins: tuple[tuple[int, int], ...]) -> int:
    return next(
        (i for i, (lo, hi) in enumerate(bins) if lo <= value < hi), len(bins) - 1
    )


def cds_span(gene: Gene) -> tuple[int, int]:
    """CDS envelope; the gene row would include UTRs the predictions lack."""
    return min(k[1] for k in gene.cds_keys), max(k[2] for k in gene.cds_keys)


def _boundary(key: CdsKey) -> tuple[str, int, int, str]:
    return key[:4]


def _ordered_cds(gene: Gene) -> list[CdsKey]:
    """CDS in transcription order, 5' to 3'."""
    return sorted(gene.cds_keys, key=lambda k: k[1], reverse=gene.strand == "-")


def _position(index: int, total: int) -> str:
    if total == 1:
        return "single"
    if index == 0:
        return "first"
    return "last" if index == total - 1 else "internal"


class CdsIndex:
    """Same-strand overlap lookup over individual CDS, returning gene IDs."""

    def __init__(self, genes: dict[str, Gene]) -> None:
        buckets: dict[tuple[str, str], list[tuple[int, int, str]]] = defaultdict(list)
        for gid, gene in genes.items():
            for _, start, end, strand, _ in gene.cds_keys:
                buckets[(gene.chrom, strand)].append((start, end, gid))
        self._items: dict[tuple[str, str], list[tuple[int, int, str]]] = {}
        self._starts: dict[tuple[str, str], list[int]] = {}
        self._max_end: dict[tuple[str, str], list[int]] = {}
        for key, items in buckets.items():
            items.sort()
            self._items[key] = items
            self._starts[key] = [item[0] for item in items]
            running, best = [], 0
            for _, end, _ in items:
                best = max(best, end)
                running.append(best)
            self._max_end[key] = running

    def overlapping(self, chrom: str, strand: str, start: int, end: int) -> set[str]:
        """IDs of genes with a CDS overlapping [start, end] on this strand."""
        items = self._items.get((chrom, strand))
        if items is None:
            return set()
        max_end = self._max_end[(chrom, strand)]
        hits: set[str] = set()
        for i in range(
            bisect.bisect_right(self._starts[(chrom, strand)], end) - 1, -1, -1
        ):
            if max_end[i] < start:
                break
            if items[i][1] >= start:
                hits.add(items[i][2])
        return hits

    def overlapping_gene(self, gene: Gene) -> set[str]:
        """IDs of indexed genes sharing coding sequence with ``gene``."""
        hits: set[str] = set()
        for _, start, end, strand, _ in gene.cds_keys:
            hits |= self.overlapping(gene.chrom, strand, start, end)
        return hits


@dataclass(frozen=True)
class LocusOutcome:
    """One reference gene, classified against the predictions."""

    gene_id: str
    chrom: str
    strand: str
    start: int
    end: int
    n_cds: int
    cds_length: int
    outcome: str
    n_overlapping: int


def classify_loci(
    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
) -> tuple[list[LocusOutcome], list[str]]:
    """Classifies every reference locus; also returns the invented predictions.

    ``correct``: its exact CDS set is predicted. Otherwise by the number of
    overlapping predictions: 0 ``missed``, 1 ``boundary``, 2+ ``fragmented``.
    Invented: a prediction overlapping no reference gene.
    """
    pred_structures = {frozenset(g.cds_keys) for g in pred_genes.values()}
    pred_index, ref_index = CdsIndex(pred_genes), CdsIndex(ref_genes)
    outcomes = []
    for gid, gene in ref_genes.items():
        if not gene.cds_keys:
            continue
        hits = pred_index.overlapping_gene(gene)
        if frozenset(gene.cds_keys) in pred_structures:
            outcome = "correct"
        else:
            outcome = (
                ("missed", "boundary")[len(hits)] if len(hits) < 2 else "fragmented"
            )
        outcomes.append(
            LocusOutcome(
                gid,
                gene.chrom,
                gene.strand,
                *cds_span(gene),
                len(gene.cds_keys),
                sum(k[2] - k[1] + 1 for k in gene.cds_keys),
                outcome,
                len(hits),
            )
        )
    invented = [
        gid
        for gid, gene in pred_genes.items()
        if gene.cds_keys and not ref_index.overlapping_gene(gene)
    ]
    return outcomes, invented


def summarise_loci(
    outcomes: list[LocusOutcome], invented: list[str], n_pred: int
) -> dict[str, float]:
    """Counts per class; precision here is correct loci over predictions."""
    counts = dict.fromkeys(LOCUS_CLASSES, 0)
    for outcome in outcomes:
        counts[outcome.outcome] += 1
    n_loci = len(outcomes)
    recall = 100.0 * counts["correct"] / n_loci if n_loci else 0.0
    precision = 100.0 * counts["correct"] / n_pred if n_pred else 0.0
    total = recall + precision
    return {
        **counts,
        "n_loci": n_loci,
        "n_pred": n_pred,
        "invented": len(invented),
        "recall": recall,
        "precision": precision,
        "gene_f1": 2 * recall * precision / total if total else 0.0,
    }


def recovery_by_covariate(outcomes: list[LocusOutcome]) -> dict[str, object]:
    """Loci and outcomes per exon-count and CDS-span bin."""
    exon_total = [0] * len(EXON_COUNT_BINS)
    exon_correct = [0] * len(EXON_COUNT_BINS)
    span_total = [0] * len(SPAN_BINS)
    span_correct = [0] * len(SPAN_BINS)
    span_outcome = {cls: [0] * len(SPAN_BINS) for cls in LOCUS_CLASSES}
    for outcome in outcomes:
        i = _bin_index(outcome.n_cds, EXON_COUNT_BINS)
        j = _half_open_bin_index(outcome.end - outcome.start + 1, SPAN_BINS)
        exon_total[i] += 1
        span_total[j] += 1
        span_outcome[outcome.outcome][j] += 1
        if outcome.outcome == "correct":
            exon_correct[i] += 1
            span_correct[j] += 1
    return {
        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
        "exon_total": exon_total,
        "exon_correct": exon_correct,
        "span_bins": [f"{lo}-{hi}" for lo, hi in SPAN_BINS],
        "span_total": span_total,
        "span_correct": span_correct,
        "span_outcome": span_outcome,
    }


def _overlapping_cds(
    pred_genes: dict[str, Gene],
    pred_index: CdsIndex,
    chrom: str,
    strand: str,
    start: int,
    end: int,
) -> list[CdsKey]:
    return [
        k
        for gid in sorted(pred_index.overlapping(chrom, strand, start, end))
        for k in pred_genes[gid].cds_keys
        if k[0] == chrom and k[3] == strand and k[1] <= end and k[2] >= start
    ]


def _best_overlap(candidates: list[CdsKey], start: int, end: int) -> CdsKey:
    """Largest overlap; ties broken by coordinate so the result is deterministic."""
    return max(
        candidates, key=lambda k: (min(end, k[2]) - max(start, k[1]), -k[1], -k[2])
    )


def exon_taxonomy(
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
    locus_outcome: dict[str, str],
) -> dict[str, object]:
    """Classifies every reference CDS (phase ignored), by position and length.

    ``exact``: some predicted CDS has both boundaries. Otherwise the most-overlapping
    same-strand predicted CDS has one (``shifted_one``) or both (``shifted_both``)
    boundaries wrong, or nothing overlaps (``missed``).
    """
    pred_index = CdsIndex(pred_genes)
    pred_boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
    by_position = {pos: dict.fromkeys(EXON_CLASSES, 0) for pos in POSITIONS}
    by_length = [dict.fromkeys(EXON_CLASSES, 0) for _ in EXON_LENGTH_BINS]
    ends = dict.fromkeys(("five_prime", "three_prime", "both"), 0)
    wrong_per_failed_gene = [0] * len(WRONG_EXON_BUCKETS)
    n_ref_cds = 0

    for gid, gene in ref_genes.items():
        ordered = _ordered_cds(gene)
        wrong = 0
        for index, key in enumerate(ordered):
            chrom, start, end, strand, _ = key
            n_ref_cds += 1
            position = _position(index, len(ordered))
            length_bin = by_length[_bin_index(end - start + 1, EXON_LENGTH_BINS)]
            if _boundary(key) in pred_boundaries:
                klass = "exact"
            else:
                wrong += 1
                candidates = _overlapping_cds(
                    pred_genes, pred_index, chrom, strand, start, end
                )
                if not candidates:
                    klass = "missed"
                else:
                    best = _best_overlap(candidates, start, end)
                    start_wrong, end_wrong = best[1] != start, best[2] != end
                    klass = (
                        "shifted_both" if start_wrong and end_wrong else "shifted_one"
                    )
                    # Transcript orientation: a minus-strand start is the 3' end.
                    if start_wrong and end_wrong:
                        ends["both"] += 1
                    elif start_wrong:
                        ends["five_prime" if strand == "+" else "three_prime"] += 1
                    else:
                        ends["three_prime" if strand == "+" else "five_prime"] += 1
            by_position[position][klass] += 1
            length_bin[klass] += 1
        if locus_outcome.get(gid) != "correct":
            wrong_per_failed_gene[_bin_index(wrong, WRONG_EXON_BUCKETS)] += 1

    ref_boundaries = {_boundary(k) for g in ref_genes.values() for k in g.cds_keys}
    ref_index = CdsIndex(ref_genes)
    pred_keys = [k for g in pred_genes.values() for k in g.cds_keys]
    invented = sum(
        1
        for k in pred_keys
        if _boundary(k) not in ref_boundaries
        and not ref_index.overlapping(k[0], k[3], k[1], k[2])
    )
    totals = {
        klass: sum(c[klass] for c in by_position.values()) for klass in EXON_CLASSES
    }
    return {
        "n_ref_cds": n_ref_cds,
        "n_pred_cds": len(pred_keys),
        "totals": totals,
        "by_position": by_position,
        "exon_length_bins": [bin_label(lo, hi) for lo, hi in EXON_LENGTH_BINS],
        "by_length": by_length,
        "ends": ends,
        "invented": invented,
        "wrong_exon_buckets": [bin_label(lo, hi) for lo, hi in WRONG_EXON_BUCKETS],
        "wrong_exons_per_failed_gene": wrong_per_failed_gene,
    }


def boundary_offsets(
    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
) -> dict[str, object]:
    """How far off shifted boundaries are, over reference CDS a prediction overlaps."""
    pred_index = CdsIndex(pred_genes)
    aligned = dict.fromkeys(POSITIONS, 0)
    shifted = dict.fromkeys(POSITIONS, 0)
    offsets = [0] * len(OFFSET_BINS)
    length_delta_mod3 = {0: 0, 1: 0, 2: 0}
    for gene in ref_genes.values():
        ordered = _ordered_cds(gene)
        for index, (chrom, start, end, strand, _) in enumerate(ordered):
            candidates = _overlapping_cds(
                pred_genes, pred_index, chrom, strand, start, end
            )
            if not candidates:
                continue
            best = _best_overlap(candidates, start, end)
            position = _position(index, len(ordered))
            aligned[position] += 1
            if best[1] == start and best[2] == end:
                continue
            shifted[position] += 1
            for offset in (abs(best[1] - start), abs(best[2] - end)):
                if offset:
                    offsets[_bin_index(offset, OFFSET_BINS)] += 1
            length_delta_mod3[((best[2] - best[1]) - (end - start)) % 3] += 1
    return {
        "aligned": aligned,
        "shifted": shifted,
        "offset_bins": [bin_label(lo, hi) for lo, hi in OFFSET_BINS],
        "offsets": offsets,
        "length_delta_mod3": length_delta_mod3,
    }


def failure_anatomy(
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
    locus_outcome: dict[str, str],
) -> dict[str, object]:
    """Wrong exons per failed gene, by gene size, and where a lone wrong exon sits."""
    boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
    grid = [[0] * len(WRONG_EXON_BUCKETS) for _ in EXON_COUNT_BINS]
    observed = dict.fromkeys(POSITIONS, 0)
    expected = dict.fromkeys(POSITIONS, 0.0)
    convertible = [0] * len(WRONG_EXON_BUCKETS)
    for gid, gene in ref_genes.items():
        if locus_outcome.get(gid) == "correct":
            continue
        ordered = _ordered_cds(gene)
        n = len(ordered)
        wrong = [
            _position(i, n)
            for i, k in enumerate(ordered)
            if _boundary(k) not in boundaries
        ]
        bucket = _bin_index(len(wrong), WRONG_EXON_BUCKETS)
        grid[_bin_index(n, EXON_COUNT_BINS)][bucket] += 1
        for i in range(bucket, len(convertible)):
            convertible[i] += 1
        if len(wrong) == 1:
            observed[wrong[0]] += 1
            if n == 1:
                expected["single"] += 1.0
            else:  # Uniform over the gene's exons.
                expected["first"] += 1.0 / n
                expected["last"] += 1.0 / n
                expected["internal"] += (n - 2) / n
    return {
        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
        "wrong_exon_buckets": [bin_label(lo, hi) for lo, hi in WRONG_EXON_BUCKETS],
        "wrong_by_gene_size": grid,
        "cumulative_convertible": convertible,
        "single_wrong_observed": observed,
        "single_wrong_expected": {k: round(v, 2) for k, v in expected.items()},
    }


def independence_check(
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
    locus_outcome: dict[str, str],
) -> dict[str, object]:
    """Observed gene recovery vs. independent exon errors, per exon-count bin.

    Predicted recovery of an n-exon gene is P(first) * P(internal)^(n-2) * P(last),
    with per-position exact rates measured on the same predictions.
    """
    boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
    hit: dict[str, int] = defaultdict(int)
    total: dict[str, int] = defaultdict(int)
    for gene in ref_genes.values():
        ordered = _ordered_cds(gene)
        for index, key in enumerate(ordered):
            position = _position(index, len(ordered))
            total[position] += 1
            hit[position] += _boundary(key) in boundaries
    rate = {p: hit[p] / total[p] for p in total}

    def predicted(n: int) -> float:
        if n == 1:
            return rate.get("single", 0.0)
        return (
            rate.get("first", 0.0)
            * rate.get("last", 0.0)
            * rate.get("internal", 0.0) ** (n - 2)
        )

    observed = [0] * len(EXON_COUNT_BINS)
    n_genes = [0] * len(EXON_COUNT_BINS)
    expected = [0.0] * len(EXON_COUNT_BINS)
    for gid, gene in ref_genes.items():
        n = len(gene.cds_keys)
        i = _bin_index(n, EXON_COUNT_BINS)
        n_genes[i] += 1
        expected[i] += predicted(n)
        observed[i] += locus_outcome.get(gid) == "correct"
    return {
        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
        "n_genes": n_genes,
        "observed": observed,
        "predicted": [round(v, 3) for v in expected],
        "position_rate": {p: round(v, 6) for p, v in sorted(rate.items())},
    }


def intron_chain(gene: Gene) -> tuple[object, ...]:
    """Internal CDS boundaries (gffcompare's intron chain); phase ignored."""
    ordered = sorted(gene.cds_keys, key=lambda k: (k[1], k[2]))
    introns = tuple((ordered[i][2], ordered[i + 1][1]) for i in range(len(ordered) - 1))
    return (gene.chrom, gene.strand, introns)


def _rates(
    tp_ref: int, n_ref: int, tp_pred: int, n_pred: int, prefix: str
) -> dict[str, float]:
    recall = 100.0 * tp_ref / n_ref if n_ref else 0.0
    precision = 100.0 * tp_pred / n_pred if n_pred else 0.0
    total = recall + precision
    return {
        f"{prefix}_recall": recall,
        f"{prefix}_precision": precision,
        f"{prefix}_f1": 2 * recall * precision / total if total else 0.0,
    }


def _chain_counts(
    refs: list[Gene], preds: list[Gene], chains: bool
) -> dict[str, int]:
    ref_exact = {frozenset(g.cds_keys) for g in refs}
    pred_exact = {frozenset(g.cds_keys) for g in preds}
    counts = {
        "n_ref": len(refs),
        "n_pred": len(preds),
        "exact_tp_ref": sum(frozenset(g.cds_keys) in pred_exact for g in refs),
        "exact_tp_pred": sum(frozenset(g.cds_keys) in ref_exact for g in preds),
    }
    if chains:
        ref_chains = {intron_chain(g) for g in refs}
        pred_chains = {intron_chain(g) for g in preds}
        counts["chain_tp_ref"] = sum(intron_chain(g) in pred_chains for g in refs)
        counts["chain_tp_pred"] = sum(intron_chain(g) in ref_chains for g in preds)
    else:  # Single-exon genes have no chain; keep the exact tally so rollups add up.
        counts["chain_tp_ref"] = counts["exact_tp_ref"]
        counts["chain_tp_pred"] = counts["exact_tp_pred"]
    return counts


def _with_rates(counts: dict[str, int]) -> dict[str, float]:
    return {
        **counts,
        **_rates(
            counts["exact_tp_ref"],
            counts["n_ref"],
            counts["exact_tp_pred"],
            counts["n_pred"],
            "exact",
        ),
        **_rates(
            counts["chain_tp_ref"],
            counts["n_ref"],
            counts["chain_tp_pred"],
            counts["n_pred"],
            "chain",
        ),
    }


def measure_chains(
    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
) -> dict[str, object]:
    """Exact-gene and intron-chain matching, for multi- and single-exon genes."""

    def split(genes: Iterable[Gene]) -> tuple[list[Gene], list[Gene]]:
        genes = [g for g in genes if g.cds_keys]
        return [g for g in genes if len(g.cds_keys) > 1], [
            g for g in genes if len(g.cds_keys) == 1
        ]

    ref_multi, ref_single = split(ref_genes.values())
    pred_multi, pred_single = split(pred_genes.values())
    multi = _chain_counts(ref_multi, pred_multi, chains=True)
    single = _chain_counts(ref_single, pred_single, chains=False)
    combined = {key: multi[key] + single[key] for key in multi}
    return {
        "multi_exon": _with_rates(multi),
        "single_exon": _with_rates(single),
        "all": _with_rates(combined),
    }


def measure_merges(
    ref_genes: dict[str, Gene],
    pred_genes: dict[str, Gene],
    outcomes: list[LocusOutcome],
) -> dict[str, object]:
    """Predictions overlapping 2+ reference loci, and what they absorb."""
    ref_index = CdsIndex(ref_genes)
    outcome_of = {o.gene_id: o.outcome for o in outcomes}
    sizes = dict.fromkeys((bin_label(lo, hi) for lo, hi in MERGE_SIZE_BINS), 0)
    merging, largest = 0, 0
    absorbed: set[str] = set()
    for gene in pred_genes.values():
        if not gene.cds_keys:
            continue
        hits = ref_index.overlapping_gene(gene)
        if len(hits) < 2:
            continue
        merging += 1
        lo, hi = MERGE_SIZE_BINS[_bin_index(len(hits), MERGE_SIZE_BINS)]
        sizes[bin_label(lo, hi)] += 1
        largest = max(largest, len(hits))
        absorbed |= hits
    absorbed_outcomes: dict[str, int] = defaultdict(int)
    for gid in absorbed:
        absorbed_outcomes[outcome_of.get(gid, "unscored")] += 1
    return {
        "n_pred": sum(1 for g in pred_genes.values() if g.cds_keys),
        "n_ref": len(outcomes),
        "merging_predictions": merging,
        "absorbed_loci": len(absorbed),
        "merge_sizes": sizes,
        "largest_merge": largest,
        "boundary_loci": sum(o.outcome == "boundary" for o in outcomes),
        "boundary_from_merge": sum(
            o.outcome == "boundary" and o.gene_id in absorbed for o in outcomes
        ),
        "absorbed_outcomes": dict(sorted(absorbed_outcomes.items())),
    }