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 2026-09-28 21:28:55 +01:00
1 parent 7bdd5dc03c
commit 95e2338897
19 files changed
+1755 -6

No files matched your search

+12
View File
@@ -11,6 +11,7 @@ import math
import numpy as np
import numpy.typing as npt
from simurg._core import core
from simurg.genotype.clusters import ClusterArrays
from simurg.genotype.gencall import _libm
@@ -22,6 +23,17 @@ NEG_NAN = np.frombuffer(np.uint32(0xFFC00000).tobytes(), dtype=np.float32)[0] #
def baf_lrr(
theta: F32Arr, r: F32Arr, raw_zero: npt.NDArray[np.bool_], c: ClusterArrays, intensity_only: npt.NDArray[np.bool_]
) -> tuple[F32Arr, F32Arr]:
if core is not None:
return core.baf_lrr( # type: ignore[no-any-return]
np.ascontiguousarray(theta, np.float32),
np.ascontiguousarray(r, np.float32),
np.ascontiguousarray(raw_zero, np.bool_),
np.ascontiguousarray(c.present, np.bool_),
np.ascontiguousarray(c.theta_mean, np.float32).ravel(),
np.ascontiguousarray(c.r_mean, np.float32).ravel(),
np.ascontiguousarray(c.mean_r, np.float32),
np.ascontiguousarray(intensity_only, np.bool_),
)
t = np.asarray(theta, np.float32)
rr = np.asarray(r, np.float32)
at, bt, ct = c.theta_mean[:, 0], c.theta_mean[:, 1], c.theta_mean[:, 2]
+29 -1
View File
@@ -15,6 +15,7 @@ from dataclasses import dataclass
import numpy as np
import numpy.typing as npt
from simurg._core import core
from simurg.genotype.clusters import ClusterArrays
F32 = np.float32
@@ -42,8 +43,17 @@ def _c_call(fn, v: float) -> float: # type: ignore[no-untyped-def]
return math.inf
_CORE_FN = {math.exp: "exp", math.log10: "log10", math.log2: "log2"}
def _libm(fn, x: F64Arr) -> F64Arr: # type: ignore[no-untyped-def]
"""Element-wise libm call (bit-identical to C)."""
"""Element-wise libm call (bit-identical to C); compiled when ``simurg_core`` is installed."""
if core is not None:
xa = np.ascontiguousarray(x, dtype=np.float64).ravel()
if (e := getattr(fn, "exponent", None)) is not None:
return core.pow(xa, e).reshape(np.shape(x)) # type: ignore[no-any-return]
if (name := _CORE_FN.get(fn)) is not None:
return getattr(core, name)(xa).reshape(np.shape(x)) # type: ignore[no-any-return]
return np.fromiter((_c_call(fn, v) for v in x.tolist()), dtype=np.float64, count=len(x))
@@ -51,6 +61,7 @@ def _pow(e: float): # type: ignore[no-untyped-def]
def f(v: float) -> float:
return math.pow(v, e)
f.exponent = e # type: ignore[attr-defined]
return f
@@ -137,6 +148,23 @@ def call_genotypes(
cutoff: float = DEFAULT_CUTOFF,
) -> Calls:
"""Bayesian calls and GenCall scores for one sample (plan §7.5)."""
if core is not None:
f32 = [np.ascontiguousarray(a, np.float32).ravel() for a in (c.theta_mean, c.theta_dev, c.r_mean, c.r_dev)]
geno, score = core.call_genotypes(
np.ascontiguousarray(theta, np.float32),
np.ascontiguousarray(r, np.float32),
np.ascontiguousarray(raw_zero, np.bool_),
np.ascontiguousarray(c.present, np.bool_),
np.ascontiguousarray(c.n, np.int32).ravel(),
*f32,
np.ascontiguousarray(c.intensity_threshold, np.float32),
np.ascontiguousarray(c.gentrain, np.float32),
np.ascontiguousarray(intensity_only, np.bool_),
F32(cutoff),
)
return Calls(
genotypes=geno, scores=score, num_calls=int((geno > 0).sum()), num_intensity_only=int(intensity_only.sum())
)
n = len(theta)
t = np.asarray(theta, np.float32)
rr = np.asarray(r, np.float32)