additional support for ms2rescore refactoring

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ms2rescore-rs"
-version = "0.4.3"
+version = "0.5.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -12,6 +12,10 @@ crate-type = ["cdylib"]
 default = []
 
 [dependencies]
-mzdata = { version = "0.59.2", features = ["thermo"] }
+mzdata = { version = "0.59.2", features = ["thermo", "parallelism"] }
 pyo3 = { version = "0.23.3", features = ["anyhow"] }
+rayon = "1.10"
 timsrust = "0.4.1"
+rustyms = "0.11"
+ordered-float = "5"
+numpy = "0.23"

src/lib.rs

@@ -3,10 +3,12 @@ mod ms2_spectrum;
 mod parse_mzdata;
 mod parse_timsrust;
 mod precursor;
+mod ms2_features;
+mod ms2pip_features;
 
 use std::collections::HashMap;
 
-use pyo3::exceptions::{PyException, PyValueError};
+use pyo3::exceptions::PyException;
 use pyo3::prelude::*;
 
 use file_types::{match_file_type, SpectrumFileType};
@@ -23,17 +25,20 @@ pub fn is_supported_file_type(spectrum_path: String) -> bool {
 
 /// Get mapping of spectrum identifiers to precursor information.
 #[pyfunction]
-pub fn get_precursor_info(spectrum_path: String) -> PyResult<HashMap<String, Precursor>> {
+pub fn get_precursor_info(py: Python<'_>, spectrum_path: String) -> PyResult<HashMap<String, Precursor>> {
     let file_type = match_file_type(&spectrum_path);
 
-    let precursors = match file_type {
+    let precursors = py.allow_threads(|| match file_type {
         SpectrumFileType::MascotGenericFormat
         | SpectrumFileType::MzML
         | SpectrumFileType::MzMLb
         | SpectrumFileType::ThermoRaw => parse_mzdata::parse_precursor_info(&spectrum_path),
         SpectrumFileType::BrukerRaw => parse_timsrust::parse_precursor_info(&spectrum_path),
-        SpectrumFileType::Unknown => return Err(PyValueError::new_err("Unsupported file type")),
-    };
+        SpectrumFileType::Unknown => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "Unsupported file type",
+        )),
+    });
 
     match precursors {
         Ok(precursors) => Ok(precursors),
@@ -43,17 +48,20 @@ pub fn get_precursor_info(spectrum_path: String) -> PyResult<HashMap<String, Pre
 
 /// Get MS2 spectra from a spectrum file.
 #[pyfunction]
-pub fn get_ms2_spectra(spectrum_path: String) -> PyResult<Vec<ms2_spectrum::MS2Spectrum>> {
+pub fn get_ms2_spectra(py: Python<'_>, spectrum_path: String) -> PyResult<Vec<ms2_spectrum::MS2Spectrum>> {
     let file_type = match_file_type(&spectrum_path);
 
-    let spectra = match file_type {
+    let spectra = py.allow_threads(|| match file_type {
         SpectrumFileType::MascotGenericFormat
         | SpectrumFileType::MzML
         | SpectrumFileType::MzMLb
         | SpectrumFileType::ThermoRaw => parse_mzdata::read_ms2_spectra(&spectrum_path),
         SpectrumFileType::BrukerRaw => parse_timsrust::read_ms2_spectra(&spectrum_path),
-        SpectrumFileType::Unknown => return Err(PyValueError::new_err("Unsupported file type")),
-    };
+        SpectrumFileType::Unknown => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "Unsupported file type",
+        )),
+    });
 
     match spectra {
         Ok(spectra) => Ok(spectra),
@@ -69,5 +77,10 @@ fn ms2rescore_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(is_supported_file_type, m)?)?;
     m.add_function(wrap_pyfunction!(get_precursor_info, m)?)?;
     m.add_function(wrap_pyfunction!(get_ms2_spectra, m)?)?;
+    m.add_function(wrap_pyfunction!(
+        ms2_features::batch_ms2_features_from_spectra,
+        m
+    )?)?;
+    m.add_function(wrap_pyfunction!(ms2pip_features::batch_ms2pip_features_numpy, m)?)?;
     Ok(())
 }