Limit gap-prior preparation to input lengths

CHANGELOG.md

@@ -3,6 +3,15 @@
 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
+
+### Changed
+
+- Gap-prior lookup preparation evaluates pure callbacks only for input lengths,
+  avoiding cubic callback work for absent lengths while preserving exact scores,
+  count reductions and lane-specific priors. Equal-length/empty axes do not
+  evaluate unused priors. The dense native lookup layout is unchanged.
+
 ## [1.3.0] — 2026-10-19
 
 ### Added

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "seqtree"
-version = "1.3.0"
+version = "1.3.1"
 description = "Fast fuzzy search over biological sequences (C++ core, Python bindings)"
 readme = "README.md"
 requires-python = ">=3.10"

python/seqtree/gapblock.py

@@ -499,19 +499,22 @@ def gapblock_score(
     return best + gap_cost(d, gap_open, gap_extend), best_i
 
 
-def _prior_cube(prior: GapPrior, width: int) -> list[int]:
+def _prior_cube(prior: GapPrior, width: int, lengths=None) -> list[int]:
     """Flatten ``prior`` to ``[m][d][i]`` with stride ``width + 1``, for the C++ kernel.
 
     Only reachable cells are filled: ``d >= 1`` (the prior never applies to equal lengths) and
     ``i <= m - d`` (the block start indexes the shorter sequence). The rest stay zero and are
-    never read.
+    never read. Input lengths restrict callback evaluation to possible full-length pairs;
+    the dense native layout is unchanged. None retains the full-width internal oracle.
     """
     w1 = width + 1
     cube = [0] * (w1 * w1 * w1)
-    for m in range(w1):
-        for d in range(1, m + 1):
+    lengths = range(w1) if lengths is None else sorted(set(lengths))
+    for row, m in enumerate(lengths):
+        for n in reversed(lengths[:row]):
+            d = m - n
             base = (m * w1 + d) * w1
-            for i in range(m - d + 1):
+            for i in range(n + 1):
                 v = int(prior(i, d, m))
                 if v < 0:
                     raise ValueError(f"gap_prior({i}, {d}, {m}) = {v} is negative")
@@ -561,7 +564,8 @@ def score_matrix(
             leaving this at 1 with a real matrix makes gaps ~14x cheaper than substitutions.
         gap_extend: Cost per additional gap column.
         gap_prior: :data:`GapPrior`, materialized once into a lookup cube and then read from C++.
-            ``None`` lets the score alone choose the block position.
+            Evaluated only for unequal input lengths and valid block starts; it must be
+            a pure nonnegative function. ``None`` lets the score alone choose the block position.
         alphabet: ``"aa"``, ``"nt"``, or ``"iupac"``. Symbols outside it raise.
         threads: Worker threads; ``0`` means one per core. Rows are disjoint, so this scales.
 
@@ -586,8 +590,9 @@ def score_matrix(
         gap_open = _default_gap_open(matrix)
     if gap_open < 0 or gap_extend < 0:
         raise ValueError("gap_open and gap_extend must be >= 0")
-    width = max((len(s) for s in q + r), default=0)
-    cube = _prior_cube(gap_prior, width) if gap_prior is not None else []
+    lengths = {len(s) for s in q + r}
+    width = max(lengths, default=0)
+    cube = _prior_cube(gap_prior, width, lengths) if gap_prior is not None else []
     return _gapblock_matrix(q, r, alphabet, matrix, gap_open, gap_extend, cube, width, threads,
                             _position_weights(position_weights_by_length, q, r),
                             *_group_options(query_group_ids, reference_group_ids, group_distances, q, r),
@@ -611,8 +616,9 @@ 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 []
+    lengths = {len(s) for axis in axes for s in axis}
+    width = max(lengths, default=0)
+    cube = _prior_cube(gap_prior, width, lengths) if gap_prior is not None else []
     return go, ge, cube, width, threads
 
 
@@ -915,11 +921,11 @@ def _additional_geometry(geometry, features):
         geometry.additional_features = terms
 
 
-def _lane_prior(geometry, prior, width):
+def _lane_prior(geometry, prior, width, queries, refs):
     if prior is not None:
         if not hasattr(geometry, "gap_prior"):
             raise RuntimeError("lane priors require native grouped gap_prior support")
-        geometry.gap_prior = _prior_cube(prior, width)
+        geometry.gap_prior = _prior_cube(prior, width, {len(s) for s in queries + refs})
 
 
 def _alignment_options(geometry, q, r, mode, scale, left, right, matrix, prior,
@@ -1099,8 +1105,8 @@ def grouped_paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresho
         reference_position_weights_beta, feature_offsets_beta, score_quantum_beta)
     _additional_geometry(alpha, additional_feature_offsets_alpha)
     _additional_geometry(beta, additional_feature_offsets_beta)
-    _lane_prior(alpha, gap_prior_alpha, width)
-    _lane_prior(beta, gap_prior_beta, width)
+    _lane_prior(alpha, gap_prior_alpha, width, qa, ra)
+    _lane_prior(beta, gap_prior_beta, width, qb, rb)
     _alignment_options(alpha, qa, ra, alignment_mode, alignment_scale, trim_left, trim_right,
         matrix, gap_prior if gap_prior_alpha is None else gap_prior_alpha, position_weights_by_length_alpha, query_position_weights_alpha,
         reference_position_weights_alpha)

tests/python/test_gapblock_prior_inputs.py

@@ -0,0 +1,58 @@
+"""Prior preparation must preserve exact scores without evaluating absent lengths."""
+import pytest
+from seqtree import gapblock as gb
+
+
+@pytest.mark.parametrize('mode', ['matrix', 'count', 'paired', 'lane-priors'])
+def test_prior_preparation_uses_input_lengths_and_matches_scalar_scores(mode):
+    qa, ra = ['AC', 'A' * 40], ['ACA', 'AC']
+    qb, rb = ['A', 'AC'], ['AC', 'A']
+    seen, seen_beta = set(), set()
+    def value(i, d, m):
+        return (i + 2*d + m) % 7
+    def prior(i, d, m):
+        seen.add((m, m-d, i))
+        return value(i, d, m)
+    def beta(i, d, m):
+        seen_beta.add((m, m-d, i))
+        return value(i, d, m) + 1
+    options = dict(gap_open=3, gap_extend=2, threads=1)
+    def scores(q, r, p):
+        return [[gb.gapblock_score(a, b, gap_open=3, gap_extend=2,
+                    gap_prior=p)[0] for b in r] for a in q]
+    expected = scores(qa, ra, value)
+    cuts = [0, 10, 100]
+    lengths = {len(s) for s in qa + ra}
+    if mode == 'matrix':
+        actual = memoryview(gb.score_matrix(qa, ra, gap_prior=prior, **options)).tolist()
+    elif mode == 'count':
+        actual = gb.count_batch(qa, ra, [cuts]*len(qa), gap_prior=prior, **options)
+        expected = [[sum(v <= c for v in row) for c in cuts] for row in expected]
+    else:
+        lanes = mode == 'lane-priors'
+        other = scores(qb, rb, lambda i, d, m: value(i, d, m) + lanes)
+        expected = [[a+b for a,b in zip(x,y)] for x,y in zip(expected,other)]
+        expected = [[[sum(v <= c for v in row) for c in cuts]]*4 for row in expected]
+        priors = dict(gap_prior_alpha=prior, gap_prior_beta=beta) if lanes else dict(gap_prior=prior)
+        chunks = list(gb.grouped_paired_sum_count_batch(qa, qb, ra, rb, cuts,
+            reference_target_offsets=[0, 1, 2], reference_target_ids=[0, 0],
+            target_peptide_ids=[0], reference_exclusive_peptide_ids=[0, 0],
+            query_gene_ids=[0, 0], reference_gene_ids=[0, 0], **options, **priors))
+        actual = [row for c in chunks for row in memoryview(c.counts).tolist()]
+        if not lanes:
+            lengths.update(len(s) for s in qb + rb)
+    assert actual == expected
+    assert seen and all(m in lengths and n in lengths and 0 <= i <= n for m, n, i in seen)
+    if seen_beta:
+        beta_lengths = {len(s) for s in qb + rb}
+        assert all(m in beta_lengths and n in beta_lengths and 0 <= i <= n for m, n, i in seen_beta)
+
+
+@pytest.mark.parametrize('queries,refs', [([], []), (['AC'], ['AC']), ([], ['AC']), (['AC'], [])])
+def test_prior_is_not_evaluated_when_no_length_difference_is_used(queries, refs):
+    def unused(i, d, m):
+        raise AssertionError('no gap-prior cell can be read')
+    matrix = gb.score_matrix(queries, refs, gap_prior=unused, threads=1)
+    assert matrix.shape == (len(queries), len(refs))
+    assert gb.count_batch(queries, refs, [[0, 1]]*len(queries), gap_prior=unused,
+                          threads=1) == [[len(refs)]*2 for _ in queries]