Merge pull request #26 from IEDB/counts-only

dmx2 · web-flow · commit 8ef2f98542c6 · 2026-06-26T20:04:49.000-07:00
Add counts_only aggregate mode
diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py
@@ -2,7 +2,7 @@
 import polars as pl
 from pathlib import Path
 from Bio import SeqIO
-from ._rs import rs_preprocess, rs_match, rs_discontinuous, rs_metadata
+from ._rs import rs_preprocess, rs_match, rs_discontinuous, rs_metadata, rs_match_counts
 
 VALID_OUTPUT_FORMATS = ['dataframe', 'csv', 'tsv', 'xlsx', 'json']
 FASTA_EXTENSIONS = {
@@ -37,6 +37,7 @@ def __init__(
     output_format='dataframe',
     output_name='',
     sequence_version=True,
+    counts_only=False,
   ):
 
     if best_match and k == 0:
@@ -57,13 +58,17 @@ def __init__(
         f'Invalid output format. Choose from: {VALID_OUTPUT_FORMATS}'
       )
 
+    if counts_only and best_match:
+      raise ValueError('counts_only is not supported with best_match.')
+
     self.proteome_file = str(proteome_file)
     self.proteome_name = str(proteome_file).split('/')[-1].split('.')[0]
     self.max_mismatches = max_mismatches
     self.preprocessed_files_path = preprocessed_files_path
     self.best_match = best_match
     self.output_format = output_format
     self.sequence_version = sequence_version
+    self.counts_only = counts_only
     self.output_name = output_name or 'PEPMatch_results'
 
     self.query = self._parse_query(query)
@@ -116,6 +121,12 @@ def match(self):
     linear_df = pl.DataFrame()
     discontinuous_df = pl.DataFrame()
 
+    if self.counts_only:
+      df = self._counts_to_dataframe(self._search_counts(self.k, self.max_mismatches))
+      if self.output_format == 'dataframe':
+        return df
+      return output_matches(df, self.output_format, self.output_name)
+
     if self.query:
       if self.best_match and self.k_specified:
         results = self._search(self.k, self.max_mismatches)
@@ -158,6 +169,30 @@ def _search(self, k, max_mismatches, peptides=None):
     print(f"Searching {len(query)} peptides against {self.proteome_name} (k={k}, max_mismatches={max_mismatches})...")
     return rs_match(pepidx_path, query, k, max_mismatches)
 
+  def _search_counts(self, k, max_mismatches):
+    pepidx_path = self._pepidx_path(k)
+    if not os.path.isfile(pepidx_path):
+      print(f"Preprocessing {self.proteome_name} with k={k}...")
+      rs_preprocess(self.proteome_file, k, pepidx_path)
+    print(f"Counting {len(self.query)} peptides against {self.proteome_name} (k={k}, max_mismatches={max_mismatches})...")
+    return rs_match_counts(pepidx_path, self.query, k, max_mismatches)
+
+  def _counts_to_dataframe(self, cols):
+    """Aggregate counts: one row per (query peptide, mismatch level) with a hit.
+    Memory is O(unique queries), independent of total hit count."""
+    qid, qseq, mm, count = cols
+    if not qid:
+      return pl.DataFrame(schema={
+        'Query ID': pl.Utf8, 'Query Sequence': pl.Utf8,
+        'Mismatches': pl.Int64, 'Count': pl.UInt64,
+      })
+    return pl.DataFrame({
+      'Query ID': qid,
+      'Query Sequence': qseq,
+      'Mismatches': pl.Series(mm, dtype=pl.Int64),
+      'Count': pl.Series(count, dtype=pl.UInt64),
+    })
+
   def best_match_search(self):
     peptides_remaining = self.query.copy()
     acc = tuple([] for _ in range(8))   # 8 columnar accumulators
diff --git a/pepmatch/rs-engine/src/lib.rs b/pepmatch/rs-engine/src/lib.rs
@@ -29,12 +29,18 @@ fn rs_metadata(pepidx_path: &str) -> matching::MetaColumns {
     matching::run_metadata(pepidx_path)
 }
 
+#[pyfunction]
+fn rs_match_counts(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> matching::CountColumns {
+    matching::run_counts(pepidx_path, peptides, k, max_mismatches)
+}
+
 #[pymodule]
 fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(rs_version, m)?)?;
     m.add_function(wrap_pyfunction!(rs_preprocess, m)?)?;
     m.add_function(wrap_pyfunction!(rs_match, m)?)?;
     m.add_function(wrap_pyfunction!(rs_discontinuous, m)?)?;
     m.add_function(wrap_pyfunction!(rs_metadata, m)?)?;
+    m.add_function(wrap_pyfunction!(rs_match_counts, m)?)?;
     Ok(())
 }
diff --git a/pepmatch/rs-engine/src/match.rs b/pepmatch/rs-engine/src/match.rs
@@ -541,3 +541,115 @@ pub(crate) fn run(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize,
         .collect();
     unzip_records(records)
 }
+
+// ── Counts-only path (aggregate; O(unique queries), no per-hit materialization) ──
+
+pub(crate) type CountColumns = (Vec<String>, Vec<String>, Vec<i64>, Vec<u64>);
+
+/// Tally accepted matches per mismatch level for one peptide, mirroring
+/// mismatch_match's dedup + validity walk exactly but without building the
+/// matched sequence, metadata, or any per-hit row.
+fn mismatch_count(peptide: &str, k: usize, max_mismatches: usize, index: &PepIndex, counts: &mut [u64]) {
+    let pep_bytes = peptide.as_bytes();
+    if pep_bytes.len() < k { return; }
+
+    let num_kmers = pep_bytes.len() - k + 1;
+    let peptide_len = pep_bytes.len();
+    let mut seen: HashSet<u64> = HashSet::new();
+
+    if peptide_len % k == 0 {
+        let mut idx = 0;
+        while idx < num_kmers {
+            let kmer = &pep_bytes[idx..idx + k];
+            if let Some(positions) = index.lookup(kmer) {
+                for kmer_hit in positions {
+                    let start = (kmer_hit as i64 - idx as i64) as u64;
+                    if seen.contains(&start) { continue; }
+                    let mismatches = check_left_neighbors(pep_bytes, idx, kmer_hit, index, k, max_mismatches, 0);
+                    let mismatches = check_right_neighbors(pep_bytes, idx, kmer_hit, index, k, max_mismatches, mismatches);
+                    if mismatches <= max_mismatches {
+                        let mut i = 0;
+                        let mut valid = true;
+                        while i < peptide_len {
+                            if index.resolve(start + i as u64).is_none() { valid = false; break; }
+                            i += k;
+                        }
+                        if valid {
+                            seen.insert(start);
+                            counts[mismatches] += 1;
+                        }
+                    }
+                }
+            }
+            idx += k;
+        }
+    } else {
+        for idx in 0..num_kmers {
+            let kmer = &pep_bytes[idx..idx + k];
+            if let Some(positions) = index.lookup(kmer) {
+                for kmer_hit in positions {
+                    let start = (kmer_hit as i64 - idx as i64) as u64;
+                    if seen.contains(&start) { continue; }
+                    let mismatches = check_left_residues(pep_bytes, idx, kmer_hit, index, max_mismatches, 0);
+                    let mismatches = check_right_residues(pep_bytes, idx, kmer_hit, index, k, max_mismatches, mismatches);
+                    if mismatches <= max_mismatches {
+                        let mut i = 0;
+                        let mut valid = true;
+                        while i < peptide_len {
+                            if i + k > peptide_len {
+                                let remaining = peptide_len - i;
+                                let back = k - remaining;
+                                if index.resolve(start + i as u64 - back as u64).is_none() { valid = false; break; }
+                            } else if index.resolve(start + i as u64).is_none() {
+                                valid = false; break;
+                            }
+                            i += k;
+                        }
+                        if valid {
+                            seen.insert(start);
+                            counts[mismatches] += 1;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn count_peptide(query_id: &str, peptide: &str, k: usize, max_mismatches: usize, index: &PepIndex) -> Vec<(String, String, i64, u64)> {
+    let mut counts = vec![0u64; max_mismatches + 1];
+    if max_mismatches == 0 {
+        counts[0] = exact_match(peptide, k, index).len() as u64;
+    } else {
+        mismatch_count(peptide, k, max_mismatches, index, &mut counts);
+    }
+    let mut out = Vec::new();
+    for (mm, &c) in counts.iter().enumerate() {
+        if c > 0 {
+            out.push((query_id.to_string(), peptide.to_string(), mm as i64, c));
+        }
+    }
+    out
+}
+
+pub(crate) fn run_counts(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> CountColumns {
+    let index = PepIndex::open(pepidx_path);
+    let rows: Vec<(String, String, i64, u64)> = peptides
+        .par_iter()
+        .flat_map_iter(|(query_id, peptide)| {
+            count_peptide(query_id, peptide, k, max_mismatches, &index).into_iter()
+        })
+        .collect();
+    let n = rows.len();
+    let mut qid = Vec::with_capacity(n);
+    let mut qseq = Vec::with_capacity(n);
+    let mut mm = Vec::with_capacity(n);
+    let mut cnt = Vec::with_capacity(n);
+    for (a, b, c, d) in rows {
+        qid.push(a);
+        qseq.push(b);
+        mm.push(c);
+        cnt.push(d);
+    }
+    (qid, qseq, mm, cnt)
+}
diff --git a/pepmatch/tests/test_counts.py b/pepmatch/tests/test_counts.py
@@ -0,0 +1,60 @@
+import pytest
+import polars as pl
+import polars.testing as plt
+from pathlib import Path
+from pepmatch import Matcher
+
+
+@pytest.fixture
+def proteome_path() -> Path:
+  return Path(__file__).parent / 'data' / 'proteome.fasta'
+
+@pytest.fixture
+def mismatch_query() -> Path:
+  return Path(__file__).parent / 'data' / 'mismatching_query.fasta'
+
+@pytest.fixture
+def exact_query() -> Path:
+  return Path(__file__).parent / 'data' / 'exact_match_query.fasta'
+
+
+def _from_full(df: pl.DataFrame) -> pl.DataFrame:
+  return (
+    df.filter(pl.col('Matched Sequence').is_not_null())
+      .group_by(['Query Sequence', 'Mismatches'])
+      .agg(pl.len().alias('Count'))
+      .with_columns(pl.col('Count').cast(pl.Int64))
+      .sort(['Query Sequence', 'Mismatches'])
+  )
+
+def _from_counts(df: pl.DataFrame) -> pl.DataFrame:
+  return (
+    df.group_by(['Query Sequence', 'Mismatches'])
+      .agg(pl.col('Count').sum().alias('Count'))
+      .with_columns(pl.col('Count').cast(pl.Int64))
+      .sort(['Query Sequence', 'Mismatches'])
+  )
+
+
+def test_counts_parity_mismatch(proteome_path, mismatch_query):
+  """counts_only must equal the full output grouped by (peptide, mismatch level)."""
+  full = Matcher(query=mismatch_query, proteome_file=proteome_path,
+                 max_mismatches=3, k=3).match()
+  counts = Matcher(query=mismatch_query, proteome_file=proteome_path,
+                   max_mismatches=3, k=3, counts_only=True).match()
+  assert counts.columns == ['Query ID', 'Query Sequence', 'Mismatches', 'Count']
+  plt.assert_frame_equal(_from_full(full), _from_counts(counts))
+
+
+def test_counts_parity_exact(proteome_path, exact_query):
+  full = Matcher(query=exact_query, proteome_file=proteome_path,
+                 max_mismatches=0, k=5).match()
+  counts = Matcher(query=exact_query, proteome_file=proteome_path,
+                   max_mismatches=0, k=5, counts_only=True).match()
+  plt.assert_frame_equal(_from_full(full), _from_counts(counts))
+
+
+def test_counts_only_rejects_best_match(proteome_path, exact_query):
+  with pytest.raises(ValueError):
+    Matcher(query=exact_query, proteome_file=proteome_path,
+            counts_only=True, best_match=True)
