Blame for src/genome_annotation/inference.py

"""Annotates a genomic region with a released NTv3 model, tiled as the pipeline does.

Each window of ``sequence_length`` bp predicts its central ``keep`` fraction; the
kept regions form a grid anchored at the contig's first base, the first window
reaching ``crop`` bp ahead of the contig (read as N) and the last one shifted back
onto its final base. Each kept region is decoded on its own and the state paths
are stitched, the overlap of the shifted last region going to the earlier window.
"""

from collections import defaultdict
from dataclasses import dataclass

import numpy as np
import torch
from transformers import AutoModel

from . import decoding
from .data import Tokenizer, reverse_complement
from .model import CODE_REVISION, predict_log_probs

DECODERS = ("argmax", "hmm", "hmm_constrained")


def load_released(clade: str, device: torch.device) -> torch.nn.Module:
    """``InstaDeepAI/NTv3_annot_<clade>`` (mammals, plants or insects), in eval mode."""
    model = AutoModel.from_pretrained(
        f"InstaDeepAI/NTv3_annot_{clade}",
        trust_remote_code=True,
        code_revision=CODE_REVISION,
    )
    return model.to(device).eval()


@dataclass(frozen=True)
class Window:
    """One model input and the contig bp its decode is kept for (0-based, half-open)."""

    start: int
    keep_start: int
    keep_stop: int


def crop_of(sequence_length: int, keep: float) -> int:
    """Context bp the head does not predict, on each side of a window."""
    return (sequence_length - int(sequence_length * keep)) // 2


def plan_windows(contig_length: int, sequence_length: int, keep: float) -> list[Window]:
    """The pipeline's windows over a whole contig, kept regions disjoint and covering it."""
    crop = crop_of(sequence_length, keep)
    stride = sequence_length - 2 * crop
    last_start = max(contig_length - stride, 0)
    windows: list[Window] = []
    covered = 0
    for grid in range(0, contig_length, stride):
        keep_start = min(grid, last_start)
        keep_stop = min(keep_start + stride, contig_length)
        windows.append(Window(keep_start - crop, max(keep_start, covered), keep_stop))
        covered = keep_stop
    return windows


def windows_for_region(
    contig_length: int, start: int, end: int, sequence_length: int, keep: float
) -> list[Window]:
    """The pipeline windows whose kept region overlaps [start, end)."""
    return [
        w
        for w in plan_windows(contig_length, sequence_length, keep)
        if w.keep_start < end and w.keep_stop > start
    ]


def window_sequence(contig: str, window: Window, sequence_length: int) -> str:
    """The window's bases, N-padded where it reaches past either contig end."""
    lo, hi = max(window.start, 0), min(window.start + sequence_length, len(contig))
    return (
        "N" * (lo - window.start)
        + contig[lo:hi]
        + "N" * (window.start + sequence_length - hi)
    )


