kemal
6524a6d1b0
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).
84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
#!/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/claude-1000/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()
|