From 757ea648aaa66fc59d718aedfb388cdb7470ac57 Mon Sep 17 00:00:00 2001 From: mikessh Date: Sat, 10 Oct 2026 13:51:50 +0300 Subject: [PATCH 1/3] Add optional group penalties to gapblock batches Co-Authored-By: Codex --- docs/gapblock.rst | 24 ++++++ include/seqtree/seqtree.hpp | 17 +++- python/seqtree/gapblock.py | 43 ++++++++-- src/_bindings.cpp | 17 ++-- src/gapblock.cpp | 60 ++++++++++++-- tests/python/test_gapblock_group_offsets.py | 91 +++++++++++++++++++++ 6 files changed, 226 insertions(+), 26 deletions(-) create mode 100644 tests/python/test_gapblock_group_offsets.py diff --git a/docs/gapblock.rst b/docs/gapblock.rst index a1a371c..01da88c 100644 --- a/docs/gapblock.rst +++ b/docs/gapblock.rst @@ -376,3 +376,27 @@ rows follow ``count_batch``. It avoids a dense score matrix but still evaluates every reference; accepted output can itself be large. Positional weights are currently unavailable in top-k and paired reducers. These operations supply geometry, without a statistical calibration or a retrieval speed guarantee. + +Additive group distances +------------------------ + +``score_matrix``, ``count_batch`` and ``ball_batch`` also accept optional +``query_group_ids``, ``reference_group_ids`` and ``group_distances``. Supply all +three together: one nonnegative integer ID per sequence on each axis, and a +nonempty rectangular matrix of nonnegative integer penalties. The total pair +score is the ordinary gapblock score plus +``group_distances[query_group_ids[i]][reference_group_ids[j]]``. The table must +use the same units as substitution, gap and position-prior penalties. Group IDs +index rows on the query axis and columns on the reference axis. No symmetry or +zero diagonal is imposed; callers wanting a distance must provide those +properties themselves. + +Dimensions, IDs and possible int32 score overflow are checked before workers +start. The small table and encoded group IDs are shared read-only, with no dense +pair-offset allocation. Bounded reducers skip a pair when its group offset alone +exceeds the requested radius, since all other penalties are nonnegative. Exact +exclusion still punctures full encoded junction identity regardless of group ID +or offset, including when a matching junction has a positive group penalty. +Omitting all three options preserves the existing scorer. Group offsets are not +exposed by top-k or paired reducers; the paired maximum-score contract is +unchanged. These additions supply geometry, not statistical calibration. diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index ffd38c9..f311107 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -223,19 +223,27 @@ std::vector align_dist_matrix(const std::vector& queries, // the diagonal zero. `matrix` may be null for unit cost. threads <= 0 => hardware_concurrency. // Optional integer weights are indexed [longer length][longer position], symmetric in q/r. // All input frames must be covered; empty tables retain the existing unweighted scorer. +// Optional group IDs index a nonnegative rectangular penalty table added per pair. +// Exact exclusion remains encoded junction identity, independent of these groups. std::vector gapblock_matrix(const std::vector& queries, const std::vector& refs, Alphabet, const SubstitutionMatrix* matrix, int32_t gap_open, int32_t gap_extend, const std::vector& prior, uint32_t prior_width, int threads = 0, - const std::vector>& weights = {}); + const std::vector>& weights = {}, + const std::vector& query_groups = {}, + const std::vector& reference_groups = {}, + const std::vector>& group_distances = {}); // Exhaustive ball rows in reference order; hit edit fields are unset (zero). std::vector> gapblock_ball_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 std::vector>& weights = {}); + const std::vector>& weights = {}, + const std::vector& query_groups = {}, + const std::vector& reference_groups = {}, + const std::vector>& group_distances = {}); // Exhaustive bounded reductions using the same gapblock cell. Top-k ties use // (penalty, ref_id); paired penalty is max(lane scores). Exact exclusion uses @@ -253,7 +261,10 @@ std::vector> gapblock_count_batch(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 std::vector>& weights = {}); + const std::vector>& weights = {}, + const std::vector& query_groups = {}, + const std::vector& reference_groups = {}, + const std::vector>& group_distances = {}); std::vector> gapblock_paired_count_batch( const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, diff --git a/python/seqtree/gapblock.py b/python/seqtree/gapblock.py index 4729643..0b91c26 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -524,6 +524,7 @@ def score_matrix( alphabet: str = "aa", threads: int = 0, position_weights_by_length=None, + query_group_ids=None, reference_group_ids=None, group_distances=None, ) -> ScoreMatrix: """Gap-block penalty of every query against every reference, in C++ with the GIL released. @@ -538,6 +539,10 @@ def score_matrix( position_weights_by_length: Optional mapping of full sequence lengths to nonnegative integer weights. Weights index the longer sequence; both axes must be covered. None or an empty mapping preserves unweighted defaults. + query_group_ids: Optional integer IDs, one per query; supply all group options together. + reference_group_ids: Optional integer IDs, one per reference. + group_distances: Nonempty rectangular nonnegative integer penalty table. + Added to each pair score in the same units; no symmetry/zero diagonal imposed. matrix: Substitution penalties; ``None`` means unit cost. gap_open: Block-opening cost. Defaults to ``2 * matrix.scale()``. See the module note: leaving this at 1 with a real matrix makes gaps ~14x cheaper than substitutions. @@ -571,7 +576,8 @@ def score_matrix( width = max((len(s) for s in q + r), default=0) cube = _prior_cube(gap_prior, width) 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)) + _position_weights(position_weights_by_length, q, r), + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) def _batch_integer(value, name, lower=0, upper=(1 << 31) - 1): @@ -596,6 +602,24 @@ def _batch_options(axes, matrix, gap_open, gap_extend, gap_prior, threads): return go, ge, cube, width, threads +def _group_options(query_ids, reference_ids, distances, queries, refs): + """Validate one small group table; native preparation rechecks its safety.""" + if query_ids is None and reference_ids is None and distances is None: + return [], [], [] + if query_ids is None or reference_ids is None or distances is None: + raise ValueError("query_group_ids, reference_group_ids and group_distances are required together") + qg = [_batch_integer(v, "query group ID") for v in query_ids] + rg = [_batch_integer(v, "reference group ID") for v in reference_ids] + if len(qg) != len(queries) or len(rg) != len(refs): + raise ValueError("group IDs must match their query/reference axes") + table = [[_batch_integer(v, "group distance") for v in row] for row in distances] + if not table or not table[0] or any(len(row) != len(table[0]) for row in table): + raise ValueError("group distances must be a nonempty rectangular matrix") + if any(v >= len(table) for v in qg) or any(v >= len(table[0]) for v in rg): + raise ValueError("group ID outside group distances") + return qg, rg, table + + def _position_weights(weights, *axes): """Validated integer tables in the longer-junction frame; never round/truncate.""" if weights is None: @@ -650,10 +674,12 @@ def topk_batch(queries, refs, k=10, matrix=None, gap_open=None, gap_extend=1, def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend=1, - gap_prior=None, alphabet="aa", threads=0, exclude_exact=False, position_weights_by_length=None): + gap_prior=None, alphabet="aa", threads=0, exclude_exact=False, position_weights_by_length=None, + query_group_ids=None, reference_group_ids=None, group_distances=None): """Count all reference scores <= each query's requested integer thresholds. - Optional positional weights index the longer full sequence, exactly as in score_matrix. + Optional positional weights and additive group distances match :func:`score_matrix`. + Exact exclusion remains full encoded identity, irrespective of group distance. ``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 @@ -664,16 +690,20 @@ def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_exten 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, - _position_weights(position_weights_by_length, q, r)) + _position_weights(position_weights_by_length, q, r), + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) def ball_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend=1, gap_prior=None, alphabet="aa", threads=0, exclude_exact=False, - position_weights_by_length=None): + position_weights_by_length=None, + query_group_ids=None, reference_group_ids=None, group_distances=None): """Exhaustive bounded ball rows, reference order, unset edit fields; no Q*N matrix. One integer cutoff per query. Negative cutoffs yield empty rows. Optional integer positional weights index the longer full sequence, symmetrically, at every length. + Optional additive group distances match :func:`score_matrix`; exact exclusion + remains full encoded identity, irrespective of group distance. This reduces output memory; it is not a trie-based candidate filter. """ q, r = list(queries), list(refs) @@ -682,7 +712,8 @@ def ball_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend raise ValueError("ball thresholds must have one value per query") go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads) return _gapblock_ball_batch(q, r, thresholds, alphabet, matrix, go, ge, cube, width, - threads, exclude_exact, _position_weights(position_weights_by_length, q, r)) + threads, exclude_exact, _position_weights(position_weights_by_length, q, r), + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) def paired_topk_batch(queries_a, queries_b, refs_a, refs_b, k=10, matrix=None, diff --git a/src/_bindings.cpp b/src/_bindings.cpp index 4c9462b..e91f09d 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -361,7 +361,7 @@ ScoreMatrix py_gapblock_matrix(const std::vector& queries, const std::optional& matrix, int32_t gap_open, int32_t gap_extend, const std::vector& prior, uint32_t prior_width, int threads, - const std::vector>& weights) { + const std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { Alphabet alph = parse_alphabet(alphabet); ScoreMatrix out; out.rows = queries.size(); @@ -369,7 +369,7 @@ ScoreMatrix py_gapblock_matrix(const std::vector& queries, { nb::gil_scoped_release release; out.data = gapblock_matrix(queries, refs, alph, matrix ? &*matrix : nullptr, gap_open, - gap_extend, prior, prior_width, threads, weights); + gap_extend, prior, prior_width, threads, weights, query_groups, reference_groups, group_distances); } return out; } @@ -801,6 +801,7 @@ NB_MODULE(_core, m) { 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("position_weights_by_length") = std::vector>{}, + nb::arg("query_group_ids") = std::vector{}, nb::arg("reference_group_ids") = std::vector{}, nb::arg("group_distances") = std::vector>{}, "Exhaustive single-gap-block penalties for every (query, ref) pair, GIL released. " "`prior` is the gap prior flattened to [m][d][i]; see seqtree.gapblock.score_matrix, " "which builds it for you."); @@ -811,17 +812,17 @@ NB_MODULE(_core, m) { 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 std::vector>& weights) { + 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 std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { 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, weights); - }, 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, nb::arg("position_weights_by_length")=std::vector>{}, "Bounded exhaustive gapblock reduction, GIL released."); + return gapblock_count_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights, query_groups, reference_groups, group_distances); + }, 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, nb::arg("position_weights_by_length")=std::vector>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, "Bounded exhaustive gapblock reduction, GIL released."); - m.def("gapblock_ball_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 std::vector>& weights) { + m.def("gapblock_ball_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 std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { const auto alph = parse_alphabet(alphabet); nb::gil_scoped_release release; - return gapblock_ball_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights); - }, 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, nb::arg("position_weights_by_length")=std::vector>{}, "Bounded exhaustive gapblock reduction, GIL released."); + return gapblock_ball_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights, query_groups, reference_groups, group_distances); + }, 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, nb::arg("position_weights_by_length")=std::vector>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, "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); diff --git a/src/gapblock.cpp b/src/gapblock.cpp index 839621c..6ba5d97 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -67,6 +67,8 @@ struct GapInput { std::vector> q, r; std::vector pen; std::vector> weights; + std::vector query_groups, reference_groups; + std::vector> group_distances; uint32_t longest = 0; uint8_t A = 0; size_t W1 = 0; @@ -75,7 +77,10 @@ struct GapInput { 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, - const std::vector>& weights = {}) { + const std::vector>& weights = {}, + const std::vector& query_groups = {}, + const std::vector& reference_groups = {}, + const std::vector>& group_distances = {}) { 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()) @@ -129,6 +134,31 @@ GapInput prepare(const std::vector& queries, const std::vector= group_distances.size()) + throw std::invalid_argument("query group ID outside group distances"); + for (const auto id : reference_groups) + if (id < 0 || size_t(id) >= columns) + throw std::invalid_argument("reference group ID outside group distances"); + in.query_groups = query_groups; + in.reference_groups = reference_groups; + in.group_distances = group_distances; + } 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"); @@ -136,19 +166,23 @@ GapInput prepare(const std::vector& queries, const std::vector uint64_t(std::numeric_limits::max()) / in.longest) throw std::invalid_argument("gapblock weighted score can overflow int32"); const uint64_t bound = uint64_t(in.longest) * weighted_pen + - uint32_t(largest_prior) + uint32_t(go) + + uint32_t(largest_prior) + uint32_t(largest_offset) + 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; } +inline int32_t group_offset(const GapInput& in, size_t i, size_t j) { + return in.group_distances.empty() ? 0 : in.group_distances[in.query_groups[i]][in.reference_groups[j]]; +} + 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, in.weights.empty() || !M ? nullptr : in.weights[M].data()); + return group_offset(in, i, j) + cell(in.q[i].data(), m, in.r[j].data(), n, in.pen.data(), in.A, go, ge, prow, scratch, in.weights.empty() || !M ? nullptr : in.weights[M].data()); } struct Ranked { @@ -220,6 +254,7 @@ std::vector> counted(const GapInput& a, const GapInput* b, 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; + if (group_offset(a, i, j) > sorted.back().first) 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, @@ -241,8 +276,10 @@ std::vector gapblock_matrix(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, - const std::vector>& weights) { - const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, std::max(threads, 0), weights); + const std::vector>& weights, + const std::vector& query_groups, const std::vector& reference_groups, + const std::vector>& group_distances) { + const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, std::max(threads, 0), weights, query_groups, reference_groups, group_distances); if (!refs.empty() && queries.size() > std::numeric_limits::max() / refs.size()) throw std::invalid_argument("gapblock matrix size overflows size_t"); std::vector out(queries.size() * refs.size()); @@ -258,16 +295,19 @@ std::vector> gapblock_ball_batch(const std::vector 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 std::vector>& weights) { + const std::vector>& weights, + const std::vector& query_groups, const std::vector& reference_groups, + const std::vector>& group_distances) { if (thresholds.size() != queries.size()) throw std::invalid_argument("ball thresholds must have one value per query"); - const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights); + const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances); std::vector> out(queries.size()); parallel_for(queries.size(), threads, [&] { return std::vector(size_t(in.longest) + 1); }, [&](size_t i, std::vector& scratch) { if (thresholds[i] < 0) return; for (size_t j = 0; j < refs.size(); ++j) { if (exclude_exact && in.q[i] == in.r[j]) continue; + if (group_offset(in, i, j) > thresholds[i]) continue; const auto score = score_cell(in, i, j, go, ge, prior, scratch.data()); if (score <= thresholds[i]) out[i].push_back(Hit{uint32_t(j), score, 0, 0, 0}); } @@ -307,8 +347,10 @@ std::vector> gapblock_count_batch(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 std::vector>& weights) { - const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights); + const std::vector>& weights, + const std::vector& query_groups, const std::vector& reference_groups, + const std::vector>& group_distances) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances); return counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact); } diff --git a/tests/python/test_gapblock_group_offsets.py b/tests/python/test_gapblock_group_offsets.py new file mode 100644 index 0000000..665b9fb --- /dev/null +++ b/tests/python/test_gapblock_group_offsets.py @@ -0,0 +1,91 @@ +"""Generic additive group penalties share the existing gapblock cell and reducers.""" +import numpy as np +import pytest +from seqtree import gapblock + + +def test_group_offsets_matrix_reducers_and_puncture(): + q=['AC','ACC','','AC']; r=['AC','CA','ACC','CA',''] + qg=[0,1,1,0]; rg=[1,0,1,0,0]; distances=[[7,2],[0,11]] + opts=dict(gap_open=3,gap_extend=2,position_weights_by_length={0:[],2:[0,2],3:[3,1,2]}) + base=np.asarray(gapblock.score_matrix(q,r,threads=1,**opts)) + expected=base+np.array([[distances[a][b] for b in rg] for a in qg]) + group=dict(query_group_ids=qg,reference_group_ids=rg,group_distances=distances) + thresholds=[[-1,0,2,5,20,5],[],[0,3,20],[20]] + for t in (1,4): + np.testing.assert_array_equal(np.asarray(gapblock.score_matrix(q,r,threads=t,**opts,**group)),expected) + for excluded in (False,True): + counts=gapblock.count_batch(q,r,thresholds,threads=t,exclude_exact=excluded,**opts,**group) + balls=gapblock.ball_batch(q,r,[5]*len(q),threads=t,exclude_exact=excluded,**opts,**group) + for i in range(len(q)): + eligible=lambda j:not(excluded and q[i]==r[j]) + assert counts[i]==[sum(eligible(j) and expected[i,j]<=c for j in range(len(r))) for c in thresholds[i]] + assert [(h.ref_id,h.score) for h in balls[i]]==[(j,int(expected[i,j])) for j in range(len(r)) if eligible(j) and expected[i,j]<=5] + # A full-junction exact match remains punctured even with a positive group penalty. + assert gapblock.count_batch(['AC'],['AC'],[[100]],exclude_exact=True,query_group_ids=[0],reference_group_ids=[1],group_distances=distances)==[[0]] + # Offset alone rejects the entire group even when its junction distance is zero. + assert gapblock.ball_batch(['AC'],['AC'],[1],query_group_ids=[0],reference_group_ids=[1],group_distances=distances)==[[]] + + +@pytest.mark.parametrize('api',['score_matrix','count_batch','ball_batch']) +@pytest.mark.parametrize('groups',[ + dict(query_group_ids=[0]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[]), + dict(query_group_ids=[],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[-1],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[1],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[0],reference_group_ids=[1],group_distances=[[0]]), + dict(query_group_ids=[0.5],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[False],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[0],[]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[-1]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[1.5]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[2**31-1]]), +]) +def test_group_validation(api,groups): + args=(['AC'],['CA']) + if api=='count_batch':args+=([[10]],) + if api=='ball_batch':args+=([10],) + with pytest.raises((ValueError,TypeError,OverflowError)): + getattr(gapblock,api)(*args,threads=1,**groups) + + +def test_defaults_and_empty_axes(): + q=['AC',''];r=['AC','CA'] + base=np.asarray(gapblock.score_matrix(q,r,threads=1)) + zero=dict(query_group_ids=[0,0],reference_group_ids=[0,0],group_distances=[[0]]) + np.testing.assert_array_equal(base,np.asarray(gapblock.score_matrix(q,r,threads=-1,**zero))) + assert gapblock.count_batch([],r,[],query_group_ids=[],reference_group_ids=[0,0],group_distances=[[5]])==[] + assert gapblock.count_batch(q,[],[[0],[5]],query_group_ids=[0,0],reference_group_ids=[],group_distances=[[5]])==[[0],[0]] + + +@pytest.mark.parametrize('api',['gapblock_matrix','gapblock_count_batch','gapblock_ball_batch']) +@pytest.mark.parametrize('groups',[ + dict(query_group_ids=[],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[-1],reference_group_ids=[0],group_distances=[[0]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[0],[]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[-1]]), + dict(query_group_ids=[0],reference_group_ids=[0],group_distances=[[2**31-1]]), +]) +def test_native_validation_before_workers(api,groups): + from seqtree import _core + args=(['AC'],['CA']) + if api=='gapblock_count_batch':args+=([[10]],) + if api=='gapblock_ball_batch':args+=([10],) + with pytest.raises(ValueError): + getattr(_core,api)(*args,threads=4,**groups) + + +def test_rectangular_group_table(): + opts=dict(query_group_ids=[0],reference_group_ids=[2,0],group_distances=[[7,4,1]]) + observed=np.asarray(gapblock.score_matrix(['AC'],['AC','AC'],threads=1,**opts)) + np.testing.assert_array_equal(observed,[[1,7]]) + + +def test_int32_boundary_group_only_score(): + maximum=2**31-1 + opts=dict(gap_open=0,gap_extend=0,query_group_ids=[0],reference_group_ids=[0],group_distances=[[maximum]],threads=4) + assert gapblock.score_matrix([''],[''],**opts)[0,0]==maximum + assert gapblock.count_batch([''],[''],[[maximum-1,maximum]],**opts)==[[0,1]] + assert [(h.ref_id,h.score) for h in gapblock.ball_batch([''],[''],[maximum],**opts)[0]]==[(0,maximum)] + assert gapblock.count_batch([''],[''],[[maximum]],exclude_exact=True,**opts)==[[0]] From 6aab661b2b51280fbf7c80a7da0c7b6493e90ba9 Mon Sep 17 00:00:00 2001 From: mikessh Date: Sat, 10 Oct 2026 16:36:29 +0300 Subject: [PATCH 2/3] Reduce linear distance mass alongside batched counts --- include/seqtree/seqtree.hpp | 8 +++ python/seqtree/gapblock.py | 14 ++++- src/_bindings.cpp | 6 ++ src/gapblock.cpp | 40 +++++++++++++- tests/python/test_gapblock_linear_mass.py | 67 +++++++++++++++++++++++ 5 files changed, 130 insertions(+), 5 deletions(-) create mode 100644 tests/python/test_gapblock_linear_mass.py diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index f311107..917e9c8 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -265,6 +265,14 @@ std::vector> gapblock_count_batch(const std::vector& query_groups = {}, const std::vector& reference_groups = {}, const std::vector>& group_distances = {}); +std::pair>, std::vector>> gapblock_count_mass_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 std::vector>& weights = {}, + const std::vector& query_groups = {}, + const std::vector& reference_groups = {}, + const std::vector>& group_distances = {}); std::vector> gapblock_paired_count_batch( const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, diff --git a/python/seqtree/gapblock.py b/python/seqtree/gapblock.py index 0b91c26..6e7a84a 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -44,6 +44,7 @@ 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_count_mass_batch as _gapblock_count_mass_batch from ._core import gapblock_ball_batch as _gapblock_ball_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 @@ -675,7 +676,7 @@ def topk_batch(queries, refs, k=10, matrix=None, gap_open=None, gap_extend=1, def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend=1, gap_prior=None, alphabet="aa", threads=0, exclude_exact=False, position_weights_by_length=None, - query_group_ids=None, reference_group_ids=None, group_distances=None): + query_group_ids=None, reference_group_ids=None, group_distances=None, linear_mass=False): """Count all reference scores <= each query's requested integer thresholds. Optional positional weights and additive group distances match :func:`score_matrix`. @@ -685,11 +686,20 @@ def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_exten 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). + With ``linear_mass=True``, return ``(counts, masses)`` from the same traversal; + each mass is ``sum(max(0, threshold - score))`` over nonexcluded references. + Boundary ties still count but contribute zero mass. Scores are nonnegative; + integer mass is uint64, with overflow bounds checked before workers. Negative + and zero thresholds have zero mass. No normalization or statistical model is + applied; default counts and paired APIs remain unchanged. """ + if not isinstance(linear_mass, bool): + raise TypeError("linear_mass must be a boolean") 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, + reduction = _gapblock_count_mass_batch if linear_mass else _gapblock_count_batch + return reduction(q, r, thresholds, alphabet, matrix, go, ge, cube, width, threads, exclude_exact, _position_weights(position_weights_by_length, q, r), *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) diff --git a/src/_bindings.cpp b/src/_bindings.cpp index e91f09d..958aeab 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -818,6 +818,12 @@ NB_MODULE(_core, m) { return gapblock_count_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights, query_groups, reference_groups, group_distances); }, 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, nb::arg("position_weights_by_length")=std::vector>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, "Bounded exhaustive gapblock reduction, GIL released."); + m.def("gapblock_count_mass_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 std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_count_mass_batch(queries, refs, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights, query_groups, reference_groups, group_distances); + }, 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, nb::arg("position_weights_by_length")=std::vector>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, "Counts and integer linear mass in one bounded traversal, GIL released."); + m.def("gapblock_ball_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 std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { const auto alph = parse_alphabet(alphabet); nb::gil_scoped_release release; diff --git a/src/gapblock.cpp b/src/gapblock.cpp index 6ba5d97..9de30e0 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -236,12 +236,23 @@ std::vector> ranked(const GapInput& a, const GapInput* b, ui 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) { + const std::vector& prior, int threads, bool exclude_exact, + std::vector>* masses = nullptr) { 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()); + if (masses) { + masses->resize(a.q.size()); + if (std::any_of(a.pen.begin(), a.pen.end(), [](int32_t cost) { return cost < 0; })) + throw std::invalid_argument("linear mass requires nonnegative scores"); + // Every score admitted to a row is <= its largest positive threshold. + // Bound both score sums and threshold*count before workers start. + for (const auto& row : thresholds) for (auto threshold : row) + if (threshold > 0 && a.r.size() > std::numeric_limits::max() / uint64_t(threshold)) + throw std::overflow_error("gapblock linear mass exceeds uint64"); + } 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) { @@ -250,7 +261,9 @@ std::vector> counted(const GapInput& a, const GapInput* b, 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()); + std::vector score_bins(masses ? sorted.size() : 0); out[i].resize(sorted.size()); + if (masses) (*masses)[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; @@ -259,12 +272,20 @@ std::vector> counted(const GapInput& a, const GapInput* b, 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())]; + if (pos != sorted.end()) { + const auto bin = size_t(pos - sorted.begin()); + ++bins[bin]; + if (masses) score_bins[bin] += uint64_t(score); + } } - uint64_t cumulative = 0; + uint64_t cumulative = 0, score_sum = 0; for (size_t t = 0; t < sorted.size(); ++t) { cumulative += bins[t]; out[i][sorted[t].second] = cumulative; + if (masses) { + score_sum += score_bins[t]; + (*masses)[i][sorted[t].second] = uint64_t(std::max(sorted[t].first, 0)) * cumulative - score_sum; + } } }); return out; @@ -354,6 +375,19 @@ std::vector> gapblock_count_batch(const std::vector>, std::vector>> gapblock_count_mass_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 std::vector>& weights, + const std::vector& query_groups, const std::vector& reference_groups, + const std::vector>& group_distances) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances); + std::vector> masses; + auto counts = counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact, &masses); + return {std::move(counts), std::move(masses)}; +} + std::vector> gapblock_paired_count_batch( const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, diff --git a/tests/python/test_gapblock_linear_mass.py b/tests/python/test_gapblock_linear_mass.py new file mode 100644 index 0000000..4239dfc --- /dev/null +++ b/tests/python/test_gapblock_linear_mass.py @@ -0,0 +1,67 @@ +"""Integer linear taper shares exact counts, scoring and full-identity puncture.""" +import numpy as np +import pytest +from seqtree import gapblock, SubstitutionMatrix + + +def scalar(scores, thresholds, queries, refs, exclude_exact): + counts=[];masses=[] + for i,row in enumerate(thresholds): + eligible=[int(scores[i,j]) for j in range(len(refs)) + if not(exclude_exact and queries[i]==refs[j])] + counts.append([sum(d<=t for d in eligible) for t in row]) + masses.append([sum(max(0,t-d) for d in eligible) for t in row]) + return counts,masses + + +@pytest.mark.parametrize('exclude_exact',[False,True]) +def test_linear_mass_scalar_oracle_and_threads(exclude_exact): + q=['AC','ACC','','AC'];r=['AC','CA','ACC','CA',''] + thresholds=[[-1,0,2,5,20,5],[],[0,3,20],[20,0]] + options=dict(gap_open=3,gap_extend=2,position_weights_by_length={0:[],2:[0,2],3:[3,1,2]}, + query_group_ids=[0,1,1,0],reference_group_ids=[1,0,1,0,0],group_distances=[[7,2],[0,11]]) + scores=np.asarray(gapblock.score_matrix(q,r,threads=1,**options)) + expected=scalar(scores,thresholds,q,r,exclude_exact) + for threads in (1,4): + actual=gapblock.count_batch(q,r,thresholds,threads=threads,exclude_exact=exclude_exact, + linear_mass=True,**options) + assert actual==expected + assert actual[0]==gapblock.count_batch(q,r,thresholds,threads=threads, + exclude_exact=exclude_exact,**options) + + +def test_linear_mass_boundaries_and_uint64(): + assert gapblock.count_batch(['AC'],['AC','CA'],[[0,2,1,-1,2]],linear_mass=True)==( + [[1,2,1,0,2]],[[0,2,1,0,2]]) + maximum=2**31-1 + assert gapblock.count_batch(['AC'],['AC']*4,[[maximum]],linear_mass=True)==([[4]],[[4*maximum]]) + assert gapblock.count_batch(['AC'],['AC'],[[maximum]],exclude_exact=True,linear_mass=True)==([[0]],[[0]]) + assert gapblock.count_batch(['AC'],['AC'],[[5]],linear_mass=True, + query_group_ids=[0],reference_group_ids=[0],group_distances=[[6]])==([[0]],[[0]]) + + +def test_linear_mass_empty_axes_and_prior(): + assert gapblock.count_batch([],['AC'],[],linear_mass=True)==([],[]) + assert gapblock.count_batch(['AC'],[],[[-1,0,2]],linear_mass=True)==([[0,0,0]],[[0,0,0]]) + q=['AC','ACC'];r=['ACC','AC'] + opts=dict(gap_prior=gapblock.positions_prior((1,)),matrix=SubstitutionMatrix.blosum62(),gap_open=28,gap_extend=14) + scores=np.asarray(gapblock.score_matrix(q,r,**opts)) + cut=[[0,28,50],[50,28,0]] + assert gapblock.count_batch(q,r,cut,linear_mass=True,**opts)==scalar(scores,cut,q,r,False) + + +@pytest.mark.parametrize('value',[1,None,'yes']) +def test_linear_mass_requires_explicit_boolean(value): + with pytest.raises((ValueError,TypeError)): + gapblock.count_batch(['AC'],['CA'],[[2]],linear_mass=value) + + +@pytest.mark.parametrize('options',[ + {'gap_open':-1}, {'gap_extend':-1}, + {'query_group_ids':[0],'reference_group_ids':[0],'group_distances':[[-1]]}, + {'position_weights_by_length':[[],[],[-1,1]]}, +]) +def test_linear_mass_native_validation(options): + from seqtree import _core + with pytest.raises((ValueError,TypeError,OverflowError)): + _core.gapblock_count_mass_batch(['AC'],['CA'],[[2]],**options) From 1b08c456088ea619e14068dbccf8ad99396e6ad5 Mon Sep 17 00:00:00 2001 From: mikessh Date: Sat, 10 Oct 2026 16:54:27 +0300 Subject: [PATCH 3/3] Reduce linked pair distance counts and linear mass in native batch --- include/seqtree/seqtree.hpp | 12 ++++ python/seqtree/gapblock.py | 35 ++++++++- src/_bindings.cpp | 6 ++ src/gapblock.cpp | 43 +++++++++-- tests/python/test_gapblock_paired_sum.py | 92 ++++++++++++++++++++++++ 5 files changed, 181 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_gapblock_paired_sum.py diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index 917e9c8..1912738 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -279,6 +279,18 @@ std::vector> gapblock_paired_count_batch( 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::pair>, std::vector>> gapblock_paired_sum_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 std::vector>& weights_alpha, + const std::vector>& weights_beta, + const std::vector& query_groups_alpha, const std::vector& query_groups_beta, + const std::vector& reference_groups_alpha, const std::vector& reference_groups_beta, + const std::vector>& group_distances_alpha, + const std::vector>& group_distances_beta); // --------------------------------------------------------------------------------------- // Plain string edit distances on raw characters, unit costs. Deliberately alphabet-agnostic: diff --git a/python/seqtree/gapblock.py b/python/seqtree/gapblock.py index 6e7a84a..f153baf 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -48,9 +48,10 @@ from ._core import gapblock_ball_batch as _gapblock_ball_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 +from ._core import gapblock_paired_sum_count_batch as _gapblock_paired_sum_count_batch __all__ = [ - "gapblock_score", "score_matrix", "ball_batch", "topk_batch", "count_batch", "paired_topk_batch", "paired_count_batch", "deletion_variants", "central_prior", "profile_prior", + "gapblock_score", "score_matrix", "ball_batch", "topk_batch", "count_batch", "paired_topk_batch", "paired_count_batch", "paired_sum_count_batch", "deletion_variants", "central_prior", "profile_prior", "frame_prior", "positions_prior", "embed_in_frame", "gap_cost", "GapBlockIndex", "ScoreMatrix", "IslandProfile", ] @@ -762,6 +763,38 @@ def paired_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, matrix= return _gapblock_paired_count_batch(qa, qb, ra, rb, thresholds, alphabet, matrix, go, ge, cube, width, threads, exclude_exact) +def paired_sum_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, + position_weights_by_length_alpha=None, position_weights_by_length_beta=None, + query_group_ids_alpha=None, query_group_ids_beta=None, + reference_group_ids_alpha=None, reference_group_ids_beta=None, + group_distances_alpha=None, group_distances_beta=None): + """Return (counts, linear masses) for linked rows under summed lane distances. + + Shared matrix/gap costs/prior and separate full-coordinate lane weights and + group tables use the existing kernel. Sum distance is computed in int64; + thresholds remain int32. Count every score <= cutoff and sum max(0, cutoff + - score); ties count with zero mass. Negative cutoffs return zero. Preserve + duplicate reference observations and input threshold order. Exact exclusion + removes a row only when both encoded full strings match. One GIL-released + native traversal, O(Q*S) outputs and O(max_length + max_threshold_row) worker + scratch; no matrix/edges. The existing paired-max API is unchanged. + """ + 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) + ag = _group_options(query_group_ids_alpha, reference_group_ids_alpha, group_distances_alpha, qa, ra) + bg = _group_options(query_group_ids_beta, reference_group_ids_beta, group_distances_beta, qb, rb) + return _gapblock_paired_sum_count_batch(qa, qb, ra, rb, thresholds, alphabet, matrix, go, ge, + cube, width, threads, exclude_exact, + _position_weights(position_weights_by_length_alpha, qa, ra), + _position_weights(position_weights_by_length_beta, qb, rb), + ag[0], bg[0], ag[1], bg[1], ag[2], bg[2]) + + class GapBlockIndex: """Search a reference set under the single-gap-block model, reusing the Hamming engine. diff --git a/src/_bindings.cpp b/src/_bindings.cpp index 958aeab..eed1f0a 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -842,6 +842,12 @@ NB_MODULE(_core, m) { 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("gapblock_paired_sum_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 std::vector>& weights_alpha, const std::vector>& weights_beta, const std::vector& query_groups_alpha, const std::vector& query_groups_beta, const std::vector& reference_groups_alpha, const std::vector& reference_groups_beta, const std::vector>& group_distances_alpha, const std::vector>& group_distances_beta) { + const auto alph = parse_alphabet(alphabet); + nb::gil_scoped_release release; + return gapblock_paired_sum_count_batch(qa, qb, ra, rb, thresholds, alph, matrix ? &*matrix : nullptr, go, ge, prior, width, threads, exclude_exact, weights_alpha, weights_beta, query_groups_alpha, query_groups_beta, reference_groups_alpha, reference_groups_beta, group_distances_alpha, group_distances_beta); + }, 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, nb::arg("position_weights_by_length_alpha")=std::vector>{}, nb::arg("position_weights_by_length_beta")=std::vector>{}, nb::arg("query_group_ids_alpha")=std::vector{}, nb::arg("query_group_ids_beta")=std::vector{}, nb::arg("reference_group_ids_alpha")=std::vector{}, nb::arg("reference_group_ids_beta")=std::vector{}, nb::arg("group_distances_alpha")=std::vector>{}, nb::arg("group_distances_beta")=std::vector>{}, "Linked sum counts and linear mass in one bounded traversal, 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."); diff --git a/src/gapblock.cpp b/src/gapblock.cpp index 9de30e0..b993f68 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -237,15 +237,18 @@ std::vector> ranked(const GapInput& a, const GapInput* b, ui 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, - std::vector>* masses = nullptr) { + std::vector>* masses = nullptr, bool sum_pair_distance = false) { 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"); + if (sum_pair_distance && !b) + throw std::invalid_argument("sum distance requires two linked views"); std::vector> out(a.q.size()); if (masses) { masses->resize(a.q.size()); - if (std::any_of(a.pen.begin(), a.pen.end(), [](int32_t cost) { return cost < 0; })) + if (std::any_of(a.pen.begin(), a.pen.end(), [](int32_t cost) { return cost < 0; }) || + (b && std::any_of(b->pen.begin(), b->pen.end(), [](int32_t cost) { return cost < 0; }))) throw std::invalid_argument("linear mass requires nonnegative scores"); // Every score admitted to a row is <= its largest positive threshold. // Bound both score sums and threshold*count before workers start. @@ -267,11 +270,18 @@ std::vector> counted(const GapInput& a, const GapInput* b, 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; - if (group_offset(a, i, j) > sorted.back().first) 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 int64_t offset = int64_t(group_offset(a, i, j)) + + (sum_pair_distance ? int64_t(group_offset(*b, i, j)) : 0); + if (offset > sorted.back().first) continue; + int64_t score = score_cell(a, i, j, go, ge, prior, scratch.data()); + if (sum_pair_distance) { + if (score > sorted.back().first) continue; + score += int64_t(score_cell(*b, i, j, go, ge, prior, scratch.data())); + } else 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; }); + [](const auto& threshold, int64_t value) { return threshold.first < value; }); if (pos != sorted.end()) { const auto bin = size_t(pos - sorted.begin()); ++bins[bin]; @@ -399,4 +409,25 @@ std::vector> gapblock_paired_count_batch( return counted(a, &b, thresholds, go, ge, prior, threads, exclude_exact); } +std::pair>, std::vector>> gapblock_paired_sum_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 std::vector>& weights_alpha, + const std::vector>& weights_beta, + const std::vector& query_groups_alpha, const std::vector& query_groups_beta, + const std::vector& reference_groups_alpha, const std::vector& reference_groups_beta, + const std::vector>& group_distances_alpha, + const std::vector>& group_distances_beta) { + const auto a = prepare(qa, ra, alphabet, matrix, go, ge, prior, width, threads, + weights_alpha, query_groups_alpha, reference_groups_alpha, group_distances_alpha); + const auto b = prepare(qb, rb, alphabet, matrix, go, ge, prior, width, threads, + weights_beta, query_groups_beta, reference_groups_beta, group_distances_beta); + std::vector> masses; + auto counts = counted(a, &b, thresholds, go, ge, prior, threads, exclude_exact, &masses, true); + return {std::move(counts), std::move(masses)}; +} + } // namespace seqtree diff --git a/tests/python/test_gapblock_paired_sum.py b/tests/python/test_gapblock_paired_sum.py new file mode 100644 index 0000000..eb52884 --- /dev/null +++ b/tests/python/test_gapblock_paired_sum.py @@ -0,0 +1,92 @@ +"""Linked sum reduction agrees with dense lane sums without int32 truncation.""" +import numpy as np +import pytest +from seqtree import gapblock + + +@pytest.mark.parametrize('exclude_exact',[False,True]) +def test_paired_sum_dense_oracle_and_threads(exclude_exact): + qa=['AC','ACC','','AC'];qb=['CA','AC','AC','CA'] + ra=['AC','AC','ACC','','CA'];rb=['CA','AC','AC','AC','CA'] + cut=[[-1,0,3,20,5,20],[],[0,3,20],[20,0]] + common=dict(gap_open=3,gap_extend=2,gap_prior=gapblock.positions_prior((1,))) + alpha=dict(position_weights_by_length={0:[],2:[0,2],3:[3,1,2]}, + query_group_ids=[0,1,1,0],reference_group_ids=[1,0,1,0,0],group_distances=[[7,2],[0,11]]) + beta=dict(position_weights_by_length={2:[1,3]}, + query_group_ids=[1,0,1,1],reference_group_ids=[1,0,0,1,0],group_distances=[[0,6],[4,0]]) + A=np.asarray(gapblock.score_matrix(qa,ra,**common,**alpha)).astype(np.int64) + B=np.asarray(gapblock.score_matrix(qb,rb,**common,**beta)).astype(np.int64) + expected_counts=[];expected_mass=[] + for i,row in enumerate(cut): + scores=[int(A[i,j]+B[i,j]) for j in range(len(ra)) + if not(exclude_exact and qa[i]==ra[j] and qb[i]==rb[j])] + expected_counts.append([sum(d<=t for d in scores) for t in row]) + expected_mass.append([sum(max(0,t-d) for d in scores) for t in row]) + opts={**common,**{key+'_alpha':value for key,value in alpha.items()}, + **{key+'_beta':value for key,value in beta.items()}} + for threads in (1,4): + assert gapblock.paired_sum_count_batch(qa,qb,ra,rb,cut,threads=threads, + exclude_exact=exclude_exact,**opts)==(expected_counts,expected_mass) + + +def test_paired_sum_full_identity_puncture_and_ties(): + # Alpha alone matches every reference; only the first linked full pair is exact. + assert gapblock.paired_sum_count_batch(['AC'],['CA'],['AC','AC'],['CA','AC'],[[0,2,3]], + exclude_exact=True)==([[0,1,1]],[[0,0,1]]) + # Existing max contract remains different, unchanged. + assert gapblock.paired_count_batch(['AC'],['AC'],['CA'],['CA'],[[2]])==[[1]] + assert gapblock.paired_sum_count_batch(['AC'],['AC'],['CA'],['CA'],[[2]])==([[0]],[[0]]) + + +def test_paired_sum_int32_overflow_and_empty_axes(): + maximum=2**31-1 + groups=dict(query_group_ids_alpha=[0],reference_group_ids_alpha=[0],group_distances_alpha=[[maximum]], + query_group_ids_beta=[0],reference_group_ids_beta=[0],group_distances_beta=[[maximum]]) + assert gapblock.paired_sum_count_batch([''],[''],[''],[''],[[-1,0,maximum]], + gap_open=0,gap_extend=0,**groups)==([[0,0,0]],[[0,0,0]]) + groups['group_distances_beta']=[[0]] + assert gapblock.paired_sum_count_batch([''],[''],[''],[''],[[maximum]], + gap_open=0,gap_extend=0,**groups)==([[1]],[[0]]) + assert gapblock.paired_sum_count_batch([],[],['AC'],['CA'],[])==([],[]) + assert gapblock.paired_sum_count_batch(['AC'],['CA'],[],[],[[0,3]])==([[0,0]],[[0,0]]) + + +@pytest.mark.parametrize('args,opts',[ + ((['AC'],[],['AC'],['CA'],[[2]]),{}), + ((['AC'],['CA'],['AC'],[],[[2]]),{}), + ((['AC'],['CA'],['AC'],['CA'],[]),{}), + ((['AC'],['CA'],['AC'],['CA'],[[2]]),{'position_weights_by_length_beta':{2:[-1,1]}}), + ((['AC'],['CA'],['AC'],['CA'],[[2]]),{'query_group_ids_beta':[0]}), +]) +def test_paired_sum_validation(args,opts): + with pytest.raises((ValueError,TypeError)): + gapblock.paired_sum_count_batch(*args,**opts) + + +@pytest.mark.parametrize('args,opts',[ + ((['AC'],[],['AC'],['CA'],[[2]]),{}), + ((['AC'],['CA'],['AC'],['CA'],[[2]]),{'gap_open':-1}), + ((['AC'],['CA'],['AC'],['CA'],[[2]]),{'query_group_ids_beta':[0]}), + ((['AC'],['CA'],['AC'],['CA'],[[2]]),{'position_weights_by_length_beta':[[],[],[-1,1]]}), +]) +def test_paired_sum_direct_native_validation(args,opts): + from seqtree import _core + with pytest.raises((ValueError,TypeError)): + _core.gapblock_paired_sum_count_batch(*args,**opts) + + +def test_paired_sum_mass_uint64_and_puncture_with_zero_flanks(): + maximum=2**31-1 + assert gapblock.paired_sum_count_batch(['AC'],['CA'],['AC']*4,['CA']*4,[[maximum]])==( + [[4]],[[4*maximum]]) + # Identical weighted cores are not identical full strings in both lanes. + assert gapblock.paired_sum_count_batch(['AC'],['CA'],['CC'],['CA'],[[1]], + position_weights_by_length_alpha={2:[0,1]},exclude_exact=True)==([[1]],[[1]]) + + +def test_paired_sum_cell_addition_is_int64(): + # No group bound or first-lane bound rejects: only the actual lane sum exceeds int32. + weight=2**30 + options=dict(gap_open=0,gap_extend=0,position_weights_by_length_alpha={1:[weight]}, + position_weights_by_length_beta={1:[weight]}) + assert gapblock.paired_sum_count_batch(['A'],['A'],['C'],['C'],[[2**31-1]],**options)==([[0]],[[0]])