From 98f981b59ed2dbc6e82d9b1f34a4c8f47a5801e3 Mon Sep 17 00:00:00 2001 From: Kemal Yaylali Date: Tue, 29 Sep 2026 07:56:12 +0100 Subject: [PATCH] CN: per-probe heterozygous BAF spread from the batch reference; --mosaic-prior/--switch-distance; Pi memory guard scripts/pi_run.sh runs heavy jobs in their own process group and stops them when MemAvailable drops below a floor (Raspberry Pi OS boots with cgroup_disable=memory, so a systemd MemoryMax cap is not available); a 360-sample genotype run with cohort QC had exhausted memory and the OOM killer took down the desktop session. --- scripts/pi_run.sh | 30 ++++++++++++++++++++++++++++++ src/simurg/cli/main.py | 9 ++++++++- src/simurg/cn/hmm.py | 23 ++++++++++++++--------- src/simurg/cn/model.py | 15 +++++++++++++++ src/simurg/pipeline/cn.py | 19 +++++++++++++------ 5 files changed, 80 insertions(+), 16 deletions(-) create mode 100755 scripts/pi_run.sh diff --git a/scripts/pi_run.sh b/scripts/pi_run.sh new file mode 100755 index 0000000..a39dfe9 --- /dev/null +++ b/scripts/pi_run.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Run a heavy command on the 8 GB Pi with a memory guard. Raspberry Pi OS boots with cgroup_disable=memory, so +# systemd MemoryMax cannot be used; instead the command runs in its own process group and is stopped (TERM, then +# KILL) as soon as the system's MemAvailable drops below MIN_FREE_MB -- before the kernel OOM killer would pick +# victims such as the desktop session. Prints the peak resident memory of the command's process tree. +# MIN_FREE_MB=1200 bash scripts/pi_run.sh +set -uo pipefail +MIN_FREE_MB=${MIN_FREE_MB:-1200} +setsid "$@" & +pid=$! +peak=0 +tripped=0 +tree_rss() { # kB, summed over the process group + ps -o rss= -g "$pid" 2>/dev/null | awk '{s+=$1} END {print s+0}' +} +while kill -0 "$pid" 2>/dev/null; do + rss=$(tree_rss); (( rss > peak )) && peak=$rss + avail=$(awk '/^MemAvailable:/ {print int($2/1024)}' /proc/meminfo) + if (( avail < MIN_FREE_MB )); then + echo "pi_run: MemAvailable ${avail} MiB < ${MIN_FREE_MB} MiB -- stopping the job (process group $pid)" >&2 + kill -TERM -- "-$pid" 2>/dev/null; sleep 5; kill -KILL -- "-$pid" 2>/dev/null + tripped=1 + break + fi + sleep 1 +done +wait "$pid"; rc=$? +(( tripped )) && rc=137 +echo "pi_run: exit $rc, peak RSS $((peak / 1024)) MiB, guard ${MIN_FREE_MB} MiB free" >&2 +exit $rc diff --git a/src/simurg/cli/main.py b/src/simurg/cli/main.py index b9e618e..09e1955 100644 --- a/src/simurg/cli/main.py +++ b/src/simurg/cli/main.py @@ -733,6 +733,12 @@ def cn_call_cmd( svlen_mode: Annotated[str, typer.Option("--svlen-mode", help="signed|positive")] = "signed", mask_bed: Annotated[Path | None, typer.Option("--mask-bed", help="Mask BED (default: built-in).")] = None, no_bgzip: Annotated[bool, typer.Option("--no-bgzip", help="Plain VCF (cn annotate then refuses it).")] = False, + mosaic_prior: Annotated[ + float, typer.Option("--mosaic-prior", help="Advanced: mosaic states' switch mass.") + ] = 0.004, + switch_distance: Annotated[ + float, typer.Option("--switch-distance", help="Advanced: HMM switch distance D (bp).") + ] = 5e8, threads: ThreadsOpt = 0, force: ForceOpt = False, log_json: LogJson = False, @@ -761,7 +767,8 @@ def cn_call_cmd( opts = CnCallOptions( out=out, bpm=bpm, model=model, gtcs=gtcs, summary=summary, fasta=fasta, thresholds=th, smoothing=smoothing, cn_cap=cn_cap, gc_correct=gc_correct, unknown_sex_policy=unknown_sex_policy, gt_mode=gt_mode, - svlen_mode=svlen_mode, mask_bed=mask_bed, bgzip=not no_bgzip, threads=threads, force=force, + svlen_mode=svlen_mode, mask_bed=mask_bed, bgzip=not no_bgzip, mosaic_prior=mosaic_prior, + switch_d=switch_distance, threads=threads, force=force, argv=sys.argv[1:], ) # fmt: skip run_guarded(rep, lambda: run_cn_call(opts, rep)) diff --git a/src/simurg/cn/hmm.py b/src/simurg/cn/hmm.py index e039744..3ee8fb1 100644 --- a/src/simurg/cn/hmm.py +++ b/src/simurg/cn/hmm.py @@ -54,7 +54,7 @@ def _d_split(kind: str, f: float) -> float: return f / 2 -def diploid_states() -> list[State]: +def diploid_states(mosaic_prior: float = MOSAIC_PRIOR) -> list[State]: s = [ State("CN0", 0, False, "", 1.0, "DEL", LRR_MEAN[0], (), 0.01), State("CN1", 1, False, "", 1.0, "DEL", LRR_MEAN[1], (("h0", 0.0), ("h1", 1.0)), 0.06), @@ -71,11 +71,11 @@ def diploid_states() -> list[State]: d = _d_split(kind, f) cn = {"loss": 1, "gain": 3, "cnloh": 2}[kind] bands = (("aa", 0.0), ("ab-", 0.5 - d), ("ab+", 0.5 + d), ("bb", 1.0)) - s.append(State(f"M{kind}{f:.1f}", cn, kind == "cnloh", kind, f, call, mu_end * f, bands, MOSAIC_PRIOR / 12)) + s.append(State(f"M{kind}{f:.1f}", cn, kind == "cnloh", kind, f, call, mu_end * f, bands, mosaic_prior / 12)) return s -def haploid_states() -> list[State]: +def haploid_states(mosaic_prior: float = MOSAIC_PRIOR) -> list[State]: s = [ State("H0", 0, False, "", 1.0, "DEL", LRR_MEAN[0], (), 0.02), State("H1", 1, False, "", 1.0, "", 0.0, (("h0", 0.0), ("h1", 1.0)), 0.80), @@ -83,9 +83,9 @@ def haploid_states() -> list[State]: ] for f in FRACTIONS: s.append(State(f"Hgain{f:.1f}", 2, False, "gain", f, "MDUP", 0.45 * f, (("h0", 0.0), ("h1", 1.0)), - MOSAIC_PRIOR / 8)) # fmt: skip + mosaic_prior / 8)) # fmt: skip s.append(State(f"Hloss{f:.1f}", 0, False, "loss", f, "MDEL", -1.0 * f, (("h0", 0.0), ("h1", 1.0)), - MOSAIC_PRIOR / 8)) # fmt: skip + mosaic_prior / 8)) # fmt: skip return s @@ -100,7 +100,7 @@ def _geno_weights(key: str, p: F64) -> F64: return table[key] -def _norm(x: F64, mu: float, sd: float) -> F64: +def _norm(x: F64, mu: float, sd: float | F64) -> F64: return np.exp(-0.5 * ((x - mu) / sd) ** 2) / (sd * math.sqrt(2 * math.pi)) @@ -114,8 +114,9 @@ class Params: extra: dict[str, float] = field(default_factory=dict) -def emissions(states: list[State], lrr: F64, baf: F64, pfb: F64, par: Params) -> F64: - """Per-probe, per-state emission likelihood (n × S); missing LRR/BAF contribute 1.""" +def emissions(states: list[State], lrr: F64, baf: F64, pfb: F64, par: Params, het_sd: F64 | None = None) -> F64: + """Per-probe, per-state emission likelihood (n × S); missing LRR/BAF contribute 1. ``het_sd``: per-probe spread + of heterozygous BAF in the reference; heterozygous components use max(σ_b, het_sd).""" n = len(lrr) out = np.ones((n, len(states)), dtype=np.float64) lok = np.isfinite(lrr) @@ -123,13 +124,17 @@ def emissions(states: list[State], lrr: F64, baf: F64, pfb: F64, par: Params) -> x = np.where(lok, lrr, 0.0) b = np.clip(np.where(bok, baf, 0.5), 0.0, 1.0) p = np.clip(np.where(np.isfinite(pfb), pfb, 0.5), 0.01, 0.99) + sd_het: float | F64 = ( + np.maximum(par.baf_sd, np.nan_to_num(het_sd, nan=par.baf_sd)) if het_sd is not None else par.baf_sd + ) for j, st in enumerate(states): sd = 2.0 if st.cn == 0 and not st.mosaic else par.lrr_sd le = (1 - par.eps) * _norm(x, st.lrr_mu + par.lrr_offset, sd) + par.eps / 8.0 if st.baf: be = np.zeros(n) for key, mean in st.baf: - be += _geno_weights(key, p) * _norm(b, mean, par.baf_sd) + bsd = par.baf_sd if mean in (0.0, 1.0) or het_sd is None else sd_het + be += _geno_weights(key, p) * _norm(b, mean, bsd) be = (1 - par.eps) * be + par.eps else: be = np.ones(n) diff --git a/src/simurg/cn/model.py b/src/simurg/cn/model.py index 15f7d39..cc205a5 100644 --- a/src/simurg/cn/model.py +++ b/src/simurg/cn/model.py @@ -71,3 +71,18 @@ def reference_stats( np.nan_to_num(het_c, nan=0.5).astype(np.float32), n_het, ) + + +def reference_het_spread(baf: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: + """Per-probe robust SD (1.4826·MAD) of heterozygous BAF (0.2 < BAF < 0.8) across reference samples; NaN when + fewer than 5 heterozygous samples.""" + het = np.where((baf > 0.2) & (baf < 0.8), baf, np.nan) + with np.errstate(all="ignore"): + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + med = np.nanmedian(het, axis=0) + mad = np.nanmedian(np.abs(het - med), axis=0) + n = np.isfinite(het).sum(axis=0) + return np.where(n >= 5, 1.4826 * mad, np.nan).astype(np.float32) diff --git a/src/simurg/pipeline/cn.py b/src/simurg/pipeline/cn.py index 6b05a22..1aa9699 100644 --- a/src/simurg/pipeline/cn.py +++ b/src/simurg/pipeline/cn.py @@ -46,6 +46,7 @@ class CytoModel: ref_lrr_f: npt.NDArray[np.float32] | None = None # same, X/Y from female reference samples ref_lrr_m: npt.NDArray[np.float32] | None = None # same, X/Y from male reference samples ref_het: npt.NDArray[np.float32] | None = None # per-probe median heterozygous BAF + ref_het_sd: npt.NDArray[np.float32] | None = None # per-probe robust SD of heterozygous BAF def run_build_model( @@ -119,8 +120,10 @@ def run_build_model( ref_by_sex[sex] = col n_f, n_m = sexes.count("F"), sexes.count("M") rep.info(f"batch reference sexes: {n_f} F, {n_m} M, {len(sexes) - n_f - n_m} U") + from simurg.cn.model import reference_het_spread + ref_cols = {"ref_lrr": ref_lrr, "ref_lrr_f": ref_by_sex["F"], "ref_lrr_m": ref_by_sex["M"], - "ref_het_baf": ref_het, "ref_n_het": n_het} # fmt: skip + "ref_het_baf": ref_het, "ref_het_sd": reference_het_spread(baf), "ref_n_het": n_het} # fmt: skip n_ref = len(paths) rep.info(f"batch reference from {n_ref} GTC(s)") header = { @@ -173,7 +176,8 @@ def read_model(path: Path, manifest: Manifest) -> CytoModel: pfb = col("pfb") assert pfb is not None return CytoModel(header, table.column("name").to_pylist(), gc.astype(np.float32), pfb, ref_lrr=col("ref_lrr"), - ref_lrr_f=col("ref_lrr_f"), ref_lrr_m=col("ref_lrr_m"), ref_het=col("ref_het_baf")) # fmt: skip + ref_lrr_f=col("ref_lrr_f"), ref_lrr_m=col("ref_lrr_m"), ref_het=col("ref_het_baf"), + ref_het_sd=col("ref_het_sd")) # fmt: skip # ================================================================================================ cn call @@ -202,6 +206,8 @@ class CnCallOptions: svlen_mode: str = "signed" mask_bed: Path | None = None bgzip: bool = True + mosaic_prior: float = 0.004 # [OURS] tuning knobs (hmm.MOSAIC_PRIOR, hmm.SWITCH_D) + switch_d: float = 5e8 threads: int = 0 force: bool = False argv: list[str] = field(default_factory=list) @@ -340,7 +346,7 @@ def call_sample(si: SampleInput, ctx: dict[str, Any], opts: CnCallOptions) -> di seg_sd = _robust_sd(lrr[auto], 0.05) b_het = baf[auto & (baf > 0.3) & (baf < 0.7)] baf_sd = _robust_sd(b_het, 0.02) - dip, hap = hmm.diploid_states(), hmm.haploid_states() + dip, hap = hmm.diploid_states(opts.mosaic_prior), hmm.haploid_states(opts.mosaic_prior) counts: dict[str, int] = {} events: list[Any] = [] segments: list[dict[str, Any]] = [] @@ -378,9 +384,10 @@ def call_sample(si: SampleInput, ctx: dict[str, Any], opts: CnCallOptions) -> di offset = float(np.nanmedian(x_raw)) if haploid else 0.0 if not math.isfinite(offset): offset = 0.0 - prm = hmm.Params(seg_sd, baf_sd, lrr_offset=offset) - e = hmm.emissions(states, x_seg, np.full(len(b), np.nan) if haploid else b, pfb[idx], prm) - sw = hmm.switch_probs(p) + prm = hmm.Params(seg_sd, baf_sd, lrr_offset=offset, switch_d=opts.switch_d) + het_sd = model.ref_het_sd[idx].astype(np.float64) if model.ref_het_sd is not None else None + e = hmm.emissions(states, x_seg, np.full(len(b), np.nan) if haploid else b, pfb[idx], prm, het_sd) + sw = hmm.switch_probs(p, opts.switch_d) path = hmm.viterbi(e, sw, prior) post = hmm.forward_backward(e, sw, prior) evs = extract(chrom, path, states, p)