From ba2593a3e4e07295f880965b29f09de14f0d0e43 Mon Sep 17 00:00:00 2001 From: Roman Young Date: Wed, 14 Oct 2026 17:23:11 +0000 Subject: [PATCH 1/2] indel search: label each group's hits from the index that found them indel_search splits a batch by query length (short queries on a 2-mer table, the rest on their own k) but joined every group's hits to the last group's index for metadata. Protein numbers are each index's own FASTA ordinals, so when the per-k .pepidx files come from different versions of the FASTA (they're reused by name, eg. after the FASTA is edited in place) a short query's hit got named after whichever protein holds that number in the other file. Each group's frame is now built from its own index and the frames are stacked with vertical_relaxed. With consistent index files nothing changes. An all-miss indel search also returned Matched Sequence and Indel Positions as null-typed columns; they're cast to String so every search has the same schema. Recall is unchanged: 100% on all 11 local indel benchmark sets with outputs identical to master, and 9,000 random mixed batches (a third with stale tables) match the brute-force oracle. New tests cover per-index labeling with stale tables, a batch equal to its short and long queries searched separately, and the all-miss types. test_metadata_rows' stale-index test now checks the error hint directly, since indel search no longer reads another group's index. --- pepmatch/matcher.py | 19 +++++---- pepmatch/tests/test_indel_search.py | 63 +++++++++++++++++++++++++++- pepmatch/tests/test_metadata_rows.py | 24 +++-------- 3 files changed, 79 insertions(+), 27 deletions(-) diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 95ae25c..00adb65 100755 --- a/pepmatch/matcher.py +++ b/pepmatch/matcher.py @@ -309,8 +309,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 +326,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 +336,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]) From ddaba9d47e433b1a9a2242629d16a081b3b5080f Mon Sep 17 00:00:00 2001 From: Roman Young Date: Wed, 14 Oct 2026 17:23:11 +0000 Subject: [PATCH 2/2] perf: read one-edit indel placements off the common prefix and suffix With one edit, a placement is valid exactly where the two strings agree before it and after it, so the valid positions are one run between the common prefix and the common suffix. _indel_placements now reads that run off the two lengths instead of building agreement runs for every shift and walking them. Two-edit pairs keep the general path. Placements are identical and in the same order: test_indel_positions_equivalence checks them against the original brute-force implementation exhaustively over small alphabets, and 1,000,000 random pairs agree as well. Annotation time halves, and output building for 50k 1-indel queries drops from 0.91 s to 0.65 s. --- pepmatch/matcher.py | 55 +++++++++++++++++++++++++++------------------ 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 00adb65..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]