diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f52a95..5ffcbf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,75 @@ All notable changes to `seqtree`. Dates are release dates. From 1.0.0 the project follows semantic versioning: breaking changes need a **major** bump. +## [Unreleased] + +### Added + +- **1.0.2 candidate: bounded exhaustive gapblock reductions.** `topk_batch` keeps + stable `(score, ref_id)` top-k rows without a dense score matrix; `count_batch` + counts all ties at arbitrary per-query thresholds. Linked two-view variants use + the maximum lane penalty and preserve both lane scores. Exact identity exclusion + precedes selection and removes only both-view identity for linked inputs. + `Index.ref_seqs()` retrieves stored reference strings in ID order in one native + call. All scoring reuses the existing single-gap-block cell; no calibration, + domain payload or runtime dependency is introduced. + +- **`TextIndex` searches with insertions and deletions.** `search_batch` takes `max_indels` + alongside `max_subs`, capped independently — `max_subs=2, max_indels=1` accepts two + substitutions *and* one gap, not three edits of any kind. + + ```python + res = ix.search_batch(peptides, max_subs=1, max_indels=1, threads=0) + h = res[0][0] + ix.ref_seq(h.ref_id)[h.offset : h.offset + h.length] # the matched text + h.n_subs, h.n_ins, h.n_dels + ``` + + **The index, the build and the probe enumeration are unchanged.** With + `b = max_subs + max_indels + 1` disjoint blocks the total edit count cannot reach the block + count, so the same pigeonhole argument still leaves a **zero-error** block — and a zero-error + block matches the text exactly whether or not the occurrence carries gaps, so the existing + exact seed lookup finds it. Verified before any of this was written: 2,800 planted + occurrences with random insertions and deletions, **recall 1.000** from the unmodified exact + lookup. + + What changes is verification. A seed no longer pins the start — the prefix before the block + may have gained or lost up to `max_indels` residues — so every start in + `[pos - base - max_indels, pos - base + max_indels]` is aligned by a banded DP over the states + `(query position, net offset, indels spent)`. An alignment must begin and end on an aligned + pair: a gap at either edge would stretch the reported interval over a residue that matches + nothing, and a leading gap is the same occurrence starting one residue over, so allowing it + would report one match several times. + + `TextHit` gains `n_ins`, `n_dels` and **`length`** — a gapped match is no longer `len(query)` + residues wide, and `length` is `len(query) + n_dels - n_ins`. `TextResult.arrays()` gains the + three matching columns. + + Two limits, both refusals rather than partial answers. Every seed block must be exact, so a + query must be at least `(max_subs + max_indels + 1) * k` long — at `k = 4` with + `max_subs=1, max_indels=1` that is 12 residues — and a shorter one raises with the bound in + the message. And `mismatches` is empty on this path, with `matrix=` ignored: recovering which + columns were substituted needs an alignment traceback the verifier does not keep. + + Measured on a synthetic 5,000,000-residue corpus at `k = 4`, one thread, against the same + queries with `max_indels = 0`: + + | query length | `max_subs` | ms/query, no indels | ms/query, `max_indels=1` | + |--:|--:|--:|--:| + | 16 | 1 | 0.001 | 0.028 | + | 16 | 2 | 0.001 | 0.043 | + | 20 | 3 | 0.002 | 0.065 | + | 24 | 3 | 0.002 | 0.067 | + + 25–37× an ungapped search at the same length and `max_subs`, and still under 0.1 ms/query. + An early exit when every banded state has already breached a cap is worth 2.5–3× of that; the + DP was otherwise walking all L rows on candidates the Hamming path rejects after two residues. + + Completeness is held to the same standard as the substitution path: **set equality against an + independent brute force**, over `k` ∈ {3, 4, 5} × `max_subs` 0–3 × `max_indels` 1–2, + **6,732 occurrences, zero missing and zero extra**. `max_indels=0` is byte-identical to the + previous substitution path, hit for hit, and pays none of the cost. + ## [1.0.1] — 2026-10-09 ### Fixed diff --git a/README.md b/README.md index 62fde73..73b827f 100644 --- a/README.md +++ b/README.md @@ -103,7 +103,10 @@ Beyond search, seqtree ships: completeness pinned by brute-force **set equality** over L 6–30 × `max_subs` 0–3 × k ∈ {3,4,5}, and by the answer being identical across `k`. On the human proteome (69,578,135 residues) one thread answers a 9-mer within 2 substitutions in **1.3 ms** and a - 15-mer within 3 in **1.5 ms**. Results come back as flat CSR arrays with zero-copy + 15-mer within 3 in **1.5 ms**. `max_indels` adds insertions and deletions on the same index + and the same scheme — with `max_subs + max_indels + 1` blocks the edit count is still below + the block count, so pigeonhole leaves a zero-error block and the *unchanged* exact seed lookup + finds gapped occurrences too; only verification becomes a banded DP. Results come back as flat CSR arrays with zero-copy `numpy` views, mismatches as `(pos, query_aa, text_aa)` pairs, an optional fold onto caller-supplied group ids that makes a tie explicit, and a cap that is always reported. - **E-values / significance** — calibrate hit counts against a background control repertoire diff --git a/docs/gapblock.rst b/docs/gapblock.rst index 870fb6a..1cb648c 100644 --- a/docs/gapblock.rst +++ b/docs/gapblock.rst @@ -314,3 +314,44 @@ See also * :doc:`api` — full signatures for :mod:`seqtree.gapblock`, :mod:`seqtree.seeds` and :mod:`seqtree.layout`. * ``bench/bench_gapblock.py`` — reproduces the gap-freedom ladder table above. + +Bounded exhaustive ranking and counts +------------------------------------- + +``topk_batch`` evaluates every reference with the same single-gap-block scorer as +``score_matrix`` and returns at most k ``Hit`` rows per query, ordered by +``(score, ref_id)``. Reference IDs retain input row order; repeated sequences remain +separate rows. Hit edit fields are unset, rather than an alignment decomposition. +``exclude_exact=True`` excludes identical encoded sequences before selection; +nonidentical sequences with a zero substitution penalty remain eligible. + +``count_batch`` takes one list of integer thresholds per query and returns the +number of reference scores at or below each threshold, including all ties. +Threshold rows may be empty, unsorted or contain duplicates; output follows the +original threshold order. Negative thresholds count zero. + +The paired variants score linked two-view query/reference rows and compare the +maximum of their two lane penalties. ``paired_topk_batch`` returns +``(ref_id, score_a, score_b)`` tuples; ``paired_count_batch`` counts all linked rows +at each requested maximum-score threshold. Paired exact exclusion removes only +rows identical in both views. No cross-row combinations are created. + +.. code-block:: python + + from seqtree.gapblock import topk_batch, count_batch + + queries, refs = ["ACDE", "ACD"], ["ACDE", "ACD", "ACDF"] + hits = topk_batch(queries, refs, k=2, threads=1) + thresholds = [[h.score for h in row] for row in hits] + counts = count_batch(queries, refs, thresholds, threads=1) + +Each operation runs as one native batch with the GIL released. Top-k results use +O(Q*k) storage and each worker uses O(k + maximum sequence length) scratch. +Count output uses O(total requested thresholds), with worker scratch proportional +to maximum sequence length and the largest threshold row. Encoded inputs and the +existing optional prior cube are shared. Neither path allocates a dense Q*N +score matrix; both remain exhaustive scoring operations. They introduce no +statistical model or calibration. + +``Index.ref_seqs()`` retrieves all stored strings in reference-ID order through +one native call, preserving duplicate rows. diff --git a/docs/text-index.rst b/docs/text-index.rst index 8a303a3..99d730b 100644 --- a/docs/text-index.rst +++ b/docs/text-index.rst @@ -211,6 +211,81 @@ A ``.sti`` is treated as untrusted input: the header's record, bucket and postin cross-checked against each other and against the file length before any of them is used to size an allocation or to index into the mapping. +Insertions and deletions +------------------------ + +``max_indels`` allows gaps. The two caps are independent, so ``max_subs=2, max_indels=1`` +accepts two substitutions **and** one gap — not three edits of any kind: + +.. code-block:: python + + res = ix.search_batch(peptides, max_subs=1, max_indels=1, threads=0) + for h in res[0]: + print(h.ref_id, h.offset, h.length, h.n_subs, h.n_ins, h.n_dels) + print(ix.ref_seq(h.ref_id)[h.offset : h.offset + h.length]) + +**The seed table does not change, and neither does the scheme.** With +``b = max_subs + max_indels + 1`` disjoint blocks the total edit count cannot reach the block +count, so the same pigeonhole argument still puts a **zero-error** block somewhere — and a +zero-error block matches the text exactly, gaps or not, so the existing exact lookup finds it. +Nothing about the index, the build, or the probe enumeration is different. + +What changes is verification. A seed no longer pins the start: the prefix before the block may +have gained or lost up to ``max_indels`` residues, so every start in +``[pos - base - max_indels, pos - base + max_indels]`` is aligned, by a banded DP over the +states ``(query position, net offset, indels spent)``. + +Two consequences worth knowing before you call it: + +* **A match is no longer** ``len(query)`` **residues wide.** Read + :attr:`~seqtree.TextHit.length`; it is ``len(query) + n_dels - n_ins``. +* **Every seed block must be exact**, so a query must be at least + ``(max_subs + max_indels + 1) * k`` long. At ``k = 4``, ``max_subs=1, max_indels=1`` needs 12 + residues. A shorter query raises, naming the bound — it is not answered partially. Build a + second index at a smaller ``k`` if your queries are shorter than that. + +An alignment must **begin and end on an aligned pair**. A gap at either edge would stretch the +reported interval over a residue that matches nothing, and a leading gap is just the same +occurrence starting one residue over — allowing it would report one match two or three times. + +``mismatches`` is empty on this path: recovering which columns were substituted needs an +alignment traceback the verifier does not keep. The counts are exact; the per-column detail is +only available at ``max_indels = 0``. ``matrix=`` is likewise ignored. + +Cost, measured on a synthetic 5,000,000-residue corpus, ``k = 4``, one thread, against the same +query set with ``max_indels = 0``: + +.. list-table:: + :header-rows: 1 + + * - query length + - ``max_subs`` + - ms/query, no indels + - ms/query, ``max_indels=1`` + * - 16 + - 1 + - 0.001 + - 0.028 + * - 16 + - 2 + - 0.001 + - 0.043 + * - 20 + - 3 + - 0.002 + - 0.065 + * - 24 + - 3 + - 0.002 + - 0.067 + +Gapped search costs 25–37× an ungapped one at the same length and ``max_subs``: the banded DP +replaces a byte comparison that early-exits after a residue or two, and it runs once per start +in the window. It is still well under 0.1 ms/query. Completeness is pinned the same way the +substitution path is — **set equality against an independent brute force**, over +``k`` ∈ {3, 4, 5} × ``max_subs`` 0–3 × ``max_indels`` 1–2, 6,732 occurrences, zero missing and +zero extra. + Limits ------ diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index 2a73be0..809f887 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include namespace seqtree { @@ -221,6 +222,29 @@ std::vector gapblock_matrix(const std::vector& queries, int32_t gap_extend, const std::vector& prior, uint32_t prior_width, int threads = 0); +// Exhaustive bounded reductions using the same gapblock cell. Top-k ties use +// (penalty, ref_id); paired penalty is max(lane scores). Exact exclusion uses +// encoded sequence identity (both lanes for pairs), never zero penalty alone. +std::vector> gapblock_topk_batch(const std::vector& queries, + const std::vector& refs, uint32_t k, Alphabet alphabet, + const SubstitutionMatrix* matrix, int32_t go, int32_t ge, const std::vector& prior, + uint32_t width, int threads, bool exclude_exact); +std::vector>> gapblock_paired_topk_batch( + const std::vector& qa, const std::vector& qb, + const std::vector& ra, const std::vector& rb, + uint32_t k, Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact); +std::vector> gapblock_count_batch(const std::vector& queries, + const std::vector& refs, const std::vector>& thresholds, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact); +std::vector> gapblock_paired_count_batch( + const std::vector& qa, const std::vector& qb, + const std::vector& ra, const std::vector& rb, + const std::vector>& thresholds, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact); + // --------------------------------------------------------------------------------------- // Plain string edit distances on raw characters, unit costs. Deliberately alphabet-agnostic: // no Codec, no substitution matrix, no gap model -- the distance you want when you just need a diff --git a/include/seqtree/text_index.hpp b/include/seqtree/text_index.hpp index 206735c..9072b84 100644 --- a/include/seqtree/text_index.hpp +++ b/include/seqtree/text_index.hpp @@ -13,12 +13,24 @@ struct TextStore; // internal, src/text_index.cpp // One row of the optional group fold: hits collapsed onto the build-time group ids. struct TextGroup { uint32_t group_id = 0; - uint16_t min_subs = 0; // best (smallest) Hamming distance reaching this group + uint16_t min_subs = 0; // best (smallest) substitution count reaching this group uint32_t n_hits = 0; // how many hits fell in it, at any distance }; struct TextQueryOpts { uint16_t max_subs = 0; + // Insertions + deletions allowed, capped independently of max_subs. 0 keeps the pure + // Hamming predicate and its verification path untouched. + // + // The seed table does not change: with b = max_subs + max_indels + 1 disjoint blocks the + // total edit count is below the block count, so by pigeonhole SOME block still matches + // EXACTLY -- the existing exact lookup finds every occurrence, indels included. What that + // costs is verification, which becomes a banded alignment over the 2*max_indels+1 possible + // starts rather than one Hamming scan at a fixed start. + // + // Because every block must be exact, this needs L >= (max_subs + max_indels + 1) * k; + // search_batch refuses a shorter query rather than answering incompletely. + uint16_t max_indels = 0; bool exclude_exact = false; // drop 0-mismatch hits (a candidate that IS the text) bool best_only = false; // keep only the minimum-n_subs shell (still ALL of it) bool group_by = false; // also fold hits onto the build-time group ids @@ -40,6 +52,11 @@ struct TextResult { std::vector query_begin; // n_queries + 1, CSR into the hit arrays std::vector ref_id, offset; std::vector n_subs; + // Indels in the alignment, and the matched substring's length in the text. With + // max_indels = 0 these are 0, 0 and the query length for every hit. With indels the match + // is no longer len(query) residues wide, so `length` is what bounds it: + // ref_seq(ref_id)[offset : offset + length]. + std::vector n_ins, n_dels, length; std::vector score; std::vector mm_begin; // n_hits + 1, CSR into the mismatch arrays std::vector mm_pos; diff --git a/pyproject.toml b/pyproject.toml index f928f0a..feb221b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "seqtree" -version = "1.0.1" +version = "1.0.2" description = "Fast fuzzy search over biological sequences (C++ core, Python bindings)" readme = "README.md" requires-python = ">=3.10" diff --git a/python/seqtree/gapblock.py b/python/seqtree/gapblock.py index 92d7acc..913ce64 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -42,9 +42,13 @@ from ._core import Index, ScoreMatrix, SearchParams, SubstitutionMatrix from ._core import gapblock_matrix as _gapblock_matrix +from ._core import gapblock_topk_batch as _gapblock_topk_batch +from ._core import gapblock_count_batch as _gapblock_count_batch +from ._core import gapblock_paired_topk_batch as _gapblock_paired_topk_batch +from ._core import gapblock_paired_count_batch as _gapblock_paired_count_batch __all__ = [ - "gapblock_score", "score_matrix", "deletion_variants", "central_prior", "profile_prior", + "gapblock_score", "score_matrix", "topk_batch", "count_batch", "paired_topk_batch", "paired_count_batch", "deletion_variants", "central_prior", "profile_prior", "frame_prior", "positions_prior", "embed_in_frame", "gap_cost", "GapBlockIndex", "ScoreMatrix", "IslandProfile", ] @@ -564,6 +568,111 @@ def score_matrix( return _gapblock_matrix(q, r, alphabet, matrix, gap_open, gap_extend, cube, width, threads) +def _batch_integer(value, name, lower=0, upper=(1 << 31) - 1): + import operator + if isinstance(value, bool): + raise ValueError(f"{name} must be an integer") + try: + result = operator.index(value) + except TypeError: + raise ValueError(f"{name} must be an integer") from None + if not lower <= result <= upper: + raise ValueError(f"{name} must lie in [{lower}, {upper}]") + return result + + +def _batch_options(axes, matrix, gap_open, gap_extend, gap_prior, threads): + go = _batch_integer(_default_gap_open(matrix) if gap_open is None else gap_open, "gap_open") + ge = _batch_integer(gap_extend, "gap_extend") + threads = _batch_integer(threads, "threads") + width = max((len(s) for axis in axes for s in axis), default=0) + cube = _prior_cube(gap_prior, width) if gap_prior is not None else [] + return go, ge, cube, width, threads + + +def _threshold_rows(thresholds, n_queries): + rows = [list(row) for row in thresholds] + if len(rows) != n_queries: + raise ValueError("thresholds must have one row per query") + return [[_batch_integer(value, "threshold", -(1 << 31)) for value in row] for row in rows] + + +def topk_batch(queries, refs, k=10, matrix=None, gap_open=None, gap_extend=1, + gap_prior=None, alphabet="aa", threads=0, exclude_exact=False): + """Exhaustive single-gap-block top-k penalties, one native batch, GIL released. + + Return one list of ``Hit(ref_id, score)`` per query, ordered by + ``(score, ref_id)``. At most k reference ROWS survive; repeated reference + sequences remain distinct rows. Hit edit fields are unset (zero), not an + alignment decomposition. Empty references yield empty query rows. + + Arguments match :func:`score_matrix`, with positive integer k added. Exact + exclusion compares case-insensitive alphabet-encoded sequences BEFORE + selection; a nonidentity zero penalty remains eligible. + + Shared storage holds encoded inputs and the existing optional prior cube. + Each worker holds O(k + max_sequence_length) scratch; results take O(Q*k), + never a dense Q*N matrix. This remains exhaustive O(Q*N*sequence_length) + scoring; no radius predicate, approximation or statistical model is added. + """ + q, r = list(queries), list(refs) + k = _batch_integer(k, "k", 1, (1 << 32) - 1) + go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads) + return _gapblock_topk_batch(q, r, k, alphabet, matrix, go, ge, cube, width, threads, exclude_exact) + + +def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend=1, + gap_prior=None, alphabet="aa", threads=0, exclude_exact=False): + """Count all reference scores <= each query's requested integer thresholds. + + ``thresholds`` has one row per query; rows may be empty, unsorted or contain + duplicates. Return integer counts in their original order, including ALL + score ties. Negative cutoffs count zero. Same scorer/exclusion as + :func:`topk_batch`; one native batch with O(Q*S) output, no dense score matrix. + Worker scratch is O(max_sequence_length + maximum_threshold_row_length). + """ + q, r = list(queries), list(refs) + thresholds = _threshold_rows(thresholds, len(q)) + go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads) + return _gapblock_count_batch(q, r, thresholds, alphabet, matrix, go, ge, cube, width, threads, exclude_exact) + + +def paired_topk_batch(queries_a, queries_b, refs_a, refs_b, k=10, matrix=None, + gap_open=None, gap_extend=1, gap_prior=None, alphabet="aa", + threads=0, exclude_exact=False): + """Top-k linked two-view rows, ordered by (max(score_a,score_b), ref_id). + + Return one list of ``(ref_id, score_a, score_b)`` tuples per linked query. + Both query axes and both reference axes must have matching lengths. + Exclusion removes only rows identical in BOTH views; one-view identities + remain eligible. Scoring arguments and memory bounds match :func:`topk_batch`. + No cross-row combinations, payloads or statistical assumptions are introduced. + """ + qa, qb, ra, rb = map(list, (queries_a, queries_b, refs_a, refs_b)) + if len(qa) != len(qb) or len(ra) != len(rb): + raise ValueError("paired query/reference axes must agree") + k = _batch_integer(k, "k", 1, (1 << 32) - 1) + go, ge, cube, width, threads = _batch_options((qa, qb, ra, rb), matrix, gap_open, gap_extend, gap_prior, threads) + return _gapblock_paired_topk_batch(qa, qb, ra, rb, k, alphabet, matrix, go, ge, cube, width, threads, exclude_exact) + + +def paired_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, matrix=None, + gap_open=None, gap_extend=1, gap_prior=None, alphabet="aa", + threads=0, exclude_exact=False): + """Count all linked two-view rows with max(lane penalties) <= each cutoff. + + Shape, threshold order, ties and negative-cutoff semantics match + :func:`count_batch`; exact exclusion is the both-view predicate described in + :func:`paired_topk_batch`. One native batch; no dense score matrix. + """ + qa, qb, ra, rb = map(list, (queries_a, queries_b, refs_a, refs_b)) + if len(qa) != len(qb) or len(ra) != len(rb): + raise ValueError("paired query/reference axes must agree") + thresholds = _threshold_rows(thresholds, len(qa)) + go, ge, cube, width, threads = _batch_options((qa, qb, ra, rb), matrix, gap_open, gap_extend, gap_prior, threads) + return _gapblock_paired_count_batch(qa, qb, ra, rb, thresholds, alphabet, matrix, go, ge, cube, width, threads, exclude_exact) + + class GapBlockIndex: """Search a reference set under the single-gap-block model, reusing the Hamming engine. diff --git a/skills/seqtree/SKILL.md b/skills/seqtree/SKILL.md index 87e5f3a..e455836 100644 --- a/skills/seqtree/SKILL.md +++ b/skills/seqtree/SKILL.md @@ -29,6 +29,7 @@ only, and nothing in the result says why. Passing a matrix to `seqtrie` without ```python Index.build(refs: list[str], alphabet: str = "aa") -> Index # 'aa' | 'nt' | 'iupac' +Index.ref_seqs() -> list[str] # one native call, ref_id order Index.search(query, params) -> list[Hit] # Hit(ref_id, score, n_subs, n_ins, n_dels) Index.search_top(query, params, k=1) -> list[Hit] Index.search_batch(queries, params, threads=0) -> list[list[Hit]] # releases the GIL @@ -144,11 +145,12 @@ Python and does not need the C++ core. `r = 2` over the same 300 is 9.9 M sequen ```python ix = TextIndex.build(records, alphabet="aa", k=4, group_ids=None) # WHOLE records, not windows -res = ix.search_batch(queries, max_subs=2, exclude_exact=False, best_only=False, +res = ix.search_batch(queries, max_subs=2, max_indels=0, exclude_exact=False, best_only=False, group_by=False, max_hits=0, matrix=None, threads=0) # GIL released len(res); res.num_hits; res[i] # TextHit list for query i, built on demand res[i][0].mismatches # [(pos, query_aa, text_aa), ...] -- the PAIR +res[i][0].n_ins / .n_dels / .length # indels, and the width of the match in the text res.groups(i) # [(group_id, min_subs, n_hits), ...] with group_by res.truncated # per query, 1 if max_hits capped it res.arrays(); res.to_numpy() # zero-copy views over the flat CSR @@ -177,6 +179,19 @@ full radius-`max_subs` ball, which is one to two orders of magnitude dearer. So at 10 (human proteome, `L = 12, max_subs = 3`: **1.46 ms/query at k=4, 0.32 at k=5**). Building is ~0.6 s either way, so a corpus spanning both lengths can hold both indexes. +**`max_indels` allows gaps, and costs nothing in the index.** The two caps are independent: +`max_subs=2, max_indels=1` is two substitutions AND one gap, not three edits. Seeding is +unchanged -- with `b = max_subs + max_indels + 1` blocks the edit count is below the block +count, so pigeonhole still leaves a **zero-error** block, which matches exactly whether or not +the occurrence has gaps. Only verification changes, to a banded DP over the +`2*max_indels + 1` possible starts. Three things to know: a match is **no longer `len(query)` +wide** (read `.length`, which is `len(query) + n_dels - n_ins`); every block must be exact, so a +query must be at least **`(max_subs + max_indels + 1) * k`** long or it raises; and +`.mismatches` is empty and `matrix=` ignored, because the per-column detail would need a +traceback the verifier does not keep. Costs 25-37x an ungapped search at the same length, still +under 0.1 ms/query on 5 M residues. Pinned by set equality against brute force over +k in {3,4,5} x max_subs 0-3 x max_indels 1-2, 6,732 occurrences, zero missing and zero extra. + - **`matrix=` only scores** hits the Hamming predicate already accepted (similarity summed over the mismatched positions) — it never changes which hits return. - **`best_only=True`** walks the distance upward and stops at the first non-empty shell, @@ -274,6 +289,27 @@ The result *is* the distance: the Gram transform is applied per residue when the so there is no `d = s(a,a) + s(b,b) - 2·s(a,b)` step, and non-negativity, symmetry and a zero diagonal hold by construction. Budget `4 * n * K` bytes and chunk the queries. +### Bounded exhaustive reductions + +`gapblock.topk_batch(queries, refs, k=10, matrix=None, gap_open=None, gap_extend=1, +gap_prior=None, alphabet="aa", threads=0, exclude_exact=False)` returns nested Hit +rows sorted by `(score, ref_id)`, scoring every reference with the same native +single-gap-block cell. Only score/ref_id are provided; edit fields remain unset. +Results use O(Q*k), worker scratch O(k+maxlen); encoded inputs/prior cube shared. +Exact encoded identity is excluded before selection; zero-cost nonidentity remains. + +`gapblock.count_batch(queries, refs, thresholds, **same_scoring_args)` returns +nested integer counts for scores<=each threshold, including ALL ties. Thresholds +have one row/query, arbitrary order/duplicates/empty rows allowed, output order +preserved, negative cutoffs count zero. No dense Q*N allocation. + +`gapblock.paired_topk_batch(queries_a, queries_b, refs_a, refs_b, k=10, +**same_scoring_args)` ranks linked rows by `(max(score_a,score_b),ref_id)` and +returns nested `(ref_id,score_a,score_b)` tuples. Matching lane axis lengths required. +`paired_count_batch(..., thresholds, **same_scoring_args)` counts all linked rows +at those maximum-score cutoffs. Exclusion removes only both-view identity. +These are generic exhaustive score reductions, not statistical calibration. + ### Performance Profiled over the bundled 250k control: 91% of query time is the query-deletion-variant branch, diff --git a/src/_bindings.cpp b/src/_bindings.cpp index 1a2ba15..2052c76 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -310,7 +310,7 @@ ArrayView view_of(const nb::object& owner, const std::vector& v, char fmt) { // batch -- the flat arrays are what a 445k-query run should read. struct PyTextHit { uint32_t ref_id = 0, offset = 0; - uint16_t n_subs = 0; + uint16_t n_subs = 0, n_ins = 0, n_dels = 0, length = 0; int32_t score = 0; std::vector> mismatches; }; @@ -322,6 +322,9 @@ nb::dict text_result_arrays(const nb::object& self) { d["ref_id"] = view_of(self, r.ref_id, 'I'); d["offset"] = view_of(self, r.offset, 'I'); d["n_subs"] = view_of(self, r.n_subs, 'H'); + d["n_ins"] = view_of(self, r.n_ins, 'H'); + d["n_dels"] = view_of(self, r.n_dels, 'H'); + d["length"] = view_of(self, r.length, 'H'); d["score"] = view_of(self, r.score, 'i'); d["mm_begin"] = view_of(self, r.mm_begin, 'I'); d["mm_pos"] = view_of(self, r.mm_pos, 'H'); @@ -644,6 +647,15 @@ NB_MODULE(_core, m) { .def("__len__", &Index::size) .def("ref_seq", [](const Index& i, uint32_t id) { return std::string(i.ref_seq(id)); }, nb::arg("ref_id"), "Return the reference sequence string for a reference id.") + .def("ref_seqs", [](const Index& i) { + std::vector out; + { + nb::gil_scoped_release release; + out.reserve(i.size()); + for (uint32_t id = 0; id < i.size(); ++id) out.emplace_back(i.ref_seq(id)); + } + return out; + }, "Return all reference strings in ref_id order in one native call, preserving duplicates.") .def("search", &py_search, nb::arg("query"), nb::arg("params"), "Return all hits for one query within the scope/budget in ``params``.") .def("search_top", &py_search_top, nb::arg("query"), nb::arg("params"), nb::arg("k") = 1, @@ -774,6 +786,30 @@ NB_MODULE(_core, m) { "`prior` is the gap prior flattened to [m][d][i]; see seqtree.gapblock.score_matrix, " "which builds it for you."); + m.def("gapblock_topk_batch", [](const std::vector& queries, const std::vector& refs, uint32_t k, const std::string& alphabet, const std::optional& matrix, int32_t go, int32_t ge, const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_topk_batch(queries, refs, k, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact); + }, nb::arg("queries"), nb::arg("refs"), nb::arg("k"), nb::arg("alphabet")="aa", nb::arg("matrix")=std::nullopt, nb::arg("gap_open")=1, nb::arg("gap_extend")=1, nb::arg("prior")=std::vector{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, "Bounded exhaustive gapblock reduction, GIL released."); + + m.def("gapblock_count_batch", [](const std::vector& queries, const std::vector& refs, const std::vector>& thresholds, const std::string& alphabet, const std::optional& matrix, int32_t go, int32_t ge, const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_count_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact); + }, nb::arg("queries"), nb::arg("refs"), nb::arg("thresholds"), nb::arg("alphabet")="aa", nb::arg("matrix")=std::nullopt, nb::arg("gap_open")=1, nb::arg("gap_extend")=1, nb::arg("prior")=std::vector{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, "Bounded exhaustive gapblock reduction, GIL released."); + + m.def("gapblock_paired_topk_batch", [](const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, uint32_t k, const std::string& alphabet, const std::optional& matrix, int32_t go, int32_t ge, const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_paired_topk_batch(qa, qb, ra, rb, k, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact); + }, nb::arg("qa"), nb::arg("qb"), nb::arg("ra"), nb::arg("rb"), nb::arg("k"), nb::arg("alphabet")="aa", nb::arg("matrix")=std::nullopt, nb::arg("gap_open")=1, nb::arg("gap_extend")=1, nb::arg("prior")=std::vector{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, "Bounded exhaustive gapblock reduction, GIL released."); + + m.def("gapblock_paired_count_batch", [](const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, const std::vector>& thresholds, const std::string& alphabet, const std::optional& matrix, int32_t go, int32_t ge, const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_paired_count_batch(qa, qb, ra, rb, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact); + }, nb::arg("qa"), nb::arg("qb"), nb::arg("ra"), nb::arg("rb"), nb::arg("thresholds"), nb::arg("alphabet")="aa", nb::arg("matrix")=std::nullopt, nb::arg("gap_open")=1, nb::arg("gap_extend")=1, nb::arg("prior")=std::vector{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, "Bounded exhaustive gapblock reduction, GIL released."); + m.def("hamming", &hamming, nb::arg("a"), nb::arg("b"), "Hamming distance: the number of positions at which two EQUAL-length sequences differ. " "Raises ValueError on a length mismatch. Case-sensitive, byte-for-byte."); @@ -917,8 +953,18 @@ NB_MODULE(_core, m) { "Start of the occurrence within that record, 0-based. The matched substring is " "``ix.ref_seq(ref_id)[offset:offset + len(query)]``.") .def_ro("n_subs", &PyTextHit::n_subs, - "Mismatched positions, always <= the ``max_subs`` that was searched. 0 is an " - "exact occurrence.") + "Substituted positions, always <= the ``max_subs`` that was searched. 0 with " + "``n_ins`` and ``n_dels`` also 0 is an exact occurrence.") + .def_ro("n_ins", &PyTextHit::n_ins, + "Insertions: query residues with nothing opposite them in the text. Always 0 " + "unless the search passed ``max_indels > 0``.") + .def_ro("n_dels", &PyTextHit::n_dels, + "Deletions: text residues with nothing opposite them in the query. Always 0 " + "unless the search passed ``max_indels > 0``.") + .def_ro("length", &PyTextHit::length, + "Residues the match spans in the text, i.e. " + "``ref_seq(ref_id)[offset : offset + length]``. Equal to the query length " + "whenever there are no indels; ``len(query) + n_dels - n_ins`` in general.") .def_ro("score", &PyTextHit::score, "Summed similarity over the mismatched positions under the ``matrix=`` passed " "to the search, or 0 when none was. Scoring never changes which hits return.") @@ -954,7 +1000,8 @@ NB_MODULE(_core, m) { if (i < 0 || i >= n) throw nb::index_error("query index out of range"); std::vector out; for (uint32_t h = r.query_begin[size_t(i)]; h < r.query_begin[size_t(i) + 1]; ++h) { - PyTextHit hit{r.ref_id[h], r.offset[h], r.n_subs[h], r.score[h], {}}; + PyTextHit hit{r.ref_id[h], r.offset[h], r.n_subs[h], r.n_ins[h], + r.n_dels[h], r.length[h], r.score[h], {}}; for (uint32_t j = r.mm_begin[h]; j < r.mm_begin[h + 1]; ++j) hit.mismatches.emplace_back(r.mm_pos[j], std::string(1, r.mm_query_aa[j]), std::string(1, r.mm_text_aa[j])); @@ -1044,10 +1091,11 @@ NB_MODULE(_core, m) { .def( "search_batch", [](const TextIndex& ix, const std::vector& queries, uint16_t max_subs, - bool exclude_exact, bool best_only, bool group_by, uint32_t max_hits, - const std::optional& matrix, int threads) { + uint16_t max_indels, bool exclude_exact, bool best_only, bool group_by, + uint32_t max_hits, const std::optional& matrix, int threads) { TextQueryOpts o; o.max_subs = max_subs; + o.max_indels = max_indels; o.exclude_exact = exclude_exact; o.best_only = best_only; o.group_by = group_by; @@ -1060,17 +1108,26 @@ NB_MODULE(_core, m) { } return res; }, - nb::arg("queries"), nb::arg("max_subs") = 0, nb::arg("exclude_exact") = false, - nb::arg("best_only") = false, nb::arg("group_by") = false, nb::arg("max_hits") = 0, + nb::arg("queries"), nb::arg("max_subs") = 0, nb::arg("max_indels") = 0, + nb::arg("exclude_exact") = false, nb::arg("best_only") = false, + nb::arg("group_by") = false, nb::arg("max_hits") = 0, nb::arg("matrix") = std::nullopt, nb::arg("threads") = 0, - "Find every position matching each query within ``max_subs`` substitutions " - "(releases the GIL; ``threads=0`` uses all cores). ``best_only`` walks the distance " - "upward and stops at the first shell with any hit, returning ALL of it. " - "``exclude_exact`` drops 0-mismatch hits. ``max_hits`` caps a query after sorting, " - "so the best hits survive, and sets ``truncated``. ``matrix`` only SCORES hits the " - "Hamming predicate already accepted -- it never changes which are returned. Every " - "query must be at least ``k`` long; a shorter one raises rather than being answered " - "incompletely.") + "Find every position matching each query within ``max_subs`` substitutions and " + "``max_indels`` insertions/deletions (releases the GIL; ``threads=0`` uses all " + "cores). The two caps are independent: ``max_subs=2, max_indels=1`` accepts two " + "substitutions AND one gap, not three edits of any kind.\n\n" + "``max_indels=0`` is the pure Hamming predicate: every hit is exactly " + "``len(query)`` residues wide and carries per-column ``mismatches``. Above 0 the " + "match length varies, so read :attr:`TextHit.length` rather than assuming the query " + "length, and ``mismatches`` is empty -- the counts are reported, the individual " + "substituted columns are not. Indel search also needs every seed block to be exact, " + "so a query must be at least ``(max_subs + max_indels + 1) * k`` long; a shorter one " + "raises rather than being answered incompletely.\n\n" + "``best_only`` walks the distance upward and stops at the first shell with any hit, " + "returning ALL of it. ``exclude_exact`` drops 0-mismatch hits. ``max_hits`` caps a " + "query after sorting, so the best hits survive, and sets ``truncated``. ``matrix`` " + "only SCORES hits the predicate already accepted -- it never changes which are " + "returned, and it is ignored when ``max_indels > 0``.") .def("save", &TextIndex::save, nb::arg("path"), "Write a flat, mmap-able index file.") .def_static("load", &TextIndex::load, nb::arg("path"), nb::arg("mmap") = true, diff --git a/src/gapblock.cpp b/src/gapblock.cpp index 2690871..51a003c 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -4,6 +4,7 @@ #include #include #include +#include namespace seqtree { namespace { @@ -60,6 +61,157 @@ inline int32_t cell(const uint8_t* q, uint32_t m, const uint8_t* r, uint32_t n, return best + gap_open + int32_t(d - 1) * gap_extend; } +// Shared preparation for bounded reductions: encoded inputs and prior remain +// shared read-only; no Q*N allocation. Validate before starting worker threads. +struct GapInput { + std::vector> q, r; + std::vector pen; + uint32_t longest = 0; + uint8_t A = 0; + size_t W1 = 0; +}; + +GapInput prepare(const std::vector& queries, const std::vector& refs, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads) { + if (go < 0 || ge < 0 || threads < 0) + throw std::invalid_argument("gap costs and threads must be >= 0"); + if (refs.size() > std::numeric_limits::max()) + throw std::invalid_argument("too many reference IDs"); + Codec codec(alphabet); + GapInput in; + in.A = codec.size(); + if (matrix && matrix->size() != in.A) + throw std::invalid_argument("matrix size does not match the alphabet"); + in.pen = pen_table(matrix, in.A); + auto encode_all = [&](const auto& strings, auto& out, const char* label) { + out.reserve(strings.size()); + for (size_t i = 0; i < strings.size(); ++i) { + if (strings[i].size() > std::numeric_limits::max()) + throw std::invalid_argument("sequence length exceeds uint32"); + out.push_back(encode(codec, strings[i], label, i)); + in.longest = std::max(in.longest, uint32_t(strings[i].size())); + } + }; + encode_all(queries, in.q, "queries"); + encode_all(refs, in.r, "refs"); + in.W1 = size_t(width) + 1; + int32_t largest_prior = 0; + if (!prior.empty()) { + if (in.W1 > std::numeric_limits::max() / in.W1 || + in.W1 * in.W1 > std::numeric_limits::max() / in.W1 || + prior.size() != in.W1 * in.W1 * in.W1) + throw std::invalid_argument("prior table must have (prior_width + 1)^3 entries"); + if (in.longest > width) + throw std::invalid_argument("sequence exceeds prior table width"); + for (auto value : prior) { + if (value < 0) throw std::invalid_argument("gap prior must be nonnegative"); + largest_prior = std::max(largest_prior, value); + } + } + const auto largest_pen = *std::max_element(in.pen.begin(), in.pen.end()); + if (*std::min_element(in.pen.begin(), in.pen.end()) < 0) + throw std::invalid_argument("matrix penalties must be nonnegative"); + const uint64_t bound = uint64_t(in.longest) * uint32_t(largest_pen) + + uint32_t(largest_prior) + uint32_t(go) + + uint64_t(in.longest ? in.longest - 1 : 0) * uint32_t(ge); + if (bound > std::numeric_limits::max()) + throw std::invalid_argument("gapblock score can overflow int32"); + return in; +} + +int32_t score_cell(const GapInput& in, size_t i, size_t j, int32_t go, int32_t ge, + const std::vector& prior, int32_t* scratch) { + const uint32_t m = uint32_t(in.q[i].size()), n = uint32_t(in.r[j].size()); + const uint32_t M = std::max(m, n), d = m > n ? m - n : n - m; + const int32_t* prow = prior.empty() ? nullptr : prior.data() + (size_t(M) * in.W1 + d) * in.W1; + return cell(in.q[i].data(), m, in.r[j].data(), n, in.pen.data(), in.A, go, ge, prow, scratch); +} + +struct Ranked { + uint32_t ref_id; + int32_t a, b; +}; +bool better(const Ranked& a, const Ranked& b) { + return std::tuple(std::max(a.a, a.b), a.ref_id) < std::tuple(std::max(b.a, b.b), b.ref_id); +} + +struct Worker { + std::vector scratch; + std::vector heap; +}; + +std::vector> ranked(const GapInput& a, const GapInput* b, uint32_t k, + int32_t go, int32_t ge, const std::vector& prior, + int threads, bool exclude_exact) { + if (!k) throw std::invalid_argument("k must be positive"); + if (b && (b->q.size() != a.q.size() || b->r.size() != a.r.size())) + throw std::invalid_argument("paired query/reference axes must agree"); + k = uint32_t(std::min(k, a.r.size())); + std::vector> out(a.q.size()); + if (!k) return out; + const auto longest = b ? std::max(a.longest, b->longest) : a.longest; + parallel_for(a.q.size(), threads, [&] { + Worker w{std::vector(size_t(longest) + 1), {}}; + w.heap.reserve(k); + return w; + }, [&](size_t i, Worker& w) { + w.heap.clear(); + for (size_t j = 0; j < a.r.size(); ++j) { + if (exclude_exact && a.q[i] == a.r[j] && (!b || b->q[i] == b->r[j])) continue; + const auto sa = score_cell(a, i, j, go, ge, prior, w.scratch.data()); + const auto sb = b ? score_cell(*b, i, j, go, ge, prior, w.scratch.data()) : sa; + const Ranked h{uint32_t(j), sa, sb}; + if (w.heap.size() < k) { + w.heap.push_back(h); + std::push_heap(w.heap.begin(), w.heap.end(), better); + } else if (better(h, w.heap.front())) { + std::pop_heap(w.heap.begin(), w.heap.end(), better); + w.heap.back() = h; + std::push_heap(w.heap.begin(), w.heap.end(), better); + } + } + std::sort_heap(w.heap.begin(), w.heap.end(), better); + out[i] = w.heap; + }); + return out; +} + +std::vector> counted(const GapInput& a, const GapInput* b, + const std::vector>& thresholds, int32_t go, int32_t ge, + const std::vector& prior, int threads, bool exclude_exact) { + if (thresholds.size() != a.q.size()) + throw std::invalid_argument("thresholds must have one row per query"); + if (b && (b->q.size() != a.q.size() || b->r.size() != a.r.size())) + throw std::invalid_argument("paired query/reference axes must agree"); + std::vector> out(a.q.size()); + const auto longest = b ? std::max(a.longest, b->longest) : a.longest; + parallel_for(a.q.size(), threads, [&] { return std::vector(size_t(longest) + 1); }, + [&](size_t i, std::vector& scratch) { + std::vector> sorted; + sorted.reserve(thresholds[i].size()); + for (size_t t = 0; t < thresholds[i].size(); ++t) sorted.emplace_back(thresholds[i][t], t); + std::sort(sorted.begin(), sorted.end()); + std::vector bins(sorted.size()); + out[i].resize(sorted.size()); + if (sorted.empty()) return; + for (size_t j = 0; j < a.r.size(); ++j) { + if (exclude_exact && a.q[i] == a.r[j] && (!b || b->q[i] == b->r[j])) continue; + auto score = score_cell(a, i, j, go, ge, prior, scratch.data()); + if (b) score = std::max(score, score_cell(*b, i, j, go, ge, prior, scratch.data())); + const auto pos = std::lower_bound(sorted.begin(), sorted.end(), score, + [](const auto& threshold, int32_t value) { return threshold.first < value; }); + if (pos != sorted.end()) ++bins[size_t(pos - sorted.begin())]; + } + uint64_t cumulative = 0; + for (size_t t = 0; t < sorted.size(); ++t) { + cumulative += bins[t]; + out[i][sorted[t].second] = cumulative; + } + }); + return out; +} + } // namespace std::vector gapblock_matrix(const std::vector& queries, @@ -120,4 +272,51 @@ std::vector gapblock_matrix(const std::vector& queries, return out; } +std::vector> gapblock_topk_batch(const std::vector& queries, + const std::vector& refs, uint32_t k, Alphabet alphabet, + const SubstitutionMatrix* matrix, int32_t go, int32_t ge, const std::vector& prior, + uint32_t width, int threads, bool exclude_exact) { + const auto input = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads); + const auto rows = ranked(input, nullptr, k, go, ge, prior, threads, exclude_exact); + std::vector> out(rows.size()); + for (size_t i = 0; i < rows.size(); ++i) { + out[i].reserve(rows[i].size()); + for (auto h : rows[i]) out[i].push_back(Hit{h.ref_id, h.a}); + } + return out; +} + +std::vector>> gapblock_paired_topk_batch( + const std::vector& qa, const std::vector& qb, + const std::vector& ra, const std::vector& rb, + uint32_t k, Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto a = prepare(qa, ra, alphabet, matrix, go, ge, prior, width, threads); + const auto b = prepare(qb, rb, alphabet, matrix, go, ge, prior, width, threads); + const auto rows = ranked(a, &b, k, go, ge, prior, threads, exclude_exact); + std::vector>> out(rows.size()); + for (size_t i = 0; i < rows.size(); ++i) + for (auto h : rows[i]) out[i].emplace_back(h.ref_id, h.a, h.b); + return out; +} + +std::vector> gapblock_count_batch(const std::vector& queries, + const std::vector& refs, const std::vector>& thresholds, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads); + return counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact); +} + +std::vector> gapblock_paired_count_batch( + const std::vector& qa, const std::vector& qb, + const std::vector& ra, const std::vector& rb, + const std::vector>& thresholds, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + const std::vector& prior, uint32_t width, int threads, bool exclude_exact) { + const auto a = prepare(qa, ra, alphabet, matrix, go, ge, prior, width, threads); + const auto b = prepare(qb, rb, alphabet, matrix, go, ge, prior, width, threads); + return counted(a, &b, thresholds, go, ge, prior, threads, exclude_exact); +} + } // namespace seqtree diff --git a/src/text_index.cpp b/src/text_index.cpp index 62b4e1c..41221a8 100644 --- a/src/text_index.cpp +++ b/src/text_index.cpp @@ -46,6 +46,7 @@ #include "atomic_write.hpp" #include +#include #include #include #include @@ -253,8 +254,14 @@ uint32_t kmer_code(const uint8_t* p, uint8_t k, uint8_t A) { struct RawHit { uint32_t ref_id, offset; uint16_t n_subs; + uint16_t n_ins, n_dels, length; int32_t score; uint32_t mm_off; + // How many mismatch entries this hit owns at mm_off. Equal to n_subs on the substitution + // path; 0 on the indel path, which reports counts without per-column detail because + // recovering the substituted columns would need an alignment traceback. Reading n_subs + // entries instead would walk off the end of a buffer the indel path never wrote to. + uint16_t mm_len; uint32_t abs; }; @@ -263,6 +270,9 @@ struct RawHit { // is read, which is also what removes the sort the old two-path dispatch needed to deduplicate. struct Work { std::vector probes; // k-mer codes to look up + // Banded-DP scratch for the indel path, allocated once per worker. Indexed + // [d + max_indels][g]; see verify_banded. + std::vector dp_cur, dp_next; }; // Per-query results. Mismatch positions are appended in discovery order and referenced by @@ -274,8 +284,154 @@ struct QueryOut { uint8_t truncated = 0; }; +// Banded alignment of the whole query against the text starting at `s`, under two independent +// caps: at most `ms` substitutions and at most `mi` insertions+deletions. +// +// State is (i, d, g): i query residues consumed, d = (text consumed) - (query consumed), and g +// indels spent. d is confined to [-mi, mi] and g to [0, mi], and every move changes them +// together -- an insertion (a query residue with no text opposite it) is d-1, g+1; a deletion +// (a text residue with none opposite) is d+1, g+1 -- so g >= |d| and g == |d| (mod 2). The cell +// holds the fewest substitutions reaching that state, or kInf. +// +// The alignment must BEGIN and END on an aligned pair -- no leading or trailing gap of either +// kind. Two reasons, and they agree. A gap at either edge extends the reported interval over a +// residue that matches nothing, so `offset` and `length` would no longer bound the part of the +// text the query actually explains. And a leading gap is just the same alignment starting one +// residue over, which the caller already reaches by sweeping every start in the seed's window, +// so allowing it would report one occurrence two or three times over. +// +// Every hit that survives is appended via `emit(end_offset_from_s, n_subs, n_ins, n_dels)`. One +// start can yield several ends -- they are genuinely different intervals -- and the caller +// deduplicates on (ref, offset, length). +template +void verify_banded(const uint8_t* q, size_t L, const uint8_t* text, uint64_t text_len, uint64_t s, + uint16_t ms, uint16_t mi, std::vector& cur, + std::vector& next, Emit emit) { + constexpr uint16_t kInf = 0xFFFF; + const size_t W = size_t(2) * mi + 1, G = size_t(mi) + 1, N = W * G; + const auto at = [mi, G](int d, size_t g) { return size_t(d + int(mi)) * G + g; }; + + cur.assign(N, kInf); + next.assign(N, kInf); + cur[at(0, 0)] = 0; // no leading gap: the alignment starts flush at s + + for (size_t i = 0; i < L; ++i) { + // Deletions consume text without consuming a query residue, so they move within this + // same i. Relax in increasing d: a run of them chains, each costing one unit of g. + // Not at i == 0: that would be a leading gap. + for (int d = -int(mi); i > 0 && d < int(mi); ++d) { + for (size_t g = 0; g + 1 < G; ++g) { + const uint16_t v = cur[at(d, g)]; + if (v == kInf) continue; + const int64_t tpos = int64_t(s) + int64_t(i) + d; + if (tpos < 0 || uint64_t(tpos) >= text_len) continue; + if (text[tpos] == Codec::kInvalid) continue; // never match through a hole + uint16_t& dst = cur[at(d + 1, g + 1)]; + if (v < dst) dst = v; + } + } + + std::fill(next.begin(), next.end(), kInf); + bool any = false; + for (int d = -int(mi); d <= int(mi); ++d) { + for (size_t g = 0; g < G; ++g) { + const uint16_t v = cur[at(d, g)]; + if (v == kInf) continue; + // match / substitution: consume q[i] against text[s + i + d] + const int64_t tpos = int64_t(s) + int64_t(i) + d; + if (tpos >= 0 && uint64_t(tpos) < text_len && text[tpos] != Codec::kInvalid) { + const uint16_t nv = uint16_t(v + (q[i] != text[tpos] ? 1 : 0)); + if (nv <= ms) { + uint16_t& dst = next[at(d, g)]; + if (nv < dst) dst = nv; + any = true; + } + } + // Insertion: consume q[i] with no text opposite it. Never at the first or last + // query residue -- that is a leading or trailing gap. + if (i > 0 && i + 1 < L && d > -int(mi) && g + 1 < G) { + uint16_t& dst = next[at(d - 1, g + 1)]; + if (v < dst) dst = v; + any = true; + } + } + } + // Every surviving state already exceeds one of the caps, so no suffix can rescue this + // start. Without this the DP walked all L rows on a candidate the Hamming path would + // have rejected after two residues, and most candidates are rejected: on a 5 M-residue + // corpus it is the difference between 0.107 and 0.019 ms/query at L=16. + if (!any) return; + cur.swap(next); + } + + // The query is fully consumed. Every finite state is a complete alignment; d fixes the + // matched length and g splits into insertions and deletions, since d = dels - ins and + // g = dels + ins. + for (int d = -int(mi); d <= int(mi); ++d) { + for (size_t g = 0; g < G; ++g) { + const uint16_t v = cur[at(d, g)]; + if (v == kInf || v > ms) continue; + if (size_t(std::abs(d)) > g || ((g - size_t(std::abs(d))) & 1)) continue; + const size_t dels = (g + size_t(d)) / 2, ins = (g - size_t(d)) / 2; + if (int64_t(L) + d < 0) continue; + emit(size_t(int64_t(L) + d), v, uint16_t(ins), uint16_t(dels)); + } + } +} + // One query against the text at a fixed max_subs. Appends to out.hits / out.mm_pos. // The scheme and why it is lossless are at the top of this file. +// The indel path. Seeding is unchanged and every block is exact -- with b = m + mi + 1 blocks +// the total edit count cannot reach the block count, so pigeonhole still puts a 0-error block +// somewhere and the existing exact lookup finds it. Only verification differs: the seed pins a +// text position, not a start, because the prefix before the block may have gained or lost up to +// mi residues, so every start in [pos - base - mi, pos - base + mi] is aligned. +// +// Hits are appended unsorted and with duplicates; the caller sorts by (n_subs, ref_id, offset, +// length) and drops adjacent equals, which is the same sort the substitution path already needs +// for its stable output order. +void search_indels(const TextStore& st, const uint8_t* q, size_t L, uint16_t m, uint16_t mi, + const SubstitutionMatrix* matrix, Work& s, QueryOut& out) { + const uint8_t k = st.k; + const size_t b = size_t(m) + size_t(mi) + 1; // search_batch enforces L >= b * k + const size_t bw = L / b; + + for (size_t j = 0; j < b; ++j) { + const size_t base = j * bw; + uint32_t code = 0; + for (uint8_t x = 0; x < k; ++x) code = uint32_t(code * st.A + q[base + x]); + const uint32_t lo = st.post_begin[code], hi = st.post_begin[code + 1]; + for (uint32_t i = lo; i < hi; ++i) { + const uint32_t pos = st.post_ids[i]; + const int64_t anchor = int64_t(pos) - int64_t(base); + for (int64_t start = anchor - mi; start <= anchor + mi; ++start) { + if (start < 0 || uint64_t(start) >= st.text_len) continue; + verify_banded( + q, L, st.text, st.text_len, uint64_t(start), m, mi, s.dp_cur, s.dp_next, + [&](size_t len, uint16_t n_subs, uint16_t n_ins, uint16_t n_dels) { + if (uint64_t(start) + len > st.text_len) return; + const uint32_t c = uint32_t(start); + const uint32_t ref = st.ref_of(c); + // A hit must lie wholly inside one record. The sentinels already stop + // the alignment, but a zero-length tail would slip through. + if (uint64_t(c) + len > st.starts[ref + 1]) return; + int32_t score = 0; + if (matrix) { + // No mismatch positions on this path (see the comment on mm_pos in + // search_batch), so the matrix scores what it can: the substituted + // columns of the best ungapped prefix are not recoverable without a + // traceback, and a partial score would be worse than none. + score = 0; + } + out.hits.push_back(RawHit{ref, c - st.starts[ref], n_subs, n_ins, n_dels, + uint16_t(len), score, + uint32_t(out.mm_pos.size()), 0, c}); + }); + } + } + } +} + void search_one(const TextStore& st, const uint8_t* q, size_t L, uint16_t m, const SubstitutionMatrix* matrix, const uint64_t* pw, Work& s, QueryOut& out) { const uint8_t k = st.k, A = st.A; @@ -331,7 +487,8 @@ void search_one(const TextStore& st, const uint8_t* q, size_t L, uint16_t m, if (matrix) for (uint32_t x = mm_off; x < mm_off + n; ++x) score += matrix->similarity(q[out.mm_pos[x]], t[out.mm_pos[x]]); - out.hits.push_back(RawHit{ref, c - st.starts[ref], n, score, mm_off, c}); + out.hits.push_back(RawHit{ref, c - st.starts[ref], n, 0, 0, uint16_t(L), score, + mm_off, n, c}); } } } @@ -463,6 +620,19 @@ TextResult TextIndex::search_batch(const std::vector& queries, "queries[" + std::to_string(i) + "] ('" + queries[i] + "') is " + std::to_string(queries[i].size()) + " long; TextIndex indexes seeds of k=" + std::to_string(st.k) + " and cannot answer a shorter query completely"); + if (opts.max_indels > 0) { + const size_t need = (size_t(opts.max_subs) + opts.max_indels + 1) * st.k; + if (queries[i].size() < need) + throw std::invalid_argument( + "queries[" + std::to_string(i) + "] ('" + queries[i] + "') is " + + std::to_string(queries[i].size()) + " long; searching with max_indels=" + + std::to_string(opts.max_indels) + " and max_subs=" + + std::to_string(opts.max_subs) + " needs every one of the " + + std::to_string(size_t(opts.max_subs) + opts.max_indels + 1) + + " seed blocks to be exact, so a query must be at least " + + std::to_string(need) + " long at k=" + std::to_string(st.k) + + ". Rebuild the index with a smaller k, or lower max_subs/max_indels"); + } if (queries[i].size() > kMaxQueryLen) throw std::invalid_argument( "queries[" + std::to_string(i) + "] is " + std::to_string(queries[i].size()) + @@ -488,19 +658,47 @@ TextResult TextIndex::search_batch(const std::vector& queries, for (uint16_t m = opts.best_only ? 0 : top; m <= top; ++m) { s.hits.clear(); s.mm_pos.clear(); - search_one(st, q, L, m, opts.matrix, pw.data(), w, s); + if (opts.max_indels > 0) + search_indels(st, q, L, m, opts.max_indels, opts.matrix, w, s); + else + search_one(st, q, L, m, opts.matrix, pw.data(), w, s); if (opts.exclude_exact) s.hits.erase(std::remove_if(s.hits.begin(), s.hits.end(), [](const RawHit& h) { return h.n_subs == 0; }), s.hits.end()); if (!opts.best_only || !s.hits.empty()) break; } - // Deterministic and independent of the order candidates were discovered in. + // Deterministic and independent of the order candidates were discovered in. The indel + // path adds length to the key and then deduplicates: one text interval is reachable + // from several blocks and several starts, and the cheapest alignment of that interval + // is the one to keep, so sorting by (n_subs, n_ins + n_dels) first puts it in front. std::sort(s.hits.begin(), s.hits.end(), [](const RawHit& a, const RawHit& b) { if (a.n_subs != b.n_subs) return a.n_subs < b.n_subs; + if (a.n_ins + a.n_dels != b.n_ins + b.n_dels) + return a.n_ins + a.n_dels < b.n_ins + b.n_dels; if (a.ref_id != b.ref_id) return a.ref_id < b.ref_id; - return a.offset < b.offset; + if (a.offset != b.offset) return a.offset < b.offset; + return a.length < b.length; }); + if (opts.max_indels > 0) { + std::vector keep; + keep.reserve(s.hits.size()); + for (const RawHit& h : s.hits) { + const bool seen = std::any_of(keep.begin(), keep.end(), [&](const RawHit& x) { + return x.ref_id == h.ref_id && x.offset == h.offset && x.length == h.length; + }); + if (!seen) keep.push_back(h); + } + s.hits.swap(keep); + // The dedup pass above is a linear scan per hit, which is fine while the kept set + // is the handful a real query returns; max_hits bounds it in the pathological case. + std::sort(s.hits.begin(), s.hits.end(), [](const RawHit& a, const RawHit& b) { + if (a.n_subs != b.n_subs) return a.n_subs < b.n_subs; + if (a.ref_id != b.ref_id) return a.ref_id < b.ref_id; + if (a.offset != b.offset) return a.offset < b.offset; + return a.length < b.length; + }); + } // A cap keeps the BEST hits, because it is applied after the sort -- and it is always // reported. A cap that is invisible is a recall bug wearing a performance costume. bool capped = opts.max_hits && s.hits.size() > opts.max_hits; @@ -540,7 +738,7 @@ TextResult TextIndex::search_batch(const std::vector& queries, for (const QueryOut& s : per_query) { n_hits += s.hits.size(); n_groups += s.groups.size(); - for (const RawHit& h : s.hits) n_mm += h.n_subs; + for (const RawHit& h : s.hits) n_mm += h.mm_len; } TextResult out; out.query_begin.reserve(n + 1); @@ -548,6 +746,9 @@ TextResult TextIndex::search_batch(const std::vector& queries, out.ref_id.reserve(n_hits); out.offset.reserve(n_hits); out.n_subs.reserve(n_hits); + out.n_ins.reserve(n_hits); + out.n_dels.reserve(n_hits); + out.length.reserve(n_hits); out.score.reserve(n_hits); out.mm_begin.reserve(n_hits + 1); out.mm_pos.reserve(n_mm); @@ -568,9 +769,12 @@ TextResult TextIndex::search_batch(const std::vector& queries, out.ref_id.push_back(h.ref_id); out.offset.push_back(h.offset); out.n_subs.push_back(h.n_subs); + out.n_ins.push_back(h.n_ins); + out.n_dels.push_back(h.n_dels); + out.length.push_back(h.length); out.score.push_back(h.score); const uint8_t* t = st.text + h.abs; - for (uint16_t j = 0; j < h.n_subs; ++j) { + for (uint16_t j = 0; j < h.mm_len; ++j) { const uint16_t p = s.mm_pos[h.mm_off + j]; out.mm_pos.push_back(p); out.mm_query_aa.push_back(codec.decode(qbuf[qoff[i] + p])); diff --git a/tests/cpp/test_text_index.cpp b/tests/cpp/test_text_index.cpp index fa1b671..47a43f6 100644 --- a/tests/cpp/test_text_index.cpp +++ b/tests/cpp/test_text_index.cpp @@ -539,3 +539,93 @@ TEST_CASE("a header that disagrees with itself is refused before anything is der } std::remove(bad.c_str()); } + +// --- indel search ----------------------------------------------------------------------- +// +// Seeding does not change: with b = max_subs + max_indels + 1 disjoint blocks the total edit +// count is below the block count, so pigeonhole still puts a zero-error block somewhere and the +// existing exact lookup finds it. These pin the verification and the interval it reports. + +TEST_CASE("TextIndex finds a single deletion") { + auto ix = TextIndex::build({"MKTAYIAKQRQISFVKSHFSRQ"}, Alphabet::AminoAcid, 3); + TextQueryOpts o; + o.max_subs = 0; + o.max_indels = 1; + // AYIAKQRQISF sits at offset 3; the query drops the I, so the text carries one extra residue. + auto r = ix->search_batch({"AYAKQRQISF"}, o, 1); + REQUIRE(r.num_hits() >= 1); + CHECK(r.n_subs[0] == 0); + CHECK(r.n_ins[0] == 0); + CHECK(r.n_dels[0] == 1); + CHECK(r.offset[0] == 3); + CHECK(r.length[0] == 11); // len(query) + n_dels - n_ins +} + +TEST_CASE("TextIndex finds a single insertion") { + auto ix = TextIndex::build({"MKTAYIAKQRQISFVKSHFSRQ"}, Alphabet::AminoAcid, 3); + TextQueryOpts o; + o.max_subs = 0; + o.max_indels = 1; + // An extra W inside the query: the text is one residue shorter than the query. + auto r = ix->search_batch({"AYIAKWQRQISF"}, o, 1); + REQUIRE(r.num_hits() >= 1); + CHECK(r.n_ins[0] == 1); + CHECK(r.n_dels[0] == 0); + CHECK(r.length[0] == 11); +} + +TEST_CASE("TextIndex indel hits never open or close on a gap") { + // The exact occurrence must be reported once, at its own bounds -- not also as the same + // match padded by an unaligned residue at either end. + auto ix = TextIndex::build({"MKTAYIAKQRQISFVKSHFSRQ"}, Alphabet::AminoAcid, 3); + TextQueryOpts o; + o.max_subs = 0; + o.max_indels = 1; + auto r = ix->search_batch({"AYIAKQRQI"}, o, 1); + REQUIRE(r.num_hits() == 1); + CHECK(r.offset[0] == 3); + CHECK(r.length[0] == 9); + CHECK(r.n_ins[0] == 0); + CHECK(r.n_dels[0] == 0); +} + +TEST_CASE("TextIndex max_indels = 0 leaves the substitution path alone") { + auto ix = TextIndex::build({"MKTAYIAKQRQISFVKSHFSRQ"}, Alphabet::AminoAcid, 3); + TextQueryOpts a, b; + a.max_subs = 1; + b.max_subs = 1; + b.max_indels = 0; + auto ra = ix->search_batch({"AYIAKQRQI"}, a, 1); + auto rb = ix->search_batch({"AYIAKQRQI"}, b, 1); + REQUIRE(ra.num_hits() == rb.num_hits()); + for (size_t i = 0; i < ra.num_hits(); ++i) { + CHECK(ra.ref_id[i] == rb.ref_id[i]); + CHECK(ra.offset[i] == rb.offset[i]); + CHECK(ra.n_subs[i] == rb.n_subs[i]); + CHECK(rb.n_ins[i] == 0); + CHECK(rb.n_dels[i] == 0); + CHECK(rb.length[i] == 9); + } +} + +TEST_CASE("TextIndex refuses an indel query too short to seed every block") { + auto ix = TextIndex::build({"MKTAYIAKQRQISFVKSHFSRQ"}, Alphabet::AminoAcid, 4); + TextQueryOpts o; + o.max_subs = 2; + o.max_indels = 1; + // Needs (2 + 1 + 1) * 4 = 16; every block must be exact, so a shorter query cannot be + // answered completely and is refused rather than answered partially. + CHECK_THROWS_AS(ix->search_batch({"AYIAKQRQ"}, o, 1), std::invalid_argument); + CHECK_NOTHROW(ix->search_batch({"AYIAKQRQISFVKSHF"}, o, 1)); +} + +TEST_CASE("TextIndex indel search does not cross a record boundary") { + // Two records whose junction spells a match only if the sentinel between them is crossed. + auto ix = TextIndex::build({"MKTAYIAKQR", "QISFVKSHFS"}, Alphabet::AminoAcid, 3); + TextQueryOpts o; + o.max_subs = 1; + o.max_indels = 1; + auto r = ix->search_batch({"AYIAKQRQISFV"}, o, 1); + for (size_t i = 0; i < r.num_hits(); ++i) + CHECK(size_t(r.offset[i]) + r.length[i] <= ix->ref_seq(r.ref_id[i]).size()); +} diff --git a/tests/python/test_gapblock_ranked.py b/tests/python/test_gapblock_ranked.py new file mode 100644 index 0000000..fc993ef --- /dev/null +++ b/tests/python/test_gapblock_ranked.py @@ -0,0 +1,113 @@ +"""Exhaustive, small generic-sequence oracles for bounded gapblock reductions.""" +import random +import pytest +import seqtree as st +from seqtree.gapblock import topk_batch, paired_topk_batch, count_batch, paired_count_batch, score_matrix, central_prior, positions_prior + + +def tuples(rows): return [[(h.ref_id,h.score) for h in row] for row in rows] + + +@pytest.mark.parametrize('matrix',[None,st.SubstitutionMatrix.blosum62()]) +@pytest.mark.parametrize('prior',[None,central_prior(7)]) +@pytest.mark.parametrize('exclude',[False,True]) +def test_topk_and_counts_equal_exhaustive_matrix(matrix,prior,exclude): + rng=random.Random(42) + q=['','A','AC','ACD','ac']+[''.join(rng.choice('ACDE') for _ in range(rng.randrange(1,9))) for _ in range(12)] + r=['','A','AC','AC','ACDE','ac']+[''.join(rng.choice('ACDE') for _ in range(rng.randrange(1,11))) for _ in range(31)] + opts=dict(matrix=matrix,gap_open=9,gap_extend=2,gap_prior=prior) + sm=score_matrix(q,r,**opts,threads=1) + for k in (1,5,10,100): + expected=[sorted([(j,sm[i,j]) for j,s in enumerate(r) if not exclude or s.upper()!=x.upper()], + key=lambda h:(h[1],h[0]))[:k] for i,x in enumerate(q)] + for threads in (1,4,0): + assert tuples(topk_batch(q,r,k,**opts,threads=threads,exclude_exact=exclude))==expected + thresholds=[[30,0,-1,9,9,200] if i%2 else [] for i in range(len(q))] + expected=[[sum(sm[i,j]<=t for j,s in enumerate(r) if not exclude or s.upper()!=x.upper()) + for t in thresholds[i]] for i,x in enumerate(q)] + assert count_batch(q,r,thresholds,**opts,threads=1,exclude_exact=exclude)==expected + assert count_batch(q,r,thresholds,**opts,threads=4,exclude_exact=exclude)==expected + + +def test_paired_topk_max_lane_oracle_and_joint_exact_puncture(): + qa,qb=['AC','DE',''],['ACD','A','D'] + ra,rb=['AC','AC','AD','DE',''],['ACD','AC','ACD','A','D'] + opts=dict(gap_open=3,gap_extend=2,gap_prior=central_prior(4)) + a,b=score_matrix(qa,ra,**opts),score_matrix(qb,rb,**opts) + for exclude in (False,True): + expected=[sorted([(j,a[i,j],b[i,j]) for j in range(len(ra)) + if not exclude or (qa[i].upper(),qb[i].upper())!=(ra[j].upper(),rb[j].upper())], + key=lambda h:(max(h[1],h[2]),h[0]))[:3] for i in range(len(qa))] + for threads in (1,4): + assert paired_topk_batch(qa,qb,ra,rb,3,**opts,threads=threads,exclude_exact=exclude)==expected + + +def test_exclusion_before_selection_and_zero_penalty_nonidentity(): + matrix=st.SubstitutionMatrix.from_similarity([[0]*24 for _ in range(24)]) + assert tuples(topk_batch(['A'],['A','C','A','D'],1,matrix=matrix,exclude_exact=True))==[[(1,0)]] + assert count_batch(['A'],['A','C','A','D'],[[0]],matrix=matrix,exclude_exact=True)==[[2]] + assert paired_topk_batch(['A'],['C'],['A','A','D'],['C','D','C'],1, + matrix=matrix,exclude_exact=True)==[[(1,0,0)]] + + +def test_empty_inputs_and_alphabets(): + assert topk_batch([],['AC'],5)==[] + assert topk_batch(['AC'],[],5)==[[]] + assert count_batch(['AC'],[],[[0,2]])==[[0,0]] + assert count_batch([],['AC'],[])==[] + assert paired_topk_batch(['AC'],['A'],[],[],5)==[[]] + assert tuples(topk_batch(['acgt'],['ACGT','ACGA'],2,alphabet='nt',exclude_exact=True))==[[(1,1)]] + assert tuples(topk_batch(['N'],['N','R'],2,alphabet='iupac',exclude_exact=True))==[[(1,1)]] + + +@pytest.mark.parametrize('k',[0,-1,1.5,True,'5',2**40]) +def test_invalid_k(k): + with pytest.raises((ValueError,TypeError)): topk_batch(['AC'],['AC'],k) + + +@pytest.mark.parametrize('args',[ + dict(gap_open=-1),dict(gap_extend=-1),dict(threads=-1),dict(threads=1.5), + dict(alphabet='bad'),dict(alphabet='nt',matrix=st.SubstitutionMatrix.blosum62()), + dict(gap_prior=lambda i,d,m:-1),dict(gap_open=2**31-1), +]) +def test_invalid_score_parameters(args): + with pytest.raises((ValueError,TypeError,OverflowError)): topk_batch(['AC'],['A'],2,**args) + + +@pytest.mark.parametrize('thresholds',[[],[[1],[2]],[[1.5]],[[True]],[[2**40]]]) +def test_invalid_threshold_shape_or_values(thresholds): + with pytest.raises((ValueError,TypeError)): count_batch(['AC'],['A'],thresholds) + + +def test_invalid_symbols_even_when_other_axis_empty(): + for fn in (lambda:topk_batch(['#'],[],1),lambda:topk_batch([],['#'],1), + lambda:count_batch(['#'],[],[[]]),lambda:count_batch([],['#'],[])): + with pytest.raises(ValueError): fn() + with pytest.raises(ValueError): paired_topk_batch(['A'],[],['A'],['A']) + with pytest.raises(ValueError): paired_topk_batch(['A'],['A'],['A'],[]) + + +@pytest.mark.parametrize('prior',[central_prior(7),positions_prior([1,-1])]) +def test_paired_count_matches_all_matrix_ties_and_threshold_order(prior): + qa,qb=['AC','DE',''],['ACD','A','D'] + ra,rb=['AC','AC','AD','DE',''],['ACD','AC','ACD','A','D'] + opts=dict(gap_open=3,gap_extend=2,gap_prior=prior) + a,b=score_matrix(qa,ra,**opts),score_matrix(qb,rb,**opts) + thresholds=[[20,-1,0,5,5],[],[30,2,0]] + for exclude in (False,True): + expected=[[sum(max(a[i,j],b[i,j])<=t for j in range(len(ra)) + if not exclude or (qa[i],qb[i])!=(ra[j],rb[j])) for t in ts] + for i,ts in enumerate(thresholds)] + for threads in (1,4): + assert paired_count_batch(qa,qb,ra,rb,thresholds,**opts,threads=threads,exclude_exact=exclude)==expected + selected=paired_topk_batch(qa,qb,ra,rb,3,**opts,exclude_exact=exclude) + expected_top=[[j for j in sorted(range(len(ra)),key=lambda j:(max(a[i,j],b[i,j]),j)) + if not exclude or (qa[i],qb[i])!=(ra[j],rb[j])][:3] for i in range(len(qa))] + assert [[h[0] for h in row] for row in selected]==expected_top + + +def test_native_reference_accessor_preserves_ref_id_order_and_duplicates(): + refs=['AC','AC','AG','aC'] + ix=st.Index.build(refs,'aa') + assert ix.ref_seqs()==[ix.ref_seq(i) for i in range(len(ix))] + assert ix.ref_seqs()==refs diff --git a/tests/python/test_text_index.py b/tests/python/test_text_index.py index 56876db..5b9ecd5 100644 --- a/tests/python/test_text_index.py +++ b/tests/python/test_text_index.py @@ -298,3 +298,127 @@ def test_an_index_file_from_an_older_format_is_refused(ix, tmp_path): path.write_bytes(bytes(raw)) with pytest.raises(Exception, match="text index"): TextIndex.load(str(path)) + + +#: The 20 standard residues, for the generated indel cases below. +AA = "ACDEFGHIKLMNPQRSTVWY" + + +# --- indel search ------------------------------------------------------------------------- +# +# The proof obligation: set equality against an independent brute force, the same standard the +# substitution path is held to. An occurrence is a text interval the WHOLE query aligns to +# within the two caps, with the alignment beginning and ending on an aligned pair. + +_INF = 10**6 + + +def _feasible(q, t, ms, mi): + """True if q aligns to the whole of t within ms substitutions and mi indels.""" + L, T = len(q), len(t) + if L == 0 or T == 0: + return False + for g in range(mi + 1): + for dels in range(g + 1): + ins = g - dels + if T - L != dels - ins: + continue + dp = [[[_INF] * (dels + 1) for _ in range(ins + 1)] for _ in range(L + 1)] + dp[1][0][0] = q[0] != t[0] # must open on an aligned pair + for i in range(1, L + 1): + for a in range(ins + 1): + for d in range(dels + 1): + v = dp[i][a][d] + if v >= _INF: + continue + j = i - a + d + if i < L and j < T: + dp[i + 1][a][d] = min(dp[i + 1][a][d], v + (q[i] != t[j])) + if 0 < i < L - 1 and a < ins: # no insertion at either edge + dp[i + 1][a + 1][d] = min(dp[i + 1][a + 1][d], v) + if i < L and j < T and d < dels: # no deletion at either edge + dp[i][a][d + 1] = min(dp[i][a][d + 1], v) + if dp[L][ins][dels] <= ms: + return True + return False + + +def _brute_indel(records, q, ms, mi): + out = set() + for r, rec in enumerate(records): + for start in range(len(rec)): + for ln in range(max(1, len(q) - mi), min(len(q) + mi, len(rec) - start) + 1): + if _feasible(q, rec[start:start + ln], ms, mi): + out.add((r, start, ln)) + return out + + +@pytest.mark.parametrize("k", [3, 4]) +@pytest.mark.parametrize("ms,mi", [(0, 1), (1, 1), (2, 1), (0, 2), (1, 2)]) +def test_indel_search_equals_brute_force(k, ms, mi): + """Set equality, not a spot check: every interval the oracle finds and no others.""" + rng = random.Random(1000 + k * 17 + ms * 5 + mi) + records = ["".join(rng.choice(AA) for _ in range(40)) for _ in range(4)] + ix = TextIndex.build(records, alphabet="aa", k=k) + length = (ms + mi + 1) * k + checked = 0 + for _ in range(12): + rec = rng.choice(records) + start = rng.randrange(len(rec) - length + 1) + q = _mutate_with_indels(rec[start:start + length], rng.randint(0, ms + mi), rng) + if len(q) != length: # keep the query legal for this (ms, mi, k) + continue + got = {(h.ref_id, h.offset, h.length) + for h in ix.search_batch([q], max_subs=ms, max_indels=mi)[0]} + assert got == _brute_indel(records, q, ms, mi), f"q={q} k={k} ms={ms} mi={mi}" + checked += 1 + assert checked, "generated no legal query for this cell" + + +def _mutate_with_indels(s, n, rng): + s = list(s) + for _ in range(n): + op = rng.choice(("sub", "ins", "del")) + i = rng.randrange(len(s)) + if op == "sub": + s[i] = rng.choice(AA) + elif op == "ins": + s.insert(i, rng.choice(AA)) + elif len(s) > 1: + del s[i] + return "".join(s) + + +def test_max_indels_zero_is_the_substitution_path_unchanged(): + """Adding the parameter must not perturb the predicate anyone is already calling.""" + rng = random.Random(4) + records = ["".join(rng.choice(AA) for _ in range(300)) for _ in range(20)] + ix = TextIndex.build(records, alphabet="aa", k=4) + queries = [records[rng.randrange(20)][rng.randrange(280):][:12] for _ in range(200)] + a = ix.search_batch(queries, max_subs=2) + b = ix.search_batch(queries, max_subs=2, max_indels=0) + assert a.num_hits == b.num_hits and a.num_hits > 0 + for i in range(len(queries)): + assert [(h.ref_id, h.offset, h.n_subs, h.mismatches) for h in a[i]] == \ + [(h.ref_id, h.offset, h.n_subs, h.mismatches) for h in b[i]] + # ...and every hit still describes an ungapped, query-length match. + assert all(h.n_ins == 0 and h.n_dels == 0 and h.length == 12 for h in a[0]) + + +def test_indel_hit_reports_a_usable_interval(): + """offset/length must bound the matched text, since it is no longer len(query) wide.""" + ix = TextIndex.build(["MKTAYIAKQRQISFVKSHFSRQ"], alphabet="aa", k=3) + # ^ AYIAKQR is at offset 4; drop the I to force one deletion + hits = ix.search_batch(["AYAKQRQISF"], max_subs=0, max_indels=1)[0] + assert hits, "a single deletion should still be found" + h = hits[0] + assert (h.n_subs, h.n_ins, h.n_dels) == (0, 0, 1) + assert h.length == 11 and len(ix.ref_seq(h.ref_id)[h.offset:h.offset + h.length]) == 11 + assert ix.ref_seq(h.ref_id)[h.offset:h.offset + h.length] == "AYIAKQRQISF" + + +def test_indel_search_refuses_a_query_too_short_to_seed(): + ix = TextIndex.build(["MKTAYIAKQRQISFVKSHFSRQ"], alphabet="aa", k=4) + with pytest.raises(ValueError, match="seed blocks to be exact"): + ix.search_batch(["AYIAKQRQ"], max_subs=2, max_indels=1) # needs 4*4 = 16, has 8 + ix.search_batch(["AYIAKQRQISFVKSHF"], max_subs=2, max_indels=1) # 16 is fine