Add group offsets to batched gap-block counts

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.

include/seqtree/seqtree.hpp

@@ -223,19 +223,27 @@ std::vector<int32_t> align_dist_matrix(const std::vector<std::string>& 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<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,
-                                     const std::vector<std::vector<int32_t>>& weights = {});
+                                     const std::vector<std::vector<int32_t>>& weights = {},
+                                     const std::vector<int32_t>& query_groups = {},
+                                     const std::vector<int32_t>& reference_groups = {},
+                                     const std::vector<std::vector<int32_t>>& group_distances = {});
 
 // 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 = {});
+        const std::vector<std::vector<int32_t>>& weights = {},
+        const std::vector<int32_t>& query_groups = {},
+        const std::vector<int32_t>& reference_groups = {},
+        const std::vector<std::vector<int32_t>>& 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,13 +261,36 @@ std::vector<std::vector<uint64_t>> gapblock_count_batch(const std::vector<std::s
         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<std::vector<int32_t>>& weights = {});
+        const std::vector<std::vector<int32_t>>& weights = {},
+        const std::vector<int32_t>& query_groups = {},
+        const std::vector<int32_t>& reference_groups = {},
+        const std::vector<std::vector<int32_t>>& group_distances = {});
+std::pair<std::vector<std::vector<uint64_t>>, std::vector<std::vector<uint64_t>>> gapblock_count_mass_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<std::vector<int32_t>>& weights = {},
+        const std::vector<int32_t>& query_groups = {},
+        const std::vector<int32_t>& reference_groups = {},
+        const std::vector<std::vector<int32_t>>& group_distances = {});
 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,
         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);
+std::pair<std::vector<std::vector<uint64_t>>, std::vector<std::vector<uint64_t>>> gapblock_paired_sum_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,
+        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<std::vector<int32_t>>& weights_alpha,
+        const std::vector<std::vector<int32_t>>& weights_beta,
+        const std::vector<int32_t>& query_groups_alpha, const std::vector<int32_t>& query_groups_beta,
+        const std::vector<int32_t>& reference_groups_alpha, const std::vector<int32_t>& reference_groups_beta,
+        const std::vector<std::vector<int32_t>>& group_distances_alpha,
+        const std::vector<std::vector<int32_t>>& group_distances_beta);
 
 // ---------------------------------------------------------------------------------------
 // Plain string edit distances on raw characters, unit costs. Deliberately alphabet-agnostic:

python/seqtree/gapblock.py

@@ -44,12 +44,14 @@
 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
+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",
 ]
@@ -524,6 +526,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 +541,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 +578,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 +604,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,30 +676,45 @@ 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, linear_mass=False):
     """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
     :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,
-                                 _position_weights(position_weights_by_length, q, r))
+    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))
 
 
 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 +723,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,
@@ -721,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.