kemalandClaude Opus 5.5
5bdd41b048
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). Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv
358 lines
13 KiB
Rust
358 lines
13 KiB
Rust
//! 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)))
|
|
}
|