From d557ed0e0dff833929017d94a7549864478442db Mon Sep 17 00:00:00 2001 From: mikessh Date: Sat, 10 Oct 2026 21:20:44 +0300 Subject: [PATCH] Support symmetric per-input positional weights in native reductions Co-Authored-By: Codex --- docs/gapblock.rst | 14 +++++ include/seqtree/seqtree.hpp | 25 ++++++-- python/seqtree/gapblock.py | 46 +++++++++++--- src/_bindings.cpp | 30 ++++----- src/gapblock.cpp | 63 ++++++++++++++----- .../python/test_gapblock_sequence_weights.py | 54 ++++++++++++++++ 6 files changed, 188 insertions(+), 44 deletions(-) create mode 100644 tests/python/test_gapblock_sequence_weights.py diff --git a/docs/gapblock.rst b/docs/gapblock.rst index 01da88c..3630a60 100644 --- a/docs/gapblock.rst +++ b/docs/gapblock.rst @@ -400,3 +400,17 @@ 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. + +Per-input positional geometry +----------------------------- + +``score_matrix``, ``ball_batch`` and ``count_batch`` accept +``query_position_weights`` and ``reference_position_weights``: nonnegative integer +rows matching each full input sequence. Each aligned substitution uses the larger +of its query and reference weights. This rule is symmetric when the substitution +matrix is symmetric; gap costs and exact identity exclusion remain unchanged. +Supply both axes and do not combine them with ``position_weights_by_length``. + +``paired_sum_count_batch`` accepts separate ``_alpha`` and ``_beta`` versions of +these arguments. Each lane retains its own coordinates; distances sum against the +same linked reference row. These are geometry options, not confidence estimates. diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index 1912738..dd7d6c4 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -223,6 +223,9 @@ 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. +// Alternatively, per-input rows index each sequence's full coordinates. Aligned +// substitution costs use max(query weight, reference weight); gaps are unchanged. +// Both axes are required; these rows cannot combine with length profiles. // 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, @@ -233,7 +236,9 @@ std::vector gapblock_matrix(const std::vector& queries, const std::vector>& weights = {}, const std::vector& query_groups = {}, const std::vector& reference_groups = {}, - const std::vector>& group_distances = {}); + const std::vector>& group_distances = {}, + const std::vector>& query_weights = {}, + const std::vector>& reference_weights = {}); // Exhaustive ball rows in reference order; hit edit fields are unset (zero). std::vector> gapblock_ball_batch(const std::vector& queries, @@ -243,7 +248,9 @@ std::vector> gapblock_ball_batch(const std::vector const std::vector>& weights = {}, const std::vector& query_groups = {}, const std::vector& reference_groups = {}, - const std::vector>& group_distances = {}); + const std::vector>& group_distances = {}, + const std::vector>& query_weights = {}, + const std::vector>& reference_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 @@ -264,7 +271,9 @@ std::vector> gapblock_count_batch(const std::vector>& weights = {}, const std::vector& query_groups = {}, const std::vector& reference_groups = {}, - const std::vector>& group_distances = {}); + const std::vector>& group_distances = {}, + const std::vector>& query_weights = {}, + const std::vector>& reference_weights = {}); 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, @@ -272,7 +281,9 @@ std::pair>, std::vector> const std::vector>& weights = {}, const std::vector& query_groups = {}, const std::vector& reference_groups = {}, - const std::vector>& group_distances = {}); + const std::vector>& group_distances = {}, + const std::vector>& query_weights = {}, + const std::vector>& reference_weights = {}); std::vector> gapblock_paired_count_batch( const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, @@ -290,7 +301,11 @@ std::pair>, std::vector> 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 std::vector>& group_distances_beta, + const std::vector>& query_weights_alpha = {}, + const std::vector>& query_weights_beta = {}, + const std::vector>& reference_weights_alpha = {}, + const std::vector>& reference_weights_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 f153baf..35ccf73 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -527,6 +527,7 @@ def score_matrix( threads: int = 0, position_weights_by_length=None, query_group_ids=None, reference_group_ids=None, group_distances=None, + query_position_weights=None, reference_position_weights=None, ) -> ScoreMatrix: """Gap-block penalty of every query against every reference, in C++ with the GIL released. @@ -541,6 +542,11 @@ 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_position_weights: Optional nonnegative integer rows, one per query and residue. + Supply with reference_position_weights; incompatible with length profiles. + Substitutions use max(query_weight, reference_weight) at aligned coordinates. + Gap costs remain unchanged. + reference_position_weights: Per-reference rows with the same contract. 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. @@ -579,7 +585,8 @@ def score_matrix( 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), - *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r), + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r)) def _batch_integer(value, name, lower=0, upper=(1 << 31) - 1): @@ -644,6 +651,22 @@ def _position_weights(weights, *axes): return [tables.get(length, []) for length in range(max(tables) + 1)] +def _sequence_weights(query_weights, reference_weights, by_length, queries, refs): + """Integer full-coordinate weights; aligned substitutions use the larger weight.""" + if query_weights is None and reference_weights is None: + return [], [] + if query_weights is None or reference_weights is None: + raise ValueError("both query and reference positional weights are required") + if by_length: + raise ValueError("length and per-input weights cannot be combined") + def validate(rows, axis): + rows = [list(row) for row in rows] + if len(rows) != len(axis) or any(len(row) != len(seq) for row, seq in zip(rows, axis)): + raise ValueError("per-input weights must match their sequence axis and lengths") + return [[_batch_integer(value, "position weight") for value in row] for row in rows] + return validate(query_weights, queries), validate(reference_weights, refs) + + def _threshold_rows(thresholds, n_queries): rows = [list(row) for row in thresholds] if len(rows) != n_queries: @@ -677,7 +700,8 @@ 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, linear_mass=False): + query_group_ids=None, reference_group_ids=None, group_distances=None, linear_mass=False, + query_position_weights=None, reference_position_weights=None): """Count all reference scores <= each query's requested integer thresholds. Optional positional weights and additive group distances match :func:`score_matrix`. @@ -702,13 +726,15 @@ def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_exten 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)) + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r), + *_sequence_weights(query_position_weights, reference_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, - query_group_ids=None, reference_group_ids=None, group_distances=None): + query_group_ids=None, reference_group_ids=None, group_distances=None, + query_position_weights=None, reference_position_weights=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 @@ -724,7 +750,8 @@ def ball_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend 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), - *_group_options(query_group_ids, reference_group_ids, group_distances, q, r)) + *_group_options(query_group_ids, reference_group_ids, group_distances, q, r), + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r)) def paired_topk_batch(queries_a, queries_b, refs_a, refs_b, k=10, matrix=None, @@ -769,7 +796,9 @@ def paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, mat 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): + group_distances_alpha=None, group_distances_beta=None, + query_position_weights_alpha=None, query_position_weights_beta=None, + reference_position_weights_alpha=None, reference_position_weights_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 @@ -788,11 +817,14 @@ def paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, mat 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) + aw = _sequence_weights(query_position_weights_alpha, reference_position_weights_alpha, position_weights_by_length_alpha, qa, ra) + bw = _sequence_weights(query_position_weights_beta, reference_position_weights_beta, position_weights_by_length_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]) + ag[0], bg[0], ag[1], bg[1], ag[2], bg[2], + aw[0], bw[0], aw[1], bw[1]) class GapBlockIndex: diff --git a/src/_bindings.cpp b/src/_bindings.cpp index eed1f0a..d91575c 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& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { + const std::vector>& weights, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights) { 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, query_groups, reference_groups, group_distances); + gap_extend, prior, prior_width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights); } return out; } @@ -801,7 +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>{}, + nb::arg("query_group_ids") = std::vector{}, nb::arg("reference_group_ids") = std::vector{}, nb::arg("group_distances") = std::vector>{}, nb::arg("query_position_weights") = std::vector>{}, nb::arg("reference_position_weights") = 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."); @@ -812,23 +812,23 @@ 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, const std::vector& query_groups, const std::vector& reference_groups, const std::vector>& group_distances) { + 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 std::vector>& query_weights, const std::vector>& reference_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, 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."); + 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, query_weights, reference_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>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, nb::arg("query_position_weights")=std::vector>{}, nb::arg("reference_position_weights")=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) { + 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 std::vector>& query_weights, const std::vector>& reference_weights) { 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."); + 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, query_weights, reference_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>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, nb::arg("query_position_weights")=std::vector>{}, nb::arg("reference_position_weights")=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) { + 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 std::vector>& query_weights, const std::vector>& reference_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, 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."); + 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, query_weights, reference_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>{}, nb::arg("query_group_ids")=std::vector{}, nb::arg("reference_group_ids")=std::vector{}, nb::arg("group_distances")=std::vector>{}, nb::arg("query_position_weights")=std::vector>{}, nb::arg("reference_position_weights")=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); @@ -842,11 +842,11 @@ 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) { + 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 std::vector>& query_weights_alpha, const std::vector>& query_weights_beta, const std::vector>& reference_weights_alpha, const std::vector>& reference_weights_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."); + 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, query_weights_alpha, query_weights_beta, reference_weights_alpha, reference_weights_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>{}, nb::arg("query_position_weights_alpha")=std::vector>{}, nb::arg("query_position_weights_beta")=std::vector>{}, nb::arg("reference_position_weights_alpha")=std::vector>{}, nb::arg("reference_position_weights_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. " diff --git a/src/gapblock.cpp b/src/gapblock.cpp index b993f68..1e67733 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -38,7 +38,8 @@ 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* weights = nullptr) { + const int32_t* prior_row, int32_t* suf, const int32_t* weights = nullptr, + const int32_t* qw = nullptr, const int32_t* rw = 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 +48,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] * (weights ? weights[j + d] : 1); + suf[j] = suf[j + 1] + pen[size_t(a) * A + b] * (qw && rw ? std::max(qw[j + (q_longer ? d : 0)], rw[j + (q_longer ? 0 : d)]) : (weights ? weights[j + d] : 1)); } if (d == 0) return suf[0]; @@ -56,7 +57,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]] * (weights ? weights[i] : 1); + if (i < L) pre += pen[size_t(q[i]) * A + r[i]] * (qw && rw ? std::max(qw[i], rw[i]) : (weights ? weights[i] : 1)); } return best + gap_open + int32_t(d - 1) * gap_extend; } @@ -66,7 +67,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; + std::vector> weights, query_weights, reference_weights; std::vector query_groups, reference_groups; std::vector> group_distances; uint32_t longest = 0; @@ -80,7 +81,9 @@ GapInput prepare(const std::vector& queries, const std::vector>& weights = {}, const std::vector& query_groups = {}, const std::vector& reference_groups = {}, - const std::vector>& group_distances = {}) { + const std::vector>& group_distances = {}, + const std::vector>& query_weights = {}, + const std::vector>& reference_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()) @@ -134,6 +137,26 @@ GapInput prepare(const std::vector& queries, const std::vector n ? m - n : n - m; const int32_t* prow = prior.empty() ? nullptr : prior.data() + (size_t(M) * in.W1 + d) * in.W1; - 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()); + 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(), + in.query_weights.empty() ? nullptr : in.query_weights[i].data(), + in.reference_weights.empty() ? nullptr : in.reference_weights[j].data()); } struct Ranked { @@ -309,8 +334,8 @@ std::vector gapblock_matrix(const std::vector& queries, const std::vector& prior, uint32_t width, int threads, 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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights) { + const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, std::max(threads, 0), weights, query_groups, reference_groups, group_distances, query_weights, reference_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()); @@ -328,10 +353,10 @@ std::vector> gapblock_ball_batch(const std::vector 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 std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_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, query_groups, reference_groups, group_distances); + const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_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) { @@ -380,8 +405,8 @@ std::vector> gapblock_count_batch(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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights); return counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact); } @@ -391,8 +416,8 @@ std::pair>, std::vector> 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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights); std::vector> masses; auto counts = counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact, &masses); return {std::move(counts), std::move(masses)}; @@ -420,11 +445,15 @@ std::pair>, std::vector> 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 std::vector>& group_distances_beta, + const std::vector>& query_weights_alpha, + const std::vector>& query_weights_beta, + const std::vector>& reference_weights_alpha, + const std::vector>& reference_weights_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); + weights_alpha, query_groups_alpha, reference_groups_alpha, group_distances_alpha, query_weights_alpha, reference_weights_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); + weights_beta, query_groups_beta, reference_groups_beta, group_distances_beta, query_weights_beta, reference_weights_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)}; diff --git a/tests/python/test_gapblock_sequence_weights.py b/tests/python/test_gapblock_sequence_weights.py new file mode 100644 index 0000000..f032dec --- /dev/null +++ b/tests/python/test_gapblock_sequence_weights.py @@ -0,0 +1,54 @@ +"""Per-input positional geometry: exhaustive symmetric oracle and shared reductions.""" +import numpy as np +import pytest +from seqtree import gapblock + + +def oracle(q, r, qw, rw, go=7, ge=3): + d = abs(len(q)-len(r)) + scores = [] + for start in range(min(len(q), len(r))+1): + score = go+(d-1)*ge if d else 0 + for i in range(min(len(q), len(r))): + qi = i+(d if len(q)>len(r) and i>=start else 0) + ri = i+(d if len(r)>len(q) and i>=start else 0) + score += (q[qi]!=r[ri])*max(qw[qi], rw[ri]) + scores.append(score) + return min(scores) + + +def test_sequence_weights_oracle_and_linked_reductions(): + q = ['', 'AC', 'ACA', 'CAAA'] + r = ['', 'CA', 'ACA', 'ACAAA', 'CA'] + qw = [[], [10,100], [100,10,10], [10,100,100,10]] + rw = [[], [100,10], [10,100,10], [10,10,100,100,10], [10,100]] + expected = np.array([[oracle(a,b,aw,bw) for b,bw in zip(r,rw)] for a,aw in zip(q,qw)]) + cuts = [[-1,0,10,100,250,100] for _ in q] + opts = dict(gap_open=7, gap_extend=3, query_position_weights=qw, reference_position_weights=rw) + for threads in (1,4): + np.testing.assert_array_equal(np.asarray(gapblock.score_matrix(q,r,threads=threads,**opts)), expected) + reverse = gapblock.score_matrix(r,q,gap_open=7,gap_extend=3,threads=threads, + query_position_weights=rw,reference_position_weights=qw) + np.testing.assert_array_equal(np.asarray(reverse).T,expected) + counts,masses = gapblock.count_batch(q,r,cuts,threads=threads,linear_mass=True,exclude_exact=True,**opts) + paircounts,pairmasses = gapblock.paired_sum_count_batch(q,q,r,r,cuts,gap_open=7,gap_extend=3, + threads=threads,exclude_exact=True,query_position_weights_alpha=qw, + query_position_weights_beta=qw,reference_position_weights_alpha=rw,reference_position_weights_beta=rw) + for i,a in enumerate(q): + eligible = np.array([a!=b for b in r]) + assert counts[i] == [int(np.sum(eligible & (expected[i]<=c))) for c in cuts[i]] + assert masses[i] == [int(np.sum(np.maximum(0,c-expected[i][eligible]))) for c in cuts[i]] + assert paircounts[i] == [int(np.sum(eligible & (2*expected[i]<=c))) for c in cuts[i]] + assert pairmasses[i] == [int(np.sum(np.maximum(0,c-2*expected[i][eligible]))) for c in cuts[i]] + uniform = dict(query_position_weights=[[100]*len(a) for a in q],reference_position_weights=[[100]*len(b) for b in r]) + old = gapblock.score_matrix(q,r,gap_open=7,gap_extend=3,position_weights_by_length={n:[100]*n for n in range(6)}) + np.testing.assert_array_equal(np.asarray(old),np.asarray(gapblock.score_matrix(q,r,gap_open=7,gap_extend=3,**uniform))) + + +@pytest.mark.parametrize('qw,rw,by_length', [([[1]],[[1,1]],None), ([[1,1]],None,None), + ([[1,-1]],[[1,1]],None), ([[True,1]],[[1,1]],None), ([[1,1.5]],[[1,1]],None), + ([[2**31-1]*2],[[1,1]],None), ([[1,1]],[[1,1]],{2:[1,1]})]) +def test_sequence_weight_validation(qw,rw,by_length): + with pytest.raises((ValueError,TypeError,OverflowError)): + gapblock.count_batch(['AA'],['AC'],[[100]],query_position_weights=qw, + reference_position_weights=rw,position_weights_by_length=by_length)