P6: gtc-to-vcf (VCF/BCF + indexes), harmonisation, exports (plink, bedgraph, finalreport)
- 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
This commit is contained in:
kemal committed
1 parent
c1fcb088c5
commit
c1699b4465
19 files changed
+3181
-9
No files matched your search
@@ -0,0 +1,156 @@
|
||||
"""BCF 2.2 writer and CSI index (spec OUT-040, OUT-038).
|
||||
|
||||
Records are built from typed values; the file is BGZF-compressed and indexed with a CSI index (min_shift 14,
|
||||
depth 5, the pseudo-bin carries record counts). [DOC: VCF/BCF specification v4.3 §6 (BCF2), CSI index format —
|
||||
layouts only]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
from collections.abc import Sequence
|
||||
|
||||
from simurg.formats.bgzf import PSEUDO_BIN, BgzfWriter, bgzf_compress, reg2bin
|
||||
|
||||
INT8, INT16, INT32, FLOAT, CHAR = 1, 2, 3, 5, 7
|
||||
FLOAT_MISSING = struct.pack("<I", 0x7F800001)
|
||||
INT8_VECTOR_END = -127
|
||||
INT8_MISSING = -128
|
||||
|
||||
|
||||
def _type_byte(n: int, t: int) -> bytes:
|
||||
if n < 15:
|
||||
return bytes([(n << 4) | t])
|
||||
return bytes([(15 << 4) | t]) + typed_ints([n])
|
||||
|
||||
|
||||
def typed_ints(values: Sequence[int]) -> bytes:
|
||||
"""Smallest integer type that holds every value (missing/vector-end sentinels excluded from the range)."""
|
||||
lo, hi = (min(values), max(values)) if values else (0, 0)
|
||||
if lo >= -120 and hi <= 127:
|
||||
return _type_byte(len(values), INT8) + struct.pack(f"<{len(values)}b", *values)
|
||||
if lo >= -32760 and hi <= 32767:
|
||||
return _type_byte(len(values), INT16) + struct.pack(f"<{len(values)}h", *values)
|
||||
return _type_byte(len(values), INT32) + struct.pack(f"<{len(values)}i", *values)
|
||||
|
||||
|
||||
def typed_string(s: str) -> bytes:
|
||||
b = s.encode()
|
||||
return _type_byte(len(b), CHAR) + b
|
||||
|
||||
|
||||
def float_bytes(values: Sequence[float | None]) -> bytes:
|
||||
out = bytearray()
|
||||
for v in values:
|
||||
out += FLOAT_MISSING if v is None or v != v else struct.pack("<f", v)
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def gt_bytes(genotypes: Sequence[tuple[int, ...] | None], ploidies: Sequence[int]) -> tuple[int, bytes]:
|
||||
"""FORMAT/GT for all samples: width = max ploidy; unphased (a+1)<<1; missing 0; haploid padded with END."""
|
||||
width = max(1, max(ploidies))
|
||||
vals: list[int] = []
|
||||
for g, p in zip(genotypes, ploidies, strict=True):
|
||||
p = max(p, 1)
|
||||
cells = [0] * p if g is None else [(a + 1) << 1 for a in g]
|
||||
cells += [INT8_VECTOR_END] * (width - len(cells))
|
||||
vals += cells
|
||||
return width, bytes(vals_to_int8(vals))
|
||||
|
||||
|
||||
def vals_to_int8(vals: Sequence[int]) -> bytes:
|
||||
return struct.pack(f"<{len(vals)}b", *vals)
|
||||
|
||||
|
||||
def encode_record(
|
||||
contig_id: int,
|
||||
pos0: int,
|
||||
rlen: int,
|
||||
vid: str,
|
||||
alleles: Sequence[str],
|
||||
filter_id: int,
|
||||
n_sample: int,
|
||||
fmt: Sequence[tuple[int, bytes, bytes]],
|
||||
) -> bytes:
|
||||
"""``fmt``: (format dictionary index, type-descriptor bytes, per-sample payload) in FORMAT order."""
|
||||
shared = bytearray()
|
||||
shared += struct.pack("<iii", contig_id, pos0, rlen)
|
||||
shared += FLOAT_MISSING # QUAL
|
||||
shared += struct.pack("<I", (len(alleles) << 16) | 0) # n_allele << 16 | n_info
|
||||
shared += struct.pack("<I", (len(fmt) << 24) | n_sample)
|
||||
shared += typed_string(vid if vid else ".")
|
||||
for a in alleles:
|
||||
shared += typed_string(a)
|
||||
shared += typed_ints([filter_id])
|
||||
indiv = bytearray()
|
||||
for key, desc, payload in fmt:
|
||||
indiv += typed_ints([key]) + desc + payload
|
||||
return struct.pack("<II", len(shared), len(indiv)) + bytes(shared) + bytes(indiv)
|
||||
|
||||
|
||||
class CsiBuilder:
|
||||
"""CSI index (min_shift 14, depth 5) with per-bin loffset and the pseudo-bin."""
|
||||
|
||||
def __init__(self, n_ref: int) -> None:
|
||||
self.refs: list[dict[str, object]] = [{} for _ in range(n_ref)]
|
||||
|
||||
def add(self, rid: int, beg: int, end: int, vstart: int, vend: int) -> None:
|
||||
r = self.refs[rid]
|
||||
if not r:
|
||||
r.update(bins={}, n=0, first=vstart, last=vend)
|
||||
bins: dict[int, list[list[int]]] = r["bins"] # type: ignore[assignment]
|
||||
b = reg2bin(beg, end)
|
||||
chunks = bins.setdefault(b, [])
|
||||
if chunks and chunks[-1][1] == vstart:
|
||||
chunks[-1][1] = vend
|
||||
else:
|
||||
chunks.append([vstart, vend])
|
||||
r["n"] = int(r["n"]) + 1 # type: ignore[call-overload]
|
||||
r["last"] = vend
|
||||
|
||||
def encode(self) -> bytes:
|
||||
out = bytearray(b"CSI\x01") + struct.pack("<iii", 14, 5, 0) + struct.pack("<i", len(self.refs))
|
||||
for r in self.refs:
|
||||
if not r:
|
||||
out += struct.pack("<i", 0)
|
||||
continue
|
||||
bins: dict[int, list[list[int]]] = r["bins"] # type: ignore[assignment]
|
||||
out += struct.pack("<i", len(bins) + 1)
|
||||
for b in sorted(bins):
|
||||
ch = bins[b]
|
||||
out += struct.pack("<IQi", b, ch[0][0], len(ch))
|
||||
for s, e in ch:
|
||||
out += struct.pack("<QQ", s, e)
|
||||
out += struct.pack("<IQi", PSEUDO_BIN, 0, 2) + struct.pack("<QQQQ", r["first"], r["last"], r["n"], 0)
|
||||
out += struct.pack("<Q", 0)
|
||||
return bgzf_compress(bytes(out))
|
||||
|
||||
|
||||
class BcfWriter:
|
||||
"""Stream BCF records into BGZF; header text must follow VCF 4.3 with PASS as the first FILTER line."""
|
||||
|
||||
def __init__(self, fh: object, header_text: str) -> None:
|
||||
self.w = BgzfWriter(fh) # type: ignore[arg-type]
|
||||
text = header_text.encode() + b"\0"
|
||||
self.w.write_record(b"BCF\x02\x02" + struct.pack("<I", len(text)) + text)
|
||||
|
||||
def write(self, record: bytes) -> tuple[int, int]:
|
||||
return self.w.write_record(record)
|
||||
|
||||
def close(self) -> None:
|
||||
self.w.close()
|
||||
|
||||
|
||||
def header_dictionaries(header_text: str) -> tuple[dict[str, int], dict[str, int]]:
|
||||
"""(string dictionary: FILTER/INFO/FORMAT IDs → index, contig dictionary: name → index) as htslib builds them."""
|
||||
strings: dict[str, int] = {"PASS": 0}
|
||||
contigs: dict[str, int] = {}
|
||||
for line in header_text.splitlines():
|
||||
for kind in ("FILTER", "INFO", "FORMAT"):
|
||||
if line.startswith(f"##{kind}=<ID="):
|
||||
key = line[len(kind) + 7 :].split(",", 1)[0].split(">", 1)[0]
|
||||
strings.setdefault(key, len(strings))
|
||||
if line.startswith("##contig=<ID="):
|
||||
name = line[13:].split(",", 1)[0].split(">", 1)[0]
|
||||
contigs.setdefault(name, len(contigs))
|
||||
return strings, contigs
|
||||
@@ -0,0 +1,142 @@
|
||||
"""BGZF writer and tabix (``.tbi``) index for position-sorted text files (spec OUT-010, OUT-038).
|
||||
|
||||
Pure Python, deterministic (fixed compression level, no timestamps). Records never straddle a block when they fit,
|
||||
and the index carries the htslib pseudo-bin with mapped/unmapped counts so ``bcftools index --nrecords`` works.
|
||||
[DOC: SAM/BAM format specification §4.1 (BGZF) and §5 (binning, linear index); tabix index format — layouts only]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
|
||||
BLOCK_DATA = 0xFF00 # uncompressed bytes per block (as htslib)
|
||||
EOF_BLOCK = bytes.fromhex("1f8b08040000000000ff0600424302001b0003000000000000000000")
|
||||
TBI_MIN_SHIFT, TBI_DEPTH = 14, 5
|
||||
PSEUDO_BIN = 37450
|
||||
|
||||
|
||||
def _block(data: bytes, level: int) -> bytes:
|
||||
c = zlib.compressobj(level, zlib.DEFLATED, -15)
|
||||
comp = c.compress(data) + c.flush()
|
||||
bsize = len(comp) + 25
|
||||
head = struct.pack("<4BIBBHBBHH", 0x1F, 0x8B, 8, 4, 0, 0, 0xFF, 6, ord("B"), ord("C"), 2, bsize)
|
||||
return head + comp + struct.pack("<II", zlib.crc32(data) & 0xFFFFFFFF, len(data))
|
||||
|
||||
|
||||
class BgzfWriter:
|
||||
"""Write text into BGZF blocks; ``tell()`` returns the virtual offset of the next byte."""
|
||||
|
||||
def __init__(self, fh: BinaryIO, level: int = 6) -> None:
|
||||
self.fh, self.level = fh, level
|
||||
self.buf = bytearray()
|
||||
self.coffset = 0
|
||||
|
||||
def tell(self) -> int:
|
||||
return (self.coffset << 16) | len(self.buf)
|
||||
|
||||
def _flush_block(self) -> None:
|
||||
if self.buf:
|
||||
blk = _block(bytes(self.buf[:BLOCK_DATA]), self.level)
|
||||
self.fh.write(blk)
|
||||
self.coffset += len(blk)
|
||||
del self.buf[:BLOCK_DATA]
|
||||
|
||||
def write_record(self, data: bytes) -> tuple[int, int]:
|
||||
"""Append ``data``; start a new block first when it would not fit. Returns (start, end) virtual offsets."""
|
||||
if self.buf and len(self.buf) + len(data) > BLOCK_DATA:
|
||||
self._flush_block()
|
||||
start = self.tell()
|
||||
self.buf += data
|
||||
while len(self.buf) >= BLOCK_DATA:
|
||||
self._flush_block()
|
||||
return start, self.tell()
|
||||
|
||||
def close(self) -> None:
|
||||
while self.buf:
|
||||
self._flush_block()
|
||||
self.fh.write(EOF_BLOCK)
|
||||
|
||||
|
||||
def bgzf_compress(data: bytes, level: int = 6) -> bytes:
|
||||
out = bytearray()
|
||||
for i in range(0, len(data), BLOCK_DATA):
|
||||
out += _block(data[i : i + BLOCK_DATA], level)
|
||||
return bytes(out + EOF_BLOCK)
|
||||
|
||||
|
||||
def reg2bin(beg: int, end: int) -> int:
|
||||
end -= 1
|
||||
if beg >> 14 == end >> 14:
|
||||
return ((1 << 15) - 1) // 7 + (beg >> 14)
|
||||
if beg >> 17 == end >> 17:
|
||||
return ((1 << 12) - 1) // 7 + (beg >> 17)
|
||||
if beg >> 20 == end >> 20:
|
||||
return ((1 << 9) - 1) // 7 + (beg >> 20)
|
||||
if beg >> 23 == end >> 23:
|
||||
return ((1 << 6) - 1) // 7 + (beg >> 23)
|
||||
if beg >> 26 == end >> 26:
|
||||
return ((1 << 3) - 1) // 7 + (beg >> 26)
|
||||
return 0
|
||||
|
||||
|
||||
class TabixBuilder:
|
||||
"""Collect (contig, beg0, end0, start voffset, end voffset) in file order and emit a VCF ``.tbi``."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.names: list[str] = []
|
||||
self.refs: dict[str, dict[str, object]] = {}
|
||||
|
||||
def add(self, contig: str, beg: int, end: int, vstart: int, vend: int) -> None:
|
||||
r = self.refs.get(contig)
|
||||
if r is None:
|
||||
self.names.append(contig)
|
||||
r = self.refs[contig] = {"bins": {}, "lin": [], "n": 0, "first": vstart, "last": vend}
|
||||
bins: dict[int, list[list[int]]] = r["bins"] # type: ignore[assignment]
|
||||
b = reg2bin(beg, end)
|
||||
chunks = bins.setdefault(b, [])
|
||||
if chunks and chunks[-1][1] == vstart:
|
||||
chunks[-1][1] = vend
|
||||
else:
|
||||
chunks.append([vstart, vend])
|
||||
lin: list[int] = r["lin"] # type: ignore[assignment]
|
||||
for w in range(beg >> TBI_MIN_SHIFT, ((end - 1) >> TBI_MIN_SHIFT) + 1):
|
||||
while len(lin) <= w:
|
||||
lin.append(-1)
|
||||
if lin[w] == -1:
|
||||
lin[w] = vstart
|
||||
r["n"] = int(r["n"]) + 1 # type: ignore[call-overload]
|
||||
r["last"] = vend
|
||||
|
||||
def encode(self) -> bytes:
|
||||
names = b"".join(n.encode() + b"\0" for n in self.names)
|
||||
out = bytearray(b"TBI\x01")
|
||||
out += struct.pack("<8i", len(self.names), 2, 1, 2, 0, ord("#"), 0, len(names)) + names
|
||||
for n in self.names:
|
||||
r = self.refs[n]
|
||||
bins: dict[int, list[list[int]]] = r["bins"] # type: ignore[assignment]
|
||||
out += struct.pack("<i", len(bins) + 1)
|
||||
for b in sorted(bins):
|
||||
out += struct.pack("<Ii", b, len(bins[b]))
|
||||
for s, e in bins[b]:
|
||||
out += struct.pack("<QQ", s, e)
|
||||
out += struct.pack("<Ii", PSEUDO_BIN, 2) + struct.pack("<QQQQ", r["first"], r["last"], r["n"], 0)
|
||||
lin: list[int] = r["lin"] # type: ignore[assignment]
|
||||
fill = 0
|
||||
filled = []
|
||||
for v in lin: # htslib back-fills empty windows with the next non-empty offset
|
||||
filled.append(v)
|
||||
for k in range(len(filled) - 1, -1, -1):
|
||||
if filled[k] == -1:
|
||||
filled[k] = fill
|
||||
else:
|
||||
fill = filled[k]
|
||||
out += struct.pack("<i", len(filled)) + b"".join(struct.pack("<Q", v) for v in filled)
|
||||
out += struct.pack("<Q", 0) # n_no_coor
|
||||
return bgzf_compress(bytes(out))
|
||||
|
||||
|
||||
def write_tbi(path: Path, builder: TabixBuilder) -> None:
|
||||
path.write_bytes(builder.encode())
|
||||
@@ -0,0 +1,350 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user