seqtree 1.1.1__tar.gz → 1.2.0__tar.gz

{seqtree-1.1.1 → seqtree-1.2.0}/CHANGELOG.md

@@ -3,6 +3,24 @@
 All notable changes to `seqtree`. Dates are release dates. From 1.0.0 the project follows semantic
 versioning: breaking changes need a **major** bump.
 
+## [1.2.0] — 2026-10-17
+
+### Fixed
+
+- Exact exclusion masks matching hits without shrinking model populations or
+  supplied annotation counters, including withheld-label models.
+
+### Added
+
+- Fixed query/reference partitions for native group audits.
+- Bounded sparse grouped radius-hit output with deterministic ordering and explicit memory-cap failures.
+- Optional grouped geometry compatible with the supported nanobind 2.5 floor.
+
+- Generic float64 global-alignment target maxima with numeric annotation CSR,
+  group offsets, explicit matrix symbols and independent candidate edit budgets.
+- Optional aggregate insertion/deletion limits for batched edit histograms;
+  omitted and redundant limits preserve the existing count paths.
+
 ## [1.1.0] — 2026-10-11
 
 ### Added

{seqtree-1.1.1 → seqtree-1.2.0}/CMakeLists.txt

@@ -30,6 +30,7 @@ add_library(seqtree_core STATIC
   src/searcher.cpp
   src/gapblock.cpp
   src/pairwise.cpp
+  src/alignment_max.cpp
   src/distance.cpp
 )
 target_include_directories(seqtree_core PUBLIC include PRIVATE src)

