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 Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv
This commit is contained in:
kemalandClaude Opus 5.5 committed
1 parent
60a05d5048
commit
b56fdebaaf
19 files changed
+3181
-9
No files matched your search
@@ -370,6 +370,153 @@ def genotype_cmd(
|
||||
run_guarded(rep, lambda: run_genotype(opts, rep))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------------
|
||||
@app.command("gtc-to-vcf")
|
||||
def gtc_to_vcf_cmd(
|
||||
bpm: Annotated[Path, typer.Option("--bpm", help="BPM manifest (required).")],
|
||||
fasta: Annotated[Path, typer.Option("--fasta", help="Reference FASTA (plain or bgzip; required).")],
|
||||
out: Annotated[Path, typer.Option("--out", help="Output folder (required).")],
|
||||
gtcs: Annotated[Path | None, typer.Option("--gtcs", help="Folder searched recursively for GTC files.")] = None,
|
||||
sheet: Annotated[Path | None, typer.Option("--sheet", help="Sample sheet (CSV).")] = None,
|
||||
csv: Annotated[Path | None, typer.Option("--csv", help="CSV manifest (SourceSeq; required for indels).")] = None,
|
||||
no_bgzip: Annotated[bool, typer.Option("--no-bgzip", help="Plain VCF, no index.")] = False,
|
||||
genome_cache: Annotated[bool, typer.Option("--genome-cache/--no-genome-cache", help="mmap genome cache.")] = True,
|
||||
exclude_loci: Annotated[
|
||||
Path | None, typer.Option("--exclude-loci", "--filter-loci", help="Probe names to drop (one per line).")
|
||||
] = None,
|
||||
include_loci: Annotated[Path | None, typer.Option("--include-loci", help="Only these probe names.")] = None,
|
||||
unsquash_duplicates: Annotated[bool, typer.Option("--unsquash-duplicates", help="One record per probe.")] = False,
|
||||
inf1_nocall_informative: Annotated[
|
||||
bool, typer.Option("--inf1-nocall-informative", help="Read Infinium I no-calls as allele absence.")
|
||||
] = False,
|
||||
drop_ambiguous: Annotated[bool, typer.Option("--drop-ambiguous", help="Drop low-confidence strand loci.")] = False,
|
||||
filter_gs: Annotated[float | None, typer.Option("--filter-gs", help="FILTER LowGS when GS < X.")] = None,
|
||||
contig_alias: Annotated[str, typer.Option("--contig-alias", help="strict|auto|FILE")] = "strict",
|
||||
ploidy_policy: Annotated[str, typer.Option("--ploidy-policy", help="sex-aware|diploid")] = "sex-aware",
|
||||
threads: Annotated[int, typer.Option("--threads", min=0, max=1024, help="Worker processes (0 = auto).")] = 0,
|
||||
merge: Annotated[
|
||||
bool, typer.Option("--merge", help="One multi-sample BCF (+ .csi) instead of per-sample VCFs.")
|
||||
] = False,
|
||||
force: Annotated[bool, typer.Option("--force", help="Replace existing outputs.")] = False,
|
||||
log_json: Annotated[bool, typer.Option("--log-json", "--json-log", help="JSON Lines log on stderr.")] = False,
|
||||
debug: Annotated[bool, typer.Option("--debug", help="Debug logging with tracebacks.")] = False,
|
||||
quiet: Annotated[bool, typer.Option("--quiet", "--no-progress", help="Suppress info messages.")] = False,
|
||||
) -> None:
|
||||
"""GTC files → one bgzipped, indexed SNV VCF per sample (plus-strand alleles, merged multi-probe sites)."""
|
||||
from simurg.pipeline.vcf import VcfOptions, run_gtc_to_vcf
|
||||
|
||||
rep = Reporter(
|
||||
command="gtc-to-vcf",
|
||||
level="debug" if debug else "info",
|
||||
json_stderr=log_json,
|
||||
debug=debug,
|
||||
quiet=quiet,
|
||||
log_file=out / "logs" / "gtc-to-vcf.log",
|
||||
)
|
||||
opts = VcfOptions(
|
||||
out=out, bpm=bpm, fasta=fasta, gtcs=gtcs, sheet=sheet, csv=csv, bgzip=not no_bgzip, genome_cache=genome_cache,
|
||||
exclude_loci=exclude_loci, include_loci=include_loci, unsquash_duplicates=unsquash_duplicates,
|
||||
inf1_nocall_informative=inf1_nocall_informative, drop_ambiguous=drop_ambiguous, filter_gs=filter_gs,
|
||||
contig_alias=contig_alias, ploidy_policy=ploidy_policy, threads=threads, merge=merge, force=force,
|
||||
argv=sys.argv[1:],
|
||||
) # fmt: skip
|
||||
run_guarded(rep, lambda: run_gtc_to_vcf(opts, rep))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------------
|
||||
export_app = typer.Typer(help="Exports from GTCs or VCFs: bedgraph, finalreport, plink, locus-summary.")
|
||||
app.add_typer(export_app, name="export")
|
||||
|
||||
|
||||
def _plain_reporter(command: str, log_json: bool, debug: bool, quiet: bool) -> Reporter:
|
||||
return Reporter(command=command, level="debug" if debug else "info", json_stderr=log_json, debug=debug, quiet=quiet)
|
||||
|
||||
|
||||
@export_app.command("bedgraph")
|
||||
def export_bedgraph_cmd(
|
||||
bpm: Annotated[Path, typer.Option("--bpm", help="BPM manifest (required).")],
|
||||
out: Annotated[Path, typer.Option("--out", help="Output folder (required).")],
|
||||
gtcs: Annotated[Path | None, typer.Option("--gtcs", help="Folder searched recursively for GTC files.")] = None,
|
||||
sheet: Annotated[Path | None, typer.Option("--sheet", help="Sample sheet (CSV).")] = None,
|
||||
no_gzip: Annotated[bool, typer.Option("--no-gzip", help="Plain .bedgraph files.")] = False,
|
||||
chrom_style: Annotated[str, typer.Option("--chrom-style", help="manifest|ucsc")] = "manifest",
|
||||
smoothed_lrr: Annotated[int, typer.Option("--smoothed-lrr", min=0, help="Extra moving-median LRR track, K.")] = 0,
|
||||
force: Annotated[bool, typer.Option("--force", help="Replace existing outputs.")] = False,
|
||||
log_json: Annotated[bool, typer.Option("--log-json", "--json-log")] = False,
|
||||
debug: Annotated[bool, typer.Option("--debug")] = False,
|
||||
quiet: Annotated[bool, typer.Option("--quiet", "--no-progress")] = False,
|
||||
) -> None:
|
||||
"""BAF and LRR bedGraph tracks per sample (OUT-050 … OUT-052)."""
|
||||
from simurg.pipeline.export import BedgraphOptions, run_export_bedgraph
|
||||
|
||||
rep = _plain_reporter("export bedgraph", log_json, debug, quiet)
|
||||
opts = BedgraphOptions(out=out, bpm=bpm, gtcs=gtcs, sheet=sheet, gzip=not no_gzip, chrom_style=chrom_style,
|
||||
smoothed_lrr=smoothed_lrr, force=force, argv=sys.argv[1:]) # fmt: skip
|
||||
run_guarded(rep, lambda: run_export_bedgraph(opts, rep))
|
||||
|
||||
|
||||
@export_app.command("finalreport")
|
||||
def export_finalreport_cmd(
|
||||
bpm: Annotated[Path, typer.Option("--bpm", help="BPM manifest (required).")],
|
||||
out: Annotated[Path, typer.Option("--out", help="Output folder (required).")],
|
||||
gtcs: Annotated[Path | None, typer.Option("--gtcs", help="Folder searched recursively for GTC files.")] = None,
|
||||
sheet: Annotated[Path | None, typer.Option("--sheet", help="Sample sheet (CSV).")] = None,
|
||||
csv: Annotated[Path | None, typer.Option("--csv", help="CSV manifest (source/reference strands).")] = None,
|
||||
egt: Annotated[Path | None, typer.Option("--egt", help="Cluster file (GT Score column).")] = None,
|
||||
matrix: Annotated[bool, typer.Option("--matrix", help="Loci × samples genotype matrix.")] = False,
|
||||
force: Annotated[bool, typer.Option("--force", help="Replace existing outputs.")] = False,
|
||||
log_json: Annotated[bool, typer.Option("--log-json", "--json-log")] = False,
|
||||
debug: Annotated[bool, typer.Option("--debug")] = False,
|
||||
quiet: Annotated[bool, typer.Option("--quiet", "--no-progress")] = False,
|
||||
) -> None:
|
||||
"""Long-format final report (one row per locus × sample), or a genotype matrix (OUT-071)."""
|
||||
from simurg.pipeline.export import FinalReportOptions, run_export_finalreport
|
||||
|
||||
rep = _plain_reporter("export finalreport", log_json, debug, quiet)
|
||||
opts = FinalReportOptions(out=out, bpm=bpm, gtcs=gtcs, sheet=sheet, csv=csv, egt=egt, matrix=matrix, force=force,
|
||||
argv=sys.argv[1:]) # fmt: skip
|
||||
run_guarded(rep, lambda: run_export_finalreport(opts, rep))
|
||||
|
||||
|
||||
@export_app.command("plink")
|
||||
def export_plink_cmd(
|
||||
vcfs: Annotated[Path, typer.Option("--vcfs", help="Folder with SimurgArray per-sample *.snv.vcf.gz.")],
|
||||
out: Annotated[Path, typer.Option("--out", help="Output folder (required).")],
|
||||
sheet: Annotated[Path | None, typer.Option("--sheet", help="Sample sheet (reported sex).")] = None,
|
||||
gtcs: Annotated[Path | None, typer.Option("--gtcs", help="GTC folder (estimated sex).")] = None,
|
||||
multiallelic: Annotated[str, typer.Option("--multiallelic", help="split|drop")] = "split",
|
||||
prefix: Annotated[str, typer.Option("--prefix", help="Output file prefix.")] = "cohort",
|
||||
force: Annotated[bool, typer.Option("--force", help="Replace existing outputs.")] = False,
|
||||
log_json: Annotated[bool, typer.Option("--log-json", "--json-log")] = False,
|
||||
debug: Annotated[bool, typer.Option("--debug")] = False,
|
||||
quiet: Annotated[bool, typer.Option("--quiet", "--no-progress")] = False,
|
||||
) -> None:
|
||||
"""PLINK 1 .bed/.bim/.fam from SimurgArray VCFs (OUT-043)."""
|
||||
from simurg.pipeline.export import PlinkOptions, run_export_plink
|
||||
|
||||
rep = _plain_reporter("export plink", log_json, debug, quiet)
|
||||
opts = PlinkOptions(out=out, vcfs=vcfs, sheet=sheet, gtcs=gtcs, multiallelic=multiallelic, prefix=prefix,
|
||||
force=force, argv=sys.argv[1:]) # fmt: skip
|
||||
run_guarded(rep, lambda: run_export_plink(opts, rep))
|
||||
|
||||
|
||||
@export_app.command("locus-summary")
|
||||
def export_locus_summary_cmd(
|
||||
bpm: Annotated[Path, typer.Option("--bpm", help="BPM manifest (required).")],
|
||||
egt: Annotated[Path, typer.Option("--egt", help="EGT cluster file (required).")],
|
||||
out: Annotated[Path, typer.Option("--out", help="Output folder (required).")],
|
||||
gtcs: Annotated[Path | None, typer.Option("--gtcs", help="Folder searched recursively for GTC files.")] = None,
|
||||
sheet: Annotated[Path | None, typer.Option("--sheet", help="Sample sheet (CSV).")] = None,
|
||||
threads: Annotated[int, typer.Option("--threads", min=0, max=1024)] = 0,
|
||||
force: Annotated[bool, typer.Option("--force", help="Replace existing outputs.")] = False,
|
||||
log_json: Annotated[bool, typer.Option("--log-json", "--json-log")] = False,
|
||||
debug: Annotated[bool, typer.Option("--debug")] = False,
|
||||
quiet: Annotated[bool, typer.Option("--quiet", "--no-progress")] = False,
|
||||
) -> None:
|
||||
"""Per-locus summary statistics (OUT-072); same as ``simurg qc loci``."""
|
||||
_qc_cohort("loci", out, bpm, gtcs, sheet, egt, threads, force, log_json, debug, quiet)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------------------------
|
||||
qc_app = typer.Typer(help="Quality control: controls, samples, contamination, relatedness, loci, thresholds.")
|
||||
app.add_typer(qc_app, name="qc")
|
||||
|
||||
@@ -89,6 +89,28 @@ CATALOGUE: dict[str, MessageDef] = dict(
|
||||
_m("E_EGT_LOAD", "error", "Cluster (EGT) file cannot be loaded", "IN-044"),
|
||||
_m("W_EGT_NAME", "warning", "Cluster file name is not portable", "IN-045"),
|
||||
_m("W_MANIFEST_NO_REFSTRAND", "warning", "Manifest lacks RefStrand annotation", "IN-040"),
|
||||
# --- reference genome and VCF export (IN-060 .. IN-065, OUT-010 .. OUT-043) --------------------------
|
||||
_m("E_FASTA_FORMAT", "error", "FASTA file is malformed", "IN-060"),
|
||||
_m("E_FAI_MISMATCH", "error", "FASTA index does not match the FASTA", "IN-060"),
|
||||
_m("E_BUILD_MISMATCH", "error", "Manifest genome build differs from the FASTA build", "IN-063"),
|
||||
_m("W_SKIP_UNMAPPED", "warning", "Locus without a genomic position was skipped", "OUT-019"),
|
||||
_m("W_SKIP_INTENSITY_ONLY", "warning", "Intensity-only locus was skipped", "OUT-019"),
|
||||
_m("W_SKIP_INDEL_CONTEXT", "warning", "Indel sequence cannot be determined (no CSV SourceSeq)", "OUT-019"),
|
||||
_m("W_SKIP_INDEL_REF", "warning", "Indel reference allele cannot be determined from the FASTA", "OUT-019"),
|
||||
_m("W_INDEL_PARTIAL_MATCH", "warning", "Indel source sequence matches the genome only partially", "OUT-019"),
|
||||
_m("W_SKIP_EXCLUDED", "warning", "Locus excluded by --exclude-loci/--include-loci", "OUT-019"),
|
||||
_m("W_SKIP_CONTIG", "warning", "Manifest chromosome absent from the FASTA", "OUT-019"),
|
||||
_m("W_SKIP_NO_STRAND", "warning", "Locus has neither RefStrand nor a sequence to align", "OUT-042"),
|
||||
_m("W_SKIP_AMBIGUOUS", "warning", "Ambiguous-strand locus dropped by --drop-ambiguous", "OUT-042"),
|
||||
_m("W_REF_NOT_QUERIED", "warning", "The reference allele is not among the probe's alleles", "OUT-015"),
|
||||
_m("W_MERGE_INCONSISTENT", "warning", "Probes at a site disagree; genotype set missing", "OUT-031"),
|
||||
_m("W_MERGE_AMBIGUOUS", "warning", "Probes at a site allow several genotypes; set missing", "OUT-031"),
|
||||
_m("W_MERGE_ALL_NC", "warning", "Every probe at a merged site is a no-call", "OUT-031"),
|
||||
_m("W_INF1_NC_CONFLICT", "warning", "Infinium I no-call conflicts with other assays", "OUT-037"),
|
||||
_m("W_NO_BGZIP", "warning", "Uncompressed VCF written; downstream commands need bgzip + index", "OUT-010"),
|
||||
_m("W_UNKNOWN_LOCI", "warning", "Names in a locus list are not in the manifest", "OUT-036"),
|
||||
_m("E_AUX_INVALID", "error", "Auxiliary-loci record is invalid", "OUT-035"),
|
||||
_m("W_LARGE_OUTPUT", "warning", "Output will be very large", "EDGE-036"),
|
||||
# --- QC (QC-001 .. QC-064) -----------------------------------------------------------------
|
||||
_m("W_CTRL_MISSING", "warning", "Controls needed by a metric are absent; metric not evaluable", "QC-007"),
|
||||
_m("W_CTRL_ZERO_DIVISOR", "warning", "Control metric has a zero denominator (NaN)", "QC-008"),
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,121 @@
|
||||
"""Genotypes of VCF sites from probe calls (spec OUT-018, OUT-031, OUT-032, OUT-037; plan §8.3).
|
||||
|
||||
Single-probe sites map the A/B call straight to allele indices. Sites with several probes use the exact candidate
|
||||
algorithm of OUT-031: enumerate every genotype over {REF} ∪ ALT with the site's ploidy, predict each probe's call
|
||||
from the alleles' signal classes (Infinium II: dye group of the extended base; Infinium I and indels: exact allele,
|
||||
no signal otherwise) and keep the candidates consistent with every called probe.
|
||||
[OSS: Illumina GTCtoVCF ``GenotypeFormat.RecordCombiner`` (Apache-2.0) — the consistency idea; the class model and
|
||||
the ploidy-aware enumeration are ours.]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from itertools import combinations_with_replacement
|
||||
|
||||
from simurg.harmonize.sites import GREEN, RED, ProbeSite, Site
|
||||
|
||||
NC, AA, AB, BB = 0, 1, 2, 3
|
||||
CALL_TEXT = {NC: "NC", AA: "AA", AB: "AB", BB: "BB"}
|
||||
|
||||
|
||||
def signal_class(p: ProbeSite, allele: str) -> str | None:
|
||||
"""OUT-031 step 1."""
|
||||
a, b = p.alleles
|
||||
if p.kind == "snv" and p.infinium == 2 and len(allele) == 1:
|
||||
grp = RED if allele in RED else GREEN
|
||||
if a in grp:
|
||||
return "A"
|
||||
if b in grp:
|
||||
return "B"
|
||||
return None
|
||||
if allele == a:
|
||||
return "A"
|
||||
if allele == b:
|
||||
return "B"
|
||||
return None
|
||||
|
||||
|
||||
def predicted_call(p: ProbeSite, genotype: tuple[str, ...]) -> int:
|
||||
classes = sorted(c for c in (signal_class(p, z) for z in genotype) if c is not None)
|
||||
if not classes:
|
||||
return NC
|
||||
u = set(classes)
|
||||
if u == {"A"}:
|
||||
return AA
|
||||
if u == {"B"}:
|
||||
return BB
|
||||
return AB
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Outcome:
|
||||
alleles: tuple[int, ...] | None # allele indices (sorted) or None for missing
|
||||
reason: str # "" | "het_in_haploid" | W_MERGE_* | W_INF1_NC_CONFLICT
|
||||
|
||||
|
||||
def single(site: Site, call: int, ploidy: int) -> Outcome:
|
||||
p = site.probes[0]
|
||||
if call == NC:
|
||||
return Outcome(None, "")
|
||||
ia, ib = site.allele_index(p.alleles[0]), site.allele_index(p.alleles[1])
|
||||
if ploidy == 1:
|
||||
if call == AB:
|
||||
return Outcome(None, "het_in_haploid")
|
||||
return Outcome((ia if call == AA else ib,), "")
|
||||
pair = {AA: (ia, ia), AB: (ia, ib), BB: (ib, ib)}[call]
|
||||
return Outcome(tuple(sorted(pair)), "")
|
||||
|
||||
|
||||
def merged(site: Site, calls: tuple[int, ...], ploidy: int, inf1_nc_informative: bool = False) -> Outcome:
|
||||
"""OUT-031 steps 3–6."""
|
||||
if all(c == NC for c in calls):
|
||||
return Outcome(None, "W_MERGE_ALL_NC")
|
||||
alleles = [site.ref, *site.alts]
|
||||
cands = list(combinations_with_replacement(range(len(alleles)), ploidy))
|
||||
keep = [
|
||||
g
|
||||
for g in cands
|
||||
if all(
|
||||
c == NC or predicted_call(p, tuple(alleles[i] for i in g)) == c
|
||||
for p, c in zip(site.probes, calls, strict=True)
|
||||
)
|
||||
]
|
||||
reason = ""
|
||||
if inf1_nc_informative:
|
||||
for p, c in zip(site.probes, calls, strict=True):
|
||||
if c == NC and p.infinium == 1:
|
||||
sub = [g for g in keep if all(signal_class(p, alleles[i]) is None for i in g)]
|
||||
if sub:
|
||||
keep = sub
|
||||
else:
|
||||
reason = "W_INF1_NC_CONFLICT"
|
||||
if len(keep) == 1:
|
||||
return Outcome(keep[0], reason)
|
||||
return Outcome(None, "W_MERGE_INCONSISTENT" if not keep else "W_MERGE_AMBIGUOUS")
|
||||
|
||||
|
||||
def genotype(site: Site, calls: tuple[int, ...], ploidy: int, inf1_nc_informative: bool = False) -> Outcome:
|
||||
if len(site.probes) == 1:
|
||||
return single(site, calls[0], ploidy)
|
||||
key = (calls, ploidy, inf1_nc_informative)
|
||||
hit = site.memo.get(key)
|
||||
if hit is None:
|
||||
hit = site.memo[key] = merged(site, calls, ploidy, inf1_nc_informative)
|
||||
assert isinstance(hit, Outcome)
|
||||
return hit
|
||||
|
||||
|
||||
def gt_text(o: Outcome, ploidy: int) -> str:
|
||||
if o.alleles is None:
|
||||
return "." if ploidy == 1 else "./."
|
||||
return "/".join(str(i) for i in o.alleles)
|
||||
|
||||
|
||||
def failure_message(site: Site, calls: tuple[int, ...], reason: str) -> str:
|
||||
what = {"W_MERGE_INCONSISTENT": "inconsistent", "W_MERGE_AMBIGUOUS": "ambiguous", "W_MERGE_ALL_NC": "all no-calls"}
|
||||
probes = ", ".join(
|
||||
f"{p.name} ({'InfiniumI' if p.infinium == 1 else 'InfiniumII'}): {CALL_TEXT[c]}"
|
||||
for p, c in zip(site.probes, calls, strict=True)
|
||||
)
|
||||
return f"Failed to combine genotypes ({what.get(reason, reason)}) - {probes}"
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Manifest loci → plus-strand VCF sites, with an audit row per locus (spec IN-046, OUT-014 … OUT-021, OUT-030,
|
||||
OUT-042; plan §8.1–§8.2).
|
||||
|
||||
[DOC: Illumina "TOP/BOT strand and A/B allele" technical note — strand conventions]
|
||||
[OSS: Illumina GTCtoVCF ``BPMRecord.py`` (Apache-2.0) — indel source-sequence handling: design-strand orientation,
|
||||
left shift, IUPAC-aware flank matching and the deletion-versus-insertion context score; ported below.]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
COMPLEMENT = dict(zip("ABCDGHKMRTVYNID", "TVGHCDMKYABRNID", strict=True))
|
||||
DEGENERACY = {"A": "A", "C": "C", "G": "G", "T": "T", "R": "AG", "Y": "CT", "S": "GC", "W": "AT", "K": "GT",
|
||||
"M": "AC", "B": "CGT", "D": "AGT", "H": "ACT", "V": "ACG", "N": "ACGT"} # fmt: skip
|
||||
RED, GREEN = set("AT"), set("CG")
|
||||
SEQ_RE = re.compile(r"^([^\[]*)\[([^/\]]*)/([^\]]*)\](.*)$")
|
||||
|
||||
Fetch = Callable[[str, int, int], str] # (contig, start0, end0) -> upper-case bases
|
||||
|
||||
|
||||
def revcomp(s: str) -> str:
|
||||
return "".join(COMPLEMENT.get(c, "N") for c in reversed(s))
|
||||
|
||||
|
||||
def complement_base(b: str) -> str:
|
||||
return COMPLEMENT.get(b, "N")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeSite:
|
||||
"""One manifest locus mapped to the plus strand."""
|
||||
|
||||
index: int
|
||||
name: str
|
||||
contig: str
|
||||
pos: int # VCF POS (1-based)
|
||||
ref: str
|
||||
alleles: tuple[str, str] # VCF allele strings of the probe's A and B alleles
|
||||
kind: str # "snv" | "indel"
|
||||
infinium: int # 1 or 2
|
||||
manifest_chrom: str
|
||||
par: bool = False
|
||||
confidence: str = "high"
|
||||
flipped: bool = False # design strand is minus
|
||||
reason: str = "ok"
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditRow:
|
||||
locus: str
|
||||
original_alleles: str
|
||||
strand_in: str
|
||||
transformed_alleles: str
|
||||
strand_out: str
|
||||
ref: str
|
||||
alt: str
|
||||
flipped: bool
|
||||
swapped: bool
|
||||
reason: str
|
||||
build: str
|
||||
confidence: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Site:
|
||||
"""A VCF record: one or more probes that share CHROM/POS/REF (and the indel span)."""
|
||||
|
||||
contig: str
|
||||
pos: int
|
||||
ref: str
|
||||
alts: list[str]
|
||||
probes: list[ProbeSite] = field(default_factory=list)
|
||||
kind: str = "snv"
|
||||
memo: dict[object, object] = field(default_factory=dict, repr=False, compare=False) # merged-genotype cache
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return (";" if self.kind == "indel" else ",").join(p.name for p in self.probes)
|
||||
|
||||
def allele_index(self, allele: str) -> int:
|
||||
return 0 if allele == self.ref else 1 + self.alts.index(allele)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ indel helpers
|
||||
def split_source(seq: str) -> tuple[str, str, str, str] | None:
|
||||
"""``left[a/b]right`` → (left, a, b, right); None when malformed."""
|
||||
m = SEQ_RE.match(seq.strip().upper())
|
||||
return (m.group(1), m.group(2), m.group(3), m.group(4)) if m else None
|
||||
|
||||
|
||||
def left_shift(five: str, indel: str, three: str) -> tuple[str, str]:
|
||||
"""Shift an indel fully 5′ within its context (GTCtoVCF ``determine_left_shift``)."""
|
||||
while indel and five.endswith(indel):
|
||||
five, three = five[: -len(indel)], indel + three
|
||||
while indel and five and len(indel) * five[-1] == indel:
|
||||
three, five = five[-1] + three, five[:-1]
|
||||
return five, three
|
||||
|
||||
|
||||
def _match(a: str, b: str) -> bool:
|
||||
return a in DEGENERACY.get(b, "")
|
||||
|
||||
|
||||
def max_suffix_match(genome: str, design: str) -> int:
|
||||
n = 0
|
||||
for x, y in zip(genome[::-1], design[::-1], strict=False):
|
||||
if not _match(x, y):
|
||||
break
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def max_prefix_match(genome: str, design: str) -> int:
|
||||
n = 0
|
||||
for x, y in zip(genome, design, strict=False):
|
||||
if not _match(x, y):
|
||||
break
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
def classify_indel(fetch: Fetch, contig: str, pos: int, five: str, indel: str, three: str) -> tuple[bool | None, float]:
|
||||
"""Is the plus-strand ``indel`` sequence present in the reference at ``pos`` (a reference deletion)?
|
||||
Returns (is_deletion or None when undecidable, best context score). Context scoring as GTCtoVCF."""
|
||||
start = pos - 1
|
||||
L = len(indel)
|
||||
seq_match = fetch(contig, start, start + L) == indel
|
||||
d5, d3 = left_shift(fetch(contig, start - len(five), start), indel,
|
||||
fetch(contig, start + L, start + L + len(three))) # fmt: skip
|
||||
i5, i3 = left_shift(fetch(contig, start - len(five) + 1, start + 1), indel,
|
||||
fetch(contig, start + 1, start + len(three) + 1)) # fmt: skip
|
||||
dm = (max_suffix_match(d5, five), max_prefix_match(d3, three))
|
||||
dmax = min(len(d5), len(five)) + min(len(d3), len(three)) + L
|
||||
dscore = (sum(dm) + L if seq_match else 0) / dmax if dmax else 0.0
|
||||
im = (max_suffix_match(i5, five), max_prefix_match(i3, three))
|
||||
imax = min(len(i5), len(five)) + min(len(i3), len(three))
|
||||
iscore = sum(im) / imax if imax else 0.0
|
||||
is_del = seq_match and dscore > iscore and min(dm) >= 1
|
||||
is_ins = iscore > dscore and min(im) >= 1
|
||||
if is_del == is_ins:
|
||||
return None, max(dscore, iscore)
|
||||
return is_del, dscore if is_del else iscore
|
||||
|
||||
|
||||
def normalize_indel(fetch: Fetch, contig: str, pos: int, indel: str, is_deletion: bool) -> tuple[int, str, str]:
|
||||
"""Left-normalised, anchored VCF representation. ``pos`` is the manifest position: the first deleted base for
|
||||
a deletion, the base 5′ of the insertion for an insertion (IN-046)."""
|
||||
if is_deletion:
|
||||
p0 = pos - 1 # 0-based first deleted base
|
||||
while p0 > 0 and fetch(contig, p0 - 1, p0) == indel[-1]:
|
||||
indel = indel[-1] + indel[:-1]
|
||||
p0 -= 1
|
||||
anchor = fetch(contig, p0 - 1, p0)
|
||||
return p0, anchor + indel, anchor # VCF POS = p0 (1-based anchor)
|
||||
a0 = pos - 1 # 0-based anchor base
|
||||
while a0 > 0 and fetch(contig, a0, a0 + 1) == indel[-1]:
|
||||
indel = indel[-1] + indel[:-1]
|
||||
a0 -= 1
|
||||
anchor = fetch(contig, a0, a0 + 1)
|
||||
return a0 + 1, anchor, anchor + indel
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ strand
|
||||
def strand_from_flanks(fetch: Fetch, contig: str, pos: int, left: str, right: str) -> tuple[str, float]:
|
||||
"""'+' when the design flanks match the reference as written, '-' when reverse-complemented; with a score."""
|
||||
left, right = left[-30:], right[:30]
|
||||
g_left = fetch(contig, pos - 1 - len(left), pos - 1)
|
||||
g_right = fetch(contig, pos, pos + len(right))
|
||||
|
||||
def score(lf: str, rt: str) -> float:
|
||||
n = len(lf) + len(rt)
|
||||
if not n:
|
||||
return 0.0
|
||||
hits = sum(_match(a, b) for a, b in zip(g_left[::-1], lf[::-1], strict=False))
|
||||
hits += sum(_match(a, b) for a, b in zip(g_right, rt, strict=False))
|
||||
return hits / n
|
||||
|
||||
fwd, rev = score(left, right), score(revcomp(right), revcomp(left))
|
||||
return ("+", fwd) if fwd >= rev else ("-", rev)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ mapping
|
||||
@dataclass
|
||||
class ManifestView:
|
||||
"""The manifest columns the mapper needs (BPM and/or CSV)."""
|
||||
|
||||
names: Sequence[str]
|
||||
snps: Sequence[str]
|
||||
chroms: Sequence[str]
|
||||
positions: Sequence[int]
|
||||
ilmn_strands: Sequence[str]
|
||||
source_strands: Sequence[str]
|
||||
ref_strands: Sequence[str]
|
||||
source_seqs: Sequence[str]
|
||||
assay_types: Sequence[int] # 0 Infinium II, 1/2 Infinium I
|
||||
intensity_only: Sequence[bool]
|
||||
|
||||
|
||||
def map_locus(
|
||||
i: int, mv: ManifestView, fetch: Fetch, contig_of: Callable[[str], str | None], build: str = ""
|
||||
) -> tuple[ProbeSite | None, AuditRow]:
|
||||
name, snp, chrom = mv.names[i], mv.snps[i], mv.chroms[i]
|
||||
pos = int(mv.positions[i])
|
||||
ilmn, ref_strand = mv.ilmn_strands[i] or "", (mv.ref_strands[i] or "").strip()
|
||||
audit = AuditRow(name, snp, f"{ilmn}/{ref_strand or '?'}", "", "+", "", "", False, False, "ok", build, "high")
|
||||
|
||||
def skip(code: str) -> tuple[None, AuditRow]:
|
||||
audit.reason, audit.strand_out = code, ""
|
||||
return None, audit
|
||||
|
||||
if chrom in ("0", "") or pos <= 0:
|
||||
return skip("W_SKIP_UNMAPPED")
|
||||
if mv.intensity_only[i]:
|
||||
return skip("W_SKIP_INTENSITY_ONLY")
|
||||
contig = contig_of(chrom)
|
||||
if contig is None:
|
||||
return skip("W_SKIP_CONTIG")
|
||||
m = re.match(r"^\[([^/\]]+)/([^\]]+)\]$", snp.strip())
|
||||
if not m:
|
||||
return skip("W_SKIP_UNMAPPED")
|
||||
a, b = m.group(1).upper(), m.group(2).upper()
|
||||
infinium = 2 if int(mv.assay_types[i]) == 0 else 1
|
||||
par = chrom == "XY"
|
||||
src = mv.source_seqs[i] if i < len(mv.source_seqs) else ""
|
||||
parts = split_source(src) if src else None
|
||||
if a in ("I", "D") or b in ("I", "D"):
|
||||
# ---------------------------------------------------------------- indel
|
||||
if parts is None or "-" not in (parts[1], parts[2]):
|
||||
return skip("W_SKIP_INDEL_CONTEXT")
|
||||
five, x, y, three = parts
|
||||
indel_seq = y if x == "-" else x
|
||||
s_strand = (mv.source_strands[i] or "")[:1].upper()
|
||||
d_strand = ilmn[:1].upper()
|
||||
if not ref_strand or s_strand in ("U", "") or d_strand in ("U", ""):
|
||||
return skip("W_SKIP_INDEL_CONTEXT")
|
||||
on_design = s_strand == d_strand
|
||||
if on_design != (ref_strand == "+"):
|
||||
five, indel_seq, three = revcomp(three), revcomp(indel_seq), revcomp(five)
|
||||
five, three = left_shift(five, indel_seq, three)
|
||||
is_del, score = classify_indel(fetch, contig, pos, five, indel_seq, three)
|
||||
if is_del is None:
|
||||
return skip("W_SKIP_INDEL_REF")
|
||||
vpos, vref, valt = normalize_indel(fetch, contig, pos, indel_seq, is_del)
|
||||
if not vref or "N" in vref:
|
||||
return skip("W_SKIP_INDEL_REF")
|
||||
# I (insertion allele) carries the sequence; D lacks it
|
||||
allele_of = {"I": vref if is_del else valt, "D": valt if is_del else vref}
|
||||
alleles = (allele_of[a], allele_of[b])
|
||||
audit.reason = "ok" if score >= 1.0 else "W_INDEL_PARTIAL_MATCH"
|
||||
audit.transformed_alleles = f"{alleles[0]}/{alleles[1]}"
|
||||
audit.ref, audit.alt = vref, valt
|
||||
audit.swapped = alleles[0] != vref
|
||||
return ProbeSite(i, name, contig, vpos, vref, alleles, "indel", infinium, chrom, par,
|
||||
"high", ref_strand == "-", audit.reason), audit # fmt: skip
|
||||
# -------------------------------------------------------------------- SNV
|
||||
confidence = "high"
|
||||
strand = ref_strand
|
||||
if strand not in ("+", "-"):
|
||||
if parts is None:
|
||||
return skip("W_SKIP_NO_STRAND") # no RefStrand and no sequence to align
|
||||
five, x, y, three = parts
|
||||
src_plus, score = strand_from_flanks(fetch, contig, pos, five, three)
|
||||
same = (mv.source_strands[i] or "")[:1].upper() == ilmn[:1].upper()
|
||||
strand = src_plus if same else ("-" if src_plus == "+" else "+")
|
||||
if score < 0.9:
|
||||
confidence = "low"
|
||||
if {a, b} in ({"A", "T"}, {"C", "G"}):
|
||||
confidence = "low" # plan §8.2: ambiguous without RefStrand
|
||||
pa, pb = (a, b) if strand == "+" else (complement_base(a), complement_base(b))
|
||||
ref = fetch(contig, pos - 1, pos)
|
||||
if not ref or ref not in "ACGT":
|
||||
return skip("W_SKIP_UNMAPPED")
|
||||
audit.transformed_alleles = f"{pa}/{pb}"
|
||||
audit.ref = ref
|
||||
audit.alt = ",".join(x for x in (pa, pb) if x != ref)
|
||||
audit.flipped = strand == "-"
|
||||
audit.swapped = pa != ref
|
||||
audit.confidence = confidence
|
||||
if ref not in (pa, pb):
|
||||
audit.reason = "W_REF_NOT_QUERIED"
|
||||
return ProbeSite(i, name, contig, pos, ref, (pa, pb), "snv", infinium, chrom, par, confidence,
|
||||
strand == "-", audit.reason), audit # fmt: skip
|
||||
|
||||
|
||||
def group_sites(probes: Sequence[ProbeSite], contig_order: dict[str, int], *, unsquash: bool = False) -> list[Site]:
|
||||
"""OUT-030/OUT-033/OUT-034: probes sharing CHROM/POS/REF (and, for indels, the ALT) form one site; ALT alleles in
|
||||
order of first appearance over BPM order; records sorted by contig (.fai order), POS, REF, ALT, ID."""
|
||||
groups: dict[tuple[str, int, str, str, str], list[ProbeSite]] = {}
|
||||
for p in sorted(probes, key=lambda p: p.index):
|
||||
extra = ",".join(sorted(x for x in p.alleles if x != p.ref)) if p.kind == "indel" else ""
|
||||
groups.setdefault((p.contig, p.pos, p.ref, p.kind, extra), []).append(p)
|
||||
sites: list[Site] = []
|
||||
for (contig, pos, ref, kind, _), ps in groups.items():
|
||||
alts: list[str] = []
|
||||
for p in ps:
|
||||
for x in p.alleles:
|
||||
if x != ref and x not in alts:
|
||||
alts.append(x)
|
||||
if unsquash and len(ps) > 1 and len(alts) <= 1:
|
||||
sites += [Site(contig, pos, ref, [x for x in p.alleles if x != ref], [p], kind) for p in ps]
|
||||
else:
|
||||
sites.append(Site(contig, pos, ref, alts, ps, kind))
|
||||
sites.sort(key=lambda s: (contig_order.get(s.contig, 1 << 30), s.pos, s.ref, ",".join(s.alts), s.id))
|
||||
return sites
|
||||
@@ -0,0 +1,422 @@
|
||||
"""``simurg export bedgraph | finalreport | plink | locus-summary`` (spec CLI-023, OUT-043, OUT-050 … OUT-052,
|
||||
OUT-071, OUT-072). None of them re-genotypes: they read GTCs (or SimurgArray VCFs) plus the manifest.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import io
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import IO, Any, cast
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from simurg import PRODUCT_NAME, __version__
|
||||
from simurg.core.atomic import atomic_open, check_outputs_absent
|
||||
from simurg.core.errors import ExitCode, SimurgError
|
||||
from simurg.core.log import Reporter
|
||||
from simurg.core.provenance import RunManifest
|
||||
from simurg.formats.bpm import Manifest
|
||||
from simurg.formats.samplesheet import output_stem
|
||||
|
||||
_COMP = bytes.maketrans(b"ACGTDI", b"TGCADI")
|
||||
LARGE_OUTPUT = 100_000_000
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ inputs
|
||||
@dataclass
|
||||
class GtcSample:
|
||||
sample_id: str
|
||||
sample_name: str
|
||||
path: str
|
||||
reported_sex: str = "U"
|
||||
|
||||
|
||||
def gtc_samples(gtcs: Path | None, sheet_path: Path | None, rep: Reporter) -> list[GtcSample]:
|
||||
"""Sheet order when a sheet is given, else sorted sample IDs (OUT-076)."""
|
||||
from simurg.formats.discovery import discover_gtcs
|
||||
from simurg.formats.samplesheet import parse_samplesheet
|
||||
|
||||
sheet = parse_samplesheet(sheet_path) if sheet_path else None
|
||||
root = gtcs or (sheet.root_folder if sheet else None)
|
||||
if root is None:
|
||||
raise SimurgError("E_NO_INPUT", "give --gtcs or a sample sheet with a RootFolder")
|
||||
found = discover_gtcs(root)
|
||||
if sheet is None:
|
||||
return [GtcSample(k, k, str(v)) for k, v in sorted(found.items())]
|
||||
out = []
|
||||
for r in sheet.rows:
|
||||
if r.sample_id in found:
|
||||
out.append(GtcSample(r.sample_id, r.sample_name or r.sample_id, str(found[r.sample_id]), r.reported_sex))
|
||||
else:
|
||||
rep.warn("W_QC_UNMATCHED_SAMPLES", f"{r.sample_id}: no GTC found", sample_id=r.sample_id)
|
||||
if not out:
|
||||
raise SimurgError("E_NO_INPUT", "none of the sheet's samples has a GTC")
|
||||
return out
|
||||
|
||||
|
||||
def chrom_rank(chrom: str) -> tuple[int, str]:
|
||||
c = chrom[3:] if chrom.lower().startswith("chr") else chrom
|
||||
if c.isdigit():
|
||||
return (int(c) if c != "0" else 1000, "")
|
||||
return ({"X": 23, "XY": 24, "Y": 25, "MT": 26, "M": 26}.get(c, 100), c)
|
||||
|
||||
|
||||
def genome_order(m: Manifest) -> npt.NDArray[np.int64]:
|
||||
"""Loci by (chromosome, position, name); unmapped (Chr 0) last."""
|
||||
keys = [
|
||||
((1000, "") if c in ("0", "") or p <= 0 else chrom_rank(c), int(p), n)
|
||||
for c, p, n in zip(m.chroms, m.map_infos.tolist(), m.names, strict=True)
|
||||
]
|
||||
return np.array(sorted(range(len(keys)), key=keys.__getitem__), dtype=np.int64)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ bedGraph
|
||||
def _chrom_label(chrom: str, style: str) -> str:
|
||||
c = "X" if chrom == "XY" else chrom
|
||||
if style == "ucsc":
|
||||
c = "M" if c == "MT" else c
|
||||
return c if c.startswith("chr") else f"chr{c}"
|
||||
return c
|
||||
|
||||
|
||||
def moving_median(v: npt.NDArray[np.float64], k: int) -> npt.NDArray[np.float64]:
|
||||
if k <= 0 or len(v) == 0:
|
||||
return v.copy()
|
||||
pad = np.pad(v, k, mode="edge")
|
||||
win = np.lib.stride_tricks.sliding_window_view(pad, 2 * k + 1)
|
||||
return np.nanmedian(win, axis=1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BedgraphLayout:
|
||||
"""Sample-independent part of OUT-050: loci in output order with their labels (computed once per manifest)."""
|
||||
|
||||
order: npt.NDArray[np.int64] # mapped loci, genome order
|
||||
chrom: list[str] # output chromosome label per ordered locus
|
||||
pos: npt.NDArray[np.int64]
|
||||
|
||||
|
||||
def bedgraph_layout(m: Manifest, style: str) -> BedgraphLayout:
|
||||
order = genome_order(m)
|
||||
mapped = np.array([c not in ("0", "") for c in m.chroms]) & (m.map_infos > 0)
|
||||
order = order[mapped[order]]
|
||||
return BedgraphLayout(order, [_chrom_label(m.chroms[i], style) for i in order.tolist()], m.map_infos[order])
|
||||
|
||||
|
||||
def bedgraph_rows(m: Manifest, g: Any, layout: BedgraphLayout, which: str) -> list[tuple[str, int, float]]:
|
||||
"""OUT-050 rows (chrom, pos, value) for one track, in output order, same-position probes averaged."""
|
||||
vals = g.b_allele_freqs if which == "baf" else g.logr_ratios
|
||||
if len(vals) != m.num_loci:
|
||||
raise SimurgError("E_GTC_UNSUITABLE", "GTC has no BAF/LRR (GTC v3?)")
|
||||
o = layout.order
|
||||
v = vals[o].astype(np.float64)
|
||||
keep = ~((g.raw_x[o] == 0) & (g.raw_y[o] == 0)) & np.isfinite(v)
|
||||
if which == "baf":
|
||||
keep &= ~m.intensity_only[o]
|
||||
idx = np.flatnonzero(keep)
|
||||
if not len(idx):
|
||||
return []
|
||||
chrom = [layout.chrom[i] for i in idx.tolist()]
|
||||
pos = layout.pos[idx]
|
||||
new = np.ones(len(idx), np.bool_)
|
||||
new[1:] = (pos[1:] != pos[:-1]) | np.array([chrom[k] != chrom[k - 1] for k in range(1, len(chrom))])
|
||||
starts = np.flatnonzero(new)
|
||||
sums = np.add.reduceat(v[idx], starts)
|
||||
counts = np.diff(np.append(starts, len(idx)))
|
||||
means = sums / counts
|
||||
return [(chrom[k], int(pos[k]), float(x)) for k, x in zip(starts.tolist(), means.tolist(), strict=True)]
|
||||
|
||||
|
||||
def write_bedgraph(path: Path, name: str, which: str, rows: Sequence[tuple[str, int, float]], gz: bool) -> None:
|
||||
limits = "0:1" if which == "baf" else "-2:2"
|
||||
head = (
|
||||
f'track type=bedGraph name="{name} {which.upper()}" description="{PRODUCT_NAME} {__version__}" '
|
||||
f"visibility=full autoScale=off viewLimits={limits}\n"
|
||||
)
|
||||
body = "".join(f"{c}\t{p - 1}\t{p}\t{v:.6g}\n" for c, p, v in rows)
|
||||
data = (head + body).encode()
|
||||
with atomic_open(path, "wb") as fh:
|
||||
cast(IO[bytes], fh).write(gzip.compress(data, mtime=0) if gz else data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BedgraphOptions:
|
||||
out: Path
|
||||
bpm: Path
|
||||
gtcs: Path | None = None
|
||||
sheet: Path | None = None
|
||||
gzip: bool = True
|
||||
chrom_style: str = "manifest"
|
||||
smoothed_lrr: int = 0
|
||||
force: bool = False
|
||||
argv: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def run_export_bedgraph(opts: BedgraphOptions, rep: Reporter) -> int:
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
if opts.chrom_style not in ("manifest", "ucsc"):
|
||||
raise SimurgError("E_CONFIG_KEY", "--chrom-style must be manifest or ucsc")
|
||||
rm = RunManifest(command="export bedgraph", argv=opts.argv)
|
||||
samples = gtc_samples(opts.gtcs, opts.sheet, rep)
|
||||
ext = ".bedgraph.gz" if opts.gzip else ".bedgraph"
|
||||
outs = []
|
||||
for s in samples:
|
||||
stem = output_stem(s.sample_name)
|
||||
outs += [opts.out / f"{stem}.baf{ext}", opts.out / f"{stem}.lrr{ext}"]
|
||||
if opts.smoothed_lrr:
|
||||
outs.append(opts.out / f"{stem}.lrr_smoothed{ext}")
|
||||
check_outputs_absent(outs, force=opts.force)
|
||||
m = read_bpm(opts.bpm)
|
||||
layout = bedgraph_layout(m, opts.chrom_style)
|
||||
opts.out.mkdir(parents=True, exist_ok=True)
|
||||
failed = 0
|
||||
for s in samples:
|
||||
try:
|
||||
g = read_gtc(s.path)
|
||||
stem = output_stem(s.sample_name)
|
||||
for which in ("baf", "lrr"):
|
||||
rows = bedgraph_rows(m, g, layout, which)
|
||||
write_bedgraph(opts.out / f"{stem}.{which}{ext}", s.sample_name, which, rows, opts.gzip)
|
||||
if which == "lrr" and opts.smoothed_lrr:
|
||||
sm: list[tuple[str, int, float]] = []
|
||||
for chrom in dict.fromkeys(c for c, _, _ in rows):
|
||||
sub = [r for r in rows if r[0] == chrom]
|
||||
med = moving_median(np.array([r[2] for r in sub]), opts.smoothed_lrr)
|
||||
sm += [(c, p, float(v)) for (c, p, _), v in zip(sub, med, strict=True)]
|
||||
write_bedgraph(opts.out / f"{stem}.lrr_smoothed{ext}", s.sample_name, "lrr", sm, opts.gzip)
|
||||
except SimurgError as e:
|
||||
failed += 1
|
||||
rep.error(e.code, f"{s.sample_id}: {e.message}", sample_id=s.sample_id)
|
||||
rm.status = "completed"
|
||||
rm.summary = {"samples": len(samples), "failed": failed}
|
||||
rm.write(opts.out / "run_manifest.export_bedgraph.json")
|
||||
return ExitCode.OK if not failed else (ExitCode.PARTIAL if failed < len(samples) else ExitCode.ALL_FAILED)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ final report
|
||||
FR_COLUMNS = ["SNP Name", "SNP", "Sample ID", "Allele1 - Top", "Allele2 - Top", "Allele1 - Forward",
|
||||
"Allele2 - Forward", "Allele1 - Plus", "Allele2 - Plus", "GC Score", "GT Score", "Log R Ratio",
|
||||
"B Allele Freq", "Chr", "Position"] # fmt: skip
|
||||
|
||||
|
||||
def allele_letters(m: Manifest) -> dict[str, list[tuple[str, str]]]:
|
||||
"""Per locus (A, B) letters on the TOP, Forward (source) and Plus (reference) strands."""
|
||||
top, fwd, plus = [], [], []
|
||||
for snp, ilmn, src, ref in zip(m.snps, m.ilmn_strands, m.source_strands, m.ref_strands, strict=True):
|
||||
mm = re.match(r"^\[([^/\]]+)/([^\]]+)\]$", snp.strip())
|
||||
a, b = (mm.group(1), mm.group(2)) if mm else ("-", "-")
|
||||
|
||||
def comp(x: str) -> str:
|
||||
return x.encode().translate(_COMP).decode() if len(x) == 1 else x
|
||||
|
||||
d = (ilmn or "")[:1].upper()
|
||||
top.append((a, b) if d in ("T", "P", "") or len(a) > 1 else (comp(a), comp(b)))
|
||||
same_src = (src or "")[:1].upper() == d or not src
|
||||
fwd.append((a, b) if same_src else (comp(a), comp(b)))
|
||||
plus.append((a, b) if (ref or "+").strip() == "+" else (comp(a), comp(b)))
|
||||
return {"top": top, "fwd": fwd, "plus": plus}
|
||||
|
||||
|
||||
def _pair(call: int, ab: tuple[str, str]) -> tuple[str, str]:
|
||||
a, b = ab
|
||||
return {1: (a, a), 2: (a, b), 3: (b, b)}.get(call, ("-", "-"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class FinalReportOptions:
|
||||
out: Path
|
||||
bpm: Path
|
||||
gtcs: Path | None = None
|
||||
sheet: Path | None = None
|
||||
csv: Path | None = None
|
||||
egt: Path | None = None
|
||||
matrix: bool = False
|
||||
force: bool = False
|
||||
argv: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def run_export_finalreport(opts: FinalReportOptions, rep: Reporter) -> int:
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
rm = RunManifest(command="export finalreport", argv=opts.argv)
|
||||
samples = gtc_samples(opts.gtcs, opts.sheet, rep)
|
||||
target = opts.out / ("final_report_matrix.tsv.gz" if opts.matrix else "final_report.tsv.gz")
|
||||
check_outputs_absent([target], force=opts.force)
|
||||
m = read_bpm(opts.bpm)
|
||||
if opts.csv is not None:
|
||||
from simurg.formats.csv_manifest import read_csv_manifest
|
||||
|
||||
cm = read_csv_manifest(opts.csv, load_sequences=False)
|
||||
m.source_strands, m.ref_strands = cm.source_strands, cm.ref_strands
|
||||
cells = m.num_loci * len(samples)
|
||||
if cells > LARGE_OUTPUT:
|
||||
rep.warn("W_LARGE_OUTPUT", f"{cells:,} rows (~{cells * 90 / 1e9:.1f} GB uncompressed); consider --matrix")
|
||||
gt_score = np.full(m.num_loci, np.nan)
|
||||
if opts.egt is not None:
|
||||
from simurg.formats.egt import read_egt
|
||||
|
||||
egt = read_egt(opts.egt)
|
||||
idx = egt.name_index()
|
||||
for i, nm in enumerate(m.names):
|
||||
j = idx.get(nm)
|
||||
if j is not None:
|
||||
gt_score[i] = float(egt.gentrain_score[j])
|
||||
letters = allele_letters(m)
|
||||
order = genome_order(m).tolist()
|
||||
opts.out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def f4(v: float) -> str:
|
||||
return "NaN" if math.isnan(v) else f"{v:.4f}"
|
||||
|
||||
with atomic_open(target, "wb") as raw:
|
||||
gz = gzip.GzipFile(fileobj=cast(IO[bytes], raw), mode="wb", mtime=0)
|
||||
fh = io.TextIOWrapper(gz, encoding="utf-8", newline="\n")
|
||||
fh.write("#simurg_schema=simurg.finalreport/1\n")
|
||||
if opts.matrix:
|
||||
fh.write("\t".join(["SNP Name", "Chr", "Position", *[s.sample_name for s in samples]]) + "\n")
|
||||
calls = []
|
||||
for s in samples:
|
||||
calls.append(read_gtc(s.path).genotypes)
|
||||
for i in order:
|
||||
row = [m.names[i], m.chroms[i], str(int(m.map_infos[i]))]
|
||||
row += ["".join(_pair(int(c[i]), letters["top"][i])) for c in calls]
|
||||
fh.write("\t".join(row) + "\n")
|
||||
else:
|
||||
fh.write("\t".join(FR_COLUMNS) + "\n")
|
||||
for s in samples:
|
||||
g = read_gtc(s.path)
|
||||
has = len(g.b_allele_freqs) == m.num_loci
|
||||
bc = g.base_calls
|
||||
for i in order:
|
||||
c = int(g.genotypes[i])
|
||||
f, p = (_pair(c, letters[k][i]) for k in ("fwd", "plus"))
|
||||
t = (chr(bc[i, 0]), chr(bc[i, 1])) if c else ("-", "-") # TOP-strand base calls (OUT-071)
|
||||
baf = float(g.b_allele_freqs[i]) if has and not m.intensity_only[i] else math.nan
|
||||
lrr = float(g.logr_ratios[i]) if has else math.nan
|
||||
fh.write(
|
||||
f"{m.names[i]}\t{m.snps[i]}\t{s.sample_id}\t{t[0]}\t{t[1]}\t{f[0]}\t{f[1]}\t{p[0]}\t{p[1]}\t"
|
||||
f"{f4(float(g.genotype_scores[i]))}\t{f4(gt_score[i])}\t{f4(lrr)}\t{f4(baf)}\t"
|
||||
f"{m.chroms[i]}\t{int(m.map_infos[i])}\n"
|
||||
)
|
||||
fh.flush()
|
||||
fh.detach()
|
||||
gz.close()
|
||||
rm.status = "completed"
|
||||
rm.summary = {"samples": len(samples), "loci": m.num_loci}
|
||||
rm.write(opts.out / "run_manifest.export_finalreport.json")
|
||||
return ExitCode.OK
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ PLINK
|
||||
PLINK_MAGIC = bytes([0x6C, 0x1B, 0x01])
|
||||
|
||||
|
||||
def parse_vcf_records(path: Path) -> Iterator[list[str]]:
|
||||
opener = gzip.open if str(path).endswith(".gz") else open
|
||||
with opener(path, "rt") as fh:
|
||||
for line in fh:
|
||||
if not line.startswith("#"):
|
||||
yield line.rstrip("\n").split("\t")
|
||||
|
||||
|
||||
def bed_code(gt: str, alt_index: int) -> int:
|
||||
"""PLINK 1 two-bit code with A1 = the split ALT, A2 = REF: 0 hom A1, 1 missing, 2 het, 3 hom A2."""
|
||||
if gt in (".", "./."):
|
||||
return 1
|
||||
al = gt.replace("|", "/").split("/")
|
||||
if "." in al:
|
||||
return 1
|
||||
n_alt = sum(int(x) == alt_index for x in al)
|
||||
if len(al) == 1:
|
||||
n_alt *= 2
|
||||
return {0: 3, 1: 2, 2: 0}[n_alt]
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlinkOptions:
|
||||
out: Path
|
||||
vcfs: Path
|
||||
sheet: Path | None = None
|
||||
gtcs: Path | None = None
|
||||
multiallelic: str = "split"
|
||||
prefix: str = "cohort"
|
||||
force: bool = False
|
||||
argv: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def run_export_plink(opts: PlinkOptions, rep: Reporter) -> int:
|
||||
"""OUT-043: SimurgArray per-sample VCFs (identical record order) → PLINK 1 .bed/.bim/.fam."""
|
||||
if opts.multiallelic not in ("split", "drop"):
|
||||
raise SimurgError("E_CONFIG_KEY", "--multiallelic must be split or drop")
|
||||
rm = RunManifest(command="export plink", argv=opts.argv)
|
||||
vcfs = sorted(p for p in opts.vcfs.glob("*.snv.vcf*") if not p.name.endswith((".tbi", ".csi")))
|
||||
if not vcfs:
|
||||
raise SimurgError("E_NO_INPUT", f"no *.snv.vcf(.gz) files in {opts.vcfs}")
|
||||
base = opts.out / opts.prefix
|
||||
check_outputs_absent([base.with_suffix(x) for x in (".bed", ".bim", ".fam")], force=opts.force)
|
||||
names, sex = [], {}
|
||||
for v in vcfs:
|
||||
opener = gzip.open if str(v).endswith(".gz") else open
|
||||
with opener(v, "rt") as fh:
|
||||
for line in fh:
|
||||
if line.startswith("#CHROM"):
|
||||
names.append(line.rstrip("\n").split("\t")[9])
|
||||
break
|
||||
if opts.gtcs is not None:
|
||||
from simurg.formats.discovery import discover_gtcs
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
for key, gpath in discover_gtcs(opts.gtcs).items():
|
||||
sex[key] = read_gtc(gpath, check_write_complete=False).gender_char
|
||||
if opts.sheet is not None:
|
||||
from simurg.formats.samplesheet import parse_samplesheet
|
||||
|
||||
for r in parse_samplesheet(opts.sheet).rows:
|
||||
if r.reported_sex in ("M", "F"):
|
||||
sex[r.sample_name or r.sample_id] = sex[r.sample_id] = r.reported_sex
|
||||
streams = [parse_vcf_records(v) for v in vcfs]
|
||||
n = len(vcfs)
|
||||
nbytes = (n + 3) // 4
|
||||
opts.out.mkdir(parents=True, exist_ok=True)
|
||||
n_var = dropped = 0
|
||||
with (
|
||||
atomic_open(base.with_suffix(".bed"), "wb") as bed_raw,
|
||||
atomic_open(base.with_suffix(".bim"), "w", "utf-8") as bim_raw,
|
||||
):
|
||||
bed, bim = cast(IO[bytes], bed_raw), cast(IO[str], bim_raw)
|
||||
bed.write(PLINK_MAGIC)
|
||||
for recs in zip(*streams, strict=True):
|
||||
r0 = recs[0]
|
||||
if any(r[:5] != r0[:5] for r in recs[1:]):
|
||||
raise SimurgError("E_GTC_UNSUITABLE", f"VCF record order differs between samples at {r0[0]}:{r0[1]}")
|
||||
alts = [] if r0[4] == "." else r0[4].split(",")
|
||||
if len(alts) > 1 and opts.multiallelic == "drop":
|
||||
dropped += 1
|
||||
continue
|
||||
gts = [r[9].split(":", 1)[0] for r in recs]
|
||||
for k, alt in enumerate(alts, start=1):
|
||||
codes = [bed_code(g, k) for g in gts]
|
||||
buf = bytearray(nbytes)
|
||||
for j, c in enumerate(codes):
|
||||
buf[j >> 2] |= c << (2 * (j & 3))
|
||||
bed.write(bytes(buf))
|
||||
vid = r0[2] if len(alts) == 1 else f"{r0[2]}:{alt}"
|
||||
bim.write(f"{r0[0]}\t{vid}\t0\t{r0[1]}\t{alt}\t{r0[3]}\n")
|
||||
n_var += 1
|
||||
fam = "".join(f"{s}\t{s}\t0\t0\t{ {'M': 1, 'F': 2}.get(sex.get(s, 'U'), 0) }\t-9\n" for s in names)
|
||||
with atomic_open(base.with_suffix(".fam"), "w", "utf-8") as fam_raw:
|
||||
cast(IO[str], fam_raw).write(fam)
|
||||
rm.status = "completed"
|
||||
rm.summary = {"samples": n, "variants": n_var, "multiallelic_dropped": dropped}
|
||||
rm.write(opts.out / "run_manifest.export_plink.json")
|
||||
rep.info(f"PLINK fileset {base}.bed/.bim/.fam: {n} samples, {n_var} variants")
|
||||
return ExitCode.OK
|
||||
@@ -308,26 +308,27 @@ class SampleInfo:
|
||||
class CohortStore:
|
||||
"""Disk-backed (samples × loci) matrices: genotypes (uint8), BAF and GenCall scores (float32, optional)."""
|
||||
|
||||
def __init__(self, root: Path, n_samples: int, n_loci: int, *, baf: bool, scores: bool) -> None:
|
||||
def __init__(self, root: Path, n_samples: int, n_loci: int, *, baf: bool, scores: bool, lrr: bool = False) -> None:
|
||||
self.root, self.shape = root, (n_samples, n_loci)
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
self.gt = np.lib.format.open_memmap(root / "gt.npy", "w+", np.uint8, self.shape)
|
||||
self.baf = np.lib.format.open_memmap(root / "baf.npy", "w+", np.float32, self.shape) if baf else None
|
||||
self.score = np.lib.format.open_memmap(root / "score.npy", "w+", np.float32, self.shape) if scores else None
|
||||
self.lrr = np.lib.format.open_memmap(root / "lrr.npy", "w+", np.float32, self.shape) if lrr else None
|
||||
|
||||
def close(self) -> None:
|
||||
for name in ("gt", "baf", "score"):
|
||||
for name in ("gt", "baf", "score", "lrr"):
|
||||
a = getattr(self, name)
|
||||
if a is not None:
|
||||
a.flush()
|
||||
self.gt = self.baf = self.score = None # type: ignore[assignment]
|
||||
self.gt = self.baf = self.score = self.lrr = None # type: ignore[assignment]
|
||||
shutil.rmtree(self.root, ignore_errors=True)
|
||||
|
||||
|
||||
def _load_one(task: tuple[int, str, str, str, bool, bool]) -> SampleInfo:
|
||||
def _load_one(task: tuple[int, str, str, str, bool, bool, bool]) -> SampleInfo:
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
i, sid, path, root, want_baf, want_scores = task
|
||||
i, sid, path, root, want_baf, want_scores, want_lrr = task
|
||||
info = SampleInfo(sid, path)
|
||||
try:
|
||||
g = read_gtc(path)
|
||||
@@ -347,6 +348,10 @@ def _load_one(task: tuple[int, str, str, str, bool, bool]) -> SampleInfo:
|
||||
info.has_baf = len(g.b_allele_freqs) == gt.shape[1]
|
||||
baf[i] = g.b_allele_freqs if info.has_baf else np.nan
|
||||
baf.flush()
|
||||
if want_lrr:
|
||||
lr = np.load(Path(root) / "lrr.npy", mmap_mode="r+")
|
||||
lr[i] = g.logr_ratios if len(g.logr_ratios) == gt.shape[1] else np.nan
|
||||
lr.flush()
|
||||
if want_scores:
|
||||
sc = np.load(Path(root) / "score.npy", mmap_mode="r+")
|
||||
sc[i] = g.genotype_scores
|
||||
@@ -359,12 +364,19 @@ def _load_one(task: tuple[int, str, str, str, bool, bool]) -> SampleInfo:
|
||||
|
||||
|
||||
def load_cohort(
|
||||
samples: Sequence[tuple[str, str]], n_loci: int, tmp: Path, *, baf: bool, scores: bool, workers: int
|
||||
samples: Sequence[tuple[str, str]],
|
||||
n_loci: int,
|
||||
tmp: Path,
|
||||
*,
|
||||
baf: bool,
|
||||
scores: bool,
|
||||
workers: int,
|
||||
lrr: bool = False,
|
||||
) -> tuple[CohortStore, list[SampleInfo]]:
|
||||
"""Read GTCs (sample_id, path) into a :class:`CohortStore` with a process pool."""
|
||||
store = CohortStore(tmp, len(samples), n_loci, baf=baf, scores=scores)
|
||||
store = CohortStore(tmp, len(samples), n_loci, baf=baf, scores=scores, lrr=lrr)
|
||||
store.gt.flush()
|
||||
tasks = [(i, sid, p, str(tmp), baf, scores) for i, (sid, p) in enumerate(samples)]
|
||||
tasks = [(i, sid, p, str(tmp), baf, scores, lrr) for i, (sid, p) in enumerate(samples)]
|
||||
if workers <= 1 or len(tasks) <= 1:
|
||||
infos = [_load_one(t) for t in tasks]
|
||||
else:
|
||||
|
||||
@@ -0,0 +1,656 @@
|
||||
"""``simurg gtc-to-vcf``: GTC + manifest + FASTA → one bgzipped, indexed SNV VCF per sample (spec CLI-022,
|
||||
OUT-010 … OUT-042, IN-060 … IN-064; plan §8).
|
||||
|
||||
The manifest is mapped to plus-strand sites once (``simurg.harmonize``); each sample is then written by a worker
|
||||
process with vectorised genotype lookup for single-probe sites and the OUT-031 algorithm for merged sites.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import math
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
from collections import Counter
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from simurg import PRODUCT_NAME, __version__
|
||||
from simurg.core.atomic import atomic_open, check_outputs_absent
|
||||
from simurg.core.errors import ExitCode, SimurgError
|
||||
from simurg.core.log import Reporter
|
||||
from simurg.core.provenance import RunManifest, sha256_file
|
||||
from simurg.formats.bgzf import BgzfWriter, TabixBuilder, write_tbi
|
||||
from simurg.harmonize.combine import genotype as combine_genotype
|
||||
from simurg.harmonize.combine import gt_text
|
||||
from simurg.harmonize.sites import AuditRow, ManifestView, Site, group_sites, map_locus
|
||||
|
||||
PAR = { # 1-based inclusive pseudo-autosomal intervals on X (IN-065, human)
|
||||
"GRCh38": ((10_001, 2_781_479), (155_701_383, 156_030_895)),
|
||||
"GRCh37": ((60_001, 2_699_520), (154_931_044, 155_260_560)),
|
||||
}
|
||||
CLASS_AUTO, CLASS_X, CLASS_Y, CLASS_MT, CLASS_PAR = 0, 1, 2, 3, 4
|
||||
|
||||
|
||||
@dataclass
|
||||
class VcfOptions:
|
||||
out: Path
|
||||
bpm: Path
|
||||
fasta: Path
|
||||
gtcs: Path | None = None
|
||||
sheet: Path | None = None
|
||||
csv: Path | None = None
|
||||
bgzip: bool = True
|
||||
genome_cache: bool = True
|
||||
exclude_loci: Path | None = None
|
||||
include_loci: Path | None = None
|
||||
unsquash_duplicates: bool = False
|
||||
inf1_nocall_informative: bool = False
|
||||
drop_ambiguous: bool = False
|
||||
filter_gs: float | None = None
|
||||
contig_alias: str = "strict"
|
||||
ploidy_policy: str = "sex-aware"
|
||||
threads: int = 0
|
||||
merge: bool = False
|
||||
force: bool = False
|
||||
argv: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ site table
|
||||
@dataclass
|
||||
class SiteTable:
|
||||
"""Everything a worker needs, precomputed once (picklable)."""
|
||||
|
||||
prefix: list[bytes] # "CHROM\tPOS\tID\tREF\tALT\t.\t" per site
|
||||
single_rows: npt.NDArray[np.int64] # site indices of single-probe sites
|
||||
single_locus: npt.NDArray[np.int64] # their manifest locus index
|
||||
single_gt: npt.NDArray[np.object_] # (n_single, 2 ploidy, 4 calls) GT text
|
||||
single_flip: npt.NDArray[np.bool_] # BAF must be 1 − BAF
|
||||
merged_sites: list[Site]
|
||||
merged_rows: npt.NDArray[np.int64]
|
||||
site_class: npt.NDArray[np.int8]
|
||||
contigs: list[str]
|
||||
header: str
|
||||
site_beg: npt.NDArray[np.int64]
|
||||
site_end: npt.NDArray[np.int64]
|
||||
site_contig: list[str]
|
||||
site_id: list[str] = field(default_factory=list)
|
||||
site_alleles: list[list[str]] = field(default_factory=list)
|
||||
single_ab: npt.NDArray[np.int16] | None = None # (n_single, 2) allele indices of the probe's A and B alleles
|
||||
inf1_nc_informative: bool = False
|
||||
ploidy_policy: str = "sex-aware"
|
||||
filter_gs: float | None = None
|
||||
|
||||
|
||||
def _site_class(contig: str, pos: int, par_flag: bool, build: str) -> int:
|
||||
c = contig[3:] if contig.lower().startswith("chr") else contig
|
||||
if c in ("MT", "M"):
|
||||
return CLASS_MT
|
||||
if c == "Y":
|
||||
return CLASS_Y
|
||||
if c == "X":
|
||||
if par_flag or any(a <= pos <= b for a, b in PAR.get(build or "GRCh38", ())):
|
||||
return CLASS_PAR
|
||||
return CLASS_X
|
||||
return CLASS_AUTO
|
||||
|
||||
|
||||
def ploidy_for(cls: npt.NDArray[np.int8], sex: str, policy: str) -> npt.NDArray[np.int8]:
|
||||
"""OUT-018 per site; 0 marks 'haploid and always missing' (Y of female samples)."""
|
||||
p = np.full(len(cls), 2, dtype=np.int8)
|
||||
p[cls == CLASS_MT] = 1
|
||||
if policy == "diploid":
|
||||
return p
|
||||
p[cls == CLASS_Y] = 0 if sex == "F" else 1
|
||||
if sex == "M":
|
||||
p[cls == CLASS_X] = 1
|
||||
return p
|
||||
|
||||
|
||||
def read_locus_list(path: Path) -> list[str]:
|
||||
out = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.split("#", 1)[0].strip()
|
||||
if line:
|
||||
out.append(line)
|
||||
return out
|
||||
|
||||
|
||||
def build_sites(
|
||||
opts: VcfOptions, rep: Reporter, rm: RunManifest
|
||||
) -> tuple[SiteTable, list[AuditRow], Counter[str], Any]:
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.formats.fasta import ContigMap, Genome
|
||||
|
||||
bpm = read_bpm(opts.bpm)
|
||||
csvm = None
|
||||
if opts.csv is not None:
|
||||
from simurg.formats.csv_manifest import compare_manifests, read_csv_manifest
|
||||
|
||||
csvm = read_csv_manifest(opts.csv)
|
||||
problems = compare_manifests(bpm, csvm)
|
||||
if problems:
|
||||
raise SimurgError("E_BPM_CSV_MISMATCH", "BPM and CSV manifests differ:\n " + "\n ".join(problems))
|
||||
genome = Genome(opts.fasta, use_cache=opts.genome_cache)
|
||||
for n in genome.notes:
|
||||
rep.info(n)
|
||||
builds = Counter(bpm.genome_builds)
|
||||
mbuild = builds.most_common(1)[0][0] if builds else ""
|
||||
want = {"37": "GRCh37", "38": "GRCh38"}.get(mbuild.strip(), mbuild.strip())
|
||||
if genome.build and want and want != genome.build:
|
||||
raise SimurgError(
|
||||
"E_BUILD_MISMATCH", f"manifest genome build {want} but FASTA looks like {genome.build} ({opts.fasta})"
|
||||
)
|
||||
build = genome.build or want
|
||||
src = csvm or bpm
|
||||
mv = ManifestView(
|
||||
bpm.names,
|
||||
bpm.snps,
|
||||
bpm.chroms,
|
||||
bpm.map_infos.tolist(),
|
||||
src.ilmn_strands,
|
||||
src.source_strands,
|
||||
src.ref_strands if any(src.ref_strands) else bpm.ref_strands,
|
||||
csvm.source_seqs if csvm else [""] * bpm.num_loci,
|
||||
bpm.assay_types.tolist(),
|
||||
bpm.intensity_only.tolist(),
|
||||
)
|
||||
if csvm is None and any(("D" in s) for s in bpm.snps):
|
||||
rep.error("W_SKIP_INDEL_CONTEXT", "no CSV manifest: indels are skipped (give --csv with SourceSeq)")
|
||||
names = set(bpm.names)
|
||||
excl: set[str] = set()
|
||||
incl: set[str] | None = None
|
||||
for path, kind in ((opts.exclude_loci, "exclude"), (opts.include_loci, "include")):
|
||||
if path is None:
|
||||
continue
|
||||
lst = read_locus_list(path)
|
||||
unknown = [x for x in lst if x not in names]
|
||||
if unknown:
|
||||
rep.warn("W_UNKNOWN_LOCI", f"{len(unknown)} name(s) in --{kind}-loci are not in the manifest")
|
||||
if kind == "exclude":
|
||||
excl = set(lst)
|
||||
else:
|
||||
incl = set(lst)
|
||||
cmap = ContigMap([e.name for e in genome.entries], opts.contig_alias)
|
||||
probes = []
|
||||
audit: list[AuditRow] = []
|
||||
skips: Counter[str] = Counter()
|
||||
for i in range(bpm.num_loci):
|
||||
nm = bpm.names[i]
|
||||
if nm in excl or (incl is not None and nm not in incl):
|
||||
a = AuditRow(nm, bpm.snps[i], "", "", "", "", "", False, False, "W_SKIP_EXCLUDED", build, "")
|
||||
p = None
|
||||
else:
|
||||
p, a = map_locus(i, mv, genome.fetch, cmap.get, build)
|
||||
if p is not None and opts.drop_ambiguous and p.confidence == "low":
|
||||
a.reason, p = "W_SKIP_AMBIGUOUS", None
|
||||
audit.append(a)
|
||||
if p is None or a.reason.startswith("W_"):
|
||||
skips[a.reason] += 1
|
||||
if p is not None:
|
||||
probes.append(p)
|
||||
for code, count in sorted(skips.items()):
|
||||
rep.warn(code, f"{count} locus/loci: {code}")
|
||||
order = {e.name: k for k, e in enumerate(genome.entries)}
|
||||
sites = group_sites(probes, order, unsquash=opts.unsquash_duplicates)
|
||||
used = {s.contig for s in sites}
|
||||
contigs = [e.name for e in genome.entries if e.name in used]
|
||||
# static per-site text and single-probe genotype tables
|
||||
prefix, cls = [], np.zeros(len(sites), np.int8)
|
||||
single_rows, single_locus, single_flip, gts, single_ab = [], [], [], [], []
|
||||
merged_sites, merged_rows = [], []
|
||||
beg = np.zeros(len(sites), np.int64)
|
||||
end = np.zeros(len(sites), np.int64)
|
||||
for k, s in enumerate(sites):
|
||||
prefix.append(f"{s.contig}\t{s.pos}\t{s.id}\t{s.ref}\t{','.join(s.alts) if s.alts else '.'}\t.\t".encode())
|
||||
cls[k] = _site_class(s.contig, s.pos, any(p.par for p in s.probes), build)
|
||||
beg[k], end[k] = s.pos - 1, s.pos - 1 + len(s.ref)
|
||||
if len(s.probes) == 1:
|
||||
from simurg.harmonize.combine import single
|
||||
|
||||
single_rows.append(k)
|
||||
single_locus.append(s.probes[0].index)
|
||||
al_a, al_b = s.probes[0].alleles
|
||||
single_flip.append(al_a != s.ref and al_b == s.ref)
|
||||
single_ab.append((s.allele_index(al_a), s.allele_index(al_b)))
|
||||
gts.append([[gt_text(single(s, c, pl), pl) for c in range(4)] for pl in (1, 2)])
|
||||
else:
|
||||
merged_sites.append(s)
|
||||
merged_rows.append(k)
|
||||
header = vcf_header(opts, genome, build, contigs, rm)
|
||||
table = SiteTable(
|
||||
prefix=prefix,
|
||||
single_rows=np.array(single_rows, np.int64),
|
||||
single_locus=np.array(single_locus, np.int64),
|
||||
single_gt=np.array(gts, dtype=object).reshape(len(single_rows), 2, 4),
|
||||
single_flip=np.array(single_flip, np.bool_),
|
||||
merged_sites=merged_sites,
|
||||
merged_rows=np.array(merged_rows, np.int64),
|
||||
site_class=cls,
|
||||
contigs=contigs,
|
||||
header=header,
|
||||
site_beg=beg,
|
||||
site_end=end,
|
||||
site_contig=[s.contig for s in sites],
|
||||
site_id=[s.id for s in sites],
|
||||
site_alleles=[[s.ref, *s.alts] for s in sites],
|
||||
single_ab=np.array(single_ab, np.int16).reshape(-1, 2),
|
||||
inf1_nc_informative=opts.inf1_nocall_informative,
|
||||
ploidy_policy=opts.ploidy_policy,
|
||||
filter_gs=opts.filter_gs,
|
||||
)
|
||||
return table, audit, skips, genome
|
||||
|
||||
|
||||
def _sha(path: Path | None) -> str:
|
||||
return sha256_file(path) if path else ""
|
||||
|
||||
|
||||
def vcf_header(opts: VcfOptions, genome: Any, build: str, contigs: list[str], rm: RunManifest) -> str:
|
||||
"""OUT-011 (without the ##SAMPLE and #CHROM lines, which are per sample)."""
|
||||
|
||||
def lst(p: Path | None) -> str:
|
||||
return "none" if p is None else f"{p.name}:{hashlib.sha256(p.read_bytes()).hexdigest()[:16]}"
|
||||
|
||||
lines = [
|
||||
"##fileformat=VCFv4.3",
|
||||
f"##source={PRODUCT_NAME} {__version__}",
|
||||
f"##reference=file://{Path(opts.fasta).resolve()}",
|
||||
f"##simurg_genome_build={build}",
|
||||
f"##simurg_manifest={opts.bpm.name};sha256={_sha(opts.bpm)}",
|
||||
]
|
||||
if opts.csv:
|
||||
lines.append(f"##simurg_csv_manifest={opts.csv.name};sha256={_sha(opts.csv)}")
|
||||
lines.append("##simurg_cluster_file=__CLUSTER__")
|
||||
lines.append(
|
||||
"##simurg_options="
|
||||
f"unsquash_duplicates={str(opts.unsquash_duplicates).lower()};exclude_loci={lst(opts.exclude_loci)};"
|
||||
f"include_loci={lst(opts.include_loci)};inf1_nocall_informative={str(opts.inf1_nocall_informative).lower()};"
|
||||
f"auxiliary_loci=none;ploidy_policy={opts.ploidy_policy};contig_alias={opts.contig_alias}"
|
||||
+ (";drop_ambiguous=true" if opts.drop_ambiguous else "")
|
||||
+ (";unknown_sex_policy=diploid_x" if opts.ploidy_policy == "sex-aware" else "")
|
||||
)
|
||||
lines.append(f"##simurg_run={rm.run_id}")
|
||||
for c in contigs:
|
||||
lines.append(f"##contig=<ID={c},length={genome.contig_length(c)}>")
|
||||
lines.append('##FILTER=<ID=PASS,Description="All filters passed">')
|
||||
if opts.filter_gs is not None:
|
||||
lines.append(f'##FILTER=<ID=LowGS,Description="GenCall score below {opts.filter_gs}">')
|
||||
lines += [
|
||||
'##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
|
||||
'##FORMAT=<ID=GS,Number=1,Type=Float,Description="GenCall score. For records merged from several assays, '
|
||||
'the minimum GenCall score over the contributing assays">',
|
||||
'##FORMAT=<ID=BAF,Number=1,Type=Float,Description="B allele frequency, oriented so that 1 means homozygous '
|
||||
'ALT (REF=A, ALT=B relative to the reference genome)">',
|
||||
'##FORMAT=<ID=LRR,Number=1,Type=Float,Description="Log R ratio">',
|
||||
]
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ per sample
|
||||
_W: dict[str, Any] = {}
|
||||
|
||||
|
||||
def _init(table: SiteTable) -> None:
|
||||
_W["t"] = table
|
||||
|
||||
|
||||
def _f4(v: float) -> str:
|
||||
return "." if math.isnan(v) else f"{v:.4f}"
|
||||
|
||||
|
||||
def _g8(v: float) -> str:
|
||||
return "." if math.isnan(v) else f"{v:.8g}"
|
||||
|
||||
|
||||
def write_sample(task: tuple[str, str, str, str, bool]) -> dict[str, Any]:
|
||||
"""Worker: one GTC → one VCF (+ index). Returns counts for the run summary."""
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
sample_id, sample_name, gtc_path, out_path, bgzip = task
|
||||
t: SiteTable = _W["t"]
|
||||
res: dict[str, Any] = {"sample_id": sample_id, "records": 0, "reasons": Counter(), "merge_failures": []}
|
||||
try:
|
||||
g = read_gtc(gtc_path)
|
||||
n = len(t.prefix)
|
||||
sex = g.gender_char
|
||||
ploidy = ploidy_for(t.site_class, sex, t.ploidy_policy)
|
||||
gt = np.empty(n, dtype=object)
|
||||
gs = np.full(n, np.nan, np.float64)
|
||||
baf = np.full(n, np.nan, np.float64)
|
||||
lrr = np.full(n, np.nan, np.float64)
|
||||
has_baf = len(g.b_allele_freqs) == g.num_snps
|
||||
# single-probe sites (vectorised)
|
||||
rows, loc = t.single_rows, t.single_locus
|
||||
calls = g.genotypes[loc].astype(np.int64)
|
||||
pl = ploidy[rows]
|
||||
pidx = np.where(pl == 2, 1, 0)
|
||||
gt[rows] = t.single_gt[np.arange(len(rows)), pidx, calls]
|
||||
gt[rows[pl == 0]] = "."
|
||||
gs[rows] = g.genotype_scores[loc].astype(np.float64)
|
||||
if has_baf:
|
||||
b = g.b_allele_freqs[loc].astype(np.float32)
|
||||
b = np.where(t.single_flip, (np.float32(1.0) - b).astype(np.float32), b)
|
||||
baf[rows] = b.astype(np.float64)
|
||||
lrr[rows] = g.logr_ratios[loc].astype(np.float64)
|
||||
het_haploid = int(((calls == 2) & (pl == 1)).sum())
|
||||
if het_haploid:
|
||||
res["reasons"]["het_in_haploid"] += het_haploid
|
||||
# merged sites
|
||||
for s, k in zip(t.merged_sites, t.merged_rows.tolist(), strict=True):
|
||||
c = tuple(int(g.genotypes[p.index]) for p in s.probes)
|
||||
p_k = int(ploidy[k])
|
||||
gs[k] = float(min(float(g.genotype_scores[p.index]) for p in s.probes))
|
||||
if p_k == 0:
|
||||
gt[k] = "."
|
||||
continue
|
||||
o = combine_genotype(s, c, p_k, t.inf1_nc_informative)
|
||||
gt[k] = gt_text(o, p_k)
|
||||
if o.reason:
|
||||
res["reasons"][o.reason] += 1
|
||||
if o.reason.startswith("W_MERGE") and o.reason != "W_MERGE_ALL_NC":
|
||||
res["merge_failures"].append((s.id, c, o.reason))
|
||||
merged_mask = np.zeros(n, np.bool_)
|
||||
merged_mask[t.merged_rows] = True
|
||||
header = t.header.replace("__CLUSTER__", g.cluster_file or "") + (
|
||||
f"##SAMPLE=<ID={sample_name},SentrixID={sample_id}>\n"
|
||||
f"#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t{sample_name}\n"
|
||||
)
|
||||
low = t.filter_gs
|
||||
gsl, bafl, lrrl = gs.tolist(), baf.tolist(), lrr.tolist()
|
||||
tbx = TabixBuilder()
|
||||
out = Path(out_path)
|
||||
with atomic_open(out, "wb") as raw:
|
||||
if bgzip:
|
||||
w = BgzfWriter(raw) # type: ignore[arg-type]
|
||||
w.write_record(header.encode())
|
||||
else:
|
||||
raw.write(header.encode()) # type: ignore[arg-type]
|
||||
for k in range(n):
|
||||
filt = "LowGS" if low is not None and not math.isnan(gsl[k]) and gsl[k] < low else "PASS"
|
||||
if merged_mask[k]:
|
||||
body = f"{filt}\t.\tGT:GS\t{gt[k]}:{_f4(gsl[k])}\n"
|
||||
else:
|
||||
body = f"{filt}\t.\tGT:GS:BAF:LRR\t{gt[k]}:{_f4(gsl[k])}:{_g8(bafl[k])}:{_g8(lrrl[k])}\n"
|
||||
line = t.prefix[k] + body.encode()
|
||||
if bgzip:
|
||||
vs, ve = w.write_record(line)
|
||||
tbx.add(t.site_contig[k], int(t.site_beg[k]), int(t.site_end[k]), vs, ve)
|
||||
else:
|
||||
raw.write(line) # type: ignore[arg-type]
|
||||
if bgzip:
|
||||
w.close()
|
||||
if bgzip:
|
||||
write_tbi(Path(str(out) + ".tbi"), tbx)
|
||||
res["records"] = n
|
||||
res["sex"] = sex
|
||||
except SimurgError as e:
|
||||
res["error"] = f"{e.code}: {e.message}"
|
||||
except Exception as e:
|
||||
res["error"] = f"E_INTERNAL: {type(e).__name__}: {e}"
|
||||
return res
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ command
|
||||
def run_gtc_to_vcf(opts: VcfOptions, rep: Reporter) -> int:
|
||||
from simurg.formats.discovery import discover_gtcs
|
||||
from simurg.formats.samplesheet import output_stem, parse_samplesheet
|
||||
from simurg.pipeline.inputs import check_paths
|
||||
|
||||
rm = RunManifest(command="gtc-to-vcf", argv=opts.argv)
|
||||
rep.run_id = rm.run_id
|
||||
check_paths(
|
||||
[
|
||||
(opts.bpm, "BPM"),
|
||||
(opts.fasta, "FASTA"),
|
||||
(opts.csv, "CSV manifest"),
|
||||
(opts.gtcs, "GTC folder"),
|
||||
(opts.sheet, "sample sheet"),
|
||||
(opts.exclude_loci, "locus list"),
|
||||
(opts.include_loci, "locus list"),
|
||||
]
|
||||
)
|
||||
if opts.ploidy_policy not in ("sex-aware", "diploid"):
|
||||
raise SimurgError("E_CONFIG_KEY", "--ploidy-policy must be sex-aware or diploid")
|
||||
sheet = parse_samplesheet(opts.sheet) if opts.sheet else None
|
||||
root = opts.gtcs or (sheet.root_folder if sheet else None)
|
||||
if root is None:
|
||||
raise SimurgError("E_NO_INPUT", "give --gtcs or a sample sheet with a RootFolder")
|
||||
found = discover_gtcs(root)
|
||||
samples: list[tuple[str, str, str]] = [] # (sample_id, sample_name, path)
|
||||
if sheet is not None:
|
||||
for r in sheet.rows:
|
||||
if r.sample_id in found:
|
||||
samples.append((r.sample_id, r.sample_name or r.sample_id, str(found[r.sample_id])))
|
||||
else:
|
||||
rep.warn("W_QC_UNMATCHED_SAMPLES", f"{r.sample_id}: no GTC under {root}", sample_id=r.sample_id)
|
||||
else:
|
||||
samples = [(k, k, str(v)) for k, v in sorted(found.items())]
|
||||
suffix = ".snv.vcf.gz" if opts.bgzip else ".snv.vcf"
|
||||
outs = (
|
||||
[opts.out / "cohort.snv.bcf"]
|
||||
if opts.merge
|
||||
else [opts.out / f"{output_stem(name)}{suffix}" for _, name, _ in samples]
|
||||
)
|
||||
check_outputs_absent([*outs, opts.out / "audit" / "loci.parquet"], force=opts.force)
|
||||
if not opts.bgzip:
|
||||
rep.warn("W_NO_BGZIP", "writing plain VCF without index; pgx and cn commands need bgzipped, indexed VCFs")
|
||||
rep.info(f"mapping manifest {opts.bpm.name} to {opts.fasta.name}")
|
||||
table, audit, skips, _genome = build_sites(opts, rep, rm)
|
||||
for role, p in (("bpm", opts.bpm), ("csv_manifest", opts.csv), ("fasta", opts.fasta), ("sample_sheet", opts.sheet)):
|
||||
if p is not None:
|
||||
rm.add_input(role, p)
|
||||
opts.out.mkdir(parents=True, exist_ok=True)
|
||||
write_audit(opts.out / "audit" / "loci.parquet", audit)
|
||||
if opts.merge:
|
||||
n = write_cohort_bcf(table, samples, opts.out / "cohort.snv.bcf", opts, rep)
|
||||
rm.summary = {"samples": len(samples), "records": n, "skipped_loci_by_code": dict(skips)}
|
||||
rm.add_output(opts.out / "cohort.snv.bcf", "bcf/2.2")
|
||||
rm.status = "completed"
|
||||
rm.write(opts.out / "run_manifest.json")
|
||||
rep.info(f"done: cohort BCF with {len(samples)} sample(s), {n} records")
|
||||
return ExitCode.OK
|
||||
workers = max(1, min(opts.threads or (os.cpu_count() or 1), len(samples)))
|
||||
tasks = [(sid, name, path, str(o), opts.bgzip) for (sid, name, path), o in zip(samples, outs, strict=True)]
|
||||
rep.info(f"{len(table.prefix)} VCF records; writing {len(tasks)} sample(s) with {workers} worker(s)")
|
||||
if workers == 1:
|
||||
_init(table)
|
||||
results = [write_sample(t) for t in tasks]
|
||||
else:
|
||||
with ProcessPoolExecutor(
|
||||
workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(table,)
|
||||
) as ex:
|
||||
results = list(ex.map(write_sample, tasks))
|
||||
reasons: Counter[str] = Counter()
|
||||
failures = []
|
||||
for res in results:
|
||||
if "error" in res:
|
||||
rep.error(res["error"].split(":", 1)[0], f"{res['sample_id']}: {res['error']}", sample_id=res["sample_id"])
|
||||
continue
|
||||
reasons.update(res["reasons"])
|
||||
failures += [(res["sample_id"], *f) for f in res["merge_failures"]]
|
||||
for code, n in sorted(reasons.items()):
|
||||
if code.startswith("W_"):
|
||||
rep.warn(code, f"{n} genotype(s): {code}")
|
||||
write_merge_audit(opts.out / "audit" / "merges.parquet", failures)
|
||||
ok = sum("error" not in r for r in results)
|
||||
rm.summary = {"samples": len(results), "succeeded": ok, "records_per_sample": len(table.prefix),
|
||||
"skipped_loci_by_code": dict(skips), "genotype_notes": dict(reasons)} # fmt: skip
|
||||
for o in outs:
|
||||
if o.exists():
|
||||
rm.add_output(o, "vcf/4.3")
|
||||
rm.status = "completed"
|
||||
rm.write(opts.out / "run_manifest.json")
|
||||
rep.info(f"done: {ok}/{len(results)} VCF(s) in {opts.out}")
|
||||
if ok == len(results):
|
||||
return ExitCode.OK
|
||||
return ExitCode.PARTIAL if ok else ExitCode.ALL_FAILED
|
||||
|
||||
|
||||
def write_audit(path: Path, audit: list[AuditRow]) -> None:
|
||||
import pyarrow as pa
|
||||
|
||||
from simurg.core.schemas import write_parquet
|
||||
|
||||
cols = list(AuditRow.__dataclass_fields__)
|
||||
table = pa.table({c: [getattr(a, c) for a in audit] for c in cols})
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
write_parquet(table, path)
|
||||
|
||||
|
||||
def write_merge_audit(path: Path, rows: list[tuple[str, str, tuple[int, ...], str]]) -> None:
|
||||
import pyarrow as pa
|
||||
|
||||
from simurg.core.schemas import write_parquet
|
||||
|
||||
table = pa.table(
|
||||
{
|
||||
"sample_id": [r[0] for r in rows],
|
||||
"site_id": [r[1] for r in rows],
|
||||
"calls": ["/".join(str(c) for c in r[2]) for r in rows],
|
||||
"reason": [r[3] for r in rows],
|
||||
}
|
||||
)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
write_parquet(table, path)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ cohort BCF
|
||||
def _f32_payload(v: npt.NDArray[np.float32]) -> bytes:
|
||||
u = np.asarray(v, dtype="<f4").view("<u4").copy()
|
||||
u[np.isnan(np.asarray(v, dtype=np.float32))] = 0x7F800001
|
||||
return u.tobytes()
|
||||
|
||||
|
||||
def write_cohort_bcf(
|
||||
t: SiteTable, samples: list[tuple[str, str, str]], out: Path, opts: VcfOptions, rep: Reporter
|
||||
) -> int:
|
||||
"""OUT-040: one multi-sample BCF (+ .csi), same records and FORMAT rules as the per-sample VCFs."""
|
||||
from simurg.formats.bcf import (
|
||||
FLOAT,
|
||||
INT8,
|
||||
BcfWriter,
|
||||
CsiBuilder,
|
||||
_type_byte,
|
||||
encode_record,
|
||||
header_dictionaries,
|
||||
)
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.pipeline.qc import default_tmp, load_cohort
|
||||
|
||||
n_loci = read_bpm(opts.bpm).num_loci
|
||||
workers = max(1, opts.threads or (os.cpu_count() or 1))
|
||||
store, infos = load_cohort(
|
||||
[(sid, path) for sid, _, path in samples],
|
||||
n_loci,
|
||||
default_tmp(out.parent),
|
||||
baf=True,
|
||||
scores=True,
|
||||
lrr=True,
|
||||
workers=workers,
|
||||
)
|
||||
try:
|
||||
for s in infos:
|
||||
if s.error:
|
||||
raise SimurgError("E_GTC_UNSUITABLE", f"{s.sample_id}: {s.error}")
|
||||
names = [name for _, name, _ in samples]
|
||||
cluster = read_cluster_name(samples[0][2])
|
||||
header = t.header.replace("__CLUSTER__", cluster)
|
||||
header += "".join(f"##SAMPLE=<ID={name},SentrixID={sid}>\n" for sid, name, _ in samples)
|
||||
header += "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t" + "\t".join(names) + "\n"
|
||||
strings, contigs = header_dictionaries(header)
|
||||
k_gt, k_gs, k_baf, k_lrr = strings["GT"], strings["GS"], strings["BAF"], strings["LRR"]
|
||||
f_pass, f_low = strings["PASS"], strings.get("LowGS", -1)
|
||||
n = len(names)
|
||||
sexes = [s.sex_est for s in infos]
|
||||
ploidy = np.stack([ploidy_for(t.site_class, sx, t.ploidy_policy) for sx in sexes], axis=1) # sites × samples
|
||||
single_pos = np.full(len(t.prefix), -1, np.int64)
|
||||
single_pos[t.single_rows] = np.arange(len(t.single_rows))
|
||||
merged_pos = {int(k): j for j, k in enumerate(t.merged_rows.tolist())}
|
||||
csi = CsiBuilder(len(contigs))
|
||||
fdesc = _type_byte(1, FLOAT)
|
||||
written = 0
|
||||
with atomic_open(out, "wb") as raw:
|
||||
w = BcfWriter(raw, header)
|
||||
chunk = 20000
|
||||
for start in range(0, len(t.prefix), chunk):
|
||||
rows = range(start, min(len(t.prefix), start + chunk))
|
||||
need = {int(t.single_locus[single_pos[k]]) for k in rows if single_pos[k] >= 0}
|
||||
for k in rows:
|
||||
if k in merged_pos:
|
||||
need.update(p.index for p in t.merged_sites[merged_pos[k]].probes)
|
||||
loci = sorted(need)
|
||||
col = {loc: j for j, loc in enumerate(loci)}
|
||||
sel = np.array(loci, np.int64)
|
||||
gt = np.asarray(store.gt[:, sel]) if len(sel) else np.zeros((n, 0), np.uint8)
|
||||
sc = np.asarray(store.score[:, sel]) if len(sel) else np.zeros((n, 0), np.float32) # type: ignore[index]
|
||||
bf = np.asarray(store.baf[:, sel]) if len(sel) else np.zeros((n, 0), np.float32) # type: ignore[index]
|
||||
lr = np.asarray(store.lrr[:, sel]) if len(sel) else np.zeros((n, 0), np.float32) # type: ignore[index]
|
||||
for k in rows:
|
||||
pl = ploidy[k]
|
||||
if single_pos[k] >= 0:
|
||||
j = int(single_pos[k])
|
||||
c = col[int(t.single_locus[j])]
|
||||
calls = gt[:, c].astype(np.int64)
|
||||
ia, ib = (int(x) for x in t.single_ab[j]) # type: ignore[index]
|
||||
a1 = np.where(calls == 3, ib, ia)
|
||||
a2 = np.where(calls == 1, ia, ib)
|
||||
lo, hi = np.minimum(a1, a2), np.maximum(a1, a2)
|
||||
width = max(1, int(pl.max()))
|
||||
cells = np.full((n, width), -127, np.int8)
|
||||
miss = (calls == 0) | (pl == 0) | ((pl == 1) & (calls == 2))
|
||||
hap = pl <= 1
|
||||
cells[:, 0] = np.where(miss, 0, ((np.where(hap & (calls == 3), ib, lo) + 1) << 1)).astype(
|
||||
np.int8
|
||||
)
|
||||
if width == 2:
|
||||
second = np.where(miss, 0, (hi + 1) << 1).astype(np.int8)
|
||||
cells[:, 1] = np.where(hap, np.int8(-127), second)
|
||||
gs = sc[:, c]
|
||||
b = bf[:, c]
|
||||
b = np.where(t.single_flip[j], (np.float32(1.0) - b).astype(np.float32), b)
|
||||
fmt = [(k_gt, _type_byte(width, INT8), cells.tobytes()), (k_gs, fdesc, _f32_payload(gs)),
|
||||
(k_baf, fdesc, _f32_payload(b)), (k_lrr, fdesc, _f32_payload(lr[:, c]))] # fmt: skip
|
||||
gs_min = gs
|
||||
else:
|
||||
site = t.merged_sites[merged_pos[k]]
|
||||
cs = [col[p.index] for p in site.probes]
|
||||
g_all = gt[:, cs].astype(np.int64)
|
||||
gs_min = sc[:, cs].min(axis=1)
|
||||
width = max(1, int(pl.max()))
|
||||
cells = np.full((n, width), -127, np.int8)
|
||||
for si in range(n):
|
||||
p_k = int(pl[si])
|
||||
if p_k == 0:
|
||||
cells[si, 0] = 0
|
||||
continue
|
||||
o = combine_genotype(site, tuple(int(x) for x in g_all[si]), p_k, t.inf1_nc_informative)
|
||||
if o.alleles is None:
|
||||
cells[si, :p_k] = 0
|
||||
else:
|
||||
for q, a in enumerate(o.alleles):
|
||||
cells[si, q] = (a + 1) << 1
|
||||
fmt = [(k_gt, _type_byte(width, INT8), cells.tobytes()), (k_gs, fdesc, _f32_payload(gs_min))]
|
||||
low = opts.filter_gs is not None and bool(np.nanmin(gs_min) < opts.filter_gs) if n else False
|
||||
rec = encode_record(contigs[t.site_contig[k]], int(t.site_beg[k]),
|
||||
int(t.site_end[k] - t.site_beg[k]), t.site_id[k], t.site_alleles[k],
|
||||
f_low if low and f_low >= 0 else f_pass, n, fmt) # fmt: skip
|
||||
vs, ve = w.write(rec)
|
||||
csi.add(contigs[t.site_contig[k]], int(t.site_beg[k]), int(t.site_end[k]), vs, ve)
|
||||
written += 1
|
||||
w.close()
|
||||
Path(str(out) + ".csi").write_bytes(csi.encode())
|
||||
finally:
|
||||
store.close()
|
||||
return written
|
||||
|
||||
|
||||
def read_cluster_name(gtc_path: str) -> str:
|
||||
from simurg.formats.gtc import read_gtc
|
||||
|
||||
return read_gtc(gtc_path, check_write_complete=False).cluster_file or ""
|
||||
Reference in new issue
Block a user