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.

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))

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)

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)

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)