Spaces:
Running
Running
File size: 6,802 Bytes
a6af7a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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()}
@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)
|