kemal
c1699b4465
- formats/fasta: plain/BGZF FASTA, .fai/.gzi built identical to htslib, index check, build detection, contig aliases, mmap genome cache (IN-060..064) - formats/bgzf, formats/bcf: pure-Python BGZF, tabix .tbi, BCF 2.2 and CSI (OUT-010, OUT-038, OUT-040) - harmonize/sites: plus-strand mapping with audit rows, indel resolution from SourceSeq (port of GTCtoVCF, Apache-2.0) and left normalisation (OUT-014..021, OUT-030, OUT-042) - harmonize/combine: OUT-031 genotype combination (ploidy-aware, Infinium I/II signal classes, OUT-037) - pipeline/vcf: `simurg gtc-to-vcf` (per-sample VCF.gz+tbi, --merge cohort BCF+csi, audit parquet) - pipeline/export: `simurg export bedgraph|finalreport|plink|locus-summary` (OUT-043, OUT-050..052, OUT-071) - validation: 36 GSA samples vs gtc2vcf: 23.66 M identical calls, 0 unexplained differences; PLINK identical to plink2 on 637,397 variants; bcftools validates VCF/BCF and our indexes answer region queries correctly - fetch_data.sh p6 (Ensembl GRCh38 with checksum), make_oracle_vcf.sh, validate_p6_vcf.py; K-022/K-023
351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""Reference genome access (spec IN-060 … IN-064).
|
|
|
|
* Plain or BGZF-compressed FASTA with a samtools-style ``.fai`` (and ``.gzi`` for BGZF). Missing indexes are built
|
|
into the cache directory, never next to the input.
|
|
* ``Genome.fetch`` returns upper-case bases for a 0-based half-open interval. With the genome cache (default) each
|
|
needed contig is stored once as a flat uint8 file under ``$SIMURG_CACHE_DIR/genomes/<fasta-sha256>/`` and
|
|
memory-mapped; without it, bases are read through the index. Both modes return identical sequences (EDGE-010).
|
|
|
|
[DOC: samtools faidx / htslib BGZF and .gzi formats — layouts only]
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import struct
|
|
import zlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
from simurg.core.errors import SimurgError
|
|
|
|
GRCH37_CHR1 = 249_250_621
|
|
GRCH38_CHR1 = 248_956_422
|
|
_BGZF_MAGIC = b"\x1f\x8b\x08\x04"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FaiEntry:
|
|
name: str
|
|
length: int
|
|
offset: int # uncompressed byte offset of the first base
|
|
line_bases: int
|
|
line_width: int
|
|
|
|
|
|
def cache_root() -> Path:
|
|
return Path(os.environ.get("SIMURG_CACHE_DIR", Path.home() / ".cache" / "simurg"))
|
|
|
|
|
|
def is_bgzf(path: Path) -> bool:
|
|
with path.open("rb") as fh:
|
|
head = fh.read(16)
|
|
return head[:4] == _BGZF_MAGIC and head[12:14] == b"BC"
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
"""SHA-256 of a (large) file, memoised by path, size and mtime in the cache directory."""
|
|
st = path.stat()
|
|
memo = cache_root() / "sha256.json"
|
|
key = f"{path.resolve()}|{st.st_size}|{st.st_mtime_ns}"
|
|
try:
|
|
known = json.loads(memo.read_text())
|
|
except (OSError, ValueError):
|
|
known = {}
|
|
if key in known:
|
|
return str(known[key])
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as fh:
|
|
for chunk in iter(lambda: fh.read(1 << 22), b""):
|
|
h.update(chunk)
|
|
known[key] = h.hexdigest()
|
|
memo.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = memo.with_suffix(".part")
|
|
tmp.write_text(json.dumps(known, indent=0))
|
|
tmp.replace(memo)
|
|
return str(known[key])
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------ BGZF
|
|
def read_gzi(path: Path) -> list[tuple[int, int]]:
|
|
"""``.gzi``: (compressed offset, uncompressed offset) of every BGZF block after the first."""
|
|
data = path.read_bytes()
|
|
(n,) = struct.unpack_from("<Q", data, 0)
|
|
return [struct.unpack_from("<QQ", data, 8 + 16 * i) for i in range(n)]
|
|
|
|
|
|
def build_gzi(fasta: Path) -> list[tuple[int, int]]:
|
|
out = []
|
|
coff = uoff = 0
|
|
with fasta.open("rb") as fh:
|
|
while True:
|
|
head = fh.read(18)
|
|
if not head:
|
|
break
|
|
if head[:4] != _BGZF_MAGIC:
|
|
raise SimurgError("E_FASTA_FORMAT", f"{fasta}: not a BGZF file (bad block at {coff})")
|
|
bsize = struct.unpack_from("<H", head, 16)[0] + 1
|
|
body = fh.read(bsize - 18)
|
|
isize = struct.unpack_from("<I", body, len(body) - 4)[0]
|
|
if coff and isize: # htslib omits empty blocks (the EOF marker)
|
|
out.append((coff, uoff))
|
|
coff += bsize
|
|
uoff += isize
|
|
return out
|
|
|
|
|
|
def write_gzi(path: Path, entries: list[tuple[int, int]]) -> None:
|
|
buf = struct.pack("<Q", len(entries)) + b"".join(struct.pack("<QQ", c, u) for c, u in entries)
|
|
path.write_bytes(buf)
|
|
|
|
|
|
class _BgzfReader:
|
|
"""Random access to uncompressed offsets of a BGZF file via its ``.gzi``."""
|
|
|
|
def __init__(self, path: Path, gzi: list[tuple[int, int]]) -> None:
|
|
self.fh = path.open("rb")
|
|
self.blocks = [(0, 0), *gzi]
|
|
self.ustarts = np.array([u for _, u in self.blocks], dtype=np.int64)
|
|
self._cache: tuple[int, bytes] | None = None
|
|
|
|
def _block(self, i: int) -> bytes:
|
|
if self._cache and self._cache[0] == i:
|
|
return self._cache[1]
|
|
self.fh.seek(self.blocks[i][0])
|
|
head = self.fh.read(18)
|
|
bsize = struct.unpack_from("<H", head, 16)[0] + 1
|
|
body = self.fh.read(bsize - 18)
|
|
data = zlib.decompress(body[:-8], -15)
|
|
self._cache = (i, data)
|
|
return data
|
|
|
|
def read(self, uoff: int, n: int) -> bytes:
|
|
out = bytearray()
|
|
i = int(np.searchsorted(self.ustarts, uoff, side="right") - 1)
|
|
pos = uoff - int(self.ustarts[i])
|
|
while len(out) < n and i < len(self.blocks):
|
|
data = self._block(i)
|
|
out += data[pos : pos + n - len(out)]
|
|
i, pos = i + 1, 0
|
|
return bytes(out)
|
|
|
|
|
|
class _PlainReader:
|
|
def __init__(self, path: Path) -> None:
|
|
self.fh = path.open("rb")
|
|
|
|
def read(self, uoff: int, n: int) -> bytes:
|
|
self.fh.seek(uoff)
|
|
return self.fh.read(n)
|
|
|
|
|
|
def _stream_uncompressed(path: Path, bgzf: bool) -> object:
|
|
if not bgzf:
|
|
return path.open("rb")
|
|
import gzip
|
|
|
|
return gzip.open(path, "rb")
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------ FAI
|
|
def read_fai(path: Path) -> list[FaiEntry]:
|
|
out = []
|
|
for i, line in enumerate(path.read_text().splitlines(), start=1):
|
|
if not line.strip():
|
|
continue
|
|
f = line.split("\t")
|
|
try:
|
|
out.append(FaiEntry(f[0], int(f[1]), int(f[2]), int(f[3]), int(f[4])))
|
|
except (IndexError, ValueError):
|
|
raise SimurgError("E_FAI_MISMATCH", f"{path}: malformed line {i}", path=str(path)) from None
|
|
return out
|
|
|
|
|
|
def build_fai(fasta: Path, bgzf: bool) -> list[FaiEntry]:
|
|
"""samtools-compatible index; offsets are uncompressed byte offsets."""
|
|
entries: list[FaiEntry] = []
|
|
fh = _stream_uncompressed(fasta, bgzf)
|
|
off = 0
|
|
name: str | None = None
|
|
length = seq_off = lb = lw = 0
|
|
with fh: # type: ignore[attr-defined]
|
|
for raw in fh: # type: ignore[attr-defined]
|
|
if raw.startswith(b">"):
|
|
if name is not None:
|
|
entries.append(FaiEntry(name, length, seq_off, lb, lw))
|
|
name = raw[1:].split()[0].decode()
|
|
length, lb, lw = 0, 0, 0
|
|
seq_off = off + len(raw)
|
|
else:
|
|
n = len(raw.rstrip(b"\r\n"))
|
|
if lb == 0:
|
|
lb, lw = n, len(raw)
|
|
length += n
|
|
off += len(raw)
|
|
if name is not None:
|
|
entries.append(FaiEntry(name, length, seq_off, lb, lw))
|
|
if not entries:
|
|
raise SimurgError("E_FASTA_FORMAT", f"{fasta}: no sequences found", path=str(fasta))
|
|
return entries
|
|
|
|
|
|
def write_fai(path: Path, entries: list[FaiEntry]) -> None:
|
|
path.write_text("".join(f"{e.name}\t{e.length}\t{e.offset}\t{e.line_bases}\t{e.line_width}\n" for e in entries))
|
|
|
|
|
|
def detect_build(entries: list[FaiEntry]) -> str:
|
|
"""IN-063: chr1 length decides GRCh37 vs GRCh38; '' when unknown."""
|
|
for e in entries:
|
|
if e.name in ("1", "chr1"):
|
|
return {GRCH37_CHR1: "GRCh37", GRCH38_CHR1: "GRCh38"}.get(e.length, "")
|
|
return ""
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------ genome
|
|
class Genome:
|
|
"""Indexed reference genome; see module docstring."""
|
|
|
|
def __init__(self, fasta: str | Path, *, use_cache: bool = True, cache_dir: Path | None = None) -> None:
|
|
self.path = Path(fasta).resolve()
|
|
if not self.path.exists():
|
|
raise SimurgError("E_FILE_NOT_FOUND", f"FASTA not found: {fasta}", path=str(fasta))
|
|
self.bgzf = is_bgzf(self.path)
|
|
self.use_cache = use_cache
|
|
self._sha: str | None = None
|
|
self._cache_dir = cache_dir
|
|
self.notes: list[str] = []
|
|
fai = Path(str(self.path) + ".fai")
|
|
if fai.exists():
|
|
self.entries = read_fai(fai)
|
|
else:
|
|
built = self.cache_dir / "index.fai"
|
|
if not built.exists():
|
|
built.parent.mkdir(parents=True, exist_ok=True)
|
|
write_fai(built, build_fai(self.path, self.bgzf))
|
|
self.notes.append(f"built FASTA index {built}")
|
|
self.entries = read_fai(built)
|
|
self.index = {e.name: e for e in self.entries}
|
|
if self.bgzf:
|
|
gzi = Path(str(self.path) + ".gzi")
|
|
if not gzi.exists():
|
|
gzi = self.cache_dir / "index.gzi"
|
|
if not gzi.exists():
|
|
gzi.parent.mkdir(parents=True, exist_ok=True)
|
|
write_gzi(gzi, build_gzi(self.path))
|
|
self._reader: _BgzfReader | _PlainReader = _BgzfReader(self.path, read_gzi(gzi))
|
|
else:
|
|
self._reader = _PlainReader(self.path)
|
|
self._mm: dict[str, npt.NDArray[np.uint8]] = {}
|
|
self.verify()
|
|
self.build = detect_build(self.entries)
|
|
|
|
@property
|
|
def sha256(self) -> str:
|
|
if self._sha is None:
|
|
self._sha = file_sha256(self.path)
|
|
return self._sha
|
|
|
|
@property
|
|
def cache_dir(self) -> Path:
|
|
if self._cache_dir is None:
|
|
self._cache_dir = cache_root() / "genomes" / self.sha256
|
|
return self._cache_dir
|
|
|
|
def verify(self, n: int = 3) -> None:
|
|
"""IN-060: sample up to ``n`` contigs and check header position and line layout."""
|
|
picks = self.entries[:: max(1, len(self.entries) // n)][:n]
|
|
for e in picks:
|
|
name_bytes = e.name.encode()
|
|
before = self._reader.read(max(0, e.offset - len(name_bytes) - 64), min(e.offset, len(name_bytes) + 64))
|
|
first = self._reader.read(e.offset, e.line_width)
|
|
ok = before.endswith(b"\n") and b">" + name_bytes in before
|
|
if e.length >= e.line_bases:
|
|
ok = ok and len(first.rstrip(b"\r\n")) == e.line_bases and first.endswith(b"\n")
|
|
if not ok:
|
|
raise SimurgError(
|
|
"E_FAI_MISMATCH",
|
|
f"FASTA index does not match {self.path} at contig {e.name}",
|
|
hint="delete the stale .fai (it will be rebuilt in the cache) or regenerate it with samtools faidx",
|
|
path=str(self.path),
|
|
)
|
|
|
|
def contig_length(self, name: str) -> int:
|
|
return self.index[name].length
|
|
|
|
def _read_raw(self, e: FaiEntry, start: int, end: int) -> bytes:
|
|
lines_before = start // e.line_bases
|
|
uoff = e.offset + lines_before * e.line_width + start % e.line_bases
|
|
last = end - 1
|
|
uend = e.offset + (last // e.line_bases) * e.line_width + last % e.line_bases + 1
|
|
raw = self._reader.read(uoff, uend - uoff)
|
|
return raw.replace(b"\n", b"").replace(b"\r", b"")
|
|
|
|
def _contig_array(self, name: str) -> npt.NDArray[np.uint8]:
|
|
if name not in self._mm:
|
|
e = self.index[name]
|
|
f = self.cache_dir / f"{name}.u8"
|
|
if not f.exists() or f.stat().st_size != e.length:
|
|
f.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = f.with_suffix(".part")
|
|
with tmp.open("wb") as out:
|
|
step = 1 << 24
|
|
for s in range(0, e.length, step):
|
|
out.write(self._read_raw(e, s, min(e.length, s + step)).upper())
|
|
tmp.replace(f)
|
|
self._mm[name] = np.memmap(f, dtype=np.uint8, mode="r")
|
|
return self._mm[name]
|
|
|
|
def fetch(self, name: str, start: int, end: int) -> str:
|
|
"""Upper-case bases of ``name``[start:end] (0-based, half-open), clipped to the contig."""
|
|
e = self.index[name]
|
|
start, end = max(0, start), min(e.length, end)
|
|
if end <= start:
|
|
return ""
|
|
if self.use_cache:
|
|
return bytes(self._contig_array(name)[start:end]).decode("ascii")
|
|
return self._read_raw(e, start, end).decode("ascii").upper()
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------ contig names
|
|
def load_alias_file(path: Path) -> dict[str, str]:
|
|
out = {}
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
|
line = line.split("#", 1)[0].strip()
|
|
if line:
|
|
a, b = line.split("\t")[:2]
|
|
out[a] = b
|
|
return out
|
|
|
|
|
|
class ContigMap:
|
|
"""IN-061: manifest chromosome → FASTA contig (``XY`` → X). Policies: ``strict`` (verbatim), ``auto``
|
|
(``1``↔``chr1``, ``MT``↔``chrM``, ``22_KI270879v1_alt``↔``chr22_KI270879v1_alt``) or an alias TSV path."""
|
|
|
|
def __init__(self, genome_contigs: set[str] | list[str], policy: str = "strict") -> None:
|
|
self.contigs = set(genome_contigs)
|
|
self.policy = policy
|
|
self.alias = load_alias_file(Path(policy)) if policy not in ("strict", "auto") else {}
|
|
self._memo: dict[str, str | None] = {}
|
|
|
|
def get(self, chrom: str) -> str | None:
|
|
if chrom in self._memo:
|
|
return self._memo[chrom]
|
|
target = "X" if chrom == "XY" else chrom
|
|
if self.alias:
|
|
cands = [self.alias.get(chrom, ""), self.alias.get(target, ""), target]
|
|
elif self.policy == "auto":
|
|
bare = target[3:] if target.lower().startswith("chr") else target
|
|
cands = [target, bare, f"chr{bare}"]
|
|
if bare in ("MT", "M"):
|
|
cands += ["chrM", "MT", "M"]
|
|
else:
|
|
cands = [target]
|
|
hit = next((c for c in cands if c and c in self.contigs), None)
|
|
self._memo[chrom] = hit
|
|
return hit
|