diff --git a/docs/gapblock.rst b/docs/gapblock.rst index 1cb648c..a1a371c 100644 --- a/docs/gapblock.rst +++ b/docs/gapblock.rst @@ -355,3 +355,24 @@ statistical model or calibration. ``Index.ref_seqs()`` retrieves all stored strings in reference-ID order through one native call, preserving duplicate rows. + +Positional weights in bounded batches +------------------------------------- + +``score_matrix``, ``count_batch`` and ``ball_batch`` accept +``position_weights_by_length``: a mapping from full sequence length to a vector +of nonnegative integer weights of that width. Supply every input length. +The longer sequence supplies the frame for each pair; prefix positions use their +original weights and the shifted suffix uses the corresponding longer-frame +weights. Equal-length scoring is weighted Hamming scoring. This pairwise frame +does not imply a consistent embedding of a whole collection. Gap charges and +position priors remain explicit and must use the same units as the weighted +substitution penalties. Invalid widths, fractional weights and possible int32 +overflow are rejected. Omitting the mapping preserves ordinary scoring. + +``ball_batch`` takes one integer cutoff per query and returns all accepted +reference rows, including ties, in reference order. Exact exclusion and duplicate +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. diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index f812ed5..ffd38c9 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -221,11 +221,21 @@ std::vector align_dist_matrix(const std::vector& queries, // substitution cost plus `prior`. `prior` is the gap prior flattened to [m][d][i] with stride // (prior_width+1); pass it empty for no prior. It is read only when d > 0, which is what keeps // 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. 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); + uint32_t prior_width, int threads = 0, + const std::vector>& weights = {}); + +// 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 = {}); // Exhaustive bounded reductions using the same gapblock cell. Top-k ties use // (penalty, ref_id); paired penalty is max(lane scores). Exact exclusion uses @@ -242,7 +252,8 @@ std::vector>> gapblock_paired 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 std::vector& prior, uint32_t width, int threads, bool exclude_exact, + const std::vector>& weights = {}); 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 913ce64..4729643 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -44,11 +44,12 @@ 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_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 __all__ = [ - "gapblock_score", "score_matrix", "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", "deletion_variants", "central_prior", "profile_prior", "frame_prior", "positions_prior", "embed_in_frame", "gap_cost", "GapBlockIndex", "ScoreMatrix", "IslandProfile", ] @@ -522,6 +523,7 @@ def score_matrix( gap_prior: GapPrior | None = None, alphabet: str = "aa", threads: int = 0, + position_weights_by_length=None, ) -> ScoreMatrix: """Gap-block penalty of every query against every reference, in C++ with the GIL released. @@ -533,6 +535,9 @@ def score_matrix( Args: queries: Query sequences (the matrix rows). refs: Reference sequences (the matrix columns). + 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. 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. @@ -565,7 +570,8 @@ def score_matrix( 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 [] - return _gapblock_matrix(q, r, alphabet, matrix, gap_open, gap_extend, cube, width, threads) + return _gapblock_matrix(q, r, alphabet, matrix, gap_open, gap_extend, cube, width, threads, + _position_weights(position_weights_by_length, q, r)) def _batch_integer(value, name, lower=0, upper=(1 << 31) - 1): @@ -590,6 +596,28 @@ def _batch_options(axes, matrix, gap_open, gap_extend, gap_prior, threads): return go, ge, cube, width, threads +def _position_weights(weights, *axes): + """Validated integer tables in the longer-junction frame; never round/truncate.""" + if weights is None: + return [] + if not hasattr(weights, "items"): + raise ValueError("position_weights_by_length must map lengths to integer weights") + tables = {} + for length, row in weights.items(): + length = _batch_integer(length, "position frame", upper=65535) + row = list(row) + if len(row) != length: + raise ValueError("positional weight length must match its frame") + tables[length] = [_batch_integer(w, "position weight") for w in row] + if not tables: + return [] + for axis in axes: + for sequence in axis: + if len(sequence) and len(sequence) not in tables: + raise ValueError("missing positional weights for an input frame") + return [tables.get(length, []) for length in range(max(tables) + 1)] + + def _threshold_rows(thresholds, n_queries): rows = [list(row) for row in thresholds] if len(rows) != n_queries: @@ -622,9 +650,10 @@ 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): + gap_prior=None, alphabet="aa", threads=0, exclude_exact=False, position_weights_by_length=None): """Count all reference scores <= each query's requested integer thresholds. + Optional positional weights index the longer full sequence, exactly as in score_matrix. ``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 @@ -634,7 +663,26 @@ def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_exten 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) + return _gapblock_count_batch(q, r, thresholds, alphabet, matrix, go, ge, cube, width, threads, exclude_exact, + _position_weights(position_weights_by_length, 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): + """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. + This reduces output memory; it is not a trie-based candidate filter. + """ + q, r = list(queries), list(refs) + thresholds = [_batch_integer(v, "threshold", -(1 << 31)) for v in thresholds] + if len(thresholds) != len(q): + 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)) 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 0bfd965..4c9462b 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -360,7 +360,8 @@ ScoreMatrix py_gapblock_matrix(const std::vector& queries, const std::vector& refs, const std::string& alphabet, const std::optional& matrix, int32_t gap_open, int32_t gap_extend, const std::vector& prior, - uint32_t prior_width, int threads) { + uint32_t prior_width, int threads, + const std::vector>& weights) { Alphabet alph = parse_alphabet(alphabet); ScoreMatrix out; out.rows = queries.size(); @@ -368,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); + gap_extend, prior, prior_width, threads, weights); } return out; } @@ -799,7 +800,7 @@ NB_MODULE(_core, m) { m.def("gapblock_matrix", &py_gapblock_matrix, nb::arg("queries"), nb::arg("refs"), 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("prior_width") = 0, nb::arg("threads") = 0, nb::arg("position_weights_by_length") = 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."); @@ -810,11 +811,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) { + 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 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."); + 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."); + + 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 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."); 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 51a003c..839621c 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -38,7 +38,7 @@ std::vector pen_table(const SubstitutionMatrix* m, uint8_t A) { // contractually zero -- hence the early return, which also keeps s(q, q) == 0. inline int32_t cell(const uint8_t* q, uint32_t m, const uint8_t* r, uint32_t n, const int32_t* pen, uint8_t A, int32_t gap_open, int32_t gap_extend, - const int32_t* prior_row, int32_t* suf) { + const int32_t* prior_row, int32_t* suf, const int32_t* weights = nullptr) { const uint32_t L = std::min(m, n); const uint32_t d = (m > n ? m - n : n - m); const bool q_longer = m >= n; @@ -47,7 +47,7 @@ inline int32_t cell(const uint8_t* q, uint32_t m, const uint8_t* r, uint32_t n, for (uint32_t j = L; j-- > 0;) { const uint8_t a = q_longer ? q[j + d] : q[j]; const uint8_t b = q_longer ? r[j] : r[j + d]; - suf[j] = suf[j + 1] + pen[size_t(a) * A + b]; + suf[j] = suf[j + 1] + pen[size_t(a) * A + b] * (weights ? weights[j + d] : 1); } if (d == 0) return suf[0]; @@ -56,7 +56,7 @@ inline int32_t cell(const uint8_t* q, uint32_t m, const uint8_t* r, uint32_t n, for (uint32_t i = 0; i <= L; ++i) { const int32_t cand = pre + suf[i] + (prior_row ? prior_row[i] : 0); if (cand < best) best = cand; - if (i < L) pre += pen[size_t(q[i]) * A + r[i]]; + if (i < L) pre += pen[size_t(q[i]) * A + r[i]] * (weights ? weights[i] : 1); } return best + gap_open + int32_t(d - 1) * gap_extend; } @@ -66,6 +66,7 @@ inline int32_t cell(const uint8_t* q, uint32_t m, const uint8_t* r, uint32_t n, struct GapInput { std::vector> q, r; std::vector pen; + std::vector> weights; uint32_t longest = 0; uint8_t A = 0; size_t W1 = 0; @@ -73,7 +74,8 @@ 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& prior, uint32_t width, int threads, + const std::vector>& weights = {}) { 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()) @@ -109,10 +111,31 @@ GapInput prepare(const std::vector& queries, const std::vector= weights.size() || weights[length].size() != length)) + throw std::invalid_argument("missing positional weights for an input frame"); + } + } 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) + + const uint64_t weighted_pen = uint64_t(uint32_t(largest_pen)) * uint32_t(largest_weight); + if (in.longest && weighted_pen > 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) + uint64_t(in.longest ? in.longest - 1 : 0) * uint32_t(ge); if (bound > std::numeric_limits::max()) @@ -125,7 +148,7 @@ int32_t score_cell(const GapInput& in, size_t i, size_t j, int32_t go, int32_t g 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); + 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()); } struct Ranked { @@ -215,60 +238,40 @@ std::vector> counted(const GapInput& a, const GapInput* b, } // namespace std::vector gapblock_matrix(const std::vector& queries, - const std::vector& refs, Alphabet alphabet, - const SubstitutionMatrix* matrix, int32_t gap_open, - int32_t gap_extend, const std::vector& prior, - uint32_t prior_width, int threads) { - if (gap_open < 0 || gap_extend < 0) - throw std::invalid_argument("gap_open and gap_extend must be >= 0"); - - const Codec codec(alphabet); - const uint8_t A = codec.size(); - if (matrix && matrix->size() != A) - throw std::invalid_argument("matrix size does not match the alphabet"); - - const size_t N = queries.size(), K = refs.size(); - std::vector out(N * K); - if (N == 0 || K == 0) return out; - - const std::vector pen = pen_table(matrix, A); - - std::vector> qc(N), rc(K); - uint32_t longest = 0; - for (size_t i = 0; i < N; ++i) { - qc[i] = encode(codec, queries[i], "queries", i); - longest = std::max(longest, uint32_t(qc[i].size())); - } - for (size_t k = 0; k < K; ++k) { - rc[k] = encode(codec, refs[k], "refs", k); - longest = std::max(longest, uint32_t(rc[k].size())); - } - - const size_t W1 = size_t(prior_width) + 1; - if (!prior.empty()) { - if (prior.size() != W1 * W1 * W1) - throw std::invalid_argument("prior table must have (prior_width + 1)^3 entries"); - if (longest > prior_width) - throw std::invalid_argument("a sequence is longer than the prior table's width"); - } - const int32_t* P = prior.empty() ? nullptr : prior.data(); + 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); + 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()); + parallel_for(queries.size(), threads, [&] { return std::vector(size_t(in.longest) + 1); }, + [&](size_t i, std::vector& scratch) { + for (size_t j = 0; j < refs.size(); ++j) + out[i * refs.size() + j] = score_cell(in, i, j, go, ge, prior, scratch.data()); + }); + return out; +} - // Every symbol was validated above and the kernel is arithmetic, so no worker throws -- - // parallel_for's exception plumbing simply never fires here. `suf` is the per-worker - // suffix-score scratch, allocated once per thread rather than once per row. - parallel_for(N, threads, [&] { return std::vector(size_t(longest) + 1); }, - [&](size_t i, std::vector& suf) { - const uint8_t* q = qc[i].data(); - const uint32_t m = uint32_t(qc[i].size()); - int32_t* row = out.data() + i * K; - for (size_t k = 0; k < K; ++k) { - const uint32_t n = uint32_t(rc[k].size()); - const uint32_t M = std::max(m, n), d = (m > n ? m - n : n - m); - const int32_t* prow = P ? P + (size_t(M) * W1 + d) * W1 : nullptr; - row[k] = cell(q, m, rc[k].data(), n, pen.data(), A, gap_open, gap_extend, - prow, suf.data()); - } - }); +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) { + 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); + 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; + 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}); + } + }); return out; } @@ -303,8 +306,9 @@ std::vector>> gapblock_paired 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); + 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); return counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact); } diff --git a/tests/python/test_gapblock_position_weights.py b/tests/python/test_gapblock_position_weights.py new file mode 100644 index 0000000..4ec81c1 --- /dev/null +++ b/tests/python/test_gapblock_position_weights.py @@ -0,0 +1,92 @@ +"""Synthetic correctness checks for the queued positional single-block API only.""" +import numpy as np +import pytest +from seqtree import Index,SearchParams,SubstitutionMatrix,PositionalMatrix,gapblock +B=SubstitutionMatrix.blosum62() + +def exhaustive(q,r,weights,go=37,ge=11,prior=None): + length=min(len(q),len(r));longer=max(len(q),len(r));d=longer-length + scores=[] + for start in range(length+1): + value=(go+(d-1)*ge if d else 0)+(prior(start,d,longer) if prior and d else 0) + for j in range(length): + qi=j+(d if len(q)>len(r) and j>=start else 0) + ri=j+(d if len(r)>len(q) and j>=start else 0) + value+=weights[longer][max(qi,ri)]*B.penalty(q[qi],r[ri]) + scores.append(value) + return min(scores) + +@pytest.mark.parametrize('prior',[None,lambda i,d,m:3*abs(i-2)]) +def test_positions_symmetry_reducers(prior): + queries=['','A','AC','ACA','CACA','ACAAA'] + refs=queries+['ACA'] + weights={n:[1+j*17 for j in range(n)] for n in range(6)} + options=dict(matrix=B,gap_open=37,gap_extend=11,gap_prior=prior,position_weights_by_length=weights) + expected=np.array([[exhaustive(q,r,weights,prior=prior) for r in refs] for q in queries]) + for threads in (1,4): + observed=np.asarray(gapblock.score_matrix(queries,refs,threads=threads,**options)) + np.testing.assert_array_equal(observed,expected) + np.testing.assert_array_equal(np.asarray(gapblock.score_matrix(refs,queries,threads=threads,**options)).T,expected) + thresholds=[[-1,0,37,100,500,100,10000] for q in queries] + counts=gapblock.count_batch(queries,refs,thresholds,threads=threads,exclude_exact=True,**options) + balls=gapblock.ball_batch(queries,refs,[100]*len(queries),threads=threads,exclude_exact=True,**options) + for i,q in enumerate(queries): + assert counts[i]==[sum(q!=r and expected[i,j]<=cut for j,r in enumerate(refs)) for cut in thresholds[i]] + assert [(h.ref_id,h.score) for h in balls[i]]==[(j,int(expected[i,j])) for j,r in enumerate(refs) if q!=r and expected[i,j]<=100] + assert all(h.n_subs==h.n_ins==h.n_dels==0 for h in balls[i]) + + +def test_equal_length_positional_matrix(): + queries=['CASSLGQAYEQYF','CASALGQAYEQYF'] + refs=['CASSLGQAYEQYF','AASSLGQAYEQYF','CASSAGQAYEQYF','CASSLGQAYEQAF'] + length=len(queries[0]);weights={length:[117,62,56,85,117,156,153,152,98,62,71,74,98]} + p=SearchParams(max_subs=5,max_ins=0,max_dels=0,max_total_edits=5,engine='seqtm');p.pos_matrix=PositionalMatrix.from_weights(B,weights[length]) + hits=Index.build(refs,'aa').search_batch(queries,p,threads=1) + observed=np.asarray(gapblock.score_matrix(queries,refs,matrix=B,position_weights_by_length=weights,threads=1)) + for i,row in enumerate(hits): + for hit in row: + assert observed[i,hit.ref_id]==hit.score + +@pytest.mark.parametrize('weights',[{2:[1]},{2:[1,-1]},{2:[1,1.5]},{2:[True,1]},{2:[1,2**31]},{3:[1,1,1]}, {2:[2**31-1]*2}]) +@pytest.mark.parametrize('api',['score_matrix','count_batch','ball_batch']) +def test_validation(weights,api): + args=(['AA'],['AC']) + if api=='count_batch':args+=([[100]],) + if api=='ball_batch':args+=([100],) + with pytest.raises((ValueError,TypeError,OverflowError)): + getattr(gapblock,api)(*args,matrix=B,position_weights_by_length=weights,threads=1) + + +def test_defaults_zero_weight_identity_duplicates_empty(): + queries=['AA'];refs=['AA','AC','AC'] + plain=np.asarray(gapblock.score_matrix(queries,refs,matrix=B,threads=1)) + np.testing.assert_array_equal(plain,np.asarray(gapblock.score_matrix(queries,refs,matrix=B,position_weights_by_length={},threads=4))) + assert gapblock.count_batch(queries,refs,[[0]],matrix=B,position_weights_by_length={2:[0,0]},exclude_exact=True,threads=1)==[[2]] + assert gapblock.count_batch([],refs,[],matrix=B,threads=1)==[] + assert gapblock.count_batch(queries,[],[[0]],matrix=B,threads=1)==[[0]] + assert gapblock.ball_batch(queries,refs,[-1],matrix=B,threads=1)==[[]] + with pytest.raises(ValueError):gapblock.ball_batch(queries,refs,[],threads=1) + + +def test_matrix_negative_threads_preserves_automatic_concurrency(): + q=['AC','CA'];r=['AA','AC','CA'] + for weights in (None,{2:[3,7]}): + args=dict(matrix=B,position_weights_by_length=weights) + expected=np.asarray(gapblock.score_matrix(q,r,threads=1,**args)) + np.testing.assert_array_equal(expected,np.asarray(gapblock.score_matrix(q,r,threads=-1,**args))) + + +@pytest.mark.parametrize('api', ['score_matrix', 'count_batch', 'ball_batch']) +def test_zero_weights_accept_large_safe_matrix_penalties(api): + from seqtree import amino_acids + width = len(amino_acids()) + matrix = SubstitutionMatrix.from_similarity( + [[0 if i == j else -600_000_000 for j in range(width)] for i in range(width)]) + kwargs = dict(matrix=matrix, gap_open=0, gap_extend=0, + position_weights_by_length={2: [0, 0]}, threads=1) + if api == 'score_matrix': + assert np.asarray(gapblock.score_matrix(['AA'], ['AC'], **kwargs)).tolist() == [[0]] + elif api == 'count_batch': + assert gapblock.count_batch(['AA'], ['AC'], [[0]], **kwargs) == [[1]] + else: + assert [(h.ref_id, h.score) for h in gapblock.ball_batch(['AA'], ['AC'], [0], **kwargs)[0]] == [(0, 0)]