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
7bdd5dc03c
commit
95e2338897
19 files changed
+1755
-6
No files matched your search
@@ -0,0 +1,275 @@
|
||||
//! BPM LocusEntry records — port of `simurg.formats.bpm._parse_locus_entries` (format: Illumina BeadArrayFiles
|
||||
//! `BeadPoolManifest.py`, BSD-2). Same field order, checks and error messages; strings of ≤ 64 bytes are shared
|
||||
//! Python objects, as in the Python parser.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use numpy::{IntoPyArray, PyArray1, PyArray2};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyList, PyString};
|
||||
|
||||
pub const STR_COLS: usize = 16;
|
||||
// output column order (matches the Python `str_cols` list)
|
||||
const ILMN: usize = 0;
|
||||
const SNP: usize = 1;
|
||||
const CHR: usize = 2;
|
||||
const ISTR: usize = 3;
|
||||
const SSTR: usize = 4;
|
||||
const RSTR: usize = 5;
|
||||
const PLOIDY: usize = 6;
|
||||
const SPECIES: usize = 7;
|
||||
const BUILD: usize = 8;
|
||||
const SOURCE: usize = 9;
|
||||
const SVER: usize = 10;
|
||||
const CSTR: usize = 11;
|
||||
const SSEQ: usize = 12;
|
||||
const TOP: usize = 13;
|
||||
const PA: usize = 14;
|
||||
const PB: usize = 15;
|
||||
|
||||
struct Cursor<'a> {
|
||||
b: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
fn err(msg: impl Into<String>) -> PyErr {
|
||||
PyValueError::new_err(msg.into())
|
||||
}
|
||||
|
||||
impl<'a> Cursor<'a> {
|
||||
fn string(&mut self) -> PyResult<&'a [u8]> {
|
||||
let size = self.b.len();
|
||||
let mut p = self.pos;
|
||||
if p >= size {
|
||||
return Err(err("unexpected end of file in locus entry"));
|
||||
}
|
||||
let mut byte = self.b[p];
|
||||
p += 1;
|
||||
let mut ln = (byte & 0x7f) as usize;
|
||||
let mut shift = 7;
|
||||
while byte & 0x80 != 0 {
|
||||
if p >= size || shift > 28 {
|
||||
return Err(err("bad string length in locus entry"));
|
||||
}
|
||||
byte = self.b[p];
|
||||
p += 1;
|
||||
ln |= ((byte & 0x7f) as usize) << shift;
|
||||
shift += 7;
|
||||
}
|
||||
let end = p + ln;
|
||||
if end > size {
|
||||
return Err(err("string runs past end of file in locus entry"));
|
||||
}
|
||||
self.pos = end;
|
||||
Ok(&self.b[p..end])
|
||||
}
|
||||
fn i32(&mut self, what: &str) -> PyResult<i32> {
|
||||
if self.pos + 4 > self.b.len() {
|
||||
return Err(err(format!("unexpected end of file at {what}")));
|
||||
}
|
||||
let v = i32::from_le_bytes(self.b[self.pos..self.pos + 4].try_into().unwrap());
|
||||
self.pos += 4;
|
||||
Ok(v)
|
||||
}
|
||||
}
|
||||
|
||||
type Parsed<'py> = (
|
||||
Vec<Bound<'py, PyList>>,
|
||||
Bound<'py, PyArray1<i64>>,
|
||||
Bound<'py, PyArray1<i32>>,
|
||||
Bound<'py, PyArray1<i32>>,
|
||||
Bound<'py, PyArray1<i16>>,
|
||||
Bound<'py, PyArray1<i8>>,
|
||||
Bound<'py, PyArray1<bool>>,
|
||||
Bound<'py, PyArray1<i8>>,
|
||||
Option<Bound<'py, PyArray2<f32>>>,
|
||||
usize,
|
||||
);
|
||||
|
||||
/// Parse `n` LocusEntry records starting at byte 0 of `buf`. Returns the 16 string columns, map_infos,
|
||||
/// address_a, address_b, assay types (−1 = infer from sequences, entry version < 6), expected clusters,
|
||||
/// intensity-only flags, entry versions, base fractions (None unless some entry is v7+) and the bytes consumed.
|
||||
/// Errors are ValueError with the Python parser's message.
|
||||
#[pyfunction]
|
||||
pub fn bpm_locus_entries<'py>(py: Python<'py>, buf: &[u8], n: usize, names: Vec<String>) -> PyResult<Parsed<'py>> {
|
||||
let name_to_idx: HashMap<&str, usize> = names.iter().enumerate().map(|(i, s)| (s.as_str(), i)).collect();
|
||||
let mut cols: Vec<Vec<Option<Py<PyString>>>> = (0..STR_COLS).map(|_| (0..n).map(|_| None).collect()).collect();
|
||||
let mut intern: HashMap<&[u8], Py<PyString>> = HashMap::new();
|
||||
let mut map_infos = vec![0i64; n];
|
||||
let mut address_a = vec![0i32; n];
|
||||
let mut address_b = vec![0i32; n];
|
||||
let mut assay = vec![0i16; n];
|
||||
let mut exp_clusters = vec![0i8; n];
|
||||
let mut int_only = vec![false; n];
|
||||
let mut versions = vec![0i8; n];
|
||||
let mut fracs = vec![0f32; 4 * n];
|
||||
let mut seen = vec![false; n];
|
||||
let mut any_v7 = false;
|
||||
let mut c = Cursor { b: buf, pos: 0 };
|
||||
|
||||
let to_py = |raw: &'_ [u8], intern: &mut HashMap<&[u8], Py<PyString>>| -> Py<PyString> {
|
||||
if let Some(s) = intern.get(raw) {
|
||||
return s.clone_ref(py);
|
||||
}
|
||||
let s = PyString::new(py, &String::from_utf8_lossy(raw)).unbind();
|
||||
s
|
||||
};
|
||||
|
||||
for _ in 0..n {
|
||||
let ver = c.i32("locus entry")?;
|
||||
if ![4, 6, 7, 8].contains(&ver) {
|
||||
return Err(err(format!("unsupported locus entry version {ver}")));
|
||||
}
|
||||
let ilmn_id = c.string()?;
|
||||
let name = c.string()?;
|
||||
for _ in 0..3 {
|
||||
c.string()?;
|
||||
}
|
||||
c.i32("locus index")?;
|
||||
c.string()?;
|
||||
let ilmn_strand = c.string()?;
|
||||
let snp = c.string()?;
|
||||
let chrom = c.string()?;
|
||||
let ploidy = c.string()?;
|
||||
let species = c.string()?;
|
||||
let map_info = c.string()?;
|
||||
let top_seq = c.string()?;
|
||||
let cust_strand = c.string()?;
|
||||
if c.pos + 8 > buf.len() {
|
||||
return Err(err("unexpected end of file at addresses"));
|
||||
}
|
||||
let a_addr = c.i32("addresses")?;
|
||||
let b_addr = c.i32("addresses")?;
|
||||
let probe_a = c.string()?;
|
||||
let probe_b = c.string()?;
|
||||
let build = c.string()?;
|
||||
let source = c.string()?;
|
||||
let source_ver = c.string()?;
|
||||
let source_strand = c.string()?;
|
||||
let source_seq = c.string()?;
|
||||
let (mut at, mut exp_cl, mut io) = (-1i16, 0i8, 0u8);
|
||||
if ver >= 6 {
|
||||
if c.pos + 4 > buf.len() {
|
||||
return Err(err("unexpected end of file at assay type"));
|
||||
}
|
||||
exp_cl = buf[c.pos + 1] as i8;
|
||||
io = buf[c.pos + 2];
|
||||
at = buf[c.pos + 3] as i16;
|
||||
c.pos += 4;
|
||||
if at > 2 {
|
||||
return Err(err(format!("invalid assay type {at}")));
|
||||
}
|
||||
if (b_addr == 0) != (at == 0) {
|
||||
return Err(err("assay type is inconsistent with address B"));
|
||||
}
|
||||
}
|
||||
let mut f4: Option<[f32; 4]> = None;
|
||||
if ver >= 7 {
|
||||
if c.pos + 16 > buf.len() {
|
||||
return Err(err("unexpected end of file at base fractions"));
|
||||
}
|
||||
let mut f = [0f32; 4];
|
||||
for (k, v) in f.iter_mut().enumerate() {
|
||||
*v = f32::from_le_bytes(buf[c.pos + 4 * k..c.pos + 4 * k + 4].try_into().unwrap());
|
||||
}
|
||||
f4 = Some(f);
|
||||
c.pos += 16;
|
||||
}
|
||||
let ref_strand: &[u8] = if ver >= 8 { c.string()? } else { b"" };
|
||||
|
||||
let name_s = String::from_utf8_lossy(name);
|
||||
let idx = match name_to_idx.get(name_s.as_ref()) {
|
||||
Some(&i) => i,
|
||||
None => return Err(err(format!("locus entry {} not in the names table", py_repr(&name_s)))),
|
||||
};
|
||||
if seen[idx] {
|
||||
return Err(err(format!("locus entry {} appears twice", py_repr(&name_s))));
|
||||
}
|
||||
seen[idx] = true;
|
||||
for (col, raw) in [
|
||||
(ILMN, ilmn_id),
|
||||
(SNP, snp),
|
||||
(CHR, chrom),
|
||||
(ISTR, ilmn_strand),
|
||||
(SSTR, source_strand),
|
||||
(RSTR, ref_strand),
|
||||
(PLOIDY, ploidy),
|
||||
(SPECIES, species),
|
||||
(BUILD, build),
|
||||
(SOURCE, source),
|
||||
(SVER, source_ver),
|
||||
(CSTR, cust_strand),
|
||||
(SSEQ, source_seq),
|
||||
(TOP, top_seq),
|
||||
(PA, probe_a),
|
||||
(PB, probe_b),
|
||||
] {
|
||||
let s = to_py(raw, &mut intern);
|
||||
if raw.len() <= 64 && !intern.contains_key(raw) {
|
||||
intern.insert(raw, s.clone_ref(py));
|
||||
}
|
||||
cols[col][idx] = Some(s);
|
||||
}
|
||||
map_infos[idx] = if map_info.is_empty() {
|
||||
0
|
||||
} else {
|
||||
let t = String::from_utf8_lossy(map_info);
|
||||
match t.trim().parse::<i64>() {
|
||||
Ok(v) => v,
|
||||
Err(_) => {
|
||||
return Err(err(format!(
|
||||
"non-integer MapInfo {} for {}",
|
||||
py_repr(&t),
|
||||
py_repr(&name_s)
|
||||
)))
|
||||
}
|
||||
}
|
||||
};
|
||||
address_a[idx] = a_addr;
|
||||
address_b[idx] = b_addr;
|
||||
assay[idx] = at;
|
||||
exp_clusters[idx] = exp_cl;
|
||||
int_only[idx] = io != 0;
|
||||
versions[idx] = ver as i8;
|
||||
if let Some(f) = f4 {
|
||||
fracs[4 * idx..4 * idx + 4].copy_from_slice(&f);
|
||||
any_v7 = true;
|
||||
}
|
||||
}
|
||||
let missing = seen.iter().filter(|s| !**s).count();
|
||||
if missing > 0 {
|
||||
return Err(err(format!("{missing} loci in the names table have no locus entry")));
|
||||
}
|
||||
let empty = PyString::new(py, "").unbind();
|
||||
let lists = cols
|
||||
.into_iter()
|
||||
.map(|col| PyList::new(py, col.into_iter().map(|s| s.unwrap_or_else(|| empty.clone_ref(py)))))
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let fr = if any_v7 {
|
||||
Some(numpy::ndarray::Array2::from_shape_vec((n, 4), fracs).map_err(|e| err(e.to_string()))?.into_pyarray(py))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok((
|
||||
lists,
|
||||
map_infos.into_pyarray(py),
|
||||
address_a.into_pyarray(py),
|
||||
address_b.into_pyarray(py),
|
||||
assay.into_pyarray(py),
|
||||
exp_clusters.into_pyarray(py),
|
||||
int_only.into_pyarray(py),
|
||||
versions.into_pyarray(py),
|
||||
fr,
|
||||
c.pos,
|
||||
))
|
||||
}
|
||||
|
||||
/// Python `repr()` of a str for error messages (single quotes unless the text contains one).
|
||||
fn py_repr(s: &str) -> String {
|
||||
if s.contains('\'') && !s.contains('"') {
|
||||
format!("\"{s}\"")
|
||||
} else {
|
||||
format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'"))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user