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