P6 complete: optional VCF fields, SNV+indel merging, auxiliary MNV loci, PLINK 2 .pgen

- gtc-to-vcf --info-fields ASSAY_TYPE,NORM_BIN,GENTRAIN,CLUSTER_SEP,ILMN_STRAND,PROBE_IDS (--egt) and
  --extra-format X,Y,NORMX,NORMY,THETA,R, in per-sample VCFs and the --merge BCF (OUT-012, plan §8.3);
  equal to gtc2vcf's values (THETA/R missing where both normalized intensities are 0, K-023)
- --merge-snv-indel (OUT-021, EDGE-013): alleles extended to the longest REF; SNV probes read the base at POS
- --auxiliary-loci (OUT-035, EDGE-015): rules (a)-(e), E_AUX_INVALID; rule (d) uses the last plus-strand base
  for minus-strand designs (spec amended, as GTCtoVCF)
- export plink --format bed|pgen|both: PLINK 2 fixed-width .pgen/.pvar/.psam; plink2 reads it back to a .bed
  byte-identical to ours
- cohort store keeps raw/normalized intensities for the BCF path; tests (unit + gtc2vcf oracle); docs
This commit is contained in:
kemal committed 2026-09-28 14:29:42 +01:00
1 parent d3add008f7
commit 1001bf951e
13 files changed
+611 -50

No files matched your search

