P9 (WIP): compiled kernels (rust/simurg-core, PyO3) with bit-identical pure-Python fallback
Rust ports of the hot paths found by profiling: fused per-locus GenCall calling and scoring, BAF/LRR, libm maps with CPython math semantics, both nearest-neighbour searches of the normalization fit, and the BPM LocusEntry parser. Address coverage is memoised per IDAT layout. SIMURG_PURE_PYTHON=1 forces the pure paths. On the Pi: genotype_sample 3.4 s -> 1.1 s, 36-sample GSA run 131 s -> 46 s, BPM read 4x faster; 36 GSA + 6 Omni2.5 GTCs and QC tables byte-identical to the pure path; tests/unit/test_core_kernels.py checks every kernel on edge cases. bench/profile_genotype.py, bench/scaling.py, cloud task p9 (thread scaling).
This commit is contained in:
kemal committed
1 parent
f64bdf11f2
commit
b0494be576
19 files changed
+1755
-6
No files matched your search
@@ -0,0 +1,83 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Per-stage timing of `simurg genotype` for one sample (plan §19 measurements).
|
||||
|
||||
python bench/profile_genotype.py BPM EGT GRN_IDAT [--repeat N]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import datetime
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.formats.egt import read_egt
|
||||
from simurg.formats.gtc import write_gtc
|
||||
from simurg.formats.idat import read_idat
|
||||
from simurg.genotype.baf_lrr import baf_lrr, logr_dev
|
||||
from simurg.genotype.clusters import align_clusters
|
||||
from simurg.genotype.gencall import call_genotypes
|
||||
from simurg.genotype.sample import genotype_sample, manifest_context, params_for_preset
|
||||
from simurg.genotype.sex import estimate_sex
|
||||
from simurg.signal.normalization.apply import normalize, polar
|
||||
from simurg.signal.normalization.fit import fit_transforms
|
||||
from simurg.signal.raw_xy import LocusMapCache, assemble_raw_xy
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("bpm")
|
||||
ap.add_argument("egt")
|
||||
ap.add_argument("grn")
|
||||
ap.add_argument("--repeat", type=int, default=3)
|
||||
ap.add_argument("--out", default="/tmp/profile.gtc")
|
||||
a = ap.parse_args()
|
||||
m = read_bpm(a.bpm)
|
||||
ctx = manifest_context(m, align_clusters(read_egt(a.egt), m))
|
||||
lite = m.lite()
|
||||
ctx = replace(ctx, manifest=lite, cache=LocusMapCache(lite))
|
||||
params = params_for_preset("current")
|
||||
red_path = a.grn.replace("_Grn", "_Red")
|
||||
t: dict[str, list[float]] = {}
|
||||
|
||||
def lap(name: str, t0: float) -> float:
|
||||
now = time.perf_counter()
|
||||
t.setdefault(name, []).append(now - t0)
|
||||
return now
|
||||
|
||||
for _ in range(a.repeat):
|
||||
t0 = time.perf_counter()
|
||||
grn, red = read_idat(a.grn), read_idat(red_path)
|
||||
t0 = lap("read_idat x2", t0)
|
||||
raw = assemble_raw_xy(grn, red, lite, cache=ctx.cache)
|
||||
t0 = lap("assemble_raw_xy", t0)
|
||||
xf = fit_transforms(raw.raw_x, raw.raw_y, ctx.norm_bins, params.norm_version)
|
||||
t0 = lap("fit_transforms", t0)
|
||||
nx, ny = normalize(raw.raw_x, raw.raw_y, ctx.lookup, xf)
|
||||
th, r = polar(nx, ny)
|
||||
t0 = lap("normalize+polar", t0)
|
||||
z = (raw.raw_x == 0) & (raw.raw_y == 0)
|
||||
thc = np.where(z, np.float32(np.nan), th).astype(np.float32)
|
||||
rc = np.where(z, np.float32(np.nan), r).astype(np.float32)
|
||||
calls = call_genotypes(thc, rc, z, ctx.clusters, lite.intensity_only, params.cutoff)
|
||||
t0 = lap("call_genotypes", t0)
|
||||
estimate_sex(calls.genotypes, r, ctx.clusters.gentrain, ctx.clusters.present, ctx.sex_classes, params.sex)
|
||||
t0 = lap("estimate_sex", t0)
|
||||
_, lrr = baf_lrr(thc, rc, z, ctx.clusters, lite.intensity_only)
|
||||
logr_dev(lrr, ctx.autosomal, lite.intensity_only, ctx.clusters.present)
|
||||
t0 = lap("baf_lrr+logr_dev", t0)
|
||||
g = genotype_sample(grn, red, ctx, params, datetime.datetime(2026, 1, 1))
|
||||
t0 = lap("genotype_sample (whole)", t0)
|
||||
write_gtc(g, Path(a.out))
|
||||
t0 = lap("write_gtc", t0)
|
||||
print(f"{'stage':28s} {'median s':>9s}")
|
||||
for k, v in t.items():
|
||||
print(f"{k:28s} {sorted(v)[len(v) // 2]:9.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,109 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Thread scaling of ``simurg genotype`` (plan §18 P9 gate: ≥ 20 samples/min/core on GSA-size arrays, linear
|
||||
scaling to 32 threads).
|
||||
|
||||
For every worker count the same N samples are genotyped (``--no-cohort-qc``); the per-sample phase is timed from
|
||||
the run log ("genotyping … sample(s)" → "done"), so manifest loading is reported separately and does not blur
|
||||
the scaling. Writes a JSON summary.
|
||||
|
||||
python bench/scaling.py --bpm M.bpm --egt C.egt --idats DIR --out OUT --threads 1,2,4,8 [--samples N]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def phase_seconds(log: Path) -> tuple[float, float]:
|
||||
"""(manifest+setup seconds, per-sample phase seconds) from a genotype run log."""
|
||||
ts: dict[str, dt.datetime] = {}
|
||||
for line in log.read_text().splitlines():
|
||||
rec = json.loads(line)
|
||||
msg = rec.get("msg") or ""
|
||||
t = dt.datetime.fromisoformat(rec["ts"].replace("Z", "+00:00"))
|
||||
ts.setdefault("first", t)
|
||||
if msg.startswith("genotyping "):
|
||||
ts["start"] = t
|
||||
elif msg.startswith("done:"):
|
||||
ts["done"] = t
|
||||
return (ts["start"] - ts["first"]).total_seconds(), (ts["done"] - ts["start"]).total_seconds()
|
||||
|
||||
|
||||
def subset(idats: Path, n: int, dest: Path) -> Path:
|
||||
grn = sorted(p for p in idats.rglob("*_Grn.idat*"))
|
||||
if n <= 0 or n >= len(grn):
|
||||
return idats
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
for g in grn[:n]:
|
||||
d = dest / g.parent.name
|
||||
d.mkdir(exist_ok=True)
|
||||
for p in (g, g.with_name(g.name.replace("_Grn.", "_Red."))):
|
||||
link = d / p.name
|
||||
if not link.exists():
|
||||
link.symlink_to(p.resolve())
|
||||
return dest
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--bpm", required=True)
|
||||
ap.add_argument("--egt", required=True)
|
||||
ap.add_argument("--idats", required=True, type=Path)
|
||||
ap.add_argument("--out", required=True, type=Path)
|
||||
ap.add_argument("--threads", default="1,2,4")
|
||||
ap.add_argument("--samples", type=int, default=0, help="use the first N samples (0 = all)")
|
||||
ap.add_argument("--simurg", default=str(Path(sys.executable).with_name("simurg")))
|
||||
ap.add_argument("--pure-python", action="store_true", help="run with SIMURG_PURE_PYTHON=1")
|
||||
a = ap.parse_args()
|
||||
a.out.mkdir(parents=True, exist_ok=True)
|
||||
idats = subset(a.idats, a.samples, a.out / "idats")
|
||||
n = len(list(idats.rglob("*_Grn.idat*")))
|
||||
env = {**os.environ, "SIMURG_PURE_PYTHON": "1" if a.pure_python else "0"}
|
||||
runs = []
|
||||
for t in (int(x) for x in a.threads.split(",")):
|
||||
out = a.out / f"t{t}"
|
||||
t0 = time.time()
|
||||
cmd = [a.simurg, "genotype", "--bpm", a.bpm, "--egt", a.egt, "--idats", str(idats), "--out", str(out),
|
||||
"--threads", str(t), "--no-cohort-qc", "--autocall-date", "1/1/2026 12:00:00 AM",
|
||||
"--quiet", "--force"] # fmt: skip
|
||||
rc = subprocess.run(cmd, env=env).returncode
|
||||
wall = time.time() - t0
|
||||
setup, phase = phase_seconds(out / "logs" / "genotype.log")
|
||||
runs.append(
|
||||
{
|
||||
"threads": t,
|
||||
"rc": rc,
|
||||
"wall_s": round(wall, 1),
|
||||
"setup_s": round(setup, 1),
|
||||
"per_sample_phase_s": round(phase, 2),
|
||||
"samples_per_min": round(60 * n / phase, 1) if phase else None,
|
||||
"samples_per_min_per_core": round(60 * n / phase / t, 1) if phase else None,
|
||||
}
|
||||
)
|
||||
print(json.dumps(runs[-1]), flush=True)
|
||||
base = runs[0]["per_sample_phase_s"] * runs[0]["threads"]
|
||||
for r in runs:
|
||||
r["speedup_vs_first"] = round(base / r["per_sample_phase_s"], 2) if r["per_sample_phase_s"] else None
|
||||
r["efficiency"] = round(r["speedup_vs_first"] / r["threads"], 3) if r["speedup_vs_first"] else None
|
||||
from simurg._core import backend
|
||||
|
||||
summary = {
|
||||
"samples": n,
|
||||
"backend": "python" if a.pure_python else backend(),
|
||||
"cpu_count": os.cpu_count(),
|
||||
"machine": os.uname().machine,
|
||||
"runs": runs,
|
||||
}
|
||||
(a.out / "scaling.json").write_text(json.dumps(summary, indent=1))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in new issue
Block a user