diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 53d8e28..3aa4942 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 +from ._rs import rs_preprocess, rs_match, rs_discontinuous, rs_metadata VALID_OUTPUT_FORMATS = ['dataframe', 'csv', 'tsv', 'xlsx', 'json'] FASTA_EXTENSIONS = { @@ -160,7 +160,17 @@ def _search(self, k, max_mismatches, peptides=None): def best_match_search(self): peptides_remaining = self.query.copy() - all_results = [] + acc = tuple([] for _ in range(8)) # 8 columnar accumulators + + def collect_matched(cols): + matched_ids = set() + matched_col = cols[2] + for i in range(len(cols[0])): + if matched_col[i] is not None: + for j in range(8): + acc[j].append(cols[j][i]) + matched_ids.add(cols[0][i]) + return matched_ids initial_k = min(len(seq) for _, seq in peptides_remaining) ks = [initial_k] @@ -172,20 +182,11 @@ def best_match_search(self): for k in ks: if not peptides_remaining: break - min_len = min(len(seq) for _, seq in peptides_remaining) max_mm = (min_len // k) - 1 if max_mm < 0: max_mm = 0 - - results = self._search(k, max_mm, peptides_remaining) - - matched_ids = set() - for row in results: - if row[2] != '': - all_results.append(row) - matched_ids.add(row[0]) - + matched_ids = collect_matched(self._search(k, max_mm, peptides_remaining)) peptides_remaining = [ (qid, seq) for qid, seq in peptides_remaining if qid not in matched_ids ] @@ -193,31 +194,22 @@ def best_match_search(self): if peptides_remaining: max_mm = min(len(seq) for _, seq in peptides_remaining) // 2 - while peptides_remaining: shortest_len = min(len(seq) for _, seq in peptides_remaining) if max_mm >= shortest_len: break - max_mm += 1 - results = self._search(2, max_mm, peptides_remaining) - - matched_ids = set() - for row in results: - if row[2] != '': - all_results.append(row) - matched_ids.add(row[0]) - + matched_ids = collect_matched(self._search(2, max_mm, peptides_remaining)) peptides_remaining = [ (qid, seq) for qid, seq in peptides_remaining if qid not in matched_ids ] print(f" -> k=2, mismatches<={max_mm}: {len(peptides_remaining)} remaining") - if peptides_remaining: - for qid, seq in peptides_remaining: - all_results.append([qid, seq] + [''] * 14) + for qid, seq in peptides_remaining: # leftover unmatched + for j, val in enumerate((qid, seq, None, None, None, "[]", None, None)): + acc[j].append(val) - df = self._to_dataframe(all_results) + df = self._to_dataframe(acc) return self._best_match_filter(df) def _best_match_filter(self, df): @@ -274,46 +266,63 @@ def _best_match_filter(self, df): return pl.concat([matched_df, unmatched_df], how="vertical") - def _to_dataframe(self, results): - schema = { - 'Query ID': pl.Utf8, - 'Query Sequence': pl.Utf8, - 'Matched Sequence': pl.Utf8, - 'Protein ID': pl.Utf8, - 'Protein Name': pl.Utf8, - 'Species': pl.Utf8, - 'Taxon ID': pl.Utf8, - 'Gene': pl.Utf8, - 'Mismatches': pl.Utf8, - 'Mutated Positions': pl.Utf8, - 'Index start': pl.Utf8, - 'Index end': pl.Utf8, - 'Protein Existence Level': pl.Utf8, - 'Sequence Version': pl.Utf8, - 'Gene Priority': pl.Utf8, - 'SwissProt Reviewed': pl.Utf8, - } - - if not results: - return pl.DataFrame(schema=schema) - - df = pl.DataFrame(results, schema=schema, orient='row') - - df = df.with_columns( - pl.when(pl.col(col) == '').then(None).otherwise(pl.col(col)).alias(col) - for col in df.columns + def _metadata_table(self) -> pl.DataFrame: + """Per-protein metadata (built once from this proteome's index) for the edge join.""" + import glob + pattern = os.path.join(self.preprocessed_files_path, f'{self.proteome_name}_*mers.pepidx') + idx_files = glob.glob(pattern) + if not idx_files: + raise FileNotFoundError(f'No .pepidx for {self.proteome_name} in {self.preprocessed_files_path}') + pnum, pid, name, species, taxon, gene, exist, seqver, geneprio, swiss = rs_metadata(idx_files[0]) + m = pl.DataFrame({ + 'protein_num': pl.Series(pnum, dtype=pl.UInt32), + 'Protein ID': pid, 'Protein Name': name, 'Species': species, 'Taxon ID': taxon, + 'Gene': gene, 'Protein Existence Level': exist, 'Sequence Version': seqver, + 'Gene Priority': geneprio, 'SwissProt Reviewed': swiss, + }) + str_cols = ['Protein ID','Protein Name','Species','Taxon ID','Gene', + 'Protein Existence Level','Sequence Version','Gene Priority','SwissProt Reviewed'] + m = m.with_columns( + pl.when(pl.col(c) == '').then(None).otherwise(pl.col(c)).alias(c) for c in str_cols ) - - df = df.with_columns([ - pl.col('Mismatches').cast(pl.Int64), - pl.col('Index start').cast(pl.Int64), - pl.col('Index end').cast(pl.Int64), + return m.with_columns([ pl.col('Protein Existence Level').cast(pl.Int64), pl.col('Sequence Version').cast(pl.Int64), pl.col('Gene Priority').cast(pl.Int64), pl.col('SwissProt Reviewed').cast(pl.Int64).cast(pl.Boolean), ]) + FINAL_COLUMNS = [ + 'Query ID','Query Sequence','Matched Sequence','Protein ID','Protein Name','Species', + 'Taxon ID','Gene','Mismatches','Mutated Positions','Index start','Index end', + 'Protein Existence Level','Gene Priority','SwissProt Reviewed', + ] + + def _to_dataframe(self, cols): + """Build the result DataFrame from columnar Rust output, reconstructing protein + metadata via a single join (instead of cloning it into every hit row).""" + qid, qseq, matched, pnum, mm, mutated, istart, iend = cols + + if not qid: + schema = {c: pl.Utf8 for c in self.FINAL_COLUMNS} + for c in ('Mismatches','Index start','Index end','Protein Existence Level','Gene Priority'): + schema[c] = pl.Int64 + schema['SwissProt Reviewed'] = pl.Boolean + return pl.DataFrame(schema=schema) + + base = pl.DataFrame({ + 'Query ID': qid, + 'Query Sequence': qseq, + 'Matched Sequence': matched, + 'protein_num': pl.Series(pnum, dtype=pl.UInt32), + 'Mismatches': pl.Series(mm, dtype=pl.Int64), + 'Mutated Positions': mutated, + 'Index start': pl.Series(istart, dtype=pl.Int64), + 'Index end': pl.Series(iend, dtype=pl.Int64), + }) + + df = base.join(self._metadata_table(), on='protein_num', how='left').drop('protein_num') + if self.sequence_version: df = df.with_columns( pl.when(pl.col('Sequence Version').is_not_null()) @@ -322,5 +331,4 @@ def _to_dataframe(self, results): .alias('Protein ID') ) - df = df.drop('Sequence Version') - return df + return df.drop('Sequence Version').select(self.FINAL_COLUMNS) diff --git a/pepmatch/rs-engine/Cargo.lock b/pepmatch/rs-engine/Cargo.lock index e1d8689..01e4263 100644 --- a/pepmatch/rs-engine/Cargo.lock +++ b/pepmatch/rs-engine/Cargo.lock @@ -92,7 +92,7 @@ checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" [[package]] name = "pepmatch_rs_engine" -version = "1.16.1" +version = "1.16.3" dependencies = [ "memmap2", "pyo3", diff --git a/pepmatch/rs-engine/src/lib.rs b/pepmatch/rs-engine/src/lib.rs index b069201..058b630 100644 --- a/pepmatch/rs-engine/src/lib.rs +++ b/pepmatch/rs-engine/src/lib.rs @@ -15,20 +15,26 @@ fn rs_preprocess(fasta_path: &str, k: usize, output_path: &str) { } #[pyfunction] -fn rs_match(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> Vec> { +fn rs_match(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> matching::Columns { matching::run(pepidx_path, peptides, k, max_mismatches) } #[pyfunction] -fn rs_discontinuous(pepidx_path: &str, epitopes: Vec<(String, Vec<(char, usize)>)>, max_mismatches: usize) -> Vec> { +fn rs_discontinuous(pepidx_path: &str, epitopes: Vec<(String, Vec<(char, usize)>)>, max_mismatches: usize) -> matching::Columns { matching::run_discontinuous(pepidx_path, epitopes, max_mismatches) } +#[pyfunction] +fn rs_metadata(pepidx_path: &str) -> matching::MetaColumns { + matching::run_metadata(pepidx_path) +} + #[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)?)?; Ok(()) } diff --git a/pepmatch/rs-engine/src/match.rs b/pepmatch/rs-engine/src/match.rs index ac76176..270c022 100644 --- a/pepmatch/rs-engine/src/match.rs +++ b/pepmatch/rs-engine/src/match.rs @@ -326,74 +326,113 @@ fn mismatch_match(peptide: &str, k: usize, max_mismatches: usize, index: &PepInd matches } -fn search_peptide(query_id: &str, peptide: &str, k: usize, max_mismatches: usize, index: &PepIndex) -> Vec> { +struct HitRecord { + query_id: String, + query_seq: String, + matched: Option, + protein_num: Option, + mismatches: Option, + mutated: String, + idx_start: Option, + idx_end: Option, +} + +fn mutated_positions(peptide: &str, matched: &str) -> String { + let v: Vec = peptide.as_bytes().iter() + .zip(matched.as_bytes()) + .enumerate() + .filter(|(_, (a, b))| a != b) + .map(|(i, _)| (i + 1).to_string()) + .collect(); + format!("[{}]", v.join(", ")) +} + +fn hit_record(query_id: &str, peptide: &str, matched_seq: &str, mismatches: usize, encoded_start: u64) -> HitRecord { + let protein_number = (encoded_start / PROTEIN_INDEX_MULTIPLIER) as u32; + let position = (encoded_start % PROTEIN_INDEX_MULTIPLIER) as usize; + HitRecord { + query_id: query_id.to_string(), + query_seq: peptide.to_string(), + matched: Some(matched_seq.to_string()), + protein_num: Some(protein_number), + mismatches: Some(mismatches as i64), + mutated: mutated_positions(peptide, matched_seq), + idx_start: Some((position + 1) as i64), + idx_end: Some((position + peptide.len()) as i64), + } +} + +fn miss_record(query_id: &str, peptide: &str) -> HitRecord { + HitRecord { + query_id: query_id.to_string(), + query_seq: peptide.to_string(), + matched: None, + protein_num: None, + mismatches: None, + mutated: String::from("[]"), + idx_start: None, + idx_end: None, + } +} + +fn search_peptide(query_id: &str, peptide: &str, k: usize, max_mismatches: usize, index: &PepIndex) -> Vec { if max_mismatches == 0 { let hits = exact_match(peptide, k, index); if hits.is_empty() { - return vec![make_empty_row(query_id, peptide)]; + return vec![miss_record(query_id, peptide)]; } - hits.iter().map(|&hit| make_hit_row(query_id, peptide, peptide, 0, hit, index)).collect() + hits.iter().map(|&hit| hit_record(query_id, peptide, peptide, 0, hit)).collect() } else { let hits = mismatch_match(peptide, k, max_mismatches, index); if hits.is_empty() { - return vec![make_empty_row(query_id, peptide)]; + return vec![miss_record(query_id, peptide)]; } - hits.iter().map(|(matched_seq, mm, start)| make_hit_row(query_id, peptide, matched_seq, *mm, *start, index)).collect() + hits.iter().map(|(matched_seq, mm, start)| hit_record(query_id, peptide, matched_seq, *mm, *start)).collect() } } -fn make_hit_row(query_id: &str, peptide: &str, matched_seq: &str, mismatches: usize, encoded_start: u64, index: &PepIndex) -> Vec { - let protein_number = (encoded_start / PROTEIN_INDEX_MULTIPLIER) as usize; - let position = (encoded_start % PROTEIN_INDEX_MULTIPLIER) as usize; - let meta = index.get_metadata(protein_number); - - let mutated: Vec = peptide.as_bytes().iter() - .zip(matched_seq.as_bytes()) - .enumerate() - .filter(|(_, (a, b))| a != b) - .map(|(i, _)| (i + 1).to_string()) - .collect(); - - vec![ - query_id.to_string(), - peptide.to_string(), - matched_seq.to_string(), - meta[0].clone(), - meta[1].clone(), - meta[2].clone(), - meta[3].clone(), - meta[4].clone(), - mismatches.to_string(), - format!("[{}]", mutated.join(", ")), - (position + 1).to_string(), - (position + peptide.len()).to_string(), - meta[5].clone(), - meta[6].clone(), - meta[7].clone(), - meta[8].clone(), - ] -} - -fn make_empty_row(query_id: &str, peptide: &str) -> Vec { - vec![ - query_id.to_string(), - peptide.to_string(), - String::new(), String::new(), String::new(), String::new(), - String::new(), String::new(), String::new(), String::from("[]"), - String::new(), String::new(), String::new(), String::new(), - String::new(), String::new(), - ] +pub(crate) type Columns = ( + Vec, Vec, Vec>, Vec>, + Vec>, Vec, Vec>, Vec>, +); + +pub(crate) type MetaColumns = ( + Vec, Vec, Vec, Vec, Vec, + Vec, Vec, Vec, Vec, Vec, +); + +fn unzip_records(records: Vec) -> Columns { + let n = records.len(); + let mut qid = Vec::with_capacity(n); + let mut qseq = Vec::with_capacity(n); + let mut matched = Vec::with_capacity(n); + let mut pnum = Vec::with_capacity(n); + let mut mm = Vec::with_capacity(n); + let mut mutated = Vec::with_capacity(n); + let mut s = Vec::with_capacity(n); + let mut e = Vec::with_capacity(n); + for r in records { + qid.push(r.query_id); + qseq.push(r.query_seq); + matched.push(r.matched); + pnum.push(r.protein_num); + mm.push(r.mismatches); + mutated.push(r.mutated); + s.push(r.idx_start); + e.push(r.idx_end); + } + (qid, qseq, matched, pnum, mm, mutated, s, e) } pub(crate) fn run_discontinuous( pepidx_path: &str, epitopes: Vec<(String, Vec<(char, usize)>)>, max_mismatches: usize, -) -> Vec> { +) -> Columns { let index = PepIndex::open(pepidx_path); let seq = &index.mmap[index.seq_offset..index.seq_offset + index.seq_len]; - let mut all_results: Vec> = Vec::new(); + let mut records: Vec = Vec::new(); for (query_id, residues) in &epitopes { let mut found = false; @@ -430,7 +469,6 @@ pub(crate) fn run_discontinuous( if valid { found = true; - let meta = index.get_metadata(pn); let matched_str: String = residues.iter() .map(|&(_, p)| format!("{}{}", seq[base + p - 1] as char, p)) .collect::>() @@ -440,43 +478,66 @@ pub(crate) fn run_discontinuous( .map(|&(_, p)| p.to_string()) .collect(); - all_results.push(vec![ - query_id.clone(), - query_str.clone(), - matched_str, - meta[0].clone(), meta[1].clone(), meta[2].clone(), - meta[3].clone(), meta[4].clone(), - mismatches.to_string(), - format!("[{}]", mutated.join(", ")), - residues.first().map(|&(_, p)| p.to_string()).unwrap_or_default(), - residues.last().map(|&(_, p)| p.to_string()).unwrap_or_default(), - meta[5].clone(), meta[6].clone(), meta[7].clone(), meta[8].clone(), - ]); + records.push(HitRecord { + query_id: query_id.clone(), + query_seq: query_str.clone(), + matched: Some(matched_str), + protein_num: Some(pn as u32), + mismatches: Some(mismatches as i64), + mutated: format!("[{}]", mutated.join(", ")), + idx_start: residues.first().map(|&(_, p)| p as i64), + idx_end: residues.last().map(|&(_, p)| p as i64), + }); } } if !found { - let mut row = vec![query_id.clone(), query_str.clone()]; - row.extend(std::iter::repeat(String::new()).take(6)); - row.push(String::new()); - row.push(String::from("[]")); - row.extend(std::iter::repeat(String::new()).take(6)); - all_results.push(row); + records.push(HitRecord { + query_id: query_id.clone(), + query_seq: query_str.clone(), + matched: None, + protein_num: None, + mismatches: None, + mutated: String::from("[]"), + idx_start: None, + idx_end: None, + }); } } - all_results + unzip_records(records) } -pub(crate) fn run(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> Vec> { +fn metadata_columns(index: &PepIndex) -> MetaColumns { + let n = index.num_proteins; + let mut pnum = Vec::with_capacity(n); + let mut f: [Vec; 9] = Default::default(); + for v in f.iter_mut() { + v.reserve(n); + } + for pn in 1..=n { + pnum.push(pn as u32); + let m = index.get_metadata(pn); + for (i, field) in m.into_iter().enumerate() { + f[i].push(field); + } + } + let [a, b, c, d, e, g, h, i, j] = f; + (pnum, a, b, c, d, e, g, h, i, j) +} + +pub(crate) fn run_metadata(pepidx_path: &str) -> MetaColumns { let index = PepIndex::open(pepidx_path); + metadata_columns(&index) +} - let all_results: Vec>> = peptides +pub(crate) fn run(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> Columns { + let index = PepIndex::open(pepidx_path); + let records: Vec = peptides .par_iter() - .map(|(query_id, peptide)| { - search_peptide(query_id, peptide, k, max_mismatches, &index) + .flat_map_iter(|(query_id, peptide)| { + search_peptide(query_id, peptide, k, max_mismatches, &index).into_iter() }) .collect(); - - all_results.into_iter().flatten().collect() + unzip_records(records) }