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 f64bdf11f2
commit b0494be576
19 files changed
+1755 -6

No files matched your search

+51 -1
View File
@@ -32,6 +32,7 @@ from typing import Any
import numpy as np
import numpy.typing as npt
from simurg._core import core
from simurg.core.errors import FormatError
from simurg.formats.binutil import BinaryReader, load_bytes
@@ -223,7 +224,9 @@ def read_bpm(path: str | Path) -> Manifest:
def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]:
"""Hot loop: parse ``n`` LocusEntry records (optimised pure Python; Rust in phase P9)."""
"""Hot loop: parse ``n`` LocusEntry records (compiled in ``simurg_core`` when installed; same results)."""
if core is not None:
return _parse_locus_entries_core(r, n, name_to_idx)
view = memoryview(r._buf)
size = r.size
pos = r.pos
@@ -416,6 +419,53 @@ def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -
}
_STR_COLS = (
"ilmn_ids",
"snps",
"chroms",
"ilmn_strands",
"source_strands",
"ref_strands",
"ploidies",
"species",
"genome_builds",
"sources",
"source_versions",
"customer_strands",
"source_seqs",
"top_genomic_seqs",
"allele_a_probe_seqs",
"allele_b_probe_seqs",
)
def _parse_locus_entries_core(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]:
"""``_parse_locus_entries`` through ``simurg_core.bpm_locus_entries`` (identical columns and error messages)."""
assert core is not None
pos = r.pos
try:
lists, map_infos, a, b, assay, exp_cl, int_only, versions, fracs, used = core.bpm_locus_entries(
bytes(r._view[pos:]), n, list(name_to_idx)
)
except ValueError as e:
raise r.fail(str(e)) from None
out: dict[str, Any] = dict(zip(_STR_COLS, lists, strict=True))
for i in np.flatnonzero(assay < 0):
assay[i] = infer_assay_type(out["allele_a_probe_seqs"][i], out["allele_b_probe_seqs"][i], out["source_seqs"][i])
r.pos = pos + used
return {
**out,
"map_infos": map_infos,
"address_a": a,
"address_b": b,
"assay_types": assay.astype(np.uint8),
"exp_clusters": exp_cl,
"intensity_only": int_only,
"entry_versions": versions,
"base_fractions": fracs,
}
_IUPAC = {
"A": 1,
"C": 2,