diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 7830dd7..2825be1 100644 --- a/.github/workflows/tests.yml +++ b/.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 diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index f059d81..0847ba5 100755 --- a/pepmatch/matcher.py +++ b/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 diff --git a/pepmatch/tests/test_indel_positions_equivalence.py b/pepmatch/tests/test_indel_positions_equivalence.py new file mode 100644 index 0000000..8ac968d --- /dev/null +++ b/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