diff --git a/lightmotif-py/lightmotif/lib.rs b/lightmotif-py/lightmotif/lib.rs index 802d7f2dd24d5fa92a9bfc0effa1476ed626dc4c..c0a80ed19e5462f65efad2da7cc7bc9cfff69009 100644 --- a/lightmotif-py/lightmotif/lib.rs +++ b/lightmotif-py/lightmotif/lib.rs @@ -5,12 +5,15 @@ extern crate lightmotif; extern crate pyo3; use lightmotif::abc::Alphabet; +use lightmotif::abc::Dna; use lightmotif::abc::Symbol; use lightmotif::dense::DenseMatrix; +use lightmotif::num::StrictlyPositive; use lightmotif::num::Unsigned; use lightmotif::pli::platform::Backend; use lightmotif::pli::Pipeline; use lightmotif::pli::Score; +use lightmotif::pli::Threshold; use pyo3::exceptions::PyBufferError; use pyo3::exceptions::PyIndexError; @@ -286,6 +289,10 @@ impl ScoringMatrix { /// of the input sequence, stored into a striped matrix for fast /// vectorized operations. /// + /// Note: + /// This method uses the best implementation for the local platform, + /// prefering AVX2 if available. + /// pub fn calculate( slf: PyRef<'_, Self>, sequence: &mut StripedSequence, @@ -378,6 +385,37 @@ impl StripedScores { Ok(()) } + + /// Return all positions with a score greater or equal to the threshold. + /// + /// Returns: + /// `list` of `int`: The indices of the position with a score greater or + /// equal to the given threshold. Note that the indices may or may not + /// be sorted, depending on the implementation. + /// + /// Note: + /// This method uses the best implementation for the local platform, + /// prefering AVX2 if available. + /// + pub fn threshold(slf: PyRef<'_, Self>, threshold: f32) -> PyResult<Vec<usize>> { + let scores = &slf.scores; + let indices = slf.py().allow_threads(|| { + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + if let Ok(pli) = Pipeline::<Dna, _>::avx2() { + return pli.threshold(scores, threshold); + } + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + if let Ok(pli) = Pipeline::<Dna, _>::sse2() { + return pli.threshold(scores, threshold); + } + #[cfg(any(target_arch = "arm", target_arch = "aarch64"))] + if let Ok(pli) = Pipeline::<Dna, _>::neon() { + return pli.threshold(scores, threshold); + } + Pipeline::<Dna, _>::generic().threshold(scores, threshold) + }); + Ok(indices) + } } impl From<lightmotif::pli::StripedScores<C>> for StripedScores {