Spaces:
Running
Running
| """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 | |
| 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) | |