"""Training windows of (soft-masked tokens, 15-class labels) from FASTA + label BED.""" from collections.abc import Iterator from pathlib import Path import numpy as np import pandas as pd import pyfaidx import torch from genome_annotation.decoding import CLASS_INDEX IGNORE_INDEX = -1 _MASK_LABELS = ("mask_short_intron", "mask_frame_error") _LABEL_TO_CLASS = CLASS_INDEX | dict.fromkeys(_MASK_LABELS, IGNORE_INDEX) # Where rows overlap, the label later in this order wins (masks last). _PAINT_ORDER = {label: i for i, label in enumerate(_LABEL_TO_CLASS)} _CDS_P0 = CLASS_INDEX["CDS_p0"] _COMPLEMENT = str.maketrans("ACGTNacgtn", "TGCANtgcan") def reverse_complement(sequence: str) -> str: """Reverse-complements a DNA string, keeping soft-masking.""" return sequence.translate(_COMPLEMENT)[::-1] class Tokenizer: """Maps DNA strings to NTv3 token ids, one token per base. A lookup-table equivalent of the model's own (character-level) tokenizer, fast enough for long windows: ACGTN and their soft-masked lowercase get their own ids, anything else (IUPAC codes, ...) is ````. """ def __init__(self, vocab: dict[str, int]) -> None: self.lut = np.full(256, vocab[""], dtype=np.int64) for base in "ACGTNacgtn": self.lut[ord(base)] = vocab[base] def __call__(self, sequence: str) -> np.ndarray: return self.lut[np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)] class LabelTrack: """Per-base 15-class labels of one assembly, read from its training BED. Each BED row labels an interval of one strand. A CDS row carries the codon position of its 5' base, which cycles 0, 1, 2 along the transcript. """ def __init__(self, bed_path: str | Path) -> None: bed = pd.read_csv( bed_path, sep="\t", header=None, names=["chrom", "start", "end", "label", "score", "strand"], dtype={"chrom": str}, ) bed["cls"] = bed["label"].map(_LABEL_TO_CLASS).astype(np.int64) bed["order"] = bed["label"].map(_PAINT_ORDER).astype(np.int64) self.rows: dict[tuple[str, str], tuple[np.ndarray, ...]] = {} for (chrom, strand), group in bed.groupby(["chrom", "strand"]): group = group.sort_values("start") self.rows[(chrom, strand)] = tuple( group[col].to_numpy() for col in ("start", "end", "cls", "order") ) def strand_labels( self, chrom: str, strand: str, start: int, end: int ) -> np.ndarray: """(end - start,) labels of one strand, in genomic order.""" labels = np.zeros(end - start, dtype=np.int64) if (chrom, strand) not in self.rows: return labels starts, ends, classes, order = self.rows[(chrom, strand)] # Rows are sorted by start; a row overlaps the window iff it starts before # its end and ends after its start (ends are not sorted where rows overlap). hi = np.searchsorted(starts, end, side="left") candidates = np.flatnonzero(ends[:hi] > start) candidates = candidates[np.argsort(order[candidates], kind="stable")] for row_start, row_end, cls in zip( starts[candidates], ends[candidates], classes[candidates], strict=True ): s, e = max(start, row_start), min(end, row_end) if _CDS_P0 <= cls < _CDS_P0 + 3: frame = cls - _CDS_P0 if strand == "+": offsets = np.arange(s - row_start, e - row_start) else: offsets = np.arange(row_end - 1 - s, row_end - 1 - e, -1) labels[s - start : e - start] = _CDS_P0 + (frame + offsets) % 3 else: labels[s - start : e - start] = cls return labels def window(self, chrom: str, start: int, end: int, reverse: bool) -> np.ndarray: """Labels of the strand read 5'->3' in the window's orientation.""" if reverse: return self.strand_labels(chrom, "-", start, end)[::-1].copy() return self.strand_labels(chrom, "+", start, end) class Assembly: """A genome and its labels, stored as ``/genome.fa.gz`` and ``labels.bed.gz``.""" def __init__(self, directory: str | Path) -> None: self.directory = Path(directory) self.name = self.directory.name self.labels = LabelTrack(self.directory / "labels.bed.gz") self._fasta: pyfaidx.Fasta | None = None with pyfaidx.Fasta(str(self.directory / "genome.fa.gz")) as fasta: self.lengths = {name: len(record) for name, record in fasta.items()} @property def fasta(self) -> pyfaidx.Fasta: # Opened lazily so that each dataloader worker gets its own file handle. if self._fasta is None: self._fasta = pyfaidx.Fasta(str(self.directory / "genome.fa.gz")) return self._fasta def sequence(self, chrom: str, start: int, end: int) -> str: return str(self.fasta[chrom][start:end]) class TrainingWindows(torch.utils.data.IterableDataset): """Endless random windows: species uniform, position uniform within the species. Each window is reverse-complemented with probability 1/2, together with its labels, so the single-strand head learns genes on both strands. """ def __init__( self, assemblies: list[Assembly], tokenizer: Tokenizer, window: int, seed: int = 0, ) -> None: self.assemblies = assemblies self.tokenizer = tokenizer self.window = window self.seed = seed def __iter__(self) -> Iterator[tuple[torch.Tensor, torch.Tensor]]: info = torch.utils.data.get_worker_info() rng = np.random.default_rng([self.seed, info.id if info else 0]) contigs = [ [(c, n) for c, n in a.lengths.items() if n >= self.window] for a in self.assemblies ] while True: index = rng.integers(len(self.assemblies)) assembly = self.assemblies[index] sizes = np.array([n - self.window + 1 for _, n in contigs[index]]) flat = rng.integers(sizes.sum()) contig = int(np.searchsorted(np.cumsum(sizes), flat, side="right")) chrom = contigs[index][contig][0] start = int(flat - (np.cumsum(sizes)[contig] - sizes[contig])) end = start + self.window reverse = bool(rng.integers(2)) sequence = assembly.sequence(chrom, start, end) if reverse: sequence = reverse_complement(sequence) labels = assembly.labels.window(chrom, start, end, reverse) yield torch.from_numpy(self.tokenizer(sequence)), torch.from_numpy(labels)