indel annotation: annotate each distinct match once, and find indel positions without rebuilding strings

.github/workflows/tests.yml

@@ -34,6 +34,8 @@ jobs:
             test-file: pepmatch/tests/test_indel2_property.py
           - name: Indel Seed Reduction
             test-file: pepmatch/tests/test_indel_seed_reduction.py
+          - name: Indel Positions Equivalence
+            test-file: pepmatch/tests/test_indel_positions_equivalence.py
 
     steps:
     - name: Checkout Code

pepmatch/matcher.py

@@ -1,6 +1,5 @@
 import os
 import polars as pl
-from itertools import combinations
 from pathlib import Path
 from Bio import SeqIO
 from ._rs import rs_preprocess, rs_match, rs_discontinuous, rs_metadata, rs_match_counts, rs_indel_match
@@ -50,30 +49,44 @@ def _indel_placements(query, matched):
   """Every valid placement of the edits between `query` and `matched`, each a tuple of
   (1-based query position, residue) per edit; a repeat makes placements equivalent.
   Query-terminal deletions and boundary insertions are barred, and an inserted
-  residue's position counts the query residues preceding it."""
+  residue's position counts the query residues preceding it.
+
+  Removing n positions from the longer string splits it into n + 1 stretches, and
+  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."""
   L, M = len(query), len(matched)
   n = abs(M - L)
   if n == 0:
     return []
-  out = []
+  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,))
+
+  place(0, 0, ())
   if M < L:
-    for combo in combinations(range(1, L - 1), n):
-      if ''.join(query[i] for i in range(L) if i not in combo) == matched:
-        out.append(tuple((i + 1, query[i]) for i in combo))
-  else:
-    for combo in combinations(range(M), n):
-      if ''.join(matched[i] for i in range(M) if i not in combo) != query:
-        continue
-      placement = []
-      for k in combo:
-        p = sum(1 for x in range(k) if x not in combo) + 1
-        if p == 1 or p == L + 1:
-          placement = None
-          break
-        placement.append((p, matched[k]))
-      if placement:
-        out.append(tuple(placement))
-  return out
+    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]
 
 
 def _indel_edits(query, matched):
@@ -551,7 +564,8 @@ def _to_dataframe(self, cols, pepidx_path, is_indels=False):
     from `pepidx_path` once instead of cloning it into every hit row. 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,
-    and miss rows stay null."""
+    once per distinct pair since a matched peptide often recurs across proteins, and
+    miss rows stay null."""
     qid, qseq, matched, pnum, mm, mutated, istart, iend = cols
 
     edit_col = 'Indels' if is_indels else 'Mismatches'
@@ -566,8 +580,12 @@ def _to_dataframe(self, cols, pepidx_path, is_indels=False):
       return pl.DataFrame(schema=schema)
 
     if is_indels:
-      positions = [format_indel_positions(q, m) if m is not None else None
-                   for q, m in zip(qseq, matched)]
+      annotations = {}
+      positions = []
+      for q, m in zip(qseq, matched):
+        if m is not None and (q, m) not in annotations:
+          annotations[(q, m)] = format_indel_positions(q, m)
+        positions.append(annotations[(q, m)] if m is not None else None)
     else:
       positions = mutated
 

pepmatch/tests/test_indel_positions_equivalence.py

@@ -0,0 +1,112 @@
+import itertools
+import random
+
+import pytest
+
+from pepmatch import matcher
+from pepmatch.matcher import _indel_placements, format_indel_positions
+
+AMINO_ACIDS = 'ACDEFGHIKLMNPQRSTVWY'
+
+
+def _reference_placements(query, matched):
+  """The original implementation, kept as the reference: try every combination of edit
+  positions, rebuild the string, compare. Slow (L^n per call) but obviously correct, so
+  the fast version must agree with it exactly -- placements, and their order."""
+  L, M = len(query), len(matched)
+  n = abs(M - L)
+  if n == 0:
+    return []
+  out = []
+  if M < L:
+    for combo in itertools.combinations(range(1, L - 1), n):
+      if ''.join(query[i] for i in range(L) if i not in combo) == matched:
+        out.append(tuple((i + 1, query[i]) for i in combo))
+  else:
+    for combo in itertools.combinations(range(M), n):
+      if ''.join(matched[i] for i in range(M) if i not in combo) != query:
+        continue
+      placement = []
+      for k in combo:
+        p = sum(1 for x in range(k) if x not in combo) + 1
+        if p == 1 or p == L + 1:
+          placement = None
+          break
+        placement.append((p, matched[k]))
+      if placement:
+        out.append(tuple(placement))
+  return out
+
+
+def _edited_pairs(query, alphabet):
+  """Every (query, matched) pair one or two deletions or insertions away from `query`,
+  including terminal edits, which the annotation must reject identically."""
+  L = len(query)
+  for n in (1, 2):
+    for pos in itertools.combinations(range(L), n):
+      yield query, ''.join(c for i, c in enumerate(query) if i not in pos)
+    for pos in itertools.combinations(range(L + n), n):
+      for residues in itertools.product(alphabet, repeat=n):
+        it, extra = iter(query), iter(residues)
+        yield query, ''.join(next(extra) if i in pos else next(it) for i in range(L + n))
+
+
+def _assert_same(query, matched, monkeypatch):
+  assert _indel_placements(query, matched) == _reference_placements(query, matched), \
+    (query, matched)
+  fast = format_indel_positions(query, matched)
+  monkeypatch.setattr(matcher, '_indel_placements', _reference_placements)
+  slow = format_indel_positions(query, matched)
+  monkeypatch.undo()
+  assert fast == slow, (query, matched)
+
+
+@pytest.mark.parametrize('alphabet,max_len', [('AB', 8), ('ABC', 6)])
+def test_exhaustive_small_alphabet(alphabet, max_len, monkeypatch):
+  """A two- or three-letter alphabet makes every query a homopolymer or periodic repeat
+  somewhere, which is where placements are ambiguous and ranges/chunks get reported."""
+  pairs = set()
+  for L in range(3, max_len + 1):
+    for query in map(''.join, itertools.product(alphabet, repeat=L)):
+      pairs.update(_edited_pairs(query, alphabet))
+  for query, matched in pairs:
+    _assert_same(query, matched, monkeypatch)
+
+
+@pytest.mark.parametrize('seed', range(5))
+def test_random_peptides(seed, monkeypatch):
+  """Realistic lengths (6-25) with a share of low-complexity queries."""
+  rng = random.Random(seed)
+  for _ in range(400):
+    L = rng.randint(6, 25)
+    pool = 'AAKKL' if rng.random() < 0.3 else AMINO_ACIDS
+    query = ''.join(rng.choice(pool) for _ in range(L))
+    n = rng.choice((1, 2))
+    if rng.random() < 0.5:
+      pos = rng.sample(range(L), n)
+      matched = ''.join(c for i, c in enumerate(query) if i not in pos)
+    else:
+      chars = list(query)
+      for _ in range(n):
+        chars.insert(rng.randint(0, len(chars)), rng.choice(AMINO_ACIDS))
+      matched = ''.join(chars)
+    _assert_same(query, matched, monkeypatch)
+
+
+def test_unrelated_and_short_pairs(monkeypatch):
+  """Every pair over {A, B} one or two residues apart in length, related by indels or not,
+  down to the empty string: both versions must agree on which have no valid placement."""
+  for q_len in range(0, 7):
+    for m_len in {q_len - 2, q_len - 1, q_len + 1, q_len + 2} - {-2, -1}:
+      for query in map(''.join, itertools.product('AB', repeat=q_len)):
+        for matched in map(''.join, itertools.product('AB', repeat=m_len)):
+          _assert_same(query, matched, monkeypatch)
+
+
+@pytest.mark.parametrize('query,matched,expected', [
+  ('ABCDEF', 'ABCDEF', '[]'),            # exact match
+  ('CAAAD', 'CAAD', 'd: A[2,4]'),        # homopolymer: the range of valid positions
+  ('ABABAB', 'ABAB', 'd: BA[2], d: AB[3], d: BA[4]'),   # periodic: residues change along the range
+])
+def test_documented_examples(query, matched, expected):
+  assert format_indel_positions(query, matched) == expected