+75 -15
View File
@@ -8,6 +8,7 @@ import gzip
import io
import math
import re
import struct
from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
@@ -348,21 +349,38 @@ class PlinkOptions:
sheet: Path | None = None
gtcs: Path | None = None
multiallelic: str = "split"
format: str = "bed"
prefix: str = "cohort"
force: bool = False
argv: list[str] = field(default_factory=list)
PGEN_FROM_BED = bytes([2, 3, 1, 0]) # PLINK 1 code (A1 = ALT) → PLINK 2 hardcall (0 ref/ref, 1 het, 2 alt/alt, 3 NA)
def _pack(codes: Sequence[int], nbytes: int) -> bytes:
buf = bytearray(nbytes)
for j, c in enumerate(codes):
buf[j >> 2] |= c << (2 * (j & 3))
return bytes(buf)
def run_export_plink(opts: PlinkOptions, rep: Reporter) -> int:
"""OUT-043: SimurgArray per-sample VCFs (identical record order) → PLINK 1 .bed/.bim/.fam."""
"""OUT-043 / plan §8.4: SimurgArray per-sample VCFs (identical record order) → PLINK 1 .bed/.bim/.fam and/or
PLINK 2 .pgen/.pvar/.psam (fixed-width hardcall storage, mode 0x02 [DOC: pgen format]). Multi-allelic sites are
split (``bcftools norm -m-``) or dropped."""
if opts.multiallelic not in ("split", "drop"):
raise SimurgError("E_CONFIG_KEY", "--multiallelic must be split or drop")
if opts.format not in ("bed", "pgen", "both"):
raise SimurgError("E_CONFIG_KEY", "--format must be bed, pgen or both")
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)
want_bed, want_pgen = opts.format in ("bed", "both"), opts.format in ("pgen", "both")
targets = ([".bed", ".bim", ".fam"] if want_bed else []) + ([".pgen", ".pvar", ".psam"] if want_pgen else [])
check_outputs_absent([base.with_suffix(x) for x in targets], force=opts.force)
names, sex = [], {}
for v in vcfs:
opener = gzip.open if str(v).endswith(".gz") else open
@@ -388,12 +406,23 @@ def run_export_plink(opts: PlinkOptions, rep: Reporter) -> int:
nbytes = (n + 3) // 4
opts.out.mkdir(parents=True, exist_ok=True)
n_var = dropped = 0
pvar_header = "".join(
ln for ln in _vcf_header(vcfs[0]) if ln.startswith("##contig") or ln.startswith("##fileformat")
)
with (
atomic_open(base.with_suffix(".bed"), "wb") as bed_raw,
atomic_open(base.with_suffix(".bim"), "w", "utf-8") as bim_raw,
atomic_open(base.with_suffix(".bed"), "wb") if want_bed else _Null() as bed_raw,
atomic_open(base.with_suffix(".bim"), "w", "utf-8") if want_bed else _Null() as bim_raw,
atomic_open(base.with_suffix(".pgen"), "wb") if want_pgen else _Null() as pgen_raw,
atomic_open(base.with_suffix(".pvar"), "w", "utf-8") if want_pgen else _Null() as pvar_raw,
):
bed, bim = cast(IO[bytes], bed_raw), cast(IO[str], bim_raw)
bed.write(PLINK_MAGIC)
pgen, pvar = cast(IO[bytes], pgen_raw), cast(IO[str], pvar_raw)
if want_bed:
bed.write(PLINK_MAGIC)
if want_pgen:
# header byte 12: nonref-flag storage 1 = every REF is the reference base (from the FASTA)
pgen.write(b"l\x1b\x02" + struct.pack("<II", 0, n) + b"\x40") # variant count patched below
pvar.write(pvar_header + "#CHROM\tPOS\tID\tREF\tALT\n")
for recs in zip(*streams, strict=True):
r0 = recs[0]
if any(r[:5] != r0[:5] for r in recs[1:]):
@@ -405,18 +434,49 @@ def run_export_plink(opts: PlinkOptions, rep: Reporter) -> int:
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")
if want_bed:
bed.write(_pack(codes, nbytes))
bim.write(f"{r0[0]}\t{vid}\t0\t{r0[1]}\t{alt}\t{r0[3]}\n")
if want_pgen:
pgen.write(_pack([PGEN_FROM_BED[c] for c in codes], nbytes))
pvar.write(f"{r0[0]}\t{r0[1]}\t{vid}\t{r0[3]}\t{alt}\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)
if want_pgen:
pgen.seek(3)
pgen.write(struct.pack("<I", n_var))
sex_code = {"M": 1, "F": 2}
if want_bed:
fam = "".join(f"{s}\t{s}\t0\t0\t{sex_code.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)
if want_pgen:
psam = "#IID\tSEX\n" + "".join(f"{s}\t{ {'M': '1', 'F': '2'}.get(sex.get(s, 'U'), 'NA') }\n" for s in names)
with atomic_open(base.with_suffix(".psam"), "w", "utf-8") as psam_raw:
cast(IO[str], psam_raw).write(psam)
rm.status = "completed"
rm.summary = {"samples": n, "variants": n_var, "multiallelic_dropped": dropped}
rm.summary = {"samples": n, "variants": n_var, "multiallelic_dropped": dropped, "format": opts.format}
rm.write(opts.out / "run_manifest.export_plink.json")
rep.info(f"PLINK fileset {base}.bed/.bim/.fam: {n} samples, {n_var} variants")
rep.info(f"PLINK fileset {base} ({opts.format}): {n} samples, {n_var} variants")
return ExitCode.OK
class _Null:
"""Context manager standing in for an output that is not requested."""
def __enter__(self) -> None:
return None
def __exit__(self, *a: object) -> None:
return None
def _vcf_header(path: Path) -> list[str]:
opener = gzip.open if str(path).endswith(".gz") else open
out = []
with opener(path, "rt") as fh:
for line in fh:
if not line.startswith("##"):
break
out.append(line)
return out
+38 -7
View File
@@ -308,27 +308,45 @@ 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, lrr: bool = False) -> None:
def __init__(
self,
root: Path,
n_samples: int,
n_loci: int,
*,
baf: bool,
scores: bool,
lrr: bool = False,
intensities: 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
# raw and normalized intensities (needs root/lookup.npy: normalization-transform index per locus)
self.rx = self.ry = self.nx = self.ny = None
if intensities:
self.rx = np.lib.format.open_memmap(root / "rx.npy", "w+", np.uint16, self.shape)
self.ry = np.lib.format.open_memmap(root / "ry.npy", "w+", np.uint16, self.shape)
self.nx = np.lib.format.open_memmap(root / "nx.npy", "w+", np.float32, self.shape)
self.ny = np.lib.format.open_memmap(root / "ny.npy", "w+", np.float32, self.shape)
def close(self) -> None:
for name in ("gt", "baf", "score", "lrr"):
for name in ("gt", "baf", "score", "lrr", "rx", "ry", "nx", "ny"):
a = getattr(self, name)
if a is not None:
a.flush()
self.gt = self.baf = self.score = self.lrr = None # type: ignore[assignment]
self.rx = self.ry = self.nx = self.ny = None
shutil.rmtree(self.root, ignore_errors=True)
def _load_one(task: tuple[int, str, str, str, bool, bool, bool]) -> SampleInfo:
def _load_one(task: tuple[int, str, str, str, bool, bool, bool, bool]) -> SampleInfo:
from simurg.formats.gtc import read_gtc
i, sid, path, root, want_baf, want_scores, want_lrr = task
i, sid, path, root, want_baf, want_scores, want_lrr, want_int = task
info = SampleInfo(sid, path)
try:
g = read_gtc(path)
@@ -356,6 +374,15 @@ def _load_one(task: tuple[int, str, str, str, bool, bool, bool]) -> SampleInfo:
sc = np.load(Path(root) / "score.npy", mmap_mode="r+")
sc[i] = g.genotype_scores
sc.flush()
if want_int:
from simurg.signal.normalization.apply import normalize
lookup = np.load(Path(root) / "lookup.npy")
nx, ny = normalize(g.raw_x, g.raw_y, lookup, g.normalization_transforms)
for name, arr in (("rx", g.raw_x), ("ry", g.raw_y), ("nx", nx), ("ny", ny)):
mm = np.load(Path(root) / f"{name}.npy", mmap_mode="r+")
mm[i] = arr
mm.flush()
except SimurgError as e:
info.error = f"{e.code}: {e.message}"
except Exception as e:
@@ -372,11 +399,15 @@ def load_cohort(
scores: bool,
workers: int,
lrr: bool = False,
lookup: npt.NDArray[np.int32] | None = None,
) -> 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, lrr=lrr)
"""Read GTCs (sample_id, path) into a :class:`CohortStore` with a process pool. With ``lookup`` (the manifest's
normalization-transform index per locus) raw and normalized intensities are stored too."""
store = CohortStore(tmp, len(samples), n_loci, baf=baf, scores=scores, lrr=lrr, intensities=lookup is not None)
if lookup is not None:
np.save(tmp / "lookup.npy", lookup)
store.gt.flush()
tasks = [(i, sid, p, str(tmp), baf, scores, lrr) for i, (sid, p) in enumerate(samples)]
tasks = [(i, sid, p, str(tmp), baf, scores, lrr, lookup is not None) for i, (sid, p) in enumerate(samples)]
if workers <= 1 or len(tasks) <= 1:
infos = [_load_one(t) for t in tasks]
else:
+171 -7
View File
@@ -57,6 +57,11 @@ class VcfOptions:
ploidy_policy: str = "sex-aware"
threads: int = 0
merge: bool = False
merge_snv_indel: bool = False
auxiliary_loci: Path | None = None
info_fields: list[str] = field(default_factory=list)
extra_format: list[str] = field(default_factory=list)
egt: Path | None = None
force: bool = False
argv: list[str] = field(default_factory=list)
@@ -82,11 +87,33 @@ class SiteTable:
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
site_info: list[str] = field(default_factory=list) # INFO text per site ('.' without --info-fields)
site_info_typed: list[list[tuple[str, str, list[Any]]]] = field(default_factory=list) # (key, kind, values)
extra_format: list[str] = field(default_factory=list)
single_lookup: npt.NDArray[np.int32] | None = None # normalization-transform index of single-probe loci
inf1_nc_informative: bool = False
ploidy_policy: str = "sex-aware"
filter_gs: float | None = None
INFO_DEFS = { # OUT-012, one value per contributing probe (BPM order)
"ASSAY_TYPE": ("Integer", "Assay type per probe: 0 Infinium II, 1 Infinium I red, 2 Infinium I green"),
"NORM_BIN": ("Integer", "Normalization bin per probe"),
"GENTRAIN": ("Float", "GenTrain score per probe (cluster file)"),
"CLUSTER_SEP": ("Float", "Cluster separation score per probe (cluster file)"),
"ILMN_STRAND": ("String", "Design strand per probe (TOP, BOT, PLUS, MINUS)"),
"PROBE_IDS": ("String", "Names of the contributing probes"),
}
EXTRA_DEFS = { # plan §8.3; single-probe records only
"X": ("Integer", "Raw intensity, X (red) channel"),
"Y": ("Integer", "Raw intensity, Y (green) channel"),
"NORMX": ("Float", "Normalized X intensity"),
"NORMY": ("Float", "Normalized Y intensity"),
"THETA": ("Float", "Normalized theta = (2/pi) atan2(NORMY, NORMX)"),
"R": ("Float", "Normalized R = NORMX + NORMY"),
}
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"):
@@ -196,8 +223,21 @@ def build_sites(
probes.append(p)
for code, count in sorted(skips.items()):
rep.warn(code, f"{count} locus/loci: {code}")
if opts.auxiliary_loci is not None:
from simurg.harmonize.sites import apply_auxiliary, read_auxiliary
aux = read_auxiliary(str(opts.auxiliary_loci))
strands = mv.ref_strands
probes, errors = apply_auxiliary(probes, aux, genome.fetch, cmap.get, lambda i: (strands[i] or "").strip())
if errors:
raise SimurgError("E_AUX_INVALID", f"{len(errors)} invalid auxiliary record(s):\n " + "\n ".join(errors))
rep.info(f"{len(aux)} auxiliary locus definition(s) applied")
names = {r.name for r in aux}
for a_row in audit:
if a_row.locus in names:
a_row.reason = "aux_locus"
order = {e.name: k for k, e in enumerate(genome.entries)}
sites = group_sites(probes, order, unsquash=opts.unsquash_duplicates)
sites = group_sites(probes, order, unsquash=opts.unsquash_duplicates, merge_snv_indel=opts.merge_snv_indel)
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
@@ -222,6 +262,7 @@ def build_sites(
else:
merged_sites.append(s)
merged_rows.append(k)
site_info, site_info_typed = _site_info(opts, bpm, sites)
header = vcf_header(opts, genome, build, contigs, rm)
table = SiteTable(
prefix=prefix,
@@ -240,6 +281,10 @@ def build_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),
site_info=site_info,
site_info_typed=site_info_typed,
extra_format=list(opts.extra_format),
single_lookup=bpm.normalization_lookups()[np.array(single_locus, np.int64)] if opts.extra_format else None,
inf1_nc_informative=opts.inf1_nocall_informative,
ploidy_policy=opts.ploidy_policy,
filter_gs=opts.filter_gs,
@@ -251,6 +296,49 @@ def _sha(path: Path | None) -> str:
return sha256_file(path) if path else ""
def _site_info(
opts: VcfOptions, bpm: Any, sites: list[Site]
) -> tuple[list[str], list[list[tuple[str, str, list[Any]]]]]:
"""Static INFO per site for --info-fields (OUT-012)."""
if not opts.info_fields:
return ["."] * len(sites), [[] for _ in sites]
gentrain = sep = None
if {"GENTRAIN", "CLUSTER_SEP"} & set(opts.info_fields):
from simurg.formats.egt import read_egt
assert opts.egt is not None
egt = read_egt(opts.egt)
eidx = egt.name_index()
rows = np.array([eidx.get(n, -1) for n in bpm.names], np.int64)
gentrain = np.where(rows >= 0, egt.gentrain_score[np.maximum(rows, 0)], np.nan)
sep = np.where(rows >= 0, egt.cluster_sep[np.maximum(rows, 0)], np.nan)
bins = bpm.norm_bin_ids()
texts, typed = [], []
for s in sites:
idx = [p.index for p in s.probes]
vals: dict[str, list[Any]] = {
"ASSAY_TYPE": [int(bpm.assay_types[i]) for i in idx],
"NORM_BIN": [int(bins[i]) for i in idx],
"GENTRAIN": [float(gentrain[i]) for i in idx] if gentrain is not None else [],
"CLUSTER_SEP": [float(sep[i]) for i in idx] if sep is not None else [],
"ILMN_STRAND": [bpm.ilmn_strands[i] or "." for i in idx],
"PROBE_IDS": [p.name for p in s.probes],
}
parts, tv = [], []
for key in opts.info_fields:
v = vals[key]
kind = INFO_DEFS[key][0]
text = ",".join(
"." if (isinstance(x, float) and math.isnan(x)) else (f"{x:.6g}" if isinstance(x, float) else str(x))
for x in v
)
parts.append(f"{key}={text}")
tv.append((key, kind, v))
texts.append(";".join(parts))
typed.append(tv)
return texts, typed
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)."""
@@ -271,14 +359,18 @@ def vcf_header(opts: VcfOptions, genome: Any, build: str, contigs: list[str], rm
"##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}"
f"auxiliary_loci={lst(opts.auxiliary_loci)};ploidy_policy={opts.ploidy_policy};contig_alias={opts.contig_alias}"
+ (";drop_ambiguous=true" if opts.drop_ambiguous else "")
+ (";merge_snv_indel=true" if opts.merge_snv_indel 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">')
for key in opts.info_fields:
kind, desc = INFO_DEFS[key]
lines.append(f'##INFO=<ID={key},Number=.,Type={kind},Description="{desc}">')
if opts.filter_gs is not None:
lines.append(f'##FILTER=<ID=LowGS,Description="GenCall score below {opts.filter_gs}">')
lines += [
@@ -289,6 +381,9 @@ def vcf_header(opts: VcfOptions, genome: Any, build: str, contigs: list[str], rm
'ALT (REF=A, ALT=B relative to the reference genome)">',
'##FORMAT=<ID=LRR,Number=1,Type=Float,Description="Log R ratio">',
]
for key in opts.extra_format:
kind, desc = EXTRA_DEFS[key]
lines.append(f'##FORMAT=<ID={key},Number=1,Type={kind},Description="{desc}">')
return "\n".join(lines) + "\n"
@@ -338,6 +433,22 @@ def write_sample(task: tuple[str, str, str, str, bool]) -> dict[str, Any]:
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)
extra_cols: list[list[str]] = []
if t.extra_format:
from simurg.signal.normalization.apply import normalize, polar
rx, ry = g.raw_x[loc], g.raw_y[loc]
nx, ny = normalize(rx, ry, t.single_lookup, g.normalization_transforms) # type: ignore[arg-type]
zero = (rx == 0) & (ry == 0) # no intensity: normalized values are missing (as in the GTC semantics)
nx = np.where(zero, np.float32(np.nan), nx).astype(np.float32)
ny = np.where(zero, np.float32(np.nan), ny).astype(np.float32)
th, rr = polar(nx, ny)
src: dict[str, npt.NDArray[Any]] = {"X": rx, "Y": ry, "NORMX": nx, "NORMY": ny, "THETA": th, "R": rr}
for key in t.extra_format:
col = np.full(n, "", dtype=object)
vals = src[key]
col[rows] = [str(int(v)) for v in vals] if key in ("X", "Y") else [_g8(float(v)) for v in vals]
extra_cols.append(col.tolist())
het_haploid = int(((calls == 2) & (pl == 1)).sum())
if het_haploid:
res["reasons"]["het_in_haploid"] += het_haploid
@@ -373,10 +484,16 @@ def write_sample(task: tuple[str, str, str, str, bool]) -> dict[str, Any]:
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"
info = t.site_info[k] if t.site_info else "."
if merged_mask[k]:
body = f"{filt}\t.\tGT:GS\t{gt[k]}:{_f4(gsl[k])}\n"
body = f"{filt}\t{info}\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"
extra_keys = "".join(":" + key for key in t.extra_format)
extra_vals = "".join(":" + c[k] for c in extra_cols)
body = (
f"{filt}\t{info}\tGT:GS:BAF:LRR{extra_keys}\t"
f"{gt[k]}:{_f4(gsl[k])}:{_g8(bafl[k])}:{_g8(lrrl[k])}{extra_vals}\n"
)
line = t.prefix[k] + body.encode()
if bgzip:
vs, ve = w.write_record(line)
@@ -413,10 +530,20 @@ def run_gtc_to_vcf(opts: VcfOptions, rep: Reporter) -> int:
(opts.sheet, "sample sheet"),
(opts.exclude_loci, "locus list"),
(opts.include_loci, "locus list"),
(opts.auxiliary_loci, "auxiliary loci VCF"),
]
)
if opts.ploidy_policy not in ("sex-aware", "diploid"):
raise SimurgError("E_CONFIG_KEY", "--ploidy-policy must be sex-aware or diploid")
bad = [x for x in opts.info_fields if x not in INFO_DEFS] + [x for x in opts.extra_format if x not in EXTRA_DEFS]
if bad:
raise SimurgError(
"E_CONFIG_KEY",
f"unknown field(s) {', '.join(bad)}; --info-fields takes {','.join(INFO_DEFS)}; "
f"--extra-format takes {','.join(EXTRA_DEFS)}",
)
if {"GENTRAIN", "CLUSTER_SEP"} & set(opts.info_fields) and opts.egt is None:
raise SimurgError("E_CONFIG_KEY", "--info-fields GENTRAIN/CLUSTER_SEP need the cluster file (--egt)")
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:
@@ -534,25 +661,31 @@ def write_cohort_bcf(
from simurg.formats.bcf import (
FLOAT,
INT8,
INT32,
BcfWriter,
CsiBuilder,
_type_byte,
encode_record,
header_dictionaries,
typed_floats,
typed_ints,
typed_string,
)
from simurg.formats.bpm import read_bpm
from simurg.pipeline.qc import default_tmp, load_cohort
from simurg.signal.normalization.apply import polar
n_loci = read_bpm(opts.bpm).num_loci
bpm = read_bpm(opts.bpm)
workers = max(1, opts.threads or (os.cpu_count() or 1))
store, infos = load_cohort(
[(sid, path) for sid, _, path in samples],
n_loci,
bpm.num_loci,
default_tmp(out.parent),
baf=True,
scores=True,
lrr=True,
workers=workers,
lookup=bpm.normalization_lookups() if t.extra_format else None,
)
try:
for s in infos:
@@ -591,6 +724,9 @@ def write_cohort_bcf(
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]
if t.extra_format and len(sel):
ex_rx, ex_ry = np.asarray(store.rx[:, sel]), np.asarray(store.ry[:, sel]) # type: ignore[index]
ex_nx, ex_ny = np.asarray(store.nx[:, sel]), np.asarray(store.ny[:, sel]) # type: ignore[index]
for k in rows:
pl = ploidy[k]
if single_pos[k] >= 0:
@@ -616,6 +752,25 @@ def write_cohort_bcf(
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
if t.extra_format:
zero = (ex_rx[:, c] == 0) & (ex_ry[:, c] == 0)
enx = np.where(zero, np.float32(np.nan), ex_nx[:, c]).astype(np.float32)
eny = np.where(zero, np.float32(np.nan), ex_ny[:, c]).astype(np.float32)
th, rr = polar(enx, eny)
xsrc = {
"X": ex_rx[:, c],
"Y": ex_ry[:, c],
"NORMX": enx,
"NORMY": eny,
"THETA": th,
"R": rr,
}
for key in t.extra_format:
if key in ("X", "Y"):
payload = xsrc[key].astype("<i4").tobytes()
fmt.append((strings[key], _type_byte(1, INT32), payload))
else:
fmt.append((strings[key], fdesc, _f32_payload(xsrc[key].astype(np.float32))))
gs_min = gs
else:
site = t.merged_sites[merged_pos[k]]
@@ -637,9 +792,18 @@ def write_cohort_bcf(
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
info = []
for key, kind, vals in t.site_info_typed[k] if t.site_info_typed else []:
if kind == "Integer":
enc = typed_ints([int(v) for v in vals])
elif kind == "Float":
enc = typed_floats([float(v) for v in vals])
else:
enc = typed_string(",".join(str(v) for v in vals))
info.append((strings[key], enc))
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
f_low if low and f_low >= 0 else f_pass, n, fmt, info) # 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