[claude] Add CandidateCollection filtering, min_fragments, and z-score filter

scripts/candidate_scoring.py

@@ -276,7 +276,7 @@ def run_fdr_filtering(psm_scored_df, candidates_df, output_folder):
         FEATURE_COLUMNS,
         psm_scored_df[psm_scored_df["decoy"] == 0].copy(),
         psm_scored_df[psm_scored_df["decoy"] == 1].copy(),
-        competetive=True,
+        competitive=True,
     )
 
     psm_df = psm_df[psm_df["qval"] <= 0.01]

src/candidate/entry.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use numpy::{ndarray::Array1, IntoPyArray};
 use pyo3::prelude::*;
 use rayon::prelude::*;
@@ -119,6 +121,67 @@ impl CandidateCollection {
         Ok(Self { candidates })
     }
 
+    /// Return a new collection containing only candidates with score > cutoff.
+    pub fn filter_by_score(&self, cutoff: f32) -> CandidateCollection {
+        let filtered: Vec<Candidate> = self
+            .candidates
+            .iter()
+            .filter(|c| c.score > cutoff)
+            .map(|c| Candidate {
+                precursor_idx: c.precursor_idx,
+                rank: c.rank,
+                score: c.score,
+                scan_center: c.scan_center,
+                scan_start: c.scan_start,
+                scan_stop: c.scan_stop,
+                cycle_center: c.cycle_center,
+                cycle_start: c.cycle_start,
+                cycle_stop: c.cycle_stop,
+            })
+            .collect();
+        CandidateCollection {
+            candidates: filtered,
+        }
+    }
+
+    /// Keep only candidates whose (precursor_idx, rank) is in the provided lists.
+    pub fn filter_by_keys(
+        &self,
+        precursor_idxs: Vec<u64>,
+        ranks: Vec<u64>,
+    ) -> PyResult<CandidateCollection> {
+        if precursor_idxs.len() != ranks.len() {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "precursor_idxs and ranks must have the same length",
+            ));
+        }
+        let keys: HashSet<(usize, usize)> = precursor_idxs
+            .iter()
+            .zip(ranks.iter())
+            .map(|(&p, &r)| (p as usize, r as usize))
+            .collect();
+
+        let filtered: Vec<Candidate> = self
+            .candidates
+            .iter()
+            .filter(|c| keys.contains(&(c.precursor_idx, c.rank)))
+            .map(|c| Candidate {
+                precursor_idx: c.precursor_idx,
+                rank: c.rank,
+                score: c.score,
+                scan_center: c.scan_center,
+                scan_start: c.scan_start,
+                scan_stop: c.scan_stop,
+                cycle_center: c.cycle_center,
+                cycle_start: c.cycle_start,
+                cycle_stop: c.cycle_stop,
+            })
+            .collect();
+        Ok(CandidateCollection {
+            candidates: filtered,
+        })
+    }
+
     /// Convert the collection to separate arrays for all fields
     #[allow(clippy::type_complexity)]
     pub fn to_arrays(

src/candidate/features.rs

@@ -309,6 +309,17 @@ impl CandidateFeatureCollection {
         self.features.is_empty()
     }
 
+    /// Return a new collection containing only features with score > cutoff.
+    pub fn filter_by_score(&self, cutoff: f32) -> CandidateFeatureCollection {
+        let filtered: Vec<CandidateFeature> = self
+            .features
+            .iter()
+            .filter(|f| f.score > cutoff)
+            .cloned()
+            .collect();
+        CandidateFeatureCollection { features: filtered }
+    }
+
     pub fn to_dict_arrays(&self, py: Python) -> PyResult<PyObject> {
         let n = self.features.len();
 

src/candidate/tests.rs

@@ -432,3 +432,87 @@ fn test_candidate_feature_collection_to_dict_arrays_dtypes_and_values() {
         assert!((fwhm_rt[0] - 15.5).abs() < 1e-6);
     });
 }
+
+#[test]
+fn test_candidate_collection_filter_by_score() {
+    let candidates = vec![
+        Candidate::new(0, 0, 0.5, 0, 5, 10),
+        Candidate::new(1, 0, 1.5, 0, 5, 10),
+        Candidate::new(2, 0, 0.3, 0, 5, 10),
+        Candidate::new(3, 1, 2.0, 0, 5, 10),
+    ];
+    let collection = CandidateCollection::from_vec(candidates);
+    assert_eq!(collection.len(), 4);
+
+    let filtered = collection.filter_by_score(0.4);
+    assert_eq!(filtered.len(), 3);
+
+    let filtered_high = collection.filter_by_score(1.5);
+    assert_eq!(filtered_high.len(), 1);
+
+    let filtered_none = collection.filter_by_score(10.0);
+    assert_eq!(filtered_none.len(), 0);
+
+    let filtered_all = collection.filter_by_score(-1.0);
+    assert_eq!(filtered_all.len(), 4);
+}
+
+#[test]
+fn test_candidate_collection_filter_by_keys() {
+    let candidates = vec![
+        Candidate::new(10, 0, 0.5, 0, 5, 10),
+        Candidate::new(20, 0, 1.5, 0, 5, 10),
+        Candidate::new(10, 1, 0.3, 0, 5, 10),
+        Candidate::new(30, 0, 2.0, 0, 5, 10),
+    ];
+    let collection = CandidateCollection::from_vec(candidates);
+    assert_eq!(collection.len(), 4);
+
+    // Keep (10, 0) and (30, 0)
+    let filtered = collection.filter_by_keys(vec![10, 30], vec![0, 0]).unwrap();
+    assert_eq!(filtered.len(), 2);
+
+    // Verify the correct candidates were kept
+    let iter: Vec<_> = filtered.iter().collect();
+    assert_eq!(iter[0].precursor_idx, 10);
+    assert_eq!(iter[0].rank, 0);
+    assert_eq!(iter[1].precursor_idx, 30);
+    assert_eq!(iter[1].rank, 0);
+
+    // Empty filter returns empty collection
+    let filtered_empty = collection.filter_by_keys(vec![], vec![]).unwrap();
+    assert_eq!(filtered_empty.len(), 0);
+
+    // Non-matching keys return empty collection
+    let filtered_none = collection.filter_by_keys(vec![99], vec![99]).unwrap();
+    assert_eq!(filtered_none.len(), 0);
+
+    // Mismatched lengths return error
+    let result = collection.filter_by_keys(vec![10, 20], vec![0]);
+    assert!(result.is_err());
+}
+
+#[test]
+fn test_candidate_feature_collection_filter_by_score() {
+    // CandidateFeature::new takes 2 usize + 64 f32 = 66 args total
+    let f1 = CandidateFeature::new(
+        0, 0, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+    );
+    let f2 = CandidateFeature::new(
+        1, 0, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+    );
+    let collection = CandidateFeatureCollection::from_vec(vec![f1, f2]);
+    assert_eq!(collection.len(), 2);
+
+    let filtered = collection.filter_by_score(1.0);
+    assert_eq!(filtered.len(), 1);
+
+    let filtered_all = collection.filter_by_score(-1.0);
+    assert_eq!(filtered_all.len(), 2);
+}

src/lib.rs

@@ -1,3 +1,5 @@
+use numpy::ndarray::{s, Array1, Axis};
+use numpy::{IntoPyArray, PyReadonlyArray2};
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::PyErr;
@@ -85,6 +87,76 @@ fn get_num_threads() -> PyResult<usize> {
     Ok(threadpool::get_num_threads())
 }
 
+/// Compute z-score filter mask over a feature matrix using batched ndarray linalg.
+///
+/// The z-score sum `Σ (x_j - μ_j) / σ_j * s_j` is equivalent to `x · w - b`
+/// where `w = signs / stds` and `b = Σ(means * signs / stds)`.
+/// Processes in batches to limit memory for column extraction.
+/// NaN and infinite values are treated as 0.
+///
+/// Returns a boolean numpy array where True = passes filter (score >= threshold).
+#[pyfunction]
+fn zscore_filter_mask(
+    py: Python,
+    features: PyReadonlyArray2<'_, f64>,
+    col_indices: Vec<usize>,
+    means: Vec<f64>,
+    stds: Vec<f64>,
+    signs: Vec<f64>,
+    threshold: f64,
+) -> PyResult<PyObject> {
+    let n_zscore = col_indices.len();
+    if means.len() != n_zscore || stds.len() != n_zscore || signs.len() != n_zscore {
+        return Err(PyErr::new::<PyValueError, _>(
+            "means, stds, signs must have the same length as col_indices",
+        ));
+    }
+
+    let features = features.as_array();
+    let n_rows = features.shape()[0];
+
+    // Precompute weight vector w = signs / stds and bias b = Σ(means * w)
+    let w = Array1::from_vec(
+        signs
+            .iter()
+            .zip(&stds)
+            .map(|(&s, &st)| s / st)
+            .collect::<Vec<f64>>(),
+    );
+    let bias: f64 = means.iter().zip(w.iter()).map(|(m, wi)| m * wi).sum();
+    let adjusted_threshold = threshold + bias;
+
+    let mut mask = Array1::<bool>::from_elem(n_rows, false);
+
+    const BATCH_SIZE: usize = 500_000;
+    for batch_start in (0..n_rows).step_by(BATCH_SIZE) {
+        let batch_end = (batch_start + BATCH_SIZE).min(n_rows);
+        let batch = features.slice(s![batch_start..batch_end, ..]);
+
+        // Select z-score columns → contiguous (batch_len, n_zscore) array
+        let mut z_batch = batch.select(Axis(1), &col_indices);
+
+        // Replace NaN/inf with 0 in-place
+        z_batch.mapv_inplace(|v| {
+            if v.is_nan() || v.is_infinite() {
+                0.0
+            } else {
+                v
+            }
+        });
+
+        // Batched dot product: scores = z_batch · w (shape: batch_len)
+        let scores = z_batch.dot(&w);
+
+        // Threshold comparison
+        for (i, &s) in scores.iter().enumerate() {
+            mask[batch_start + i] = s >= adjusted_threshold;
+        }
+    }
+
+    Ok(mask.into_pyarray(py).into())
+}
+
 #[pymodule]
 fn alphadia_search_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<DIAData>()?;
@@ -105,5 +177,6 @@ fn alphadia_search_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(get_current_simd_backend, m)?)?;
     m.add_function(wrap_pyfunction!(set_num_threads, m)?)?;
     m.add_function(wrap_pyfunction!(get_num_threads, m)?)?;
+    m.add_function(wrap_pyfunction!(zscore_filter_mask, m)?)?;
     Ok(())
 }

src/peak_group_selection/mod.rs

@@ -173,6 +173,18 @@ impl PeakGroupSelection {
             let cycle_center_idx = local_maxima_indices[i];
             let score = local_maxima_values[i];
 
+            // Count non-zero fragments at the apex cycle and skip if below threshold
+            let center_col = cycle_center_idx - cycle_start_idx;
+            let n_matched = dense_xic_obs
+                .dense_xic
+                .column(center_col)
+                .iter()
+                .filter(|&&v| v > 0.0)
+                .count();
+            if n_matched < self.params.min_fragments {
+                continue;
+            }
+
             let cycle_start_idx = max(0, cycle_center_idx - self.params.peak_length);
             let cycle_stop_idx = min(
                 cycle_center_idx + self.params.peak_length + 1,

src/peak_group_selection/parameters.rs

@@ -18,6 +18,8 @@ pub struct SelectionParameters {
     pub candidate_count: usize,
     #[pyo3(get)]
     pub top_k_fragments: usize,
+    #[pyo3(get)]
+    pub min_fragments: usize,
 }
 
 #[pymethods]
@@ -39,6 +41,8 @@ impl SelectionParameters {
             candidate_count: 3,
             // maximum number of fragments to use for selecting precursors from a DIAData object.
             top_k_fragments: 12,
+            // minimum number of matched fragments at the apex cycle to keep a candidate.
+            min_fragments: 3,
         }
     }
 
@@ -64,6 +68,9 @@ impl SelectionParameters {
         if let Some(value) = config.get_item("top_k_fragments")? {
             self.top_k_fragments = value.extract::<usize>()?;
         }
+        if let Some(value) = config.get_item("min_fragments")? {
+            self.min_fragments = value.extract::<usize>()?;
+        }
         Ok(())
     }
 }