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
f64bdf11f2
commit
b0494be576
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('\'', "\\'"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,357 @@
|
||||
//! Fused per-locus ports of `simurg.genotype.gencall.call_genotypes` and `simurg.genotype.baf_lrr.baf_lrr`
|
||||
//! (which port gtc2vcf `gen_std_flair`/`modilik`/`compute_score_call_prelim`/`gencall_score_map`/`make_calls`
|
||||
//! and `compute_baf_lrr`, MIT, © Giulio Genovese). Every intermediate keeps the NumPy reference's type (f32 or
|
||||
//! f64) and evaluation order; NumPy's `fmin`/`fmax` are `f32::min`/`max`; libm calls use CPython `math`
|
||||
//! semantics. One pass per locus instead of dozens of full-array temporaries.
|
||||
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::{math_exp, math_log, math_pow};
|
||||
|
||||
const M_2_SQRTPI: f64 = 1.12837916709551257390;
|
||||
const M_SQRT1_2: f64 = 0.707106781186547524401;
|
||||
const NEG_NAN32: f32 = f32::from_bits(0xffc0_0000);
|
||||
|
||||
#[inline]
|
||||
fn smf(x: f32, a: f32, b: f32) -> f32 {
|
||||
let mid = (a + b) / 2.0;
|
||||
if a >= b {
|
||||
return if x >= mid { 1.0 } else { 0.0 };
|
||||
}
|
||||
let u = (x - a) / (b - a);
|
||||
let v = (x - b) / (b - a);
|
||||
if x <= a {
|
||||
0.0
|
||||
} else if a < x && x <= mid {
|
||||
2.0 * (u * u)
|
||||
} else if mid < x && x <= b {
|
||||
1.0 - 2.0 * (v * v)
|
||||
} else {
|
||||
1.0
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn zmf(x: f32, a: f32, b: f32) -> f32 {
|
||||
let mid = (a + b) / 2.0;
|
||||
let u = (x - a) / (b - a);
|
||||
let v = (x - b) / (b - a);
|
||||
if x <= a {
|
||||
1.0
|
||||
} else if a < x && x <= mid {
|
||||
1.0 - 2.0 * (u * u)
|
||||
} else if mid < x && x <= b {
|
||||
2.0 * (v * v)
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Normal pdf pieces of `_pdf`: (full, peak, bad) for value x, mean mu, sd sigma.
|
||||
#[inline]
|
||||
fn pdf(x: f32, mu: f32, sigma: f32) -> (f64, f64, bool) {
|
||||
let z = ((x - mu) / sigma) as f64;
|
||||
let e = math_exp(-0.5 * (z * z));
|
||||
let sd = sigma as f64;
|
||||
let full = e * 0.5 * M_2_SQRTPI * M_SQRT1_2 / sd;
|
||||
let peak = (0.5 * M_2_SQRTPI * M_SQRT1_2) / sd;
|
||||
(full, peak, sigma <= 0.0)
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Side {
|
||||
Left,
|
||||
Right,
|
||||
Full,
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get(p: (f64, f64, bool), x: f32, mu: f32, side: Side) -> f64 {
|
||||
let (full, peak, bad) = p;
|
||||
if bad {
|
||||
return f64::NAN;
|
||||
}
|
||||
match side {
|
||||
Side::Left => {
|
||||
if x < mu {
|
||||
peak
|
||||
} else {
|
||||
full
|
||||
}
|
||||
}
|
||||
Side::Right => {
|
||||
if x > mu {
|
||||
peak
|
||||
} else {
|
||||
full
|
||||
}
|
||||
}
|
||||
Side::Full => full,
|
||||
}
|
||||
}
|
||||
|
||||
fn check(len: usize, n: usize, what: &str) -> PyResult<()> {
|
||||
if len != n {
|
||||
return Err(PyValueError::new_err(format!("gencall: {what} has length {len}, expected {n}")));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Genotypes (0 NC, 1 AA, 2 AB, 3 BB) and GenCall scores; (n, 3) cluster arrays are passed flattened (C order).
|
||||
#[pyfunction]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn call_genotypes<'py>(
|
||||
py: Python<'py>,
|
||||
theta: PyReadonlyArray1<'py, f32>,
|
||||
r: PyReadonlyArray1<'py, f32>,
|
||||
raw_zero: PyReadonlyArray1<'py, bool>,
|
||||
present: PyReadonlyArray1<'py, bool>,
|
||||
n_cl: PyReadonlyArray1<'py, i32>,
|
||||
theta_mean: PyReadonlyArray1<'py, f32>,
|
||||
theta_dev: PyReadonlyArray1<'py, f32>,
|
||||
r_mean: PyReadonlyArray1<'py, f32>,
|
||||
r_dev: PyReadonlyArray1<'py, f32>,
|
||||
intensity_threshold: PyReadonlyArray1<'py, f32>,
|
||||
gentrain: PyReadonlyArray1<'py, f32>,
|
||||
intensity_only: PyReadonlyArray1<'py, bool>,
|
||||
cutoff: f32,
|
||||
) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray1<f32>>)> {
|
||||
let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?);
|
||||
let (ncl, tmv, tdv, rmv, rdv) =
|
||||
(n_cl.as_slice()?, theta_mean.as_slice()?, theta_dev.as_slice()?, r_mean.as_slice()?, r_dev.as_slice()?);
|
||||
let (ithr, gt, io) = (intensity_threshold.as_slice()?, gentrain.as_slice()?, intensity_only.as_slice()?);
|
||||
let n = th.len();
|
||||
for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (ithr.len(), "threshold"),
|
||||
(gt.len(), "gentrain"), (io.len(), "intensity_only")] {
|
||||
check(l, n, w)?;
|
||||
}
|
||||
for (l, w) in [(ncl.len(), "n"), (tmv.len(), "theta_mean"), (tdv.len(), "theta_dev"), (rmv.len(), "r_mean"),
|
||||
(rdv.len(), "r_dev")] {
|
||||
check(l, 3 * n, w)?;
|
||||
}
|
||||
let tiny_d = f64::MIN_POSITIVE * f64::EPSILON;
|
||||
let third: f32 = 1.0 / 3.0;
|
||||
let mut geno = vec![0u8; n];
|
||||
let mut scores = vec![0f32; n];
|
||||
|
||||
for i in 0..n {
|
||||
let t = th[i];
|
||||
let r = rr[i];
|
||||
let tm = [tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]];
|
||||
let td = [tdv[3 * i], tdv[3 * i + 1], tdv[3 * i + 2]];
|
||||
let rm = [rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]];
|
||||
let rd = [rdv[3 * i], rdv[3 * i + 1], rdv[3 * i + 2]];
|
||||
let nc = [ncl[3 * i], ncl[3 * i + 1], ncl[3 * i + 2]];
|
||||
|
||||
// gen_std_flair
|
||||
let z1 = 0.5f32 * (tm[1] - tm[0]) / (td[1] + td[0]);
|
||||
let z2 = 0.5f32 * (tm[2] - tm[1]) / (td[2] + td[1]);
|
||||
let mz = z1.min(z2);
|
||||
let alpha = (third * mz).max(1.0);
|
||||
let beta = (third * mz).min(1.0);
|
||||
let eta = alpha * beta;
|
||||
let mut ftd = [0f32; 3];
|
||||
for k in 0..3 {
|
||||
let v = td[k] * eta;
|
||||
ftd[k] = if v < 0.02 { 0.02 } else { v };
|
||||
}
|
||||
let med = td[0].min(td[1]).max(td[0].max(td[1]).min(td[2]));
|
||||
let min_r = if 0.2f32 < med { med } else { 0.2 };
|
||||
let mut frd = [0f32; 3];
|
||||
for k in 0..3 {
|
||||
let mrk0 = rm[k] / 7.0;
|
||||
let mrk = if mrk0 < min_r { min_r } else { mrk0 };
|
||||
frd[k] = if rd[k] < mrk { mrk } else { rd[k] };
|
||||
}
|
||||
|
||||
// modilik
|
||||
let d: [f64; 3] = std::array::from_fn(|k| 100.0 * ((t - tm[k]).abs() as f64) + ((r - rm[k]).abs() as f64));
|
||||
let (daa, dab, dbb) = (d[0], d[1], d[2]);
|
||||
let ok_t = !t.is_nan();
|
||||
let pt: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(t, tm[k], ftd[k]));
|
||||
let pr: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(r, rm[k], frd[k]));
|
||||
let pdf2 = |k: usize, st: Side, sr: Side| get(pt[k], t, tm[k], st) * get(pr[k], r, rm[k], sr);
|
||||
use Side::{Full, Left, Right};
|
||||
let mut lik = [0f64; 3];
|
||||
let mut covered = false;
|
||||
if ok_t {
|
||||
let cases: [(bool, [(usize, Side, Side); 2]); 6] = [
|
||||
(daa <= dbb && dab <= dbb && rm[0] <= rm[1], [(0, Left, Left), (1, Right, Right)]),
|
||||
(daa <= dab && dbb <= dab && rm[0] <= rm[2], [(0, Left, Left), (2, Right, Right)]),
|
||||
(dab <= daa && dbb <= daa && rm[1] <= rm[2], [(1, Left, Left), (2, Right, Right)]),
|
||||
(daa <= dbb && dab <= dbb && rm[0] > rm[1], [(0, Left, Right), (1, Right, Left)]),
|
||||
(dab <= daa && dbb <= daa && rm[1] > rm[2], [(1, Left, Right), (2, Right, Left)]),
|
||||
(daa <= dab && dbb <= dab && rm[0] > rm[2], [(0, Left, Right), (2, Right, Left)]),
|
||||
];
|
||||
for (cond, assigns) in cases {
|
||||
if cond {
|
||||
for (k, st, sr) in assigns {
|
||||
lik[k] = pdf2(k, st, sr);
|
||||
}
|
||||
covered = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !covered {
|
||||
for (k, l) in lik.iter_mut().enumerate() {
|
||||
*l = pdf2(k, Full, Full);
|
||||
}
|
||||
}
|
||||
for l in lik.iter_mut() {
|
||||
if l.is_nan() {
|
||||
*l = 0.0;
|
||||
}
|
||||
}
|
||||
if !ok_t {
|
||||
lik = [f64::NAN; 3];
|
||||
}
|
||||
|
||||
// priors and posteriors
|
||||
let ntot = nc[0] as i64 + nc[1] as i64 + nc[2] as i64;
|
||||
let p: [f64; 3] = std::array::from_fn(|k| (((nc[k] as f32) / (ntot as f32)) as f64 + 1.0) / 4.0);
|
||||
let lp: [f64; 3] = std::array::from_fn(|k| lik[k] * p[k]);
|
||||
let ev = (lp[0] + lp[1]) + lp[2];
|
||||
let (aa, ab, bb) = (lp[0] / ev, lp[1] / ev, lp[2] / ev);
|
||||
let mut iap: u8 = if aa >= ab && aa >= bb {
|
||||
1
|
||||
} else if ab >= aa && ab >= bb {
|
||||
2
|
||||
} else if bb >= aa && bb >= ab {
|
||||
3
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let mut mx = if aa > ab { aa } else { ab };
|
||||
let mut scnd = if aa > ab { ab } else { aa };
|
||||
let b_gt_mx = bb > mx;
|
||||
scnd = if b_gt_mx { mx } else if bb > scnd { bb } else { scnd };
|
||||
mx = if b_gt_mx { bb } else { mx };
|
||||
let ap_lod = math_log(mx / (tiny_d + scnd), true);
|
||||
let score_ap = smf(ap_lod as f32, 0.0, 2.0);
|
||||
|
||||
// sub-scores
|
||||
let score_r1 = smf(r, 0.0, 0.1);
|
||||
let nclus = nc.iter().filter(|&&v| v > 0).count();
|
||||
let misclust: f32 = if nclus == 1 && nc[1] == 0 {
|
||||
0.7
|
||||
} else if nclus != 3 {
|
||||
0.95
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
let g = (iap as i64 - 1).clamp(0, 2) as usize;
|
||||
let (gm_t, gd_t, gm_r, gd_r) = (tm[g], ftd[g], rm[g], frd[g]);
|
||||
let st_aa = zmf(t, gm_t + 2.0 * gd_t, gm_t + 6.0 * gd_t);
|
||||
let st_ab = zmf(((t - gm_t) / gd_t).abs(), 2.0, 6.0);
|
||||
let st_bb = smf(t, gm_t - 6.0 * gd_t, gm_t - 2.0 * gd_t);
|
||||
let mut score_t = match iap {
|
||||
1 => st_aa,
|
||||
2 => st_ab,
|
||||
3 => st_bb,
|
||||
_ => 1.0,
|
||||
};
|
||||
let r2 = smf(r, 0.0, gm_r / 10.0);
|
||||
let r3 = smf(r, gm_r - 6.0 * gd_r, gm_r - 2.0 * gd_r);
|
||||
let r4 = zmf(r, gm_r + 6.0 * gd_r, gm_r + 12.0 * gd_r);
|
||||
let mut score_r = if iap > 0 { ((score_r1 * r2) * r3) * r4 } else { 0.0 };
|
||||
let mut score_n: f32 = 1.0;
|
||||
if gt[i].is_nan() {
|
||||
score_t = f32::NAN;
|
||||
score_r = f32::NAN;
|
||||
}
|
||||
if !ok_t {
|
||||
score_t = f32::NAN;
|
||||
score_r = f32::NAN;
|
||||
score_n = f32::NAN;
|
||||
}
|
||||
let mut prelim = ((((score_ap * gt[i]) * score_t) * score_r) * score_n) * misclust;
|
||||
if r < ithr[i] {
|
||||
prelim = f32::NAN;
|
||||
iap = 0;
|
||||
}
|
||||
|
||||
// gencall_score_map
|
||||
let x = prelim as f64;
|
||||
let q = (x - 1.08) / 0.504;
|
||||
let gbell = 1.0 / (1.0 + math_pow(q * q, 1.71));
|
||||
let mut score = (math_pow(x, 0.35) * gbell) as f32;
|
||||
if score.is_nan() {
|
||||
score = 0.0;
|
||||
}
|
||||
if rz[i] || !pres[i] {
|
||||
score = 0.0;
|
||||
iap = 0;
|
||||
}
|
||||
geno[i] = if score < cutoff || io[i] { 0 } else { iap };
|
||||
scores[i] = score;
|
||||
}
|
||||
Ok((geno.into_pyarray(py), scores.into_pyarray(py)))
|
||||
}
|
||||
|
||||
/// BAF and LRR (float32, canonical NaN 0xFFC00000 where undefined).
|
||||
#[pyfunction]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn baf_lrr<'py>(
|
||||
py: Python<'py>,
|
||||
theta: PyReadonlyArray1<'py, f32>,
|
||||
r: PyReadonlyArray1<'py, f32>,
|
||||
raw_zero: PyReadonlyArray1<'py, bool>,
|
||||
present: PyReadonlyArray1<'py, bool>,
|
||||
theta_mean: PyReadonlyArray1<'py, f32>,
|
||||
r_mean: PyReadonlyArray1<'py, f32>,
|
||||
mean_r: PyReadonlyArray1<'py, f32>,
|
||||
intensity_only: PyReadonlyArray1<'py, bool>,
|
||||
) -> PyResult<(Bound<'py, PyArray1<f32>>, Bound<'py, PyArray1<f32>>)> {
|
||||
let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?);
|
||||
let (tmv, rmv, mr, io) = (theta_mean.as_slice()?, r_mean.as_slice()?, mean_r.as_slice()?, intensity_only.as_slice()?);
|
||||
let n = th.len();
|
||||
for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (mr.len(), "mean_r"), (io.len(), "intensity_only")] {
|
||||
check(l, n, w)?;
|
||||
}
|
||||
check(tmv.len(), 3 * n, "theta_mean")?;
|
||||
check(rmv.len(), 3 * n, "r_mean")?;
|
||||
let mut baf = vec![0f32; n];
|
||||
let mut lrr = vec![0f32; n];
|
||||
for i in 0..n {
|
||||
let (t, r) = (th[i], rr[i]);
|
||||
let (at, bt, ct) = (tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]);
|
||||
let (ar, br, cr) = (rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]);
|
||||
let (eq, lo, hi) = (t == bt, t < bt, t > bt);
|
||||
let (r_ref, mut b) = if eq {
|
||||
(br, 0.5f32)
|
||||
} else if lo {
|
||||
(ar + ((t - at) * (ar - br)) / (at - bt), ((t - at) / (bt - at)) * 0.5)
|
||||
} else if hi {
|
||||
(br + ((t - bt) * (cr - br)) / (ct - bt), 0.5 + ((t - bt) / (ct - bt)) * 0.5)
|
||||
} else {
|
||||
(f32::NAN, NEG_NAN32)
|
||||
};
|
||||
let denom = if io[i] { mr[i] } else { r_ref };
|
||||
let ratio = (r / denom) as f64;
|
||||
let mut l = math_log(ratio, false) as f32;
|
||||
if r == 0.0 {
|
||||
l = -f32::MAX;
|
||||
}
|
||||
if !(eq || lo || hi) {
|
||||
l = NEG_NAN32;
|
||||
}
|
||||
b = if b < 0.0 {
|
||||
0.0
|
||||
} else if b > 1.0 {
|
||||
1.0
|
||||
} else {
|
||||
b
|
||||
};
|
||||
if rz[i] || !pres[i] || b.is_nan() {
|
||||
b = NEG_NAN32;
|
||||
}
|
||||
if rz[i] || !pres[i] || l.is_nan() {
|
||||
l = NEG_NAN32;
|
||||
}
|
||||
baf[i] = b;
|
||||
lrr[i] = l;
|
||||
}
|
||||
Ok((baf.into_pyarray(py), lrr.into_pyarray(py)))
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
//! Compiled kernels for SimurgArray (plan §18 P9, §19).
|
||||
//!
|
||||
//! Every kernel is a drop-in replacement for a pure-Python/NumPy reference in the `simurg` package and returns
|
||||
//! bit-identical results (checked by `tests/unit/test_core_kernels.py`):
|
||||
//!
|
||||
//! * `exp`, `log10`, `log2`, `pow` — element-wise glibc libm calls with the error semantics of CPython's `math`
|
||||
//! module as used by `simurg.genotype.gencall._c_call` (NaN input returned as is, a domain error gives the
|
||||
//! canonical +NaN, `log10(0)` = −inf, overflow = +inf). The libm symbols are the ones CPython's `math` calls, so
|
||||
//! results equal the C reference implementations (gtc2vcf, MIT) that the Python port follows.
|
||||
//! * `closest_binned` — the per-bin candidate scan of the vendor's binned nearest-neighbour search
|
||||
//! (`simurg.signal.normalization.fit._closest_binned`): float32 squared distances, first minimum wins.
|
||||
//!
|
||||
//! * `call_genotypes`, `baf_lrr` — fused per-locus GenCall calling/scoring and BAF/LRR (`gencall.rs`).
|
||||
//! * `bpm_locus_entries` — the BPM LocusEntry parser (`simurg.formats.bpm._parse_locus_entries`).
|
||||
//!
|
||||
//! No floating-point contraction (FMA) happens: Rust never fuses `a*b+c` implicitly.
|
||||
|
||||
mod bpm;
|
||||
mod gencall;
|
||||
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[link(name = "m")]
|
||||
unsafe extern "C" {
|
||||
#[link_name = "exp"]
|
||||
fn c_exp(x: f64) -> f64;
|
||||
#[link_name = "log10"]
|
||||
fn c_log10(x: f64) -> f64;
|
||||
#[link_name = "log2"]
|
||||
fn c_log2(x: f64) -> f64;
|
||||
#[link_name = "pow"]
|
||||
fn c_pow(x: f64, y: f64) -> f64;
|
||||
}
|
||||
|
||||
/// CPython's `float('nan')`.
|
||||
const PY_NAN: f64 = f64::from_bits(0x7ff8_0000_0000_0000);
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn math_exp(v: f64) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
let r = unsafe { c_exp(v) };
|
||||
if r.is_nan() { PY_NAN } else if r.is_infinite() && v.is_finite() { f64::INFINITY } else { r }
|
||||
}
|
||||
|
||||
/// CPython `m_log10`/`m_log2` + `_c_call`: x > 0 → libm; 0 → −inf (log10) or NaN (log2); x < 0, −inf → NaN.
|
||||
#[inline]
|
||||
pub(crate) fn math_log(v: f64, ten: bool) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
if v.is_infinite() {
|
||||
return if v > 0.0 { v } else { PY_NAN };
|
||||
}
|
||||
if v > 0.0 {
|
||||
return unsafe { if ten { c_log10(v) } else { c_log2(v) } };
|
||||
}
|
||||
if v == 0.0 && ten { f64::NEG_INFINITY } else { PY_NAN }
|
||||
}
|
||||
|
||||
/// CPython `math.pow(v, e)` for a finite exponent e > 0 that is not an integer, + `_c_call`.
|
||||
#[inline]
|
||||
pub(crate) fn math_pow(v: f64, e: f64) -> f64 {
|
||||
if v.is_nan() {
|
||||
return v;
|
||||
}
|
||||
if v.is_infinite() {
|
||||
return f64::INFINITY;
|
||||
}
|
||||
if v == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let r = unsafe { c_pow(v, e) };
|
||||
if r.is_nan() { PY_NAN } else if r.is_infinite() { f64::INFINITY } else { r }
|
||||
}
|
||||
|
||||
fn map_f64<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, f: impl Fn(f64) -> f64) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
let s = x.as_slice()?;
|
||||
let out: Vec<f64> = s.iter().map(|&v| f(v)).collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn exp<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, math_exp)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn log10<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, |v| math_log(v, true))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn log2<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
map_f64(py, x, |v| math_log(v, false))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn pow<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, e: f64) -> PyResult<Bound<'py, PyArray1<f64>>> {
|
||||
if !(e.is_finite() && e > 0.0 && e.fract() != 0.0) {
|
||||
return Err(PyValueError::new_err("pow kernel: exponent must be finite, positive and not an integer"));
|
||||
}
|
||||
map_f64(py, x, |v| math_pow(v, e))
|
||||
}
|
||||
|
||||
/// For each sweep point `p` searching bin `choice[p]` (0..11 regular, 12..22 shifted; any other value: none), the
|
||||
/// index of the closest candidate among `eligible[i]` with `b1[i] == choice` (regular) or `b2[i] == choice − 12`
|
||||
/// (shifted), in ascending order; squared distance in float32; −1 if no candidate or the best is not < 1e20.
|
||||
#[pyfunction]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn closest_binned<'py>(
|
||||
py: Python<'py>,
|
||||
raw_x: PyReadonlyArray1<'py, f32>,
|
||||
raw_y: PyReadonlyArray1<'py, f32>,
|
||||
axis_x: PyReadonlyArray1<'py, f32>,
|
||||
axis_y: PyReadonlyArray1<'py, f32>,
|
||||
choice: PyReadonlyArray1<'py, i64>,
|
||||
eligible: PyReadonlyArray1<'py, i64>,
|
||||
b1: PyReadonlyArray1<'py, i64>,
|
||||
b2: PyReadonlyArray1<'py, i64>,
|
||||
) -> PyResult<Bound<'py, PyArray1<i64>>> {
|
||||
let (rx, ry) = (raw_x.as_slice()?, raw_y.as_slice()?);
|
||||
let (ax, ay, ch) = (axis_x.as_slice()?, axis_y.as_slice()?, choice.as_slice()?);
|
||||
let (el, b1, b2) = (eligible.as_slice()?, b1.as_slice()?, b2.as_slice()?);
|
||||
if ax.len() != ay.len() || ax.len() != ch.len() || el.len() != b1.len() || el.len() != b2.len() || rx.len() != ry.len() {
|
||||
return Err(PyValueError::new_err("closest_binned: length mismatch"));
|
||||
}
|
||||
let mut groups: Vec<Vec<usize>> = vec![Vec::new(); 23];
|
||||
for (i, &e) in el.iter().enumerate() {
|
||||
let e = usize::try_from(e).map_err(|_| PyValueError::new_err("closest_binned: negative index"))?;
|
||||
if e >= rx.len() {
|
||||
return Err(PyValueError::new_err("closest_binned: index out of range"));
|
||||
}
|
||||
if (0..12).contains(&b1[i]) {
|
||||
groups[b1[i] as usize].push(e);
|
||||
}
|
||||
if (0..11).contains(&b2[i]) {
|
||||
groups[12 + b2[i] as usize].push(e);
|
||||
}
|
||||
}
|
||||
let out: Vec<i64> = (0..ax.len())
|
||||
.map(|p| {
|
||||
let c = ch[p];
|
||||
if !(0..23).contains(&c) {
|
||||
return -1;
|
||||
}
|
||||
let cand = &groups[c as usize];
|
||||
let mut best_j = usize::MAX;
|
||||
let mut best = f32::INFINITY;
|
||||
for (j, &e) in cand.iter().enumerate() {
|
||||
let dx = rx[e] - ax[p];
|
||||
let dy = ry[e] - ay[p];
|
||||
let d = dx * dx + dy * dy;
|
||||
// numpy argmin: the first NaN wins, otherwise the first strict minimum
|
||||
if d.is_nan() {
|
||||
best_j = j;
|
||||
best = d;
|
||||
break;
|
||||
}
|
||||
if best_j == usize::MAX || d < best {
|
||||
best_j = j;
|
||||
best = d;
|
||||
}
|
||||
}
|
||||
if best_j == usize::MAX || !((best as f64) < 1e20) { -1 } else { cand[best_j] as i64 }
|
||||
})
|
||||
.collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
/// Exact nearest neighbour (`simurg.signal.normalization.fit._closest_slow`): for each sweep point (x, y) the
|
||||
/// index of the reference point with the smallest float32 squared distance `(x-xr)²+(y-yr)²`; first minimum wins,
|
||||
/// a NaN distance wins as in `numpy.argmin`. `xref` must not be empty.
|
||||
#[pyfunction]
|
||||
fn closest_exact<'py>(
|
||||
py: Python<'py>,
|
||||
xref: PyReadonlyArray1<'py, f32>,
|
||||
yref: PyReadonlyArray1<'py, f32>,
|
||||
x: PyReadonlyArray1<'py, f32>,
|
||||
y: PyReadonlyArray1<'py, f32>,
|
||||
) -> PyResult<Bound<'py, PyArray1<i64>>> {
|
||||
let (xr, yr, xs, ys) = (xref.as_slice()?, yref.as_slice()?, x.as_slice()?, y.as_slice()?);
|
||||
if xr.is_empty() || xr.len() != yr.len() || xs.len() != ys.len() {
|
||||
return Err(PyValueError::new_err("closest_exact: empty reference or length mismatch"));
|
||||
}
|
||||
let out: Vec<i64> = xs
|
||||
.iter()
|
||||
.zip(ys)
|
||||
.map(|(&px, &py_)| {
|
||||
let mut best_j = 0usize;
|
||||
let mut best = f32::INFINITY;
|
||||
for (j, (&rx, &ry)) in xr.iter().zip(yr).enumerate() {
|
||||
let dx = px - rx;
|
||||
let dy = py_ - ry;
|
||||
let d = dx * dx + dy * dy;
|
||||
if d.is_nan() {
|
||||
return j as i64;
|
||||
}
|
||||
if j == 0 || d < best {
|
||||
best_j = j;
|
||||
best = d;
|
||||
}
|
||||
}
|
||||
best_j as i64
|
||||
})
|
||||
.collect();
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn simurg_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
|
||||
m.add_function(wrap_pyfunction!(exp, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(log10, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(log2, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(pow, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(closest_binned, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(closest_exact, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(bpm::bpm_locus_entries, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(gencall::call_genotypes, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(gencall::baf_lrr, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user