def predict_window(
    model: torch.nn.Module,
    tokenizer: Tokenizer,
    sequence: str,
    device: torch.device,
    dtype: torch.dtype,
    one_strand_at_a_time: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """Log-probabilities of the window's kept region, for it and its reverse complement.

    ``one_strand_at_a_time`` halves peak memory (about 31 GB instead of 57 GB for a
    1 Mb window in bfloat16).

    Returns:
        (kept, 15) for the plus strand in genomic order, and (kept, 15) for the
        minus strand in reverse-complement order (5' to 3' of the minus strand).
    """
    tokens = np.stack([tokenizer(sequence), tokenizer(reverse_complement(sequence))])
    if one_strand_at_a_time:
        plus, rc = (predict_log_probs(model, t[None], device, dtype)[0] for t in tokens)
        return plus, rc
    log_probs = predict_log_probs(model, tokens, device, dtype)
    return log_probs[0], log_probs[1]


def decode(log_probs: np.ndarray, codes: np.ndarray, decoder: str) -> np.ndarray:
    """State path of one strand's kept region (both in the strand's own orientation)."""
    if decoder == "argmax":
        return log_probs.argmax(-1).astype(np.int8)
    return decoding.viterbi(log_probs, codes if decoder == "hmm_constrained" else None)


@dataclass
class RegionPrediction:
    """Model outputs of every window covering a region.

    ``per_window`` holds, for each window, the log-probabilities of its whole kept
    region: plus strand in genomic order, minus strand in its own 5'->3' order.
    """

    chrom: str
    start: int
    end: int
    windows: list[Window]
    per_window: list[tuple[np.ndarray, np.ndarray]]

    def stitched(self, strand: str) -> np.ndarray:
        """(end - start, 15) log-probabilities of one strand ("+"/"-"), genomic order."""
        out = np.zeros((self.end - self.start, decoding.NUM_CLASSES), np.float32)
        for window, (plus, rc) in zip(self.windows, self.per_window, strict=True):
            values = plus if strand == "+" else rc[::-1]
            skip = window.keep_start - (window.keep_stop - len(plus))
            out[window.keep_start - self.start : window.keep_stop - self.start] = (
                values[skip:]
            )
        return out


def predict_region(
    model: torch.nn.Module,
    tokenizer: Tokenizer,
    contig: str,
    chrom: str,
    start: int,
    end: int,
    sequence_length: int,
    keep: float,
    device: torch.device,
    dtype: torch.dtype,
    one_strand_at_a_time: bool = False,
) -> RegionPrediction:
    """Runs the model on every pipeline window covering [start, end) of a contig.

    The region is widened to whole kept regions, so decoding it reproduces the
    pipeline's decode of those bp exactly.
    """
    windows = windows_for_region(len(contig), start, end, sequence_length, keep)
    crop = crop_of(sequence_length, keep)
    per_window = []
    for window in windows:
        plus, rc = predict_window(
            model,
            tokenizer,
            window_sequence(contig, window, sequence_length),
            device,
            dtype,
            one_strand_at_a_time,
        )
        n = min(window.start + crop + len(plus), len(contig)) - (window.start + crop)
        per_window.append((plus[:n], rc[len(rc) - n :]))
    return RegionPrediction(
        chrom, windows[0].keep_start, windows[-1].keep_stop, windows, per_window
    )


def decode_region(
    prediction: RegionPrediction, contig: str, decoder: str
) -> tuple[np.ndarray, np.ndarray]:
    """Plus- and minus-strand state paths over the region, in genomic order.

    Each window is decoded over its whole kept region, as in the pipeline; the
    shifted last window then drops the bp its predecessor already decoded.
    """
    codes = decoding.encode_bases(contig)
    n = prediction.end - prediction.start
    plus, minus = np.zeros(n, np.int8), np.zeros(n, np.int8)
    for window, (lp_plus, lp_rc) in zip(
        prediction.windows, prediction.per_window, strict=True
    ):
        full_start = window.keep_stop - len(lp_plus)
        window_codes = codes[full_start : window.keep_stop]
        states_plus = decode(lp_plus, window_codes, decoder)
        states_minus = decode(
            lp_rc, decoding.reverse_complement_codes(window_codes), decoder
        )[::-1]
        skip = window.keep_start - full_start
        out = slice(
            window.keep_start - prediction.start, window.keep_stop - prediction.start
        )
        plus[out] = states_plus[skip:]
        minus[out] = states_minus[skip:]
    return plus, minus


def region_genes(
    plus: np.ndarray,
    minus: np.ndarray,
    contig: str,
    chrom: str,
    start: int,
    filters: bool = True,
    min_cds_length: int = 201,
) -> list[str]:
    """GFF3 lines of both strands, in contig coordinates.

    With ``filters``, genes with an internal in-frame stop or less than
    ``min_cds_length`` bp of CDS are dropped, as in the released annotations.
    """
    codes = (
        decoding.encode_bases(contig[start : start + len(plus)]) if filters else None
    )
    lines = decoding.find_genes(plus, "+", chrom, codes)
    lines += decoding.find_genes(minus, "-", chrom, codes)
    lines = [_shift(line, start) for line in lines]
    return drop_short_genes(lines, min_cds_length) if filters else lines


def _shift(line: str, offset: int) -> str:
    cols = line.split("\t")
    cols[3], cols[4] = str(int(cols[3]) + offset), str(int(cols[4]) + offset)
    if "Parent=" in cols[8]:  # sub-feature IDs carry their start
        parent = cols[8].split("Parent=")[1]
        cols[8] = f"ID={parent}_{cols[2]}_{cols[3]};Parent={parent}"
    return "\t".join(cols)


def drop_short_genes(lines: list[str], min_cds_length: int) -> list[str]:
    """Drops ``find_genes`` genes with less than ``min_cds_length`` bp of CDS."""
    cds_length: dict[str, int] = defaultdict(int)
    for line in lines:
        cols = line.split("\t")
        if cols[2] == "CDS":
            parent = cols[8].split("Parent=")[1]
            cds_length[parent] += int(cols[4]) - int(cols[3]) + 1

    def gene_of(line: str) -> str:
        attributes = line.split("\t")[8]
        if "Parent=" in attributes:
            return attributes.split("Parent=")[1]
        return attributes.split("ID=")[1].split(";")[0]

    return [line for line in lines if cds_length[gene_of(line)] >= min_cds_length]