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
f11b39b5a3
commit
6524a6d1b0
19 files changed
+1755
-6
No files matched your search
@@ -0,0 +1,238 @@
|
||||
"""The compiled kernels (``rust/simurg-core``, plan §18 P9) must equal the pure-Python paths bit for bit."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
import math
|
||||
from dataclasses import fields, replace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from tests.conftest import GSA_BPM, GSA_EGT, GSA_IDATS
|
||||
|
||||
import simurg.formats.bpm as bpm_mod
|
||||
import simurg.genotype.baf_lrr as baf_mod
|
||||
import simurg.genotype.gencall as gencall_mod
|
||||
import simurg.signal.normalization.fit as fit_mod
|
||||
from simurg._core import core
|
||||
from simurg.core.errors import FormatError
|
||||
from simurg.formats.bpm import read_bpm
|
||||
from simurg.synth.files import SynthLocus, encode_bpm
|
||||
|
||||
pytestmark = pytest.mark.skipif(core is None, reason="simurg_core extension not installed")
|
||||
|
||||
MODULES = (bpm_mod, gencall_mod, fit_mod, baf_mod)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pure(monkeypatch): # type: ignore[no-untyped-def]
|
||||
"""Context switch: ``with pure(): ...`` runs the pure-Python paths."""
|
||||
|
||||
class _Pure:
|
||||
def __enter__(self) -> None:
|
||||
for m in MODULES:
|
||||
monkeypatch.setattr(m, "core", None)
|
||||
|
||||
def __exit__(self, *exc: object) -> None:
|
||||
for m in MODULES:
|
||||
monkeypatch.setattr(m, "core", core)
|
||||
|
||||
return _Pure
|
||||
|
||||
|
||||
def _bits(a: np.ndarray) -> bytes:
|
||||
return np.ascontiguousarray(a).tobytes()
|
||||
|
||||
|
||||
def _edge_values() -> np.ndarray:
|
||||
neg_nan = np.frombuffer(np.uint64(0xFFF8000000000000).tobytes(), np.float64)[0]
|
||||
payload_nan = np.frombuffer(np.uint64(0x7FF8000000001234).tobytes(), np.float64)[0]
|
||||
tiny = np.finfo(np.float64).tiny
|
||||
special = [0.0, -0.0, 1.0, -1.0, math.inf, -math.inf, math.nan, neg_nan, payload_nan, tiny, tiny / 4, -tiny,
|
||||
1e-320, 709.78, 709.79, -745.2, -746.0, 1e300, -1e300, 2.0**-1074, 0.5, 10.0, 1e-5] # fmt: skip
|
||||
rng = np.random.default_rng(7)
|
||||
rand = np.concatenate(
|
||||
[rng.normal(0, 3, 20_000), rng.uniform(-800, 800, 5_000), np.exp(rng.uniform(-700, 700, 5_000))]
|
||||
)
|
||||
return np.concatenate([np.array(special, np.float64), rand, -rand[:1000]])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fn", [math.exp, math.log10, math.log2, gencall_mod._pow(1.71), gencall_mod._pow(0.35)])
|
||||
def test_libm_kernels_bit_identical(fn, pure) -> None: # type: ignore[no-untyped-def]
|
||||
x = _edge_values()
|
||||
fast = gencall_mod._libm(fn, x)
|
||||
with pure():
|
||||
ref = gencall_mod._libm(fn, x)
|
||||
assert fast.dtype == ref.dtype == np.float64
|
||||
assert _bits(fast) == _bits(ref)
|
||||
|
||||
|
||||
def test_pow_kernel_rejects_integer_exponent() -> None:
|
||||
assert core is not None
|
||||
with pytest.raises(ValueError):
|
||||
core.pow(np.ones(3), 2.0)
|
||||
|
||||
|
||||
def test_closest_binned_identical(pure) -> None: # type: ignore[no-untyped-def]
|
||||
rng = np.random.default_rng(3)
|
||||
for trial in range(20):
|
||||
n = int(rng.integers(0, 3000))
|
||||
raw_x = rng.gamma(2.0, 400.0, n).astype(np.float32)
|
||||
raw_y = rng.gamma(1.0, 60.0 if trial % 2 else 3000.0, n).astype(np.float32)
|
||||
if trial == 5 and n:
|
||||
raw_x[: n // 3] = raw_x[0] # ties: first minimum must win
|
||||
raw_y[: n // 3] = raw_y[0]
|
||||
if trial == 7 and n:
|
||||
raw_x[n // 2] = np.nan
|
||||
m = int(rng.integers(1, 400))
|
||||
axis = np.linspace(0, float(np.float32(rng.uniform(100, 5000))), m).astype(np.float32)
|
||||
zeros = np.zeros(m, np.float32)
|
||||
ax, ay = (axis, zeros) if trial % 3 else (zeros, axis)
|
||||
fast = fit_mod._closest_binned(raw_x, raw_y, ax, ay)
|
||||
with pure():
|
||||
ref = fit_mod._closest_binned(raw_x, raw_y, ax, ay)
|
||||
assert fast.dtype == ref.dtype
|
||||
assert np.array_equal(fast, ref), trial
|
||||
|
||||
|
||||
def _loci() -> list[SynthLocus]:
|
||||
return [
|
||||
SynthLocus("rs1", "1", 100, "[A/G]", 11, 0, 0, norm_id=1, source_seq="ACGT[A/G]TTGA", probe_a="ACGTA"),
|
||||
SynthLocus("rs2", "2", 200, "[A/T]", 12, 13, 1, ilmn_strand="BOT", ref_strand="-", probe_a="AC", probe_b="AT"),
|
||||
SynthLocus("rs3", "X", 300, "[C/G]", 14, 15, 2, norm_id=99),
|
||||
SynthLocus("rsü", "Y", 400, "[T/C]", 16, 0, 0, norm_id=2, intensity_only=1),
|
||||
SynthLocus("rs5", "MT", 500, "[A/C]", 17, 0, 0, norm_id=99, source_seq="A" * 200 + "[A/C]"),
|
||||
]
|
||||
|
||||
|
||||
def _manifest_dict(m) -> dict: # type: ignore[no-untyped-def]
|
||||
out = {}
|
||||
for f in fields(m):
|
||||
v = getattr(m, f.name)
|
||||
out[f.name] = (v.dtype.str, v.shape, v.tobytes()) if isinstance(v, np.ndarray) else v
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry_version", [4, 6, 7, 8])
|
||||
@pytest.mark.parametrize("order", [None, [4, 3, 2, 1, 0]])
|
||||
def test_bpm_parser_identical(tmp_path, pure, entry_version, order) -> None: # type: ignore[no-untyped-def]
|
||||
p = tmp_path / "m.bpm"
|
||||
p.write_bytes(encode_bpm(_loci(), entry_version=entry_version, entry_order=order))
|
||||
fast = _manifest_dict(read_bpm(p))
|
||||
with pure():
|
||||
ref = _manifest_dict(read_bpm(p))
|
||||
assert fast == ref
|
||||
|
||||
|
||||
def test_bpm_parser_errors_identical(tmp_path, pure) -> None: # type: ignore[no-untyped-def]
|
||||
good = encode_bpm(_loci())
|
||||
cases = {
|
||||
"truncated": good[: len(good) - 40],
|
||||
"bad_assay": encode_bpm([SynthLocus("rs1", address_b=0, assay_type=1)]),
|
||||
"bad_map": encode_bpm([SynthLocus("rs1", pos=12345)]).replace(b"\x0512345", b"\x0512x45"),
|
||||
}
|
||||
for label, data in cases.items():
|
||||
p = tmp_path / f"{label}.bpm"
|
||||
p.write_bytes(data)
|
||||
with pytest.raises(FormatError) as fast:
|
||||
read_bpm(p)
|
||||
with pure(), pytest.raises(FormatError) as ref:
|
||||
read_bpm(p)
|
||||
assert (fast.value.code, str(fast.value)) == (ref.value.code, str(ref.value)), label
|
||||
|
||||
|
||||
@pytest.mark.data
|
||||
@pytest.mark.skipif(not GSA_BPM.exists(), reason="GSA demo data not fetched")
|
||||
def test_genotype_sample_identical(pure) -> None: # type: ignore[no-untyped-def]
|
||||
from simurg.formats.egt import read_egt
|
||||
from simurg.formats.gtc import encode_gtc
|
||||
from simurg.formats.idat import read_idat
|
||||
from simurg.genotype.clusters import align_clusters
|
||||
from simurg.genotype.sample import genotype_sample, manifest_context, params_for_preset
|
||||
from simurg.signal.raw_xy import LocusMapCache
|
||||
|
||||
def run() -> bytes:
|
||||
m = read_bpm(GSA_BPM)
|
||||
ctx = manifest_context(m, align_clusters(read_egt(GSA_EGT), m))
|
||||
ctx = replace(ctx, manifest=m.lite(), cache=LocusMapCache(m.lite()))
|
||||
stem = next(iter(sorted(GSA_IDATS.rglob("*_Grn.idat")))).as_posix()[: -len("_Grn.idat")]
|
||||
g = genotype_sample(read_idat(stem + "_Grn.idat"), read_idat(stem + "_Red.idat"), ctx,
|
||||
params_for_preset("current"), dt.datetime(2026, 1, 1)) # fmt: skip
|
||||
return encode_gtc(g)
|
||||
|
||||
fast = run()
|
||||
with pure():
|
||||
ref = run()
|
||||
assert fast == ref
|
||||
|
||||
|
||||
def test_closest_exact_identical(pure) -> None: # type: ignore[no-untyped-def]
|
||||
rng = np.random.default_rng(11)
|
||||
for trial in range(10):
|
||||
n = int(rng.integers(1, 1500))
|
||||
xr = rng.normal(500, 200, n).astype(np.float32)
|
||||
yr = rng.normal(300, 100, n).astype(np.float32)
|
||||
if trial == 2:
|
||||
xr[n // 2 :] = xr[0]
|
||||
yr[n // 2 :] = yr[0]
|
||||
if trial == 3:
|
||||
xr[n // 3] = np.nan
|
||||
m = int(rng.integers(1, 300))
|
||||
x = rng.normal(500, 200, m).astype(np.float32)
|
||||
y = rng.normal(300, 100, m).astype(np.float32)
|
||||
fast = fit_mod._closest_slow(xr, yr, x, y)
|
||||
with pure():
|
||||
ref = fit_mod._closest_slow(xr, yr, x, y)
|
||||
assert np.array_equal(fast, ref), trial
|
||||
|
||||
|
||||
def _synthetic_clusters(n: int, rng: np.random.Generator): # type: ignore[no-untyped-def]
|
||||
from simurg.genotype.clusters import ClusterArrays
|
||||
|
||||
tm = np.sort(rng.uniform(0, 1, (n, 3)), axis=1).astype(np.float32)
|
||||
td = rng.uniform(0.0, 0.08, (n, 3)).astype(np.float32)
|
||||
rm = rng.uniform(0.2, 3, (n, 3)).astype(np.float32)
|
||||
rd = rng.uniform(0.0, 0.3, (n, 3)).astype(np.float32)
|
||||
counts = rng.integers(0, 200, (n, 3)).astype(np.int32)
|
||||
counts[rng.random((n, 3)) < 0.2] = 0
|
||||
counts[: n // 50] = 0 # no samples at all
|
||||
td[rng.random((n, 3)) < 0.02] = 0.0
|
||||
tm[rng.random((n, 3)) < 0.01] = np.nan
|
||||
rm[rng.random((n, 3)) < 0.01] = np.nan
|
||||
thr = rng.uniform(0, 0.3, n).astype(np.float32)
|
||||
thr[rng.random(n) < 0.3] = np.nan
|
||||
gt = rng.uniform(0, 1, n).astype(np.float32)
|
||||
gt[rng.random(n) < 0.02] = np.nan
|
||||
return ClusterArrays(
|
||||
present=rng.random(n) > 0.01, n=counts, theta_mean=tm, theta_dev=td, r_mean=rm, r_dev=rd,
|
||||
intensity_threshold=thr, gentrain=gt, mean_r=rng.uniform(0.2, 3, n).astype(np.float32),
|
||||
) # fmt: skip
|
||||
|
||||
|
||||
def test_call_genotypes_and_baf_lrr_identical(pure) -> None: # type: ignore[no-untyped-def]
|
||||
from simurg.genotype.baf_lrr import baf_lrr
|
||||
from simurg.genotype.gencall import call_genotypes
|
||||
|
||||
rng = np.random.default_rng(5)
|
||||
n = 60_000
|
||||
c = _synthetic_clusters(n, rng)
|
||||
theta = rng.uniform(-0.05, 1.05, n).astype(np.float32)
|
||||
theta[: n // 3] = c.theta_mean[: n // 3, rng.integers(0, 3)] # exactly on a cluster centre
|
||||
r = rng.gamma(2.0, 0.6, n).astype(np.float32)
|
||||
raw_zero = rng.random(n) < 0.01
|
||||
theta[raw_zero] = np.nan
|
||||
r[raw_zero] = np.nan
|
||||
r[rng.random(n) < 0.01] = 0.0
|
||||
intensity_only = rng.random(n) < 0.02
|
||||
for cutoff in (0.15, 0.0):
|
||||
fast = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff)
|
||||
with pure():
|
||||
ref = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff)
|
||||
assert _bits(fast.genotypes) == _bits(ref.genotypes)
|
||||
assert _bits(fast.scores) == _bits(ref.scores)
|
||||
assert (fast.num_calls, fast.num_intensity_only) == (ref.num_calls, ref.num_intensity_only)
|
||||
fb, fl = baf_lrr(theta, r, raw_zero, c, intensity_only)
|
||||
with pure():
|
||||
rb, rl = baf_lrr(theta, r, raw_zero, c, intensity_only)
|
||||
assert _bits(fb) == _bits(rb)
|
||||
assert _bits(fl) == _bits(rl)
|
||||
Reference in new issue
Block a user