Blame for src/genome_annotation/decoding.py

"""Constrained Viterbi decoding of 15-class annotation logits into GFF3 genes.

A simplified port of the NTv3 genome-annotation decoder: the same 15-state
Tiberius-style HMM and start / splice / stop motif constraints, without the U12
(AT..AC) intron lane.
"""

import gzip
import itertools
from collections.abc import Callable
from pathlib import Path

import numba
import numpy as np

CLASSES: tuple[str, ...] = (
    "intergenic",
    "intron_p0",
    "intron_p1",
    "intron_p2",
    "CDS_p0",
    "CDS_p1",
    "CDS_p2",
    "splice_acceptor_p0",
    "splice_acceptor_p1",
    "splice_acceptor_p2",
    "splice_donor_p0",
    "splice_donor_p1",
    "splice_donor_p2",
    "start_codon",
    "stop_codon",
)
NUM_CLASSES = len(CLASSES)
CLASS_INDEX: dict[str, int] = {name: i for i, name in enumerate(CLASSES)}

# Tiberius mean lengths: intergenic ~10 kb, intron ~4.5 kb, exon ~200 bp.
_P_IR = 1 / 10_000
_P_I = 1 / 4_500
_P_E = 1 / 200

# Allowed transitions per state; the self-loop takes the remaining mass. Start,
# stop and splice states last exactly 1 bp.
TRANSITIONS: dict[str, dict[str, float]] = {
    "intergenic": {"start_codon": _P_IR},
    "intron_p0": {"splice_acceptor_p1": _P_I},
    "intron_p1": {"splice_acceptor_p2": _P_I},
    "intron_p2": {"splice_acceptor_p0": _P_I},
    "CDS_p0": {"CDS_p1": 1 - _P_E, "splice_donor_p1": _P_E},
    "CDS_p1": {"CDS_p2": 1 - 2 * _P_E, "splice_donor_p2": _P_E, "stop_codon": _P_E},
    "CDS_p2": {"CDS_p0": 1 - _P_E, "splice_donor_p0": _P_E},
    "splice_acceptor_p0": {"CDS_p1": 1.0},
    "splice_acceptor_p1": {"CDS_p2": 1.0},
    "splice_acceptor_p2": {"CDS_p0": 1.0},
    "splice_donor_p0": {"intron_p0": 1.0},
    "splice_donor_p1": {"intron_p1": 1.0},
    "splice_donor_p2": {"intron_p2": 1.0},
    "start_codon": {"CDS_p1": 1.0},
    "stop_codon": {"intergenic": 1.0},
}


def build_log_transitions() -> np.ndarray:
    """Returns the (15, 15) log transition matrix of the HMM."""
    matrix = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.float64)
    for src, dests in TRANSITIONS.items():
        i = CLASS_INDEX[src]
        matrix[i, i] = max(0.0, 1.0 - sum(dests.values()))
        for dst, prob in dests.items():
            matrix[i, CLASS_INDEX[dst]] = prob
    with np.errstate(divide="ignore"):
        return np.log(matrix.astype(np.float32))


LOG_TRANSITIONS = build_log_transitions()

# Base codes: A=0, C=1, G=2, T=3, anything else (N, IUPAC) = 4.
A, C, G, T, N = 0, 1, 2, 3, 4
_BASE_CODE_LUT = np.full(256, N, dtype=np.uint8)
for _base, _code in zip("ACGT", (A, C, G, T), strict=True):
    _BASE_CODE_LUT[ord(_base)] = _code
    _BASE_CODE_LUT[ord(_base.lower())] = _code


def encode_bases(sequence: str) -> np.ndarray:
    """Maps a DNA string to case-insensitive base codes."""
    return _BASE_CODE_LUT[np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)]


def reverse_complement_codes(codes: np.ndarray) -> np.ndarray:
    """Reverse-complements an array of base codes."""
    return np.where(codes == N, N, T - codes)[::-1].astype(np.uint8)


# (offsets relative to the state position, bases, log-score) per accepted motif.
# Scores follow Tiberius: a slight TAG preference, and GT >> GC donors.
_MOTIFS: dict[str, list[tuple[tuple[int, ...], tuple[int, ...], float]]] = {
    "start": [((0, 1, 2), (A, T, G), 0.0)],
    "stop": [
        ((-2, -1, 0), (T, A, G), float(np.log(np.float32(0.34)))),
        ((-2, -1, 0), (T, A, A), float(np.log(np.float32(0.33)))),
        ((-2, -1, 0), (T, G, A), float(np.log(np.float32(0.33)))),
    ],
    "donor": [
        ((1, 2), (G, T), float(np.log(np.float32(0.99)))),
        ((1, 2), (G, C), float(np.log(np.float32(0.01)))),
    ],
    "acceptor": [((-2, -1), (A, G), 0.0)],
}
_MOTIF_STATES: dict[str, tuple[str, ...]] = {
    "start": ("start_codon",),
    "stop": ("stop_codon",),
    "donor": ("splice_donor_p0", "splice_donor_p1", "splice_donor_p2"),
    "acceptor": ("splice_acceptor_p0", "splice_acceptor_p1", "splice_acceptor_p2"),
}


