From 73358da9a992e548c07f91035368e2950024a76f Mon Sep 17 00:00:00 2001 From: Roman Young Date: Wed, 4 Nov 2026 21:35:57 +0000 Subject: [PATCH] fix: cap k at 16, the longest k-mer the index can store, instead of panicking --- .github/workflows/tests.yml | 2 + README.md | 4 +- pepmatch/matcher.py | 13 ++-- pepmatch/preprocessor.py | 6 ++ pepmatch/tests/test_k_limit.py | 120 +++++++++++++++++++++++++++++++++ 5 files changed, 139 insertions(+), 6 deletions(-) create mode 100644 pepmatch/tests/test_k_limit.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 4825ed4..03e02f5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -40,6 +40,8 @@ jobs: test-file: pepmatch/tests/test_metadata_rows.py - name: Indel Golden Output test-file: pepmatch/tests/test_indel_golden.py + - name: K Limit + test-file: pepmatch/tests/test_k_limit.py steps: - name: Checkout Code diff --git a/README.md b/README.md index ec87b5b..cf52f13 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,7 @@ pepmatch-preprocess -p human.fasta -k 5 #### Flags * `-p`, `--proteome` (Required): Path to the proteome FASTA file. -* `-k`, `--kmer_size` (Required): The k-mer size for indexing. +* `-k`, `--kmer_size` (Required): The k-mer size for indexing, from 2 to 16. * `-n`, `--proteome_name`: Custom name for the proteome. * `-P`, `--preprocessed_files_path`: Directory to save preprocessed files. @@ -178,7 +178,7 @@ pepmatch-match -q peptides.fasta -p human.fasta -m 0 -k 5 * `-p`, `--proteome_file` (Required): Path to the proteome FASTA file. * `-m`, `--max_mismatches`: Maximum mismatches allowed (default: 0). * `-i`, `--max_indels`: Maximum indels allowed (default: 0). Supports 1 or 2, and mutually exclusive with `-m`. At 2, a match uses two insertions or two deletions, never one of each. -* `-k`, `--kmer_size`: K-mer size (default: 5). +* `-k`, `--kmer_size`: K-mer size, from 2 to 16 (default: chosen from the query lengths, the largest that guarantees complete recall, up to 16). * `-P`, `--preprocessed_files_path`: Directory containing preprocessed files. * `-b`, `--best_match`: Enable best match mode. * `-f`, `--output_format`: Output format — `csv`, `tsv`, `xlsx`, `json` (default: `csv`). diff --git a/pepmatch/matcher.py b/pepmatch/matcher.py index 3bb1721..1a96b52 100755 --- a/pepmatch/matcher.py +++ b/pepmatch/matcher.py @@ -2,6 +2,7 @@ import polars as pl from pathlib import Path from Bio import SeqIO +from .preprocessor import MAX_INDEX_K from ._rs import rs_preprocess, rs_match, rs_discontinuous, rs_metadata, rs_match_counts, rs_indel_match VALID_OUTPUT_FORMATS = ['dataframe', 'polars', 'pandas', 'csv', 'tsv', 'xlsx', 'json'] @@ -178,6 +179,9 @@ def __init__( if k != 0 and k < 2: raise ValueError('k must be >= 2.') + if k > MAX_INDEX_K: + raise ValueError(f'k must be <= {MAX_INDEX_K}, the longest k-mer the index can store.') + if max_indels > 2: raise ValueError('max_indels > 2 is not yet supported. Only max_indels<=2 has been validated.') @@ -334,7 +338,7 @@ def indel_search(self): if short: groups.setdefault(self._clamp_k(2), []).extend(short) if rest: - rest_optimal = max(2, min(len(s) for _, s in rest) // (n + 1)) + rest_optimal = min(MAX_INDEX_K, max(2, min(len(s) for _, s in rest) // (n + 1))) groups.setdefault(self._clamp_k(rest_optimal), []).extend(rest) frames = [] @@ -376,9 +380,10 @@ def _recall_warning(self, k, edits, queries=None): def _auto_k(self, edits): """Largest k that still guarantees recall (pigeonhole): floor(L / (edits + 1)) taken on the shortest query peptide, so every peptide splits into edits+1 disjoint seeds. - Floored at 2, the preprocessor minimum.""" + Floored at 2, the preprocessor minimum, and capped at MAX_INDEX_K, the longest k-mer the + index can store: a smaller k only adds seeds, so recall still holds.""" min_len = min(len(seq) for _, seq in self.query) - return max(2, min_len // (edits + 1)) + return min(MAX_INDEX_K, max(2, min_len // (edits + 1))) def _pepidx_path(self, k): return os.path.join( @@ -439,7 +444,7 @@ def collect_matched(cols): matched_ids.add(cols[0][i]) return matched_ids - initial_k = min(len(seq) for _, seq in peptides_remaining) + initial_k = min(MAX_INDEX_K, min(len(seq) for _, seq in peptides_remaining)) ks = [initial_k] while initial_k > 2: initial_k //= 2 diff --git a/pepmatch/preprocessor.py b/pepmatch/preprocessor.py index 9261d77..a9de14b 100755 --- a/pepmatch/preprocessor.py +++ b/pepmatch/preprocessor.py @@ -2,6 +2,10 @@ from pathlib import Path from ._rs import rs_preprocess +# preprocess.rs packs each k-mer into a u128, one byte per residue, so no index can be +# built for k above 16. +MAX_INDEX_K = 16 + class Preprocessor: """Preprocesses a proteome FASTA file into a .pepidx binary index optimized for fast peptide searching. The index stores k-mer positions, @@ -29,6 +33,8 @@ def preprocess(self, k: int) -> None: """Build the k-mer index and write it to _mers.pepidx.""" if k < 2: raise ValueError('k must be >= 2.') + if k > MAX_INDEX_K: + raise ValueError(f'k must be <= {MAX_INDEX_K}, the longest k-mer the index can store.') output_path = os.path.join( self.preprocessed_files_path, f'{self.proteome_name}_{k}mers.pepidx' ) diff --git a/pepmatch/tests/test_k_limit.py b/pepmatch/tests/test_k_limit.py new file mode 100644 index 0000000..f1c4d31 --- /dev/null +++ b/pepmatch/tests/test_k_limit.py @@ -0,0 +1,120 @@ +import random +import pytest + +from pepmatch import Matcher, Preprocessor +from indel2_brute_force import brute_force_search + +# The index packs each k-mer into 16 bytes. Auto-k is the shortest query's length divided +# by (edits + 1), so before the cap any batch of long enough queries asked for a 17-mer index +# and the preprocessor panicked; best_match started from the shortest query's full length. +PROTEIN = 'GGGGMKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQGGGG' +AMINO_ACIDS = 'ACDEFGHIKLMNPQRSTVWY' + + +def window(n): + return PROTEIN[4:4 + n] + + +@pytest.fixture +def proteome(tmp_path): + path = tmp_path / 'p.fasta' + path.write_text(f'>sp|P|P P OS=Test OX=1 GN=P PE=1 SV=1\n{PROTEIN}\n') + return path + + +def search(tmp_path, proteome, **kwargs): + return Matcher(proteome_file=str(proteome), preprocessed_files_path=str(tmp_path), + output_format='dataframe', **kwargs).match() + + +@pytest.mark.parametrize('length, kwargs', [ + (17, {}), # exact: auto-k was the full 17 + (30, {}), + (34, {'max_mismatches': 1}), # 34 // 2 = 17 + (17, {'best_match': True}), # best_match started at the shortest query's length + (34, {'max_indels': 1}), # 34 // 2 = 17 + (51, {'max_indels': 2}), # 51 // 3 = 17 +]) +def test_long_query_is_found_with_k_capped_at_16(tmp_path, proteome, capsys, length, kwargs): + df = search(tmp_path, proteome, query=[window(length)], **kwargs) + assert window(length) in df['Matched Sequence'].to_list() + assert 'k=16' in capsys.readouterr().out + + +def test_counts_only_long_query(tmp_path, proteome): + df = search(tmp_path, proteome, query=[window(17)], counts_only=True) + assert df.filter(df['Mismatches'] == 0)['Count'].to_list() == [1] + + +def test_k_below_the_cap_is_unchanged(tmp_path, proteome, capsys): + search(tmp_path, proteome, query=[window(16)]) + assert '(k=16,' in capsys.readouterr().out + + +def test_explicit_k_above_16_is_rejected(tmp_path, proteome): + with pytest.raises(ValueError, match='k must be <= 16'): + search(tmp_path, proteome, query=[window(20)], k=17) + + +def test_preprocessor_rejects_k_above_16(tmp_path, proteome): + with pytest.raises(ValueError, match='k must be <= 16'): + Preprocessor(str(proteome), preprocessed_files_path=str(tmp_path)).preprocess(17) + + +def _substituted(rng, seq, n): + for _ in range(n): + i = rng.randrange(len(seq)) + seq = seq[:i] + rng.choice(AMINO_ACIDS) + seq[i + 1:] + return seq + + +@pytest.mark.parametrize('mismatches', [0, 1, 2]) +@pytest.mark.parametrize('seed', range(5)) +def test_capped_k_keeps_complete_recall_for_mismatches(tmp_path, seed, mismatches): + # Capping k only adds seeds, so long queries must still match a brute-force Hamming scan. + rng = random.Random(seed * 3 + mismatches) + proteins = {f'P{i}': ''.join(rng.choice(AMINO_ACIDS) for _ in range(rng.randint(60, 90))) for i in range(3)} + queries = [] + for _ in range(6): + p = proteins[rng.choice(list(proteins))] + length = rng.randint(17 * (mismatches + 1), min(len(p), 17 * (mismatches + 1) + 8)) + start = rng.randrange(len(p) - length + 1) + queries.append(_substituted(rng, p[start:start + length], rng.randint(0, mismatches))) + fasta = tmp_path / 'proteins.fasta' + fasta.write_text(''.join(f'>{pid}\n{seq}\n' for pid, seq in proteins.items())) + + df = search(tmp_path, fasta, query=queries, max_mismatches=mismatches) + found = {(r['Query Sequence'], r['Protein ID'], r['Matched Sequence']) + for r in df.iter_rows(named=True) if r['Matched Sequence'] is not None} + expected = { + (q, f'{pid}.1', p[s:s + len(q)]) + for q in queries for pid, p in proteins.items() for s in range(len(p) - len(q) + 1) + if sum(a != b for a, b in zip(q, p[s:s + len(q)])) <= mismatches + } + assert found == expected + + +@pytest.mark.parametrize('indels', [1, 2]) +@pytest.mark.parametrize('seed', range(3)) +def test_capped_k_keeps_complete_recall_for_indels(tmp_path, seed, indels): + rng = random.Random(100 + seed * 2 + indels) + proteins = {f'P{i}': ''.join(rng.choice(AMINO_ACIDS) for _ in range(rng.randint(70, 90))) for i in range(2)} + queries = [] + for _ in range(4): + p = proteins[rng.choice(list(proteins))] + length = rng.randint(17 * (indels + 1), 17 * (indels + 1) + 6) + start = rng.randrange(len(p) - length + 1) + q = p[start:start + length] + for _ in range(rng.randint(0, indels)): # interior deletions from the query + i = rng.randrange(1, len(q) - 1) + q = q[:i] + q[i + 1:] + queries.append(q) + fasta = tmp_path / 'proteins.fasta' + fasta.write_text(''.join(f'>{pid}\n{seq}\n' for pid, seq in proteins.items())) + + df = search(tmp_path, fasta, query=queries, max_indels=indels) + found = {(r['Query Sequence'], r['Protein ID'], r['Matched Sequence']) + for r in df.iter_rows(named=True) if r['Matched Sequence'] is not None} + expected = {(q, f'{pid}.1', w) for q in queries for pid, p in proteins.items() + for _, w in brute_force_search(q, p, indels)} + assert found == expected