indel search: label each group's hits from its own index

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

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

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])