{seqtree-1.1.1 → seqtree-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: seqtree
-Version: 1.1.1
+Version: 1.2.0
 Summary: Fast fuzzy search over biological sequences (C++ core, Python bindings)
 Keywords: sequence-search,fuzzy-matching,CDR3,immunology,bioinformatics,trie
 Author-Email: ISALGO laboratory <mikhail.shugay@gmail.com>
@@ -289,3 +289,14 @@ This repo follows **git-flow**:
 Roadmap (affine gaps, position-specific matrices, succinct memory packing) lives in
 [docs/roadmap.rst](docs/roadmap.rst). Control-set E-values already ship — see the
 [E-value guide](https://antigenomics.github.io/seqtree/evalue.html).
+
+### Floating alignment reductions
+
+`Index.alignment_max_batch` streams unit-cost search candidates through float64
+global alignment and requested numeric-target maxima without a dense pair table.
+It accepts explicit matrix symbols, small group-offset matrices and annotation CSR,
+with separate alignment-selection and final scoring coefficients.
+`Index.edit_histogram_batch(..., max_indels=...)` additionally supports an aggregate
+insertion/deletion cap; its default preserves previous behavior. See the
+[API reference](https://antigenomics.github.io/seqtree/api.html) for array contracts
+and tie, count, exclusion and unsupported-value semantics.

{seqtree-1.1.1 → seqtree-1.2.0}/README.md

@@ -257,3 +257,14 @@ This repo follows **git-flow**:
 Roadmap (affine gaps, position-specific matrices, succinct memory packing) lives in
 [docs/roadmap.rst](docs/roadmap.rst). Control-set E-values already ship — see the
 [E-value guide](https://antigenomics.github.io/seqtree/evalue.html).
+
+### Floating alignment reductions
+
+`Index.alignment_max_batch` streams unit-cost search candidates through float64
+global alignment and requested numeric-target maxima without a dense pair table.
+It accepts explicit matrix symbols, small group-offset matrices and annotation CSR,
+with separate alignment-selection and final scoring coefficients.
+`Index.edit_histogram_batch(..., max_indels=...)` additionally supports an aggregate
+insertion/deletion cap; its default preserves previous behavior. See the
+[API reference](https://antigenomics.github.io/seqtree/api.html) for array contracts
+and tie, count, exclusion and unsupported-value semantics.

seqtree-1.2.0/RELEASE_AUDIT.md

@@ -0,0 +1,12 @@
+# Release audit — 17 October 2026
+
+The requested generic batch/search APIs are implemented. No new native algorithm is needed for the audited downstream requirements.
+
+- Native batch search and edit histograms; aggregate insertion/deletion caps.
+- Float64 global-alignment target reductions, separate candidate/final scoring semantics, fixed populations under exact exclusion.
+- Grouped geometry, fixed audit partitions, bounded deterministic sparse radius hits and optional-geometry compatibility at the supported nanobind floor.
+- Shared CPU allocation checks; one-thread inline execution and explicit over-budget rejection.
+
+Verification on `0852ff5`: **379 passed, 1 Linux-only platform skip** across `test_batch_threads.py`, `test_alignment_max.py`, `test_edit_histogram.py`, `test_grouped_optional_geometry.py` and `test_grouped_sparse_hits.py`. An additional 280 checks cover grouped populations, withholding, paired lanes, offsets, score bins and partitions. Native budget was two CPUs; numerical libraries used one thread. The committed `examples/01_gapped_search.py` runs and its per-query cutoff/search output was inspected. Sphinx builds with `-W`.
+
+The local `dev` commits `0126e00`, `858056b` and `0852ff5` were three commits ahead of the remote at audit start. They are preserved and included when this audited `dev` is pushed. The changelog now covers these capabilities; public API documentation already describes the alignment reductions. [Issue 18](https://github.com/antigenomics/seqtree/issues/18) tracks integration into the release head and the remaining exact-head cross-platform/wheel gate. Publishing remains the maintainer's action; this audit creates no tag or release. Focused checks do not substitute for complete release CI.

{seqtree-1.1.1 → seqtree-1.2.0}/include/seqtree/seqtree.hpp

@@ -105,6 +105,14 @@ private:
     std::vector<uint8_t> masked_;  // [width]
 };
 
+struct AlignmentMaxResult {
+    std::vector<double> scores, raw_similarity, normalized_similarity, pfit, informativeness;
+    std::vector<uint64_t> counts, modeled_counts, aa_counts, exact_counts, candidate_counts, nw_pairs;
+    std::vector<int64_t> best_ref_ids, best_lane_ids;
+    std::vector<uint32_t> gap_counts;
+    std::vector<uint8_t> query_supported;
+};
+
 // Immutable after build. Lock-free for concurrent reads; share one Index across
 // threads, give each thread its own Searcher.
 class Index {
@@ -129,7 +137,18 @@ public:
     // Exact edit-count bins 0..max_total_edits; one reusable hit buffer per worker.
     std::vector<std::vector<uint64_t>> edit_histogram_batch(
         const std::vector<std::string>& queries, const SearchParams&,
-        int threads = 0, bool exclude_exact = false) const;
+        int threads = 0, bool exclude_exact = false, int max_indels = -1) const;
+
+    // Numeric requested-label maxima; float NW selection is independent of candidate scope.
+    AlignmentMaxResult alignment_max_batch(const std::vector<std::string>& queries,
+        const SearchParams&, const std::vector<std::vector<double>>& similarity, const std::string& symbols, int max_indels,
+        const std::vector<uint32_t>& query_target_offsets, const std::vector<uint32_t>& query_target_ids,
+        const std::vector<uint32_t>& reference_target_offsets, const std::vector<uint32_t>& reference_target_ids,
+        const std::vector<int32_t>& target_class_ids, const std::vector<uint32_t>& reference_class_offsets,
+        const std::vector<int32_t>& reference_class_ids, const std::vector<std::vector<int32_t>>& query_group_ids,
+        const std::vector<std::vector<int32_t>>& reference_group_ids, const std::vector<std::vector<std::vector<double>>>& group_offsets,
+        const std::vector<uint64_t>& reference_other_counts, const std::vector<uint64_t>& reference_multiplicities,
+        const std::vector<double>& coefficients, int threads=0, bool exclude_exact=false) const;
 
     // Per-query count of seqtm collisions: how often a reference was re-reached via a
     // different edit path during branch-and-bound (0 for the seqtrie engine / Hamming).
@@ -155,7 +174,7 @@ public:
     ~Searcher();
 
     std::vector<Hit> search(std::string_view query, const SearchParams&);
-    void search_into(std::string_view query, const SearchParams&, std::vector<Hit>& out);
+    void search_into(std::string_view query, const SearchParams&, std::vector<Hit>& out, int max_indels = -1);
     bool search_top(std::string_view query, const SearchParams&, Hit& out);
 
     // seqtm collisions from the most recent search_into on this Searcher.
@@ -285,7 +304,9 @@ gapblock_group_audit_batch(const std::vector<std::string>& queries,
         const std::vector<std::vector<int32_t>>& group_distances = {},
         const std::vector<std::vector<int32_t>>& query_weights = {},
         const std::vector<std::vector<int32_t>>& reference_weights = {},
-        const std::optional<PairedFeatureOffsets>& feature_offsets = std::nullopt);
+        const std::optional<PairedFeatureOffsets>& feature_offsets = std::nullopt,
+        const std::vector<int32_t>& query_partitions = {},
+        const std::vector<int32_t>& reference_partitions = {});
 
 std::vector<std::vector<Hit>> gapblock_topk_batch(const std::vector<std::string>& queries,
         const std::vector<std::string>& refs, uint32_t k, Alphabet alphabet,
@@ -371,6 +392,10 @@ struct GroupedGeometry {
     int32_t alignment_scale = 1;
     uint32_t trim_left = 0, trim_right = 0;
     int32_t score_quantum = 1;
+    bool return_hits = false;
+    size_t max_hit_bytes = 64*1024*1024;
+    std::optional<std::vector<int32_t>> query_target_ids;
+    std::optional<std::vector<int64_t>> query_reference_ids;
     std::vector<std::vector<int32_t>> weights, group_distances, query_weights, reference_weights;
     std::vector<int32_t> query_groups, reference_groups;
     std::optional<PairedFeatureOffsets> features;
@@ -382,6 +407,7 @@ struct CountBuffer {
 struct GroupedCountChunk {
     size_t start = 0, stop = 0;
     CountBuffer counts, sizes, exact_counts;
+    CountBuffer hits{{},{0,3}};
 };
 class GroupedCountBatch {
 public:

{seqtree-1.1.1 → seqtree-1.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "seqtree"
-version = "1.1.1"
+version = "1.2.0"
 description = "Fast fuzzy search over biological sequences (C++ core, Python bindings)"
 readme = "README.md"
 requires-python = ">=3.10"

{seqtree-1.1.1 → seqtree-1.2.0}/python/seqtree/gapblock.py

@@ -745,24 +745,34 @@ def group_audit_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap
                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_position_weights=None, reference_position_weights=None, feature_offsets=None):
+               query_position_weights=None, reference_position_weights=None, feature_offsets=None,
+               query_partition_ids=None, reference_partition_ids=None):
     """Stream component-ball count, full-score count and bounded nearest in one batch.
 
     Geometry/offset options match ball_batch. Component excludes the nonnegative
     group offset; full score includes it. For each query threshold, return three
     lists: component count, full count and minimum full score among component hits.
     Empty component balls have nearest=-1. Duplicate references count separately;
-    full encoded exact puncturing is identical. O(Q) output; no sparse hit objects.
+    full encoded exact puncturing is identical. Optional nonnegative partition IDs
+    restrict each query to references with the same ID. Both arrays are required;
+    unmatched partitions return zero counts and nearest=-1. Partition restriction
+    precedes scoring and is independent of additive group/feature distances. O(Q) output; no sparse hit objects.
     """
     q, r = list(queries), list(refs)
     thresholds = [_batch_integer(v, "threshold", -(1 << 31)) for v in thresholds]
     if len(thresholds) != len(q):
         raise ValueError("group audit thresholds must have one value per query")
+    if (query_partition_ids is None) != (reference_partition_ids is None):
+        raise ValueError("both query and reference partition IDs are required")
+    qp = [] if query_partition_ids is None else [_batch_integer(v, "query partition ID") for v in query_partition_ids]
+    rp = [] if reference_partition_ids is None else [_batch_integer(v, "reference partition ID") for v in reference_partition_ids]
+    if query_partition_ids is not None and (len(qp) != len(q) or len(rp) != len(r)):
+        raise ValueError("partition IDs must cover every query and reference")
     go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads)
     return _gapblock_group_audit_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),
-                            *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets)
+                            *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets, qp, rp)
 
 
 def ball_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend=1,
@@ -950,7 +960,8 @@ def _grouped_geometry(q, r, weights, query_groups, reference_groups, distances,
         query_groups, reference_groups, distances, q, r)
     geometry.query_weights, geometry.reference_weights = _sequence_weights(
         query_weights, reference_weights, weights, q, r)
-    geometry.features = features
+    if features is not None:
+        geometry.features = features
     return geometry
 
 
@@ -963,9 +974,21 @@ def grouped_count_batch(queries, refs, thresholds, *, reference_target_offsets,
         feature_offsets=None, query_withheld_peptide_ids=None,
         reference_peptide_offsets=None, reference_peptide_ids=None, score_quantum=1,
         alignment_mode="gapblock", alignment_scale=1, trim_left=0, trim_right=0,
-        additional_feature_offsets=None):
+        additional_feature_offsets=None, return_hits=False, query_target_ids=None,
+        query_reference_ids=None, max_hit_bytes=64*1024*1024):
     """Yield bounded label curves from one immutable prepared scoring population.
 
+    ``return_hits=True`` instead yields ``chunk.hits[H,3]`` read-only uint32 rows
+    (global query index, reference index, raw distance) at thresholds[-1], ordered
+    by query then reference; counts/sizes/exact_counts are empty. Requires one
+    query_target_ids entry per query and scans only that target's CSR references.
+    Optional query_reference_ids keeps reference indices strictly greater than
+    each query's supplied index (-1 keeps all). Exact exclusion still compares
+    full sequences. Supports the single non-withheld population API only.
+    max_hit_bytes bounds actual hit output per chunk (default64MiB); exceeding it
+    raises without dropping hits. Row assembly plus final output temporarily uses
+    at most roughly three times that cap, excluding inputs and worker DP scratch.
+
     ``alignment_mode`` defaults to gapblock. Global/local reuse the existing affine
     similarity recurrence and require a raw similarity matrix. For trimmed sequences
     q/r, junction distance is alignment_scale * max(0, sum(diag(q)) + sum(diag(r))
@@ -997,7 +1020,7 @@ def grouped_count_batch(queries, refs, thresholds, *, reference_target_offsets,
     Each chunk owns read-only uint32 buffers: counts[C,1+2*T+P,K], sizes[C,1+2*T+P],
     exact_counts[C,T], with start/stop query offsets. numpy.asarray is zero-copy;
     no NumPy import/dependency is required. Exact counts retain raw multiplicity;
-    exclude_exact punctures every curve and population size by full sequence identity.
+    exclude_exact masks hit curves by full sequence identity; model population sizes remain fixed.
     Optional query_withheld_peptide_ids is one ID/query (-1 means none) and
     requires full reference_peptide_offsets/reference_peptide_ids CSR. Each sorted
     unique row includes all raw peptide labels, including sparse/unselected IDs in
@@ -1018,6 +1041,15 @@ def grouped_count_batch(queries, refs, thresholds, *, reference_target_offsets,
     go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads)
     geometry = _grouped_geometry(q, r, position_weights_by_length, query_group_ids,
         reference_group_ids, group_distances, query_position_weights, reference_position_weights, feature_offsets, score_quantum)
+    if return_hits or query_target_ids is not None or query_reference_ids is not None:
+        if not hasattr(geometry, "return_hits"):
+            raise RuntimeError("sparse hits require native return_hits support")
+        geometry.return_hits = bool(return_hits)
+        geometry.max_hit_bytes = _batch_integer(max_hit_bytes, "max_hit_bytes", 12)
+        if query_target_ids is not None:
+            geometry.query_target_ids = [_batch_integer(x, "query target ID", 0) for x in query_target_ids]
+        if query_reference_ids is not None:
+            geometry.query_reference_ids = [_batch_integer(x, "query reference ID", -1) for x in query_reference_ids]
     _additional_geometry(geometry, additional_feature_offsets)
     _alignment_options(geometry, q, r, alignment_mode, alignment_scale, trim_left, trim_right,
         matrix, gap_prior, position_weights_by_length, query_position_weights, reference_position_weights)

{seqtree-1.1.1 → seqtree-1.2.0}/src/_bindings.cpp

@@ -188,13 +188,13 @@ nb::list py_search_batch(const Index& idx, const std::vector<std::string>& queri
 
 std::vector<std::vector<uint64_t>> py_edit_histogram_batch(
         const Index& idx, const std::vector<std::string>& queries,
-        const PyParams& pp, int threads, bool exclude_exact) {
+        const PyParams& pp, int threads, bool exclude_exact, int max_indels) {
     if (pp.max_total_edits < 1 || pp.max_total_edits > 65535)
         throw nb::value_error("edit histogram requires max_total_edits in [1, 65535]");
     auto mat = make_matrix(pp, idx.alphabet());
     SearchParams cp = to_cpp(pp, mat ? &*mat : nullptr);
     nb::gil_scoped_release release;
-    return idx.edit_histogram_batch(queries, cp, threads, exclude_exact);
+    return idx.edit_histogram_batch(queries, cp, threads, exclude_exact, max_indels);
 }
 
 std::vector<uint64_t> py_collisions_batch(const Index& idx, const std::vector<std::string>& queries,
@@ -335,6 +335,26 @@ ArrayView view_of(const nb::object& owner, const std::vector<T>& v, char fmt) {
     return ArrayView{owner, v.data(), v.size(), sizeof(T), {fmt, 0}};
 }
 
+nb::dict alignment_max_arrays(const nb::object& self) {
+    const auto& r=nb::cast<const AlignmentMaxResult&>(self); nb::dict d;
+    d["scores"]=view_of(self,r.scores,'d');
+    d["raw_similarity"]=view_of(self,r.raw_similarity,'d');
+    d["normalized_similarity"]=view_of(self,r.normalized_similarity,'d');
+    d["pfit"]=view_of(self,r.pfit,'d');
+    d["informativeness"]=view_of(self,r.informativeness,'d');
+    d["counts"]=view_of(self,r.counts,'Q');
+    d["modeled_counts"]=view_of(self,r.modeled_counts,'Q');
+    d["aa_counts"]=view_of(self,r.aa_counts,'Q');
+    d["exact_counts"]=view_of(self,r.exact_counts,'Q');
+    d["candidate_counts"]=view_of(self,r.candidate_counts,'Q');
+    d["nw_pairs"]=view_of(self,r.nw_pairs,'Q');
+    d["best_ref_ids"]=view_of(self,r.best_ref_ids,'q');
+    d["best_lane_ids"]=view_of(self,r.best_lane_ids,'q');
+    d["gap_counts"]=view_of(self,r.gap_counts,'I');
+    d["query_supported"]=view_of(self,r.query_supported,'B');
+    return d;
+}
+
 // One hit as a Python object. Built only for the query actually asked for, never for the whole
 // batch -- the flat arrays are what a 445k-query run should read.
 struct PyTextHit {
@@ -395,6 +415,7 @@ ScoreMatrix py_gapblock_matrix(const std::vector<std::string>& queries,
 }  // namespace
 
 NB_MODULE(_core, m) {
+    m.attr("exact_exclusion_hits_only") = true;
     m.doc() = "seqtree: fuzzy biological-sequence search (C++ core)";
 
     nb::class_<SubstitutionMatrix>(m, "SubstitutionMatrix",
@@ -663,6 +684,9 @@ NB_MODULE(_core, m) {
             return "Alignment(score=" + std::to_string(a.score) + ", ops='" + a.ops + "')";
         });
 
+    nb::class_<AlignmentMaxResult>(m,"AlignmentMaxResult")
+        .def("arrays",&alignment_max_arrays,"Read-only zero-copy requested-target CSR buffers; views own their result.");
+
     nb::class_<Index>(m, "Index",
                       "Immutable search index over a set of reference sequences. Build once, "
                       "then query concurrently; reference id is the position in ``refs``.")
@@ -696,11 +720,32 @@ NB_MODULE(_core, m) {
              "cores. Returns one hit list per query, in input order.")
         .def("edit_histogram_batch", &py_edit_histogram_batch,
              nb::arg("queries"), nb::arg("params"), nb::arg("threads") = 0,
-             nb::arg("exclude_exact") = false,
+             nb::arg("exclude_exact") = false, nb::arg("max_indels") = -1,
              "Exact reference counts by total n_subs+n_ins+n_dels, bins 0..max_total_edits. "
              "Requires explicit positive max_total_edits and uncapped seqtm/auto all-hit search. "
+             "max_indels>=0 also caps total insertions+deletions; -1 preserves directional-only caps. "
              "Duplicates count separately; exclude_exact retains bin 0 with count zero. threads<=0 uses the detected available CPU budget. Releases the GIL; "
              "one reusable hit buffer per worker, no all-query hit vectors.")
+        .def("alignment_max_batch", [](const Index& idx,const std::vector<std::string>& queries,const PyParams& pp,
+            const std::vector<std::vector<double>>& sim,const std::string& symbols,int max_indels,
+            const std::vector<uint32_t>& qo,const std::vector<uint32_t>& qt,const std::vector<uint32_t>& ro,
+            const std::vector<uint32_t>& rt,const std::vector<int32_t>& tc,const std::vector<uint32_t>& co,
+            const std::vector<int32_t>& ci,const std::vector<std::vector<int32_t>>& qg,
+            const std::vector<std::vector<int32_t>>& rg,const std::vector<std::vector<std::vector<double>>>& offsets,
+            const std::vector<uint64_t>& other,const std::vector<uint64_t>& multiplicities,
+            const std::vector<double>& coefficients,int threads,bool exclude_exact) {
+                auto mat=make_matrix(pp,idx.alphabet()); auto cp=to_cpp(pp,mat?&*mat:nullptr);
+                nb::gil_scoped_release release;
+                return idx.alignment_max_batch(queries,cp,sim,symbols,max_indels,qo,qt,ro,rt,tc,co,ci,qg,rg,offsets,other,multiplicities,coefficients,threads,exclude_exact);
+            }, nb::arg("queries"),nb::arg("params"),nb::arg("similarity"),nb::arg("similarity_symbols"),nb::arg("max_indels"),nb::kw_only(),
+            nb::arg("query_target_offsets"),nb::arg("query_target_ids"),nb::arg("reference_target_offsets"),nb::arg("reference_target_ids"),
+            nb::arg("target_class_ids"),nb::arg("reference_class_offsets"),nb::arg("reference_class_ids"),
+            nb::arg("query_group_ids"),nb::arg("reference_group_ids"),nb::arg("group_offsets"),nb::arg("reference_other_counts"),
+            nb::arg("reference_multiplicities"),nb::arg("coefficients"),nb::arg("threads")=0,nb::arg("exclude_exact")=false,
+            "Fused numeric requested-label max of cloglog(offset+c3*(rawNW-maxself)+gapcoef*g)*finite information. "
+            "Candidate scope is independent of float64 global NW; selection gap=max(abs(similarity)); "
+            "strict ties use M/X/Y for M and X, M/Y/X for Y, terminal M/X/Y. Group ID -1 is unsupported, not imputed. "
+            "Target filtering precedes NW. Use bounded query chunks against the same immutable Index. GIL released.")
         .def("align", &py_align, nb::arg("ref_id"), nb::arg("query"), nb::arg("params"),
              "Compute a global alignment between ``query`` and a reference, on demand.")
         .def("collisions_batch", &py_collisions_batch, nb::arg("queries"), nb::arg("params"),
@@ -839,9 +884,14 @@ NB_MODULE(_core, m) {
     nb::class_<GroupedCountChunk>(m,"GroupedCountChunk")
         .def_ro("start",&GroupedCountChunk::start).def_ro("stop",&GroupedCountChunk::stop)
         .def_ro("counts",&GroupedCountChunk::counts).def_ro("sizes",&GroupedCountChunk::sizes)
-        .def_ro("exact_counts",&GroupedCountChunk::exact_counts);
+        .def_ro("exact_counts",&GroupedCountChunk::exact_counts)
+        .def_ro("hits",&GroupedCountChunk::hits);
     nb::class_<GroupedGeometry>(m,"_GroupedGeometry")
         .def(nb::init<>()).def_rw("weights",&GroupedGeometry::weights)
+        .def_rw("return_hits",&GroupedGeometry::return_hits)
+        .def_rw("max_hit_bytes",&GroupedGeometry::max_hit_bytes)
+        .def_rw("query_target_ids",&GroupedGeometry::query_target_ids)
+        .def_rw("query_reference_ids",&GroupedGeometry::query_reference_ids)
         .def_rw("additional_features",&GroupedGeometry::additional_features)
         .def_rw("gap_prior",&GroupedGeometry::gap_prior)
         .def_rw("alignment_mode",&GroupedGeometry::alignment_mode)
@@ -895,7 +945,7 @@ NB_MODULE(_core, m) {
        nb::arg("reference_target_offsets"),nb::arg("reference_target_ids"),
        nb::arg("target_peptide_ids"),nb::arg("reference_exclusive_peptide_ids"),
        nb::arg("query_gene_ids"),nb::arg("reference_gene_ids"),nb::arg("chunk_size"),
-       nb::arg("alphabet"),nb::arg("matrix"),nb::arg("gap_open"),nb::arg("gap_extend"),
+       nb::arg("alphabet"),nb::arg("matrix").none(),nb::arg("gap_open"),nb::arg("gap_extend"),
        nb::arg("prior"),nb::arg("prior_width"),nb::arg("threads"),nb::arg("exclude_exact"),
        nb::arg("alpha"),nb::arg("beta")=std::nullopt,
        nb::arg("queries_beta")=std::vector<std::string>{},nb::arg("refs_beta")=std::vector<std::string>{},
@@ -923,11 +973,11 @@ NB_MODULE(_core, m) {
         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, feature_offsets, shared_thresholds);
     }, 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>>{}, nb::arg("query_group_ids")=std::vector<int32_t>{}, nb::arg("reference_group_ids")=std::vector<int32_t>{}, nb::arg("group_distances")=std::vector<std::vector<int32_t>>{}, nb::arg("query_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("reference_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("feature_offsets")=std::nullopt, nb::arg("shared_thresholds")=false, "Counts and integer linear mass in one bounded traversal, GIL released.");
 
-    m.def("gapblock_group_audit_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 std::vector<int32_t>& query_groups, const std::vector<int32_t>& reference_groups, const std::vector<std::vector<int32_t>>& group_distances, const std::vector<std::vector<int32_t>>& query_weights, const std::vector<std::vector<int32_t>>& reference_weights, const std::optional<PairedFeatureOffsets>& feature_offsets) {
+    m.def("gapblock_group_audit_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 std::vector<int32_t>& query_groups, const std::vector<int32_t>& reference_groups, const std::vector<std::vector<int32_t>>& group_distances, const std::vector<std::vector<int32_t>>& query_weights, const std::vector<std::vector<int32_t>>& reference_weights, const std::optional<PairedFeatureOffsets>& feature_offsets, const std::vector<int32_t>& query_partitions, const std::vector<int32_t>& reference_partitions) {
         const auto alph = parse_alphabet(alphabet);
         nb::gil_scoped_release release;
-        return gapblock_group_audit_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, feature_offsets);
-    }, 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>>{}, nb::arg("query_group_ids")=std::vector<int32_t>{}, nb::arg("reference_group_ids")=std::vector<int32_t>{}, nb::arg("group_distances")=std::vector<std::vector<int32_t>>{}, nb::arg("query_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("reference_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("feature_offsets")=std::nullopt, "Streaming component/full ball counts and minimum full score within the component ball, GIL released.");
+        return gapblock_group_audit_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, feature_offsets, query_partitions, reference_partitions);
+    }, 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>>{}, nb::arg("query_group_ids")=std::vector<int32_t>{}, nb::arg("reference_group_ids")=std::vector<int32_t>{}, nb::arg("group_distances")=std::vector<std::vector<int32_t>>{}, nb::arg("query_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("reference_position_weights")=std::vector<std::vector<int32_t>>{}, nb::arg("feature_offsets")=std::nullopt, nb::arg("query_partition_ids")=std::vector<int32_t>{}, nb::arg("reference_partition_ids")=std::vector<int32_t>{}, "Streaming component/full ball counts and minimum full score within the component ball, 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 std::vector<int32_t>& query_groups, const std::vector<int32_t>& reference_groups, const std::vector<std::vector<int32_t>>& group_distances, const std::vector<std::vector<int32_t>>& query_weights, const std::vector<std::vector<int32_t>>& reference_weights, const std::optional<PairedFeatureOffsets>& feature_offsets) {
         const auto alph = parse_alphabet(alphabet);

seqtree-1.2.0/src/alignment_max.cpp

@@ -0,0 +1,177 @@
+#include "seqtree/seqtree.hpp"
+#include "seqtree/parallel.hpp"
+#include "count_unit.hpp"
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <limits>
+#include <stdexcept>
+#include <unordered_map>
+
+namespace seqtree {
+namespace {
+struct Cell { double score=-std::numeric_limits<double>::infinity(), raw=0; uint32_t gaps=0; };
+Cell best(Cell a, const Cell& b, const Cell& c) {
+    if (b.score>a.score) a=b;
+    if (c.score>a.score) a=c;
+    return a;
+}
+using Rows=std::array<std::vector<Cell>,6>;
+// Existing strict NW ties: M/X/Y for M and X, M/Y/X for Y; terminal M/X/Y.
+Cell nw(const std::vector<uint8_t>& q,const std::vector<uint8_t>& r,
+        const std::vector<std::vector<double>>& sim,double gap,Rows& rows) {
+    for (auto& row:rows) row.assign(r.size()+1,Cell{});
+    auto& mp=rows[0]; auto& xp=rows[1]; auto& yp=rows[2];
+    auto& mc=rows[3]; auto& xc=rows[4]; auto& yc=rows[5];
+    mp[0]={0,0,0};
+    for (size_t j=1;j<=r.size();++j) yp[j]={-gap*double(j),0,uint32_t(j)};
+    for (size_t i=1;i<=q.size();++i) {
+        mc[0]=yc[0]=Cell{}; xc[0]={-gap*double(i),0,uint32_t(i)};
+        for (size_t j=1;j<=r.size();++j) {
+            mc[j]=best(mp[j-1],xp[j-1],yp[j-1]);
+            const double s=sim[q[i-1]][r[j-1]]; mc[j].score+=s; mc[j].raw+=s;
+            xc[j]=best(mp[j],xp[j],yp[j]); xc[j].score-=gap; ++xc[j].gaps;
+            yc[j]=best(mc[j-1],yc[j-1],xc[j-1]); yc[j].score-=gap; ++yc[j].gaps;
+        }
+        mp.swap(mc); xp.swap(xc); yp.swap(yc);
+    }
+    return best(mp.back(),xp.back(),yp.back());
+}
+void require(bool ok,const char* message) { if (!ok) throw std::invalid_argument(message); }
+void csr(const std::vector<uint32_t>& offsets,size_t rows,size_t entries) {
+    require(offsets.size()==rows+1 && offsets.front()==0 && offsets.back()==entries &&
+        std::is_sorted(offsets.begin(),offsets.end()),"invalid CSR offsets");
+}
+void add(uint64_t& a,uint64_t b) {
+    if (b>std::numeric_limits<uint64_t>::max()-a) throw std::overflow_error("annotation count overflow");
+    a+=b;
+}
+}
+
+AlignmentMaxResult Index::alignment_max_batch(const std::vector<std::string>& queries,
+    const SearchParams& p,const std::vector<std::vector<double>>& similarity,const std::string& symbols,int max_indels,
+    const std::vector<uint32_t>& qo,const std::vector<uint32_t>& qt,
+    const std::vector<uint32_t>& ro,const std::vector<uint32_t>& rt,const std::vector<int32_t>& tc,
+    const std::vector<uint32_t>& co,const std::vector<int32_t>& ci,
+    const std::vector<std::vector<int32_t>>& qg,const std::vector<std::vector<int32_t>>& rg,
+    const std::vector<std::vector<std::vector<double>>>& offsets,
+    const std::vector<uint64_t>& other,const std::vector<uint64_t>& multiplicities,
+    const std::vector<double>& coefficients,int threads,bool exclude_exact) const {
+    const size_t Q=queries.size(),R=size(),L=rt.size(),T=tc.size(),M=qt.size();
+    require(p.engine!=Engine::SeqTrie && p.mode==Mode::AllHits && !p.max_hits && !p.matrix && !p.pos_matrix &&
+        p.gap_open==1 && p.gap_extend==1 && !p.max_score_penalty,"candidate scope requires uncapped unit-cost seqtm");
+    require(max_indels>=0 && max_indels<=65535,"max_indels must be in [0,65535]");
+    require(coefficients.size()==3 && std::all_of(coefficients.begin(),coefficients.end(),[](double x){return std::isfinite(x);}),"three finite coefficients required");
+    require(!symbols.empty() && similarity.size()==symbols.size(),"similarity dimension does not match symbols");
+    std::array<int,256> code; code.fill(-1); double gap=0;
+    for (size_t a=0;a<symbols.size();++a) {
+        const auto c=codec().encode(symbols[a]);
+        require(c!=Codec::kInvalid && code[c]<0,"invalid or duplicate similarity symbol"); code[c]=int(a);
+        require(similarity[a].size()==symbols.size(),"similarity matrix must be square");
+        for (double x:similarity[a]) { require(std::isfinite(x),"similarity values must be finite"); gap=std::max(gap,std::abs(x)); }
+    }
+    require(gap>0,"NW selection gap max(abs(similarity)) must be positive");
+    csr(qo,Q,M); csr(ro,R,L); csr(co,R,ci.size());
+    require(other.size()==L && multiplicities.size()==L,"annotation arrays have inconsistent lengths");
+    require(qg.size()==offsets.size() && rg.size()==offsets.size(),"group offset components disagree");
+    for (size_t c=0;c<offsets.size();++c) {
+        require(qg[c].size()==Q && rg[c].size()==L,"group ID lengths disagree");
+        const size_t columns=offsets[c].empty()?0:offsets[c][0].size();
+        for (const auto& row:offsets[c]) require(row.size()==columns && std::all_of(row.begin(),row.end(),[](double x){return std::isfinite(x);}),"group offsets must be finite rectangular matrices");
+        for (int32_t x:qg[c]) require(x>=-1 && (x<0 || size_t(x)<offsets[c].size()),"query group ID out of range");
+        for (int32_t x:rg[c]) require(x>=-1 && (x<0 || size_t(x)<columns),"reference group ID out of range");
+    }
+    std::vector<size_t> seen(T,std::numeric_limits<size_t>::max());
+    for (size_t i=0;i<Q;++i) for (size_t k=qo[i];k<qo[i+1];++k) {
+        require(qt[k]<T && seen[qt[k]]!=i,"invalid or duplicate requested target"); seen[qt[k]]=i;
+    }
+    for (int32_t x:tc) require(x>=0,"target class IDs must be nonnegative");
+    std::vector<uint8_t> modeled(L,1);
+    for (size_t r=0;r<R;++r) {
+        require(co[r]<co[r+1],"each reference requires at least one class role");
+        std::vector<int32_t> labels(ci.begin()+co[r],ci.begin()+co[r+1]); std::sort(labels.begin(),labels.end());
+        require(labels.front()>=0 && std::adjacent_find(labels.begin(),labels.end())==labels.end(),"reference class roles must be unique and nonnegative");
+        for (size_t a=ro[r];a<ro[r+1];++a) {
+            require(rt[a]<T && std::binary_search(labels.begin(),labels.end(),tc[rt[a]]),"annotation target class missing from reference roles");
+            require(multiplicities[a]>0 && other[a]<ci.size() && other[a]<=R,"invalid annotation multiplicity or other-class count");
+            for (size_t c=0;c<rg.size();++c) if (rg[c][a]<0) modeled[a]=0;
+        }
+    }
+    auto encode=[&](std::string_view s) {
+        require(s.size()<=size_t(std::numeric_limits<int>::max()),"sequence exceeds native search length");
+        std::vector<uint8_t> out; out.reserve(s.size());
+        for (char a:s) { const auto c=codec().encode(a); require(c!=Codec::kInvalid && code[c]>=0,"sequence symbol absent from similarity matrix"); out.push_back(uint8_t(code[c])); }
+        return out;
+    };
+    auto self=[&](const std::vector<uint8_t>& s) {
+        double value=0; for (auto a:s) value+=similarity[a][a];
+        if (!std::isfinite(value)) throw std::overflow_error("self similarity overflow"); return value;
+    };
+    auto key=[](const std::vector<uint8_t>& s) { return std::string(s.begin(),s.end()); };
+    std::vector<std::vector<uint8_t>> refs(R),qs(Q); std::vector<double> rself(R),qself(Q);
+    std::unordered_map<std::string,uint32_t> identities;
+    for (size_t r=0;r<R;++r) { refs[r]=encode(ref_seq(uint32_t(r))); rself[r]=self(refs[r]); require(identities.emplace(key(refs[r]),uint32_t(r)).second,"alignment max index references must be unique sequences"); }
+    std::vector<int64_t> exact_ref(Q,-1);
+    for (size_t i=0;i<Q;++i) { qs[i]=encode(queries[i]); qself[i]=self(qs[i]); if (auto found=identities.find(key(qs[i]));found!=identities.end()) exact_ref[i]=found->second; }
+    AlignmentMaxResult out; const double nan=std::numeric_limits<double>::quiet_NaN();
+    out.scores.assign(M,nan); out.raw_similarity.assign(M,nan); out.normalized_similarity.assign(M,nan);
+    out.pfit.assign(M,nan); out.informativeness.assign(M,nan); out.gap_counts.assign(M,0);
+    out.best_ref_ids.assign(M,-1); out.best_lane_ids.assign(M,-1);
+    out.counts.assign(M,0); out.modeled_counts.assign(M,0); out.aa_counts.assign(M,0); out.exact_counts.assign(M,0);
+    out.candidate_counts.assign(Q,0); out.nw_pairs.assign(Q,0); out.query_supported.assign(Q,1);
+    struct Worker { Searcher search; std::vector<Hit> hits; Rows rows; std::vector<int64_t> slots; std::vector<uint64_t> aa_seen;
+        Worker(const Index& index,size_t targets):search(index),slots(targets,-1),aa_seen(targets,0) {} };
+    parallel_for(Q,threads,[&]{return Worker(*this,T);},[&](size_t i,Worker& w) {
+        for (size_t c=0;c<qg.size();++c) if (qg[c][i]<0) out.query_supported[i]=0;
+        for (size_t k=qo[i];k<qo[i+1];++k) w.slots[qt[k]]=int64_t(k);
+        if (count_unit::bounded_eligible(p,max_indels) && !queries[i].empty() && queries[i].size()<=63) {
+            count_unit::Pattern pattern(queries[i],codec()); w.hits.clear();
+            for (uint32_t r=0;r<R;++r) {
+                const int d=count_unit::bounded_distance(pattern,ref_seq(r),codec(),p,max_indels);
+                // Only ref_id is consumed below; no operation trace is reconstructed.
+                if (d<=p.max_total_edits) w.hits.push_back(Hit{r,d,0,0,0});
+            }
+        } else w.search.search_into(queries[i],p,w.hits,max_indels);
+        out.candidate_counts[i]=w.hits.size();
+        const int64_t er=exact_ref[i];
+        const uint64_t N=ci.size();
+        for (const auto& hit:w.hits) {
+            const uint32_t r=hit.ref_id; const bool exact=int64_t(r)==er; bool align=false;
+            for (size_t a=ro[r];a<ro[r+1];++a) {
+                const uint32_t t=rt[a]; const int64_t k=w.slots[t]; if (k<0) continue;
+                if (exact) add(out.exact_counts[size_t(k)],multiplicities[a]);
+                if (exclude_exact && exact) continue;
+                add(out.counts[size_t(k)],multiplicities[a]);
+                if (w.aa_seen[t]!=uint64_t(r)+1) { ++out.aa_counts[size_t(k)]; w.aa_seen[t]=uint64_t(r)+1; }
+                if (modeled[a] && out.query_supported[i]) { add(out.modeled_counts[size_t(k)],multiplicities[a]); align=true; }
+            }
+            if (!align) continue; // Filter requested targets and model availability BEFORE NW.
+            const Cell features=nw(qs[i],refs[r],similarity,gap,w.rows); ++out.nw_pairs[i];
+            const double normalized=features.raw-std::max(qself[i],rself[r]);
+            if (!std::isfinite(features.score) || !std::isfinite(normalized)) throw std::overflow_error("NW similarity overflow");
+            for (size_t a=ro[r];a<ro[r+1];++a) {
+                const uint32_t t=rt[a]; const int64_t slot=w.slots[t]; if (slot<0 || !modeled[a]) continue;
+                const size_t k=size_t(slot);
+                const uint64_t numerator=other[a]+1;
+                require(N>=numerator,"reference-role denominator is too small");
+                const double information=-std::log10(double(numerator)/double(N));
+                double eta=coefficients[0];
+                for (size_t c=0;c<offsets.size();++c) eta+=offsets[c][size_t(qg[c][i])][size_t(rg[c][a])];
+                eta+=coefficients[1]*normalized+coefficients[2]*features.gaps;
+                if (!std::isfinite(eta)) throw std::overflow_error("linear predictor overflow");
+                const double probability=-std::expm1(-std::exp(eta)),score=probability*information;
+                if (out.best_lane_ids[k]<0 || score>out.scores[k] || (score==out.scores[k] && int64_t(a)<out.best_lane_ids[k])) {
+                    out.scores[k]=score; out.best_ref_ids[k]=r; out.best_lane_ids[k]=int64_t(a);
+                    out.raw_similarity[k]=features.raw; out.normalized_similarity[k]=normalized; out.gap_counts[k]=features.gaps;
+                    out.pfit[k]=probability; out.informativeness[k]=information;
+                }
+            }
+        }
+        for (size_t k=qo[i];k<qo[i+1];++k) {
+            if (out.query_supported[i] && !out.counts[k]) out.scores[k]=0;
+            w.slots[qt[k]]=-1; w.aa_seen[qt[k]]=0;
+        }
+    });
+    return out;
+}
+} // namespace seqtree