v2 CN: unit tests (GC windows, PFB, GC-wave, smoothing, HMM DEL/LOH, effective size, thresholds, mask, CLI VCF rules)
This commit is contained in:
kemal committed
1 parent
2ad3a044a0
commit
1255f37a16
1 file changed
+173
@@ -0,0 +1,173 @@
|
||||
"""Genome-wide CN/LOH/mosaic calling (spec CYTO-001 … CYTO-040)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from simurg.cli.main import app
|
||||
from simurg.cn import hmm
|
||||
from simurg.cn.events import Event, Thresholds, apply_thresholds, effective_size, extract
|
||||
from simurg.cn.gcwave import gc_correct
|
||||
from simurg.cn.model import gc_fractions, population_bfreq
|
||||
from simurg.formats.gtc import Gtc, write_gtc
|
||||
from simurg.pipeline.cn import _mask_events, moving_median
|
||||
from simurg.synth.files import SynthLocus, encode_bpm
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ covariates
|
||||
def test_gc_fractions() -> None: # CYTO-038
|
||||
seq = np.frombuffer(b"GGGGAAAAAANNNNCCCC", dtype=np.uint8)
|
||||
out = gc_fractions(seq, np.array([1, 13, 18]), (4, 36))
|
||||
assert out[0, 0] == pytest.approx(1.0) # window [0, 2) clipped: "GG"
|
||||
assert np.isnan(out[1, 0]) # window [10, 14) = "NNNN": no called base
|
||||
assert out[2, 1] == pytest.approx(8 / 14) # whole sequence: 8 G/C over 14 called bases
|
||||
|
||||
|
||||
def test_population_bfreq() -> None:
|
||||
p = population_bfreq(np.array([[10, 0, 0], [0, 0, 10], [5, 10, 5], [0, 0, 0]], dtype=np.int32), 4)
|
||||
assert p.tolist() == pytest.approx([0.0, 1.0, 0.5, 0.5])
|
||||
|
||||
|
||||
def test_gc_correct_removes_wave() -> None: # CYTO-003
|
||||
rng = np.random.default_rng(1)
|
||||
gc = rng.uniform(0.3, 0.6, (5000, 3)).astype(np.float32)
|
||||
noise = rng.normal(0, 0.05, 5000)
|
||||
lrr = 0.8 * (gc[:, 2] - 0.45) + noise
|
||||
out, fit = gc_correct(lrr, gc, np.ones(5000, dtype=np.bool_))
|
||||
assert fit.applied and fit.gcwf > 0.02
|
||||
assert abs(np.corrcoef(out, gc[:, 2])[0, 1]) < 0.05
|
||||
same, fit2 = gc_correct(lrr[:10], gc[:10], np.ones(10, dtype=np.bool_))
|
||||
assert not fit2.applied and np.array_equal(same, lrr[:10])
|
||||
|
||||
|
||||
def test_moving_median_truncates_ends() -> None: # CYTO-004
|
||||
x = np.array([1.0, 9.0, 2.0, 3.0, np.nan, 4.0])
|
||||
m = moving_median(x, 1)
|
||||
assert m[0] == pytest.approx(5.0) # window truncated at the start: median(1, 9)
|
||||
assert m[2] == pytest.approx(3.0) and np.isnan(m[4])
|
||||
assert np.array_equal(moving_median(x, 0), x, equal_nan=True)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ HMM + events
|
||||
def _simulate(n: int = 4000, seed: int = 3) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
rng = np.random.default_rng(seed)
|
||||
pos = np.arange(1, n + 1, dtype=np.int64) * 5000
|
||||
pfb = np.full(n, 0.3)
|
||||
geno = rng.choice([0, 1, 2], n, p=[0.49, 0.42, 0.09])
|
||||
baf = np.choose(geno, [0.0, 0.5, 1.0]) + rng.normal(0, 0.03, n)
|
||||
lrr = rng.normal(0, 0.15, n)
|
||||
lrr[1000:1400] -= 0.45 # CN1 deletion: heterozygotes disappear
|
||||
baf[1000:1400] = np.where(baf[1000:1400] > 0.25, 1.0, 0.0) + rng.normal(0, 0.02, 400)
|
||||
baf[2500:3300] = np.where(baf[2500:3300] > 0.25, 1.0, 0.0) + rng.normal(0, 0.02, 800) # copy-neutral LOH
|
||||
return pos, lrr, np.clip(baf, 0, 1), pfb
|
||||
|
||||
|
||||
def test_hmm_finds_deletion_and_loh() -> None:
|
||||
pos, lrr, baf, pfb = _simulate()
|
||||
st = hmm.diploid_states()
|
||||
prior = np.array([s.prior for s in st])
|
||||
prior /= prior.sum()
|
||||
e = hmm.emissions(st, lrr, baf, pfb, hmm.Params(0.15, 0.03))
|
||||
sw = hmm.switch_probs(pos)
|
||||
path = hmm.viterbi(e, sw, prior)
|
||||
evs = extract("1", path, st, pos)
|
||||
calls = {(ev.call, ev.i0 // 100, ev.i1 // 100) for ev in evs if ev.probes >= 100}
|
||||
assert ("DEL", 10, 14) in calls
|
||||
assert ("LOH", 25, 33) in calls
|
||||
post = hmm.forward_backward(e, sw, prior)
|
||||
assert np.allclose(post.sum(axis=1), 1.0)
|
||||
assert post[1200, [0, 1]].sum() > 0.9
|
||||
|
||||
|
||||
def test_effective_size_and_thresholds() -> None: # CYTO-007/008/010
|
||||
pos = np.array([100, 200, 300, 400, 100_400, 100_500], dtype=np.int64)
|
||||
assert effective_size(pos, 100.0) == (100_500 - 100 + 1) - (100_000 - 100)
|
||||
th = Thresholds.defaults()
|
||||
ev = Event("1", 1, 30_000, 0, 12, "DEL", 1, "CN1", effective_size=30_000)
|
||||
small = Event("1", 1, 20_000, 0, 12, "DEL", 1, "CN1", effective_size=20_000)
|
||||
counts: dict[str, int] = {}
|
||||
assert apply_thresholds([ev, small], th, counts) == [ev] and counts["below_threshold"] == 1
|
||||
|
||||
|
||||
def test_mask_trimming() -> None: # CYTO-017
|
||||
mask = {"1": [(1_000, 1_999)]}
|
||||
a = Event("1", 1_500, 10_000, 0, 20, "DEL", 1, "CN1")
|
||||
b = Event("1", 1_100, 2_500, 0, 20, "DEL", 1, "CN1")
|
||||
kept = _mask_events([a, b], mask, counts := {})
|
||||
assert kept == [a] and a.start == 2_000 and "centromere_trimmed" in a.flags and counts["masked"] == 1
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------ CLI end to end
|
||||
def _dataset(tmp: Path) -> tuple[Path, Path, Path]:
|
||||
rng = np.random.default_rng(11)
|
||||
n_per = 3000
|
||||
loci = []
|
||||
for c in ("1", "2"):
|
||||
for k in range(n_per):
|
||||
loci.append(SynthLocus(f"rs{c}_{k}", c, 5000 * (k + 1), "[A/G]", len(loci) + 1, 0, 0))
|
||||
bpm = tmp / "m.bpm"
|
||||
bpm.write_bytes(encode_bpm(loci))
|
||||
fa = tmp / "g.fa"
|
||||
body = {c: "".join(rng.choice(list("ACGT"), n_per * 5000 + 10_000)) for c in ("1", "2")}
|
||||
fa.write_text("".join(f">{c}\n" + "\n".join(s[i : i + 60] for i in range(0, len(s), 60)) + "\n"
|
||||
for c, s in body.items())) # fmt: skip
|
||||
gtcs = tmp / "gtcs"
|
||||
n = len(loci)
|
||||
for sid, deletion in (("200000000001_R01C01", True), ("200000000001_R02C01", False)):
|
||||
geno = rng.choice([1, 2, 3], n, p=[0.49, 0.42, 0.09]).astype(np.uint8)
|
||||
baf = np.choose(geno - 1, [0.0, 0.5, 1.0]) + rng.normal(0, 0.03, n)
|
||||
lrr = rng.normal(0, 0.12, n)
|
||||
if deletion:
|
||||
lrr[500:900] -= 0.45
|
||||
baf[500:900] = np.where(baf[500:900] > 0.25, 1.0, 0.0)
|
||||
geno[500:900] = np.where(geno[500:900] == 2, 3, geno[500:900])
|
||||
g = Gtc(num_snps=n, genotypes=geno, genotype_scores=np.full(n, 0.9, np.float32),
|
||||
b_allele_freqs=np.clip(baf, 0, 1).astype(np.float32), logr_ratios=lrr.astype(np.float32),
|
||||
raw_x=np.full(n, 1000, np.uint16), raw_y=np.full(n, 1000, np.uint16),
|
||||
base_calls=np.zeros((n, 2), np.uint8), gender=b"F") # fmt: skip
|
||||
(gtcs / sid).mkdir(parents=True)
|
||||
write_gtc(g, gtcs / sid / f"{sid}.gtc")
|
||||
return bpm, fa, gtcs
|
||||
|
||||
|
||||
def test_cn_build_and_call(tmp_path: Path) -> None: # CYTO-028 .. CYTO-032, CYTO-037
|
||||
bpm, fa, gtcs = _dataset(tmp_path)
|
||||
model = tmp_path / "m.simurg-cyto"
|
||||
r = CliRunner().invoke(app, ["cn", "build-model", "--bpm", str(bpm), "--fasta", str(fa), "--out", str(model),
|
||||
"--quiet"]) # fmt: skip
|
||||
assert r.exit_code == 0, r.output
|
||||
out = tmp_path / "cn"
|
||||
r = CliRunner().invoke(app, ["cn", "call", "--bpm", str(bpm), "--model", str(model), "--gtcs", str(gtcs),
|
||||
"--fasta", str(fa), "--out", str(out), "--threads", "1", "--min-del-probes", "5",
|
||||
"--unknown-sex-policy", "female", "--quiet"]) # fmt: skip
|
||||
assert r.exit_code == 0, r.output
|
||||
side = json.loads((out / "cn_call_manifest.json").read_text())
|
||||
assert side["thresholds"]["min_del_probes"] == 10 # CYTO-008 floor
|
||||
with gzip.open(out / "200000000001_R01C01.cyto.cnv.vcf.gz", "rt") as fh:
|
||||
text = fh.read()
|
||||
assert (out / "200000000001_R01C01.cyto.cnv.vcf.gz.tbi").exists()
|
||||
head = [ln for ln in text.splitlines() if ln.startswith("##")]
|
||||
assert head[0] == "##fileformat=VCFv4.3" and any(h.startswith("##simurg_thresholds=") for h in head)
|
||||
recs = [ln.split("\t") for ln in text.splitlines() if not ln.startswith("#")]
|
||||
dels = [r_ for r_ in recs if r_[4] == "<DEL>"]
|
||||
assert dels, text[-2000:]
|
||||
for r_ in recs: # CYTO-029 / EDGE-025
|
||||
pos, rid, info = int(r_[1]), r_[2], r_[7]
|
||||
end = int(next(x for x in info.split(";") if x.startswith("END="))[4:])
|
||||
s, e = rid.split(":")[2].split("-")
|
||||
assert int(s) == pos + 1 and int(e) == end
|
||||
assert r_[3] in "ACGT"
|
||||
svlen = int(next(x for x in info.split(";") if x.startswith("SVLEN="))[6:])
|
||||
assert (svlen < 0) == (r_[4] == "<DEL>")
|
||||
assert r_[8] == "GT:CN:NR:LRD:MF"
|
||||
d = dels[0]
|
||||
assert d[0] == "1" and 2_300_000 < int(d[1]) < 2_700_000 and d[9].split(":")[:2] == ["0/1", "1"]
|
||||
summ = (out / "cn_sample_summary.tsv").read_text().splitlines()
|
||||
assert summ[0].startswith("sample_id\tsample_name\tstatus") and len(summ) == 3
|
||||
assert (out / "200000000001_R01C01.cn_segments.parquet").exists()
|
||||
Reference in new issue
Block a user