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:
1 parent
f64bdf11f2
commit
b0494be576
19 files changed
+1755
-6
No files matched your search
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user