diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2825be1..af56da9 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -36,6 +36,8 @@ jobs: test-file: pepmatch/tests/test_indel_seed_reduction.py - name: Indel Positions Equivalence test-file: pepmatch/tests/test_indel_positions_equivalence.py + - name: Metadata Rows + test-file: pepmatch/tests/test_metadata_rows.py steps: - name: Checkout Code diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 0847ba5..95ae25c 100755 --- a/pepmatch/matcher.py +++ b/pepmatch/matcher.py @@ -524,19 +524,28 @@ def _best_match_filter(self, df): return pl.concat([matched_df, unmatched_df], how="vertical_relaxed") - def _metadata_table(self, pepidx_path: str) -> pl.DataFrame: - """Per-protein metadata read from `pepidx_path`, the index the search just used. - Protein numbers are that index's own FASTA ordinals, so any other file -- including - one whose name merely shares this proteome's prefix -- would misattribute every hit.""" - pnum, pid, name, species, taxon, gene, exist, seqver, geneprio, swiss = rs_metadata(pepidx_path) + def _metadata_table(self, pepidx_path: str, protein_nums=None) -> pl.DataFrame: + """Metadata for `protein_nums` (every protein when None), read from `pepidx_path`, the + index the search just used. Protein numbers are that index's own FASTA ordinals, so any + other file -- including one whose name merely shares this proteome's prefix -- would + misattribute every hit. Searches pass only the proteins they hit: building the whole + proteome's table (82k rows for human) cost more than the search itself on small queries. + A hit on a protein this index lacks means the search's indexes disagree, which raises.""" + try: + pnum, pid, name, species, taxon, gene, exist, seqver, geneprio, swiss = rs_metadata(pepidx_path, protein_nums) + except ValueError as e: # .pepidx files are reused by name, so a FASTA edited in place leaves stale ones + raise ValueError( + f"{e}. This proteome's .pepidx files in {self.preprocessed_files_path} were most likely built " + f"from different versions of {self.proteome_file}; delete them to rebuild." + ) from e + str_cols = ['Protein ID','Protein Name','Species','Taxon ID','Gene', + 'Protein Existence Level','Sequence Version','Gene Priority','SwissProt Reviewed'] 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'] + }, schema_overrides={c: pl.Utf8 for c in str_cols}) # String even when nothing was hit m = m.with_columns( pl.when(pl.col(c) == '').then(None).otherwise(pl.col(c)).alias(c) for c in str_cols ) @@ -561,7 +570,8 @@ def _final_columns(self, is_indels): def _to_dataframe(self, cols, pepidx_path, is_indels=False): """Build the results frame from the columnar Rust output, joining protein metadata - from `pepidx_path` once instead of cloning it into every hit row. rs_indel_match + from `pepidx_path` once instead of cloning it into every hit row. Only the proteins + the search hit are read, and the join keeps the hits' row order. rs_indel_match reuses the mismatch slot for its edit count, so only the column name differs by mode; indel positions are derived here from (query, matched) rather than in Rust, once per distinct pair since a matched peptide often recurs across proteins, and @@ -600,7 +610,9 @@ def _to_dataframe(self, cols, pepidx_path, is_indels=False): 'Index end': pl.Series(iend, dtype=pl.Int64), }) - df = base.join(self._metadata_table(pepidx_path), on='protein_num', how='left').drop('protein_num') + hit_proteins = base.get_column('protein_num').drop_nulls().unique().to_list() + metadata = self._metadata_table(pepidx_path, hit_proteins) + df = base.join(metadata, on='protein_num', how='left', maintain_order='left').drop('protein_num') if self.sequence_version: df = df.with_columns( diff --git a/pepmatch/rs-engine/src/lib.rs b/pepmatch/rs-engine/src/lib.rs index 1655a0b..d3b3da6 100644 --- a/pepmatch/rs-engine/src/lib.rs +++ b/pepmatch/rs-engine/src/lib.rs @@ -29,9 +29,12 @@ fn rs_discontinuous(pepidx_path: &str, epitopes: Vec<(String, Vec<(char, usize)> matching::run_discontinuous(pepidx_path, epitopes, max_mismatches) } +/// Metadata for `protein_nums` (1-based FASTA ordinals in this index), or for every protein +/// when it is omitted. A number outside the index raises `ValueError`. #[pyfunction] -fn rs_metadata(pepidx_path: &str) -> matching::MetaColumns { - matching::run_metadata(pepidx_path) +#[pyo3(signature = (pepidx_path, protein_nums=None))] +fn rs_metadata(pepidx_path: &str, protein_nums: Option>) -> PyResult { + matching::run_metadata(pepidx_path, protein_nums).map_err(pyo3::exceptions::PyValueError::new_err) } #[pyfunction] diff --git a/pepmatch/rs-engine/src/match.rs b/pepmatch/rs-engine/src/match.rs index c878d8c..420f958 100644 --- a/pepmatch/rs-engine/src/match.rs +++ b/pepmatch/rs-engine/src/match.rs @@ -576,22 +576,38 @@ pub(crate) fn run_discontinuous( unzip_records(records) } -fn metadata_columns(index: &PepIndex) -> MetaColumns { +/// One row per protein in `wanted` (ascending, duplicates collapsed), or per protein in the +/// index when `wanted` is None. A search passes only the proteins it hit, so the cost follows +/// the hits rather than the proteome (82k proteins for human, most of which a search never +/// touches). Duplicates are collapsed here because a repeated row would double every hit +/// joined to it. +fn metadata_columns(index: &PepIndex, wanted: Option>) -> Result { let n = index.num_proteins; - let mut pnum = Vec::with_capacity(n); + let proteins: Vec = match wanted { + Some(mut nums) => { + nums.sort_unstable(); + nums.dedup(); + if let Some(&bad) = nums.iter().find(|&&pn| pn == 0 || pn as usize > n) { + return Err(format!("protein number {bad} is not in this index (1..={n})")); + } + nums + } + None => (1..=n as u32).collect(), + }; + let mut pnum = Vec::with_capacity(proteins.len()); let mut f: [Vec; 9] = Default::default(); for v in f.iter_mut() { - v.reserve(n); + v.reserve(proteins.len()); } - for pn in 1..=n { - pnum.push(pn as u32); - let m = index.get_metadata(pn); + for pn in proteins { + pnum.push(pn); + let m = index.get_metadata(pn as usize); 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) + Ok((pnum, a, b, c, d, e, g, h, i, j)) } /// Seeds for indel search: the n+1 disjoint query tiles with the FEWEST proteome @@ -862,9 +878,9 @@ pub(crate) fn run_indel( unzip_records(records) } -pub(crate) fn run_metadata(pepidx_path: &str) -> MetaColumns { +pub(crate) fn run_metadata(pepidx_path: &str, protein_nums: Option>) -> Result { let index = PepIndex::open(pepidx_path); - metadata_columns(&index) + metadata_columns(&index, protein_nums) } pub(crate) fn run(pepidx_path: &str, peptides: Vec<(String, String)>, k: usize, max_mismatches: usize) -> Columns { diff --git a/pepmatch/tests/test_metadata_rows.py b/pepmatch/tests/test_metadata_rows.py new file mode 100644 index 0000000..50ff73b --- /dev/null +++ b/pepmatch/tests/test_metadata_rows.py @@ -0,0 +1,160 @@ +import pytest +import polars.testing as plt + +from pathlib import Path + +import pepmatch.matcher as matcher_module +from pepmatch import Matcher, Preprocessor +from pepmatch._rs import rs_metadata + +DATA = Path(__file__).parent / 'data' + +# Three proteins. The second has no gene name (no GN=), so its Gene field is empty and must +# come back null; the third is never hit by the searches below, so its metadata is never read. +PROTEOME_FASTA = ( + '>sp|P11111|AAAA_HUMAN Alpha protein OS=Homo sapiens OX=9606 GN=AAA PE=1 SV=1\n' + 'MKVLATGGWQMKRRDEYH\n' + '>sp|P22222|BBBB_HUMAN Beta protein OS=Homo sapiens OX=9606 PE=2 SV=3\n' + 'RRHFDSAMFFLLNNMKVLAT\n' + '>tr|A0A333|A0A333_HUMAN Gamma protein OS=Homo sapiens OX=9606 GN=CCC PE=4 SV=1\n' + 'PQRSTWYVDEGHIKLMNPQ\n' +) + + +@pytest.fixture +def proteome(tmp_path) -> Path: + """The three-protein FASTA, preprocessed at k=5 next to it.""" + fasta = tmp_path / 'three.fasta' + fasta.write_text(PROTEOME_FASTA) + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(5) + return fasta + + +def index_of(proteome): + return str(proteome.with_name('three_5mers.pepidx')) + + +def rows(columns): + return list(zip(*columns)) + + +def search(proteome, query): + return Matcher( + query=query, + proteome_file=str(proteome), + k=5, + preprocessed_files_path=str(proteome.parent), + output_format='dataframe', + ).match() + + +def test_requested_rows_equal_the_full_tables_rows(proteome): + full = rows(rs_metadata(index_of(proteome))) + assert [row[0] for row in full] == [1, 2, 3] + assert rows(rs_metadata(index_of(proteome), [3, 1])) == [full[0], full[2]] # ascending + + +def test_repeated_protein_numbers_collapse(proteome): + # A repeated row would double every hit joined to it. + assert [row[0] for row in rows(rs_metadata(index_of(proteome), [2, 2, 2]))] == [2] + + +def test_no_protein_numbers_return_empty_columns(proteome): + assert rs_metadata(index_of(proteome), []) == tuple([] for _ in range(10)) + + +@pytest.mark.parametrize('protein', [0, 4]) +def test_a_protein_number_outside_the_index_raises(proteome, protein): + with pytest.raises(ValueError, match=f'protein number {protein} '): + rs_metadata(index_of(proteome), [protein]) + + +def test_hits_carry_their_own_proteins_metadata(proteome): + df = search(proteome, ['MKVLAT', 'FFLLNN']) + got = { + (row['Query Sequence'], row['Protein ID'], row['Gene'], row['Protein Existence Level'], row['SwissProt Reviewed']) + for row in df.iter_rows(named=True) + } + assert got == { + ('MKVLAT', 'P11111.1', 'AAA', 1, True), + ('MKVLAT', 'P22222.3', None, 2, True), + ('FFLLNN', 'P22222.3', None, 2, True), + } + + +METADATA_COLUMNS = [ + 'Protein ID', 'Protein Name', 'Species', 'Taxon ID', 'Gene', + 'Protein Existence Level', 'Gene Priority', 'SwissProt Reviewed', +] + + +def test_a_search_that_hits_nothing_keeps_the_metadata_column_types(proteome): + # No protein is hit, so no metadata row is read. The metadata columns must still have the + # types a search with hits gives them, not turn into untyped nulls. + hit, miss = search(proteome, ['MKVLAT']), search(proteome, ['WWWWWW']) + assert miss['Matched Sequence'].is_null().all() + assert {c: miss.schema[c] for c in METADATA_COLUMNS} == {c: hit.schema[c] for c in METADATA_COLUMNS} + + +SEARCHES = { + 'exact': dict(query=DATA / 'exact_match_query.fasta'), + 'mismatch': dict(query=DATA / 'mismatching_query.fasta', max_mismatches=3), + 'best match': dict(query=DATA / 'best_match_query.fasta', best_match=True), + 'discontinuous with linear': dict( + query=['L354, V420, G461, Q468, E486, K499, D501, M503, G509', 'T819, L822, A835, F840, S871, Y884, P886', + 'CNERFERLMKELSE', 'FAKEPEPTIDE'], + max_mismatches=1), + 'one indel': dict(query=DATA / 'indel_query.fasta', max_indels=1), + 'two indels': dict(query=DATA / 'indel_query.fasta', max_indels=2), + 'nothing hit': dict(query=['WWWWWWWWW']), +} + + +@pytest.mark.parametrize('scenario', SEARCHES) +def test_reading_only_the_hit_proteins_gives_the_frame_reading_every_protein_gives(monkeypatch, tmp_path, scenario): + # Capture the raw columns each search hands _to_dataframe, then build the frame from those + # same columns twice: reading only the hit proteins' metadata, as searches now do, and + # reading every protein's, as they did before. The input is identical, so the comparison + # is exact even for searches whose hit order changes from run to run. + captured = [] + to_dataframe = Matcher._to_dataframe + + def capture(self, cols, pepidx_path, is_indels=False): + captured.append((self, cols, pepidx_path, is_indels)) + return to_dataframe(self, cols, pepidx_path, is_indels) + + monkeypatch.setattr(Matcher, '_to_dataframe', capture) + Matcher(proteome_file=DATA / 'proteome.fasta', preprocessed_files_path=str(tmp_path), **SEARCHES[scenario]).match() + monkeypatch.undo() + assert captured + + hits = 0 + for matcher, cols, pepidx_path, is_indels in captured: + hit_proteins_only = matcher._to_dataframe(cols, pepidx_path, is_indels) + monkeypatch.setattr(matcher_module, 'rs_metadata', lambda path, protein_nums=None: rs_metadata(path)) + every_protein = matcher._to_dataframe(cols, pepidx_path, is_indels) + monkeypatch.undo() + plt.assert_frame_equal(hit_proteins_only, every_protein) + hits += hit_proteins_only['Protein ID'].is_not_null().sum() + assert (hits > 0) == (scenario != 'nothing hit') # the comparison covered real hits + + +def test_a_hit_from_a_stale_index_raises_instead_of_losing_its_metadata(tmp_path): + # .pepidx files are reused by name, so editing a FASTA in place can leave its 5-mer index + # three proteins long while a rebuilt 2-mer index has four. An indel batch of short and long + # queries searches both and reads metadata from the 5-mer index, so a hit on the fourth + # protein has no metadata row: it must fail loudly, not come back with blank metadata. + fasta = tmp_path / 'three.fasta' + fasta.write_text(PROTEOME_FASTA) + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(5) + fasta.write_text(PROTEOME_FASTA + '>sp|P44444|DDDD_HUMAN Delta protein OS=Homo sapiens OX=9606 GN=DDD PE=1 SV=1\nWWYYWWCCHH\n') + Preprocessor(str(fasta), preprocessed_files_path=str(tmp_path)).preprocess(2) + matcher = Matcher( + query=['YYWWC', 'MKVLATGGWQ'], + proteome_file=str(fasta), + max_indels=1, + preprocessed_files_path=str(tmp_path), + output_format='dataframe', + ) + with pytest.raises(ValueError, match='protein number 4 .* different versions'): + matcher.match()