def _motif_log_mask(
    codes: np.ndarray, rules: list[tuple[tuple[int, ...], tuple[int, ...], float]]
) -> np.ndarray:
    """Per-position log-score of one motif: -inf if no rule matches, 0 near an N."""
    length = len(codes)
    padded = np.pad(codes, (2, 2), constant_values=N)
    offsets = {off for rule_offsets, _, _ in rules for off in rule_offsets}

    def shifted(off: int) -> np.ndarray:
        return padded[2 + off : 2 + off + length]

    has_n = np.zeros(length, dtype=bool)
    for off in offsets:
        has_n |= shifted(off) == N
    score = np.full(length, -np.inf, dtype=np.float32)
    for rule_offsets, bases, log_score in rules:
        match = np.ones(length, dtype=bool)
        for off, base in zip(rule_offsets, bases, strict=True):
            match &= shifted(off) == base
        score[match] = log_score
    score[has_n] = 0.0
    # Motifs whose context falls outside the window are forbidden.
    pos = np.arange(length)
    score[(pos < -min(offsets)) | (pos >= length - max(offsets))] = -np.inf
    return score


def constraint_log_mask(codes: np.ndarray) -> np.ndarray:
    """(L, 15) additive log-mask vetoing start/stop/splice states off their motif."""
    mask = np.zeros((len(codes), NUM_CLASSES), dtype=np.float32)
    for motif, rules in _MOTIFS.items():
        column = _motif_log_mask(codes, rules)
        for state in _MOTIF_STATES[motif]:
            mask[:, CLASS_INDEX[state]] = column
    return mask


@numba.njit(cache=True)
def _viterbi(log_emissions: np.ndarray, log_transitions: np.ndarray) -> np.ndarray:
    length, num_states = log_emissions.shape
    backpointers = np.zeros((length, num_states), dtype=np.int8)
    scores = log_emissions[0].copy()
    stepped = np.empty(num_states, dtype=np.float32)
    for t in range(1, length):
        for j in range(num_states):
            best = -np.inf
            best_i = 0
            for i in range(num_states):
                value = scores[i] + log_transitions[i, j]
                if value > best:
                    best = value
                    best_i = i
            stepped[j] = best + log_emissions[t, j]
            backpointers[t, j] = best_i
        # Rescale by the running max so scores stay bounded; argmaxes are unchanged.
        offset = max(stepped.max(), np.float32(-1e30))
        for j in range(num_states):
            scores[j] = stepped[j] - offset
    path = np.empty(length, dtype=np.int8)
    path[-1] = np.argmax(scores)
    for t in range(length - 1, 0, -1):
        path[t - 1] = backpointers[t, path[t]]
    return path


def viterbi(
    log_emissions: np.ndarray, constrain_codes: np.ndarray | None
) -> np.ndarray:
    """Most likely state path of one strand, read 5'->3'.

    Args:
        log_emissions: (L, 15) log-probabilities in transcription direction.
        constrain_codes: Base codes in the same direction, or None to decode
            without motif constraints.

    Returns:
        (L,) int8 class indices.
    """
    log_emissions = np.asarray(log_emissions, dtype=np.float32)
    if constrain_codes is not None:
        log_emissions = log_emissions + constraint_log_mask(constrain_codes)
    return _viterbi(np.ascontiguousarray(log_emissions), LOG_TRANSITIONS)


def tile_windows(length: int, window: int) -> list[tuple[int, int]]:
    """(window start, first bp the window owns) covering [0, length) end to end.

    The last window is pulled back to end on the final base; the bp it shares
    with the previous window stay with the previous one.
    """
    if length < window:
        raise ValueError(f"Sequence of {length} bp is shorter than one window.")
    starts = list(range(0, length - window + 1, window))
    tiles = [(s, 0) for s in starts]
    if starts[-1] + window < length:
        last = length - window
        tiles.append((last, starts[-1] + window - last))
    return tiles


def annotate_sequence(
    sequence: str,
    predict: Callable[[str], tuple[np.ndarray, np.ndarray]],
    window: int,
    constrained: bool = True,
) -> tuple[np.ndarray, np.ndarray]:
    """Decodes both strands of a chromosome window by window.

    Args:
        sequence: Chromosome sequence (soft-masking allowed).
        predict: Maps a window's sequence to (L, 15) log-probabilities for the
            window itself and for its reverse complement.
        window: Window length in bp.
        constrained: Whether to apply the motif constraints.

    Returns:
        Plus- and minus-strand state paths in genomic order, each (len(sequence),).
    """
    codes = encode_bases(sequence)
    plus = np.zeros(len(sequence), dtype=np.int8)
    minus = np.zeros(len(sequence), dtype=np.int8)
    for start, keep in tile_windows(len(sequence), window):
        end = start + window
        log_probs_plus, log_probs_rc = predict(sequence[start:end])
        window_codes = codes[start:end]
        rc_codes = reverse_complement_codes(window_codes)
        states_plus = viterbi(log_probs_plus, window_codes if constrained else None)
        states_rc = viterbi(log_probs_rc, rc_codes if constrained else None)
        plus[start + keep : end] = states_plus[keep:]
        minus[start + keep : end] = states_rc[::-1][keep:]
    return plus, minus


_INTRON_STATES = {CLASS_INDEX[f"intron_p{p}"] for p in range(3)}
# Codon position at the 5' end of a CDS run, which sets its GFF3 phase.
_CODON_POSITION: dict[int, int] = {
    CLASS_INDEX[f"{kind}_p{p}"]: p
    for kind in ("CDS", "splice_acceptor", "splice_donor")
    for p in range(3)
} | {CLASS_INDEX["start_codon"]: 0, CLASS_INDEX["stop_codon"]: 0}
_STOP_CODONS = {(T, A, A), (T, A, G), (T, G, A)}


def _has_internal_stop(
    gene_states: np.ndarray, gene_codes: np.ndarray, strand: str
) -> bool:
    """Whether the spliced CDS holds an in-frame stop before its last codon."""
    is_intron = np.isin(gene_states, list(_INTRON_STATES))
    cds = gene_codes[~is_intron]
    if strand == "-":
        cds = reverse_complement_codes(cds)
    codons = cds[: len(cds) // 3 * 3].reshape(-1, 3)[:-1]
    return any(tuple(int(b) for b in codon) in _STOP_CODONS for codon in codons)


def find_genes(
    states: np.ndarray,
    strand: str,
    chrom: str,
    codes: np.ndarray | None = None,
    first_gene_id: int = 0,
) -> list[str]:
    """Turns one strand's state path into GFF3 gene/CDS/intron lines.

    A gene is a maximal run of non-intergenic states. Genes with an internal
    in-frame stop codon are dropped when ``codes`` is given.

    Args:
        states: (L,) class indices of one strand, in genomic order.
        strand: "+" or "-".
        chrom: Sequence name for the GFF3 seqid column.
        codes: Optional base codes of the chromosome, in genomic order.
        first_gene_id: Genes are numbered from first_gene_id + 1.

    Returns:
        GFF3 lines, each gene followed by its sub-features.
    """
    is_genic = np.concatenate([[False], states != CLASS_INDEX["intergenic"], [False]])
    edges = np.flatnonzero(is_genic[1:] != is_genic[:-1])
    feature_type = np.where(
        np.isin(states, list(_INTRON_STATES)), 2, np.where(states == 0, 0, 1)
    )
    type_changes = np.flatnonzero(feature_type[1:] != feature_type[:-1]) + 1
    lines: list[str] = []
    gene_id = first_gene_id
    for gene_start, gene_end in zip(edges[::2], edges[1::2], strict=True):
        gene_states = states[gene_start:gene_end]
        if codes is not None and _has_internal_stop(
            gene_states, codes[gene_start:gene_end], strand
        ):
            continue
        gene_id += 1
        gid = f"gene_{gene_id}_{'p' if strand == '+' else 'm'}"
        lines.append(
            f"{chrom}\tNTv3\tgene\t{gene_start + 1}\t{gene_end}\t.\t{strand}\t.\tID={gid}"
        )
        lo = np.searchsorted(type_changes, gene_start, "right")
        hi = np.searchsorted(type_changes, gene_end, "left")
        bounds = [gene_start, *type_changes[lo:hi].tolist(), gene_end]
        for feat_start, feat_end in itertools.pairwise(bounds):
            is_cds = feature_type[feat_start] == 1
            phase = "."
            if is_cds:
                five_prime = feat_start if strand == "+" else feat_end - 1
                phase = str((3 - _CODON_POSITION[int(states[five_prime])]) % 3)
            kind = "CDS" if is_cds else "intron"
            lines.append(
                f"{chrom}\tNTv3\t{kind}\t{feat_start + 1}\t{feat_end}\t.\t{strand}"
                f"\t{phase}\tID={gid}_{kind}_{feat_start + 1};Parent={gid}"
            )
    return lines


def write_gff3(path: str | Path, lines: list[str]) -> None:
    """Writes GFF3 lines sorted by (seqid, start), gzipped if path ends in .gz."""
    ordered = sorted(
        lines, key=lambda line: (line.split("\t")[0], int(line.split("\t")[3]))
    )
    text = "##gff-version 3\n" + "\n".join(ordered) + "\n"
    opener = gzip.open if str(path).endswith(".gz") else open
    with opener(path, "wt") as handle:
        handle.write(text)