diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 95ae25c..3bb1721 100755 --- a/pepmatch/matcher.py +++ b/pepmatch/matcher.py @@ -55,35 +55,46 @@ def _indel_placements(query, matched): stretch t must equal the shorter string shifted by t residues. Placements are read off precomputed agreement runs for each shift, left to right so they come out in lexicographic order (_indel_edits reads placements[0]), instead of rebuilding the - string for every combination of positions.""" + string for every combination of positions. With one edit that reduces to two numbers: + position c is valid when the strings agree before it (common prefix) and after it + (common suffix), so the valid positions are one run, read off those two lengths.""" L, M = len(query), len(matched) n = abs(M - L) if n == 0: return [] longer, shorter = (query, matched) if M < L else (matched, query) N = len(longer) - # runs[s][x]: consecutive positions from x where longer[y + s] == shorter[y]. - runs = [] - for s in range(n + 1): - run = [0] * (N + 1) - for x in range(len(shorter) - 1, -1, -1): - if x + s < N and longer[x + s] == shorter[x]: - run[x] = run[x + 1] + 1 - runs.append(run) - - combos = [] - - def place(t, start, combo): - if t == n: - if runs[n][start - n] >= N - start: - combos.append(combo) - return - for c in range(max(start, 1), N - 1): # never an end residue: the terminal rule - if c - start > runs[t][start - t]: - break - place(t + 1, c + 1, combo + (c,)) + if n == 1: + prefix = 0 + while prefix < N - 1 and longer[prefix] == shorter[prefix]: + prefix += 1 + suffix = 0 + while suffix < N - 1 and longer[N - 1 - suffix] == shorter[N - 2 - suffix]: + suffix += 1 + combos = [(c,) for c in range(max(1, N - 1 - suffix), min(prefix, N - 2) + 1)] + else: + # runs[s][x]: consecutive positions from x where longer[y + s] == shorter[y]. + runs = [] + for s in range(n + 1): + run = [0] * (N + 1) + for x in range(len(shorter) - 1, -1, -1): + if x + s < N and longer[x + s] == shorter[x]: + run[x] = run[x + 1] + 1 + runs.append(run) + + combos = [] + + def place(t, start, combo): + if t == n: + if runs[n][start - n] >= N - start: + combos.append(combo) + return + for c in range(max(start, 1), N - 1): # never an end residue: the terminal rule + if c - start > runs[t][start - t]: + break + place(t + 1, c + 1, combo + (c,)) - place(0, 0, ()) + place(0, 0, ()) if M < L: return [tuple((i + 1, query[i]) for i in combo) for combo in combos] return [tuple((k - u + 1, matched[k]) for u, k in enumerate(combo)) for combo in combos] @@ -309,8 +320,12 @@ def indel_search(self): Queries too short for k=3 run against a 2-mer table on their own rather than dragging the whole batch down to k=2, where 2-mer seeds explode. Every group still runs at k <= its own pigeonhole optimum, so the match set equals a single-k run. - Groups are merged as columns, not frames, so an all-miss group's null column cannot - clash with another group's String column.""" + Each group's hits are labeled from the index that found them: protein numbers are + that index's own FASTA ordinals, so labeling them from another group's index would + name the wrong protein whenever the two were built from different versions of the + proteome. vertical_relaxed lets an all-miss group's null column stack under another + group's String column, and the text columns are cast to String at the end so a batch + that hits nothing has the same schema as one that does.""" n = self.max_indels short = [(q, s) for q, s in self.query if len(s) // (n + 1) < 3] rest = [(q, s) for q, s in self.query if len(s) // (n + 1) >= 3] @@ -322,8 +337,7 @@ def indel_search(self): rest_optimal = max(2, min(len(s) for _, s in rest) // (n + 1)) groups.setdefault(self._clamp_k(rest_optimal), []).extend(rest) - combined = [[] for _ in range(8)] - metadata_path = None + frames = [] for k, queries in sorted(groups.items()): self._recall_warning(k, n, queries) pepidx_path = self._pepidx_path(k) @@ -333,10 +347,10 @@ def indel_search(self): rs_preprocess(self.proteome_file, k, pepidx_path) print(f"Searching {len(queries)} peptides against {self.proteome_name} " f"(k={k}, max_indels={n})...") - for i, column in enumerate(rs_indel_match(pepidx_path, queries, n)): - combined[i].extend(column) - metadata_path = pepidx_path - return self._to_dataframe(combined, metadata_path, is_indels=True) + results = rs_indel_match(pepidx_path, queries, n) + frames.append(self._to_dataframe(results, pepidx_path, is_indels=True)) + df = pl.concat(frames, how='vertical_relaxed') + return df.with_columns(pl.col('Matched Sequence', 'Indel Positions').cast(pl.Utf8)) def _clamp_k(self, optimal): """Resolve k for one query group: an explicit k is honored up to `optimal`, above diff --git a/pepmatch/tests/test_indel_search.py b/pepmatch/tests/test_indel_search.py index f2fad25..358ba74 100644 --- a/pepmatch/tests/test_indel_search.py +++ b/pepmatch/tests/test_indel_search.py @@ -2,7 +2,7 @@ import polars as pl import polars.testing as plt from pathlib import Path -from pepmatch import Matcher +from pepmatch import Matcher, Preprocessor from pepmatch.matcher import format_indel_positions @pytest.fixture @@ -463,3 +463,64 @@ def test_indel2_batch_warns_only_for_the_short_query(tmp_path, capsys): assert 'complete recall is not guaranteed' in out assert 'lengths: [5]' in out # only the length-5 query is unguaranteed assert 'lengths: [12]' not in out + + +def test_indel_batch_labels_each_hit_from_the_index_that_found_it(tmp_path): + # .pepidx files are reused by name, so a FASTA edited in place can leave its 5-mer + # table built from the old version and its 2-mer table from the new one. Protein + # numbers are each table's own FASTA ordinals: the short query's hit is protein 1 on + # the 2-mer table (Delta), but protein 1 on the 5-mer table is Alpha. Each hit has to + # be labeled from the table that found it. + alpha = '>sp|P11111|A_HUMAN Alpha OS=Homo sapiens OX=9606 GN=AAA PE=1 SV=1\nMKVLATGGWQMKRRDEYH\n' + delta = '>sp|P44444|D_HUMAN Delta OS=Homo sapiens OX=9606 GN=DDD PE=1 SV=1\nWWYYWWCCHH\n' + proteome_path = tmp_path / 'proteome.fasta' + proteome_path.write_text(alpha) + Preprocessor(str(proteome_path), preprocessed_files_path=str(tmp_path)).preprocess(5) + proteome_path.write_text(delta + alpha) + Preprocessor(str(proteome_path), preprocessed_files_path=str(tmp_path)).preprocess(2) + + df = Matcher(query=['YYWWC', 'MKVLATGGWQ'], proteome_file=str(proteome_path), max_indels=1, + preprocessed_files_path=str(tmp_path), output_format='dataframe').match() + hits = df.filter(pl.col('Matched Sequence').is_not_null()) + assert set(zip(hits['Query Sequence'].to_list(), hits['Protein ID'].to_list())) == { + ('YYWWC', 'P44444.1'), # found on the 2-mer table + ('MKVLATGGWQ', 'P11111.1'), # found on the 5-mer table + } + + +@pytest.mark.parametrize('max_indels, short', [ + (1, ['YYWWC', 'ABDEF']), + (2, ['YYADGY', 'ZZYDGYZ']), + (1, ['WWWWW', 'YYYYY']), # every short query misses: a null-typed group stacked on hits +]) +def test_indel_batch_returns_what_its_short_and_long_queries_return_separately(tmp_path, max_indels, short): + # Short queries run on a 2-mer table and the rest on their own k (see indel_search). + # Splitting the batch must not change a single row: the batch's frame is the short + # queries' frame stacked on the long queries' frame, each searched on its own. + proteome_path = tmp_path / 'proteome.fasta' + proteome_path.write_text('>P1\nZZYDGYZZ\n>P2\nZZABCDEFGHIKLMZZ\n>P3\nQQWWYYWWCCHHKLMNPQ\n') + long = ['ABCDEFGHIKLM', 'WWYYWWCCHHKLM', 'WWWWWWWWWWWW'] + + def search(queries, name): + query_path = tmp_path / f'{name}.fasta' + query_path.write_text(''.join(f'>{q}\n{q}\n' for q in queries)) # IDs survive the split + return Matcher(query=str(query_path), proteome_file=str(proteome_path), max_indels=max_indels, + preprocessed_files_path=str(tmp_path), output_format='dataframe').match() + + batch = search([short[0], long[0], short[1], long[1], long[2]], 'batch') + separately = pl.concat([search(short, 'short'), search(long, 'long')], how='vertical_relaxed') + plt.assert_frame_equal(batch, separately) + assert batch['Matched Sequence'].is_not_null().sum() > 0 # the comparison covered real hits + + +def test_indel_search_that_hits_nothing_keeps_the_text_column_types(tmp_path): + # With no hit anywhere, the text columns would otherwise come back null-typed, a + # different schema from any search that matches something. + proteome_path = tmp_path / 'proteome.fasta' + proteome_path.write_text('>P1\nZZYDGYZZ\n>P2\nZZABCDEFGHIKLMZZ\n') + for queries in (['WWWWWWWWWWWW'], ['WWWWW', 'WWWWWWWWWWWW']): # one table, then two + df = Matcher(query=queries, proteome_file=str(proteome_path), max_indels=1, + preprocessed_files_path=str(tmp_path), output_format='dataframe').match() + assert df['Matched Sequence'].is_null().all() + assert df.schema['Matched Sequence'] == pl.Utf8 + assert df.schema['Indel Positions'] == pl.Utf8 diff --git a/pepmatch/tests/test_metadata_rows.py b/pepmatch/tests/test_metadata_rows.py index 50ff73b..c7bdd9f 100644 --- a/pepmatch/tests/test_metadata_rows.py +++ b/pepmatch/tests/test_metadata_rows.py @@ -139,22 +139,10 @@ def capture(self, cols, pepidx_path, is_indels=False): 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', - ) +def test_a_protein_the_index_lacks_raises_with_a_rebuild_hint(proteome): + # Protein numbers are the index's own FASTA ordinals, so a number this index lacks means the + # hits came from an index built from a different version of the FASTA (best match without k + # searches several indexes and reads metadata from the last). Fail loudly and say why. + matcher = Matcher(query=['MKVLAT'], proteome_file=str(proteome), k=5, preprocessed_files_path=str(proteome.parent)) with pytest.raises(ValueError, match='protein number 4 .* different versions'): - matcher.match() + matcher._metadata_table(index_of(proteome), [1, 4])