Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv
267 lines
14 KiB
Python
267 lines
14 KiB
Python
"""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, _mf_from_baf, apply_thresholds, complete_loss, 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()
|
|
|
|
|
|
# ------------------------------------------------------------------------------------------------ annotate
|
|
def test_iscn_strings_and_genes() -> None: # CYTO-034 / CYTO-035
|
|
from simurg.cn.annotate import Bands, Genes, iscn_sample
|
|
|
|
bands = Bands.from_ucsc("chr1\t0\t2300000\tp36.33\tgneg\nchr1\t2300000\t5300000\tp36.32\tgpos25\n"
|
|
"chr9\t0\t2200000\tp24.3\tgneg\n") # fmt: skip
|
|
dele = {"chrom": "1", "start": 1_567_924, "end": 1_632_825, "call": "DEL", "call_type": "DEL", "copy_number": 1,
|
|
"mosaic": False} # fmt: skip
|
|
dup = {"chrom": "1", "start": 2_000_000, "end": 3_000_000, "call": "DUP", "call_type": "DUP", "copy_number": 3,
|
|
"mosaic": False} # fmt: skip
|
|
loh = {"chrom": "9", "start": 100, "end": 2_000_000, "call": "LOH", "call_type": "LOH", "copy_number": 2,
|
|
"mosaic": False} # fmt: skip
|
|
mos = {"chrom": "9", "start": 100, "end": 1_000, "call": "MDEL", "call_type": "DEL", "copy_number": 1,
|
|
"mosaic": True, "base_cn": 2} # fmt: skip
|
|
trad, micro = iscn_sample([loh, dup, dele], bands, "GRCh38", "F")
|
|
assert micro == ("arr[GRCh38] 1p36.33(1567924_1632825)x1,1p36.33p36.32(2000000_3000000)x3,"
|
|
"9p24.3(100_2000000)x2 hmz") # fmt: skip
|
|
assert trad == "del(1)(p36.33p36.33),dup(1)(p36.33p36.32)"
|
|
assert iscn_sample([mos], bands, "GRCh38", "F")[1] == "arr[GRCh38] 9p24.3(100_1000)x2~1"
|
|
assert iscn_sample([], bands, "GRCh38", "M")[1] == "arr[GRCh38](1-22)x2,(XY)x1"
|
|
assert iscn_sample([], bands, "GRCh37", "F")[1] == "arr[GRCh37](1-22,X)x2"
|
|
assert iscn_sample([], bands, "GRCh38", "U")[1] == "arr[GRCh38](1-22)x2"
|
|
genes = Genes.from_rows([("1", 100, 5_000, "A"), ("1", 4_000, 4_500, "B"), ("1", 200, 300, "A"),
|
|
("1", 9_000, 9_100, "C")]) # fmt: skip
|
|
assert genes.overlapping("1", 4_400, 8_999) == ["A", "B"]
|
|
|
|
|
|
def test_cn_annotate_end_to_end(tmp_path: Path) -> None: # CYTO-032 / CYTO-033
|
|
bpm, fa, gtcs = _dataset(tmp_path)
|
|
model, out = tmp_path / "m.simurg-cyto", tmp_path / "cn"
|
|
run = CliRunner().invoke
|
|
assert run(app, ["cn", "annotate", "--vcfs", str(tmp_path)]).exit_code == 2 # E_CN_SIDECAR_MISSING
|
|
assert run(app, ["cn", "build-model", "--bpm", str(bpm), "--fasta", str(fa), "--out", str(model),
|
|
"--quiet"]).exit_code == 0 # fmt: skip
|
|
r = run(app, ["cn", "call", "--bpm", str(bpm), "--model", str(model), "--gtcs", str(gtcs), "--fasta", str(fa),
|
|
"--out", str(out), "--threads", "1", "--unknown-sex-policy", "female", "--quiet"]) # fmt: skip
|
|
assert r.exit_code == 0, r.output
|
|
r = run(app, ["cn", "annotate", "--vcfs", str(out), "--min-qual", "0", "--quiet"])
|
|
assert r.exit_code == 0, r.output
|
|
doc = json.loads((out / "200000000001_R01C01.cyto.json").read_text())
|
|
assert doc["schema"] == "simurg.cyto/1" and doc["annotation_type"] == "constitutional"
|
|
assert doc["iscn_microarray"].startswith("arr[")
|
|
assert any(e["call_type"] == "DEL" and e["chrom"] == "1" for e in doc["events"])
|
|
assert {c["id"] for c in doc["chromosomes"]} == {"1", "2"}
|
|
normal = json.loads((out / "200000000001_R02C01.cyto.json").read_text())
|
|
assert normal["iscn_microarray"] in ("arr[GRCh38](1-22,X)x2", "arr[GRCh37](1-22,X)x2") or normal["events"]
|
|
|
|
|
|
def test_packaged_cn_data() -> None: # the built-in mask and cytobands ship with the package (CYTO-017, CYTO-035)
|
|
from simurg.pipeline.cn import _builtin_file, read_mask
|
|
|
|
for build in ("GRCh37", "GRCh38"):
|
|
m = read_mask(None, build)
|
|
assert len(m) >= 24 and all(s <= e for iv in m.values() for s, e in iv)
|
|
assert gzip.decompress(_builtin_file(f"cytoBand_{build}.txt.gz")).startswith(b"chr1\t0\t")
|
|
|
|
|
|
def test_mf_refinements() -> None: # CYTO-018 / CYTO-041
|
|
rng = np.random.default_rng(3)
|
|
sd = 0.04
|
|
# unsplit heterozygotes: noise alone must not look like a mosaic split
|
|
b = 0.5 + rng.normal(0, sd, 400)
|
|
mf = _mf_from_baf("MDEL", b, sd)
|
|
assert mf is not None and mf < 0.08
|
|
assert (_mf_from_baf("MDEL", b) or 0) > 0.05 # uncorrected, the old estimate was biased upwards
|
|
# a true 40 % mosaic loss is still recovered
|
|
d = 0.4 / (2 * (2 - 0.4))
|
|
b = 0.5 + rng.choice([-d, d], 400) + rng.normal(0, sd, 400)
|
|
assert abs((_mf_from_baf("MDEL", b, sd) or 0) - 0.4) < 0.08
|
|
# no heterozygous calls where ~12 are expected: a complete (constitutional) loss
|
|
hom = np.full(30, 1, dtype=np.uint8)
|
|
assert complete_loss(hom, np.full(30, 0.3))
|
|
assert not complete_loss(hom, np.full(30, 0.02)) # too few expected
|
|
mixed = hom.copy()
|
|
mixed[:8] = 2
|
|
assert not complete_loss(mixed, np.full(30, 0.3))
|
|
|
|
|
|
def test_probe_lrr_noise() -> None: # CYTO-042
|
|
from simurg.cn.model import reference_lrr_spread
|
|
from simurg.pipeline.cn import _lrr_noise
|
|
|
|
rng = np.random.default_rng(5)
|
|
sd = np.array([0.1] * 50 + [0.3] * 50)
|
|
lrr = (rng.normal(0, 1, (40, 100)) * sd + rng.normal(0, 0.2, (40, 1))).astype(np.float32) # + sample offsets
|
|
spread = reference_lrr_spread(lrr, np.zeros(100, dtype=np.float32))
|
|
assert abs(float(np.median(spread[:50])) - 0.1) < 0.02 and abs(float(np.median(spread[50:])) - 0.3) < 0.05
|
|
auto = np.ones(100, dtype=np.bool_)
|
|
s = _lrr_noise(0.2, auto, spread)
|
|
assert np.median(s[50:]) / np.median(s[:50]) > 2
|
|
assert np.all(_lrr_noise(0.2, auto, None) == 0.2)
|