Skip to content

Repository navigation

Add positional weights to native gap-block batches #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mikessh wants to merge 2 commits into from
Closed

Add positional weights to native gap-block batches #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
98bf4c9
Add positional weights to native gap-block batches
mikessh Oct 10, 2026
464a362
Allow safe zero-weight score tables (#8)
mikessh Oct 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions docs/gapblock.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
15 changes: 13 additions & 2 deletions include/seqtree/seqtree.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,11 +221,21 @@ std::vector<int32_t> align_dist_matrix(const std::vector<std::string>& 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<int32_t> gapblock_matrix(const std::vector<std::string>& queries,
const std::vector<std::string>& refs, Alphabet,
const SubstitutionMatrix* matrix, int32_t gap_open,
int32_t gap_extend, const std::vector<int32_t>& prior,
uint32_t prior_width, int threads = 0);
uint32_t prior_width, int threads = 0,
const std::vector<std::vector<int32_t>>& weights = {});

// Exhaustive ball rows in reference order; hit edit fields are unset (zero).
std::vector<std::vector<Hit>> gapblock_ball_batch(const std::vector<std::string>& queries,
const std::vector<std::string>& refs, const std::vector<int32_t>& thresholds,
Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge,
const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact,
const std::vector<std::vector<int32_t>>& 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
Expand All @@ -242,7 +252,8 @@ std::vector<std::vector<std::tuple<uint32_t, int32_t, int32_t>>> gapblock_paired
std::vector<std::vector<uint64_t>> gapblock_count_batch(const std::vector<std::string>& queries,
const std::vector<std::string>& refs, const std::vector<std::vector<int32_t>>& thresholds,
Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge,
const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact);
const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact,
const std::vector<std::vector<int32_t>>& weights = {});
std::vector<std::vector<uint64_t>> gapblock_paired_count_batch(
const std::vector<std::string>& qa, const std::vector<std::string>& qb,
const std::vector<std::string>& ra, const std::vector<std::string>& rb,
Expand Down
56 changes: 52 additions & 4 deletions python/seqtree/gapblock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]
Expand Down Expand Up @@ -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.

Expand All @@ -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.
Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down
19 changes: 13 additions & 6 deletions src/_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -360,15 +360,16 @@ ScoreMatrix py_gapblock_matrix(const std::vector<std::string>& queries,
const std::vector<std::string>& refs, const std::string& alphabet,
const std::optional<SubstitutionMatrix>& matrix, int32_t gap_open,
int32_t gap_extend, const std::vector<int32_t>& prior,
uint32_t prior_width, int threads) {
uint32_t prior_width, int threads,
const std::vector<std::vector<int32_t>>& weights) {
Alphabet alph = parse_alphabet(alphabet);
ScoreMatrix out;
out.rows = queries.size();
out.cols = refs.size();
{
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;
}
Expand Down Expand Up @@ -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<int32_t>{},
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<std::vector<int32_t>>{},
"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.");
Expand All @@ -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<int32_t>{}, 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<std::string>& queries, const std::vector<std::string>& refs, const std::vector<std::vector<int32_t>>& thresholds, const std::string& alphabet, const std::optional<SubstitutionMatrix>& matrix, int32_t go, int32_t ge, const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact) {
m.def("gapblock_count_batch", [](const std::vector<std::string>& queries, const std::vector<std::string>& refs, const std::vector<std::vector<int32_t>>& thresholds, const std::string& alphabet, const std::optional<SubstitutionMatrix>& matrix, int32_t go, int32_t ge, const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact, const std::vector<std::vector<int32_t>>& 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<int32_t>{}, 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<int32_t>{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, nb::arg("position_weights_by_length")=std::vector<std::vector<int32_t>>{}, "Bounded exhaustive gapblock reduction, GIL released.");

m.def("gapblock_ball_batch", [](const std::vector<std::string>& queries, const std::vector<std::string>& refs, const std::vector<int32_t>& thresholds, const std::string& alphabet, const std::optional<SubstitutionMatrix>& matrix, int32_t go, int32_t ge, const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact, const std::vector<std::vector<int32_t>>& 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<int32_t>{}, nb::arg("prior_width")=0, nb::arg("threads")=0, nb::arg("exclude_exact")=false, nb::arg("position_weights_by_length")=std::vector<std::vector<int32_t>>{}, "Bounded exhaustive gapblock reduction, GIL released.");

m.def("gapblock_paired_topk_batch", [](const std::vector<std::string>& qa, const std::vector<std::string>& qb, const std::vector<std::string>& ra, const std::vector<std::string>& rb, uint32_t k, const std::string& alphabet, const std::optional<SubstitutionMatrix>& matrix, int32_t go, int32_t ge, const std::vector<int32_t>& prior, uint32_t width, int threads, bool exclude_exact) {
const auto alph = parse_alphabet(alphabet);
Expand Down
Loading
Loading