diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 3aa4942..e44a23c 100755 --- 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,6 +58,9 @@ 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 @@ -64,6 +68,7 @@ def __init__( 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 index 058b630..71124dd 100644 --- a/pepmatch/rs-engine/src/lib.rs +++ b/pepmatch/rs-engine/src/lib.rs @@ -29,6 +29,11 @@ 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)?)?; @@ -36,5 +41,6 @@ fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> { 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 index 270c022..2173c77 100644 --- 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, Vec, Vec, Vec); + +/// 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 = 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 new file mode 100644 index 0000000..323e080 --- /dev/null +++ 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)