Spaces:
Running
Running
| """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 ``<unk>``. | |
| """ | |
| def __init__(self, vocab: dict[str, int]) -> None: | |
| self.lut = np.full(256, vocab["<unk>"], 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 ``<dir>/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()} | |
| 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) | |