diff --git a/docs/gapblock.rst b/docs/gapblock.rst index 88b0e8b..de9e1d7 100644 --- a/docs/gapblock.rst +++ b/docs/gapblock.rst @@ -432,3 +432,113 @@ references once with shared encoded inputs and per-worker alignment scratch; output occupies O(Q) storage without sparse hit objects or a dense matrix. The bounded minimum is not an unbounded nearest result. No statistical model or cutoff recommendation is introduced. + +Factored paired-feature offsets +------------------------------- + +``PairedFeatureOffsets(query_codes, reference_codes, costs, weight=1.0, scale=1)`` +represents an additive offset without a query-by-reference penalty table. Each +code row contains ``[head_0, tail_0, ..., head_(D-1), tail_(D-1)]`` for a positive, +shared depth ``D``. Integer codes index the query rows/reference columns of the +shared nonnegative rectangular ``costs`` table. The offset is:: + + scale * round_even(weight * sum(max(costs[q_head_i][r_head_i], + costs[q_tail_i][r_tail_i])) / D) + +``round_even`` rounds to nearest integer with ties to even. Binary64 multiplication +by ``weight`` precedes division by ``D``; ``scale`` multiplies after rounding. +Weight must be finite and nonnegative, scale a nonnegative integer. Code rows, +axes, costs and combined int32 score bounds are checked before worker execution. +Both empty axes are accepted with empty code rows; nonempty axes require positive +depth. Zero weight or scale supplies a zero offset with the same validation. +The descriptor owns its inputs and exposes no mutable Python fields; one batch +copy remains shared and immutable while the GIL is released. + +``score_matrix``, ``count_batch`` (including ``linear_mass=True``), ``ball_batch`` +and ``group_audit_batch`` accept ``feature_offsets=None``. The offset adds to any +existing ``group_distances`` penalty. The audit component excludes both additive +offsets. ``paired_sum_count_batch`` accepts ``feature_offsets_alpha=None`` and +``feature_offsets_beta=None`` independently. Exact exclusion still uses full +encoded sequence identity, including both lanes for linked rows. Omission +preserves the existing scorer. Storage is O((Q+R)*D+C), where ``C`` is the number +of code-pair costs; no Q-by-R offset allocation is made. + +Native capability can be checked by importing ``PairedFeatureOffsets`` from +``seqtree._core`` and checking the native reducer docstring for ``feature_offsets``. +The Python wrapper alone does not prove that an installed extension supports it. + +Shared batch cutoffs +-------------------- + +``count_batch`` (including ``linear_mass=True``) and ``paired_sum_count_batch`` +accept ``shared_thresholds=False``. With ``True``, ``thresholds`` is one flat +integer cutoff list for the whole query batch. Python validates that list once; +the native reducer owns and sorts one immutable cutoff vector shared by all +workers. Outputs retain their usual query-by-cutoff shape, cutoff order, +duplicate cutoffs, boundary ties, negative-cutoff behavior and exact exclusion. +There is no cross-call cache. An empty cutoff list produces an empty row per +query; empty query axes produce empty outputs. Default ``False`` preserves the +existing per-query row API. The native entry point represents shared cutoffs as +exactly one row, including for empty query axes, and checks that shape before +workers. Int32 cutoff bounds and uint64 mass overflow validation remain intact. + +Native capability is exposed by ``shared_thresholds`` in the reducer docstrings; +check both the Python signature and the installed native extension. + +Grouped label populations +------------------------- + +``grouped_count_batch`` and ``grouped_paired_sum_count_batch`` prepare one immutable +query/reference scoring population and return a native iterator over contiguous +query chunks. The geometry is exactly the existing count/summed-pair geometry, +including positional weights, additive group distances and factored feature +offsets. Each query/reference key is scored once, then its single cutoff bin +contributes to every attached target label. No statistical reducer or domain data +is added. Default count APIs are unchanged. + +Required label metadata is generic integer coding: + +* ``reference_target_offsets`` and ``reference_target_ids`` form CSR rows for + the reference axis; each row contains sorted unique target IDs. +* ``target_peptide_ids`` assigns each target to a parent label (peptide downstream). +* ``reference_exclusive_peptide_ids`` is the parent label a reference belongs to + exclusively, or -1 for shared/unselected labels. Compute it from all original + labels, before restricting requested targets. Empty target memberships are valid; + an exclusive ID may be provided even when that reference has no selected target + membership. Included memberships must agree with any supplied exclusive ID. +* ``query_gene_ids`` and ``reference_gene_ids`` code a separate canonical identity + axis for same-identity target subsets. This equality does not alter allele/group + costs. For linked pairs, encode a tuple of the two identities into one common ID. + +Thresholds are one strictly increasing int32 list, shared across the population; +negative or empty cutoffs are supported. Every chunk owns read-only uint32 buffers, +with no NumPy dependency or nested Python count lists. ``numpy.asarray(buffer)`` +and ``memoryview(buffer)`` expose the same allocation without copying. For ``C`` +chunk rows, ``T`` targets, ``P`` parent labels and ``K`` cutoffs: + +* ``chunk.start`` / ``chunk.stop`` identify its original query interval. +* ``chunk.counts`` has shape ``(C, 1+2*T+P, K)``. Channels are whole population, + targets, same-identity targets, exclusive parent populations, in that order. +* ``chunk.sizes`` has shape ``(C, 1+2*T+P)`` in the same channel order. +* ``chunk.exact_counts`` has shape ``(C, T)`` and always reports raw full-key exact + multiplicity, including when exact references are excluded from other outputs. + +With ``exclude_exact=True``, every curve and corresponding population size is +punctured by the full encoded query key, irrespective of additive groups/identities. +For paired sum geometry both full sequences must match. Subtracting an exclusive +parent channel from the whole channel therefore yields the exact complement, +including references carrying shared labels. Count overflow is excluded explicitly +by requiring at most uint32-max reference rows; addressable buffer sizes and all +metadata axes/ranges are validated. Empty query axes yield no chunks; empty +reference axes yield zero curves and sizes with ordinary shapes. + +``chunk_size=64`` controls output memory: O(C*(1+2*T+P)*(K+1)+C*T), plus O(Q+R) +encoded inputs, CSR memberships and geometry tables. Worker scratch is only +O(maximum sequence length); histogram bins live directly in disjoint chunk rows. +Keep only the current chunk for bounded memory. A retained buffer owns its chunk +and safely outlives the iterator. Workers use the existing CPU-budget scaffold; +preparation and each native chunk release the GIL. No persisted computed cache, +per-query Python search loop, global Q-by-label allocation or new distance formula +is introduced. Native capability is ``gapblock_grouped_count_batch`` plus +``GroupedCountBatch``/``GroupedCountChunk`` in ``seqtree._core``; the constructor +and chunk methods document the buffer/feature support. diff --git a/include/seqtree/seqtree.hpp b/include/seqtree/seqtree.hpp index d8ac3c0..8073111 100644 --- a/include/seqtree/seqtree.hpp +++ b/include/seqtree/seqtree.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -228,6 +229,20 @@ std::vector align_dist_matrix(const std::vector& queries, // 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. +// Factored nonnegative offsets: paired feature codes, one shared code cost table. +// Round weight * sum(max(head cost, tail cost)) / depth to nearest, ties even, +// then multiply by scale. Input descriptors are copied into immutable batch state. +struct PairedFeatureOffsets { + std::vector> query_codes, reference_codes, costs; + double weight; + int32_t scale; + PairedFeatureOffsets(std::vector> q, + std::vector> r, + std::vector> c, double w = 1.0, int32_t s = 1) + : query_codes(std::move(q)), reference_codes(std::move(r)), costs(std::move(c)), + weight(w), scale(s) {} +}; + std::vector gapblock_matrix(const std::vector& queries, const std::vector& refs, Alphabet, const SubstitutionMatrix* matrix, int32_t gap_open, @@ -238,7 +253,8 @@ std::vector gapblock_matrix(const std::vector& queries, const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}); + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt); // Exhaustive ball rows in reference order; hit edit fields are unset (zero). std::vector> gapblock_ball_batch(const std::vector& queries, @@ -250,7 +266,8 @@ std::vector> gapblock_ball_batch(const std::vector const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}); + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt); // Exhaustive bounded reductions using the same gapblock cell. Top-k ties use // (penalty, ref_id); paired penalty is max(lane scores). Exact exclusion uses @@ -267,7 +284,8 @@ gapblock_group_audit_batch(const std::vector& queries, const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}); + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt); std::vector> gapblock_topk_batch(const std::vector& queries, const std::vector& refs, uint32_t k, Alphabet alphabet, @@ -287,7 +305,8 @@ std::vector> gapblock_count_batch(const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}); + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt, bool shared_thresholds = false); 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, @@ -297,7 +316,8 @@ std::pair>, std::vector> const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}); + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt, bool shared_thresholds = false); std::vector> gapblock_paired_count_batch( const std::vector& qa, const std::vector& qb, const std::vector& ra, const std::vector& rb, @@ -319,7 +339,9 @@ std::pair>, std::vector> 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 std::vector>& reference_weights_beta = {}, + const std::optional& feature_offsets_alpha = std::nullopt, + const std::optional& feature_offsets_beta = std::nullopt, bool shared_thresholds = false); // --------------------------------------------------------------------------------------- // Plain string edit distances on raw characters, unit costs. Deliberately alphabet-agnostic: @@ -341,4 +363,48 @@ std::vector hamming_matrix(const std::vector& a, std::vector levenshtein_matrix(const std::vector& a, const std::vector& b, int threads = 0); +// Generic label histograms over one prepared native query/reference population. +struct GroupedGeometry { + std::vector additional_features; + std::optional> gap_prior; + std::string alignment_mode = "gapblock"; + int32_t alignment_scale = 1; + uint32_t trim_left = 0, trim_right = 0; + int32_t score_quantum = 1; + std::vector> weights, group_distances, query_weights, reference_weights; + std::vector query_groups, reference_groups; + std::optional features; +}; +struct CountBuffer { + std::vector data; + std::vector shape; +}; +struct GroupedCountChunk { + size_t start = 0, stop = 0; + CountBuffer counts, sizes, exact_counts; +}; +class GroupedCountBatch { +public: + GroupedCountBatch(const std::vector& qa, const std::vector& ra, + std::vector thresholds, std::vector target_offsets, + std::vector target_ids, std::vector target_peptides, + std::vector exclusive_peptides, std::vector query_genes, + std::vector reference_genes, size_t chunk_size, Alphabet alphabet, + const SubstitutionMatrix* matrix, int32_t go, int32_t ge, std::vector prior, + uint32_t width, int threads, bool exclude_exact, const GroupedGeometry& alpha, + const std::optional& beta = std::nullopt, + const std::vector& qb = {}, const std::vector& rb = {}, + std::optional> query_withheld_peptides = std::nullopt, + std::vector reference_peptide_offsets = {}, + std::vector reference_peptide_ids = {}); + std::optional next(); + size_t query_count() const; + size_t reference_count() const; + size_t target_count() const; + size_t peptide_count() const; +private: + struct Impl; + std::shared_ptr impl_; +}; + } // namespace seqtree diff --git a/pyproject.toml b/pyproject.toml index ed02978..083934d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "seqtree" -version = "1.1.0" +version = "1.1.1" description = "Fast fuzzy search over biological sequences (C++ core, Python bindings)" readme = "README.md" requires-python = ">=3.10" diff --git a/python/seqtree/gapblock.py b/python/seqtree/gapblock.py index ab2e7bf..695bcdc 100644 --- a/python/seqtree/gapblock.py +++ b/python/seqtree/gapblock.py @@ -40,7 +40,7 @@ import math from collections.abc import Callable, Iterable, Sequence -from ._core import Index, ScoreMatrix, SearchParams, SubstitutionMatrix +from ._core import Index, ScoreMatrix, SearchParams, SubstitutionMatrix, PairedFeatureOffsets 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 @@ -51,10 +51,12 @@ 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 +from ._core import gapblock_grouped_count_batch as _gapblock_grouped_count_batch, _GroupedGeometry + __all__ = [ "gapblock_score", "score_matrix", "ball_batch", "group_audit_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", + "IslandProfile", "PairedFeatureOffsets", "grouped_count_batch", "grouped_paired_sum_count_batch", ] #: A prior value large enough that no layout carrying it can ever win, yet small enough that @@ -528,7 +530,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, + query_position_weights=None, reference_position_weights=None, feature_offsets=None, ) -> ScoreMatrix: """Gap-block penalty of every query against every reference, in C++ with the GIL released. @@ -552,6 +554,8 @@ def score_matrix( 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. + feature_offsets: Optional :class:`PairedFeatureOffsets` descriptor. Adds + paired-code costs averaged at shared depth, ties-even rounded, then scaled. 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. @@ -587,7 +591,7 @@ def score_matrix( 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), - *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r)) + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets) def _batch_integer(value, name, lower=0, upper=(1 << 31) - 1): @@ -668,7 +672,11 @@ def validate(rows, axis): return validate(query_weights, queries), validate(reference_weights, refs) -def _threshold_rows(thresholds, n_queries): +def _threshold_rows(thresholds, n_queries, shared_thresholds=False): + if not isinstance(shared_thresholds, bool): + raise TypeError("shared_thresholds must be a boolean") + if shared_thresholds: + return [[_batch_integer(value, "threshold", -(1 << 31)) for value in thresholds]] rows = [list(row) for row in thresholds] if len(rows) != n_queries: raise ValueError("thresholds must have one row per query") @@ -702,13 +710,15 @@ 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_position_weights=None, reference_position_weights=None): + query_position_weights=None, reference_position_weights=None, feature_offsets=None, shared_thresholds=False): """Count all reference scores <= each query's requested integer thresholds. 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 + duplicates. With ``shared_thresholds=True``, supply one flat cutoff list; + it is validated and sorted once for the entire batch. 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). @@ -722,20 +732,20 @@ def count_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_exten 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)) + thresholds = _threshold_rows(thresholds, len(q), shared_thresholds) go, ge, cube, width, threads = _batch_options((q, r), matrix, gap_open, gap_extend, gap_prior, threads) 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), - *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r)) + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets, shared_thresholds) def group_audit_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_position_weights=None, reference_position_weights=None): + query_position_weights=None, reference_position_weights=None, feature_offsets=None): """Stream component-ball count, full-score count and bounded nearest in one batch. Geometry/offset options match ball_batch. Component excludes the nonnegative @@ -752,14 +762,14 @@ def group_audit_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap 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)) + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets) 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_position_weights=None, reference_position_weights=None): + query_position_weights=None, reference_position_weights=None, feature_offsets=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 @@ -776,7 +786,7 @@ def ball_batch(queries, refs, thresholds, matrix=None, gap_open=None, gap_extend 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), - *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r)) + *_sequence_weights(query_position_weights, reference_position_weights, position_weights_by_length, q, r), feature_offsets) def paired_topk_batch(queries_a, queries_b, refs_a, refs_b, k=10, matrix=None, @@ -823,7 +833,8 @@ def paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, mat reference_group_ids_alpha=None, reference_group_ids_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): + reference_position_weights_alpha=None, reference_position_weights_beta=None, + feature_offsets_alpha=None, feature_offsets_beta=None, shared_thresholds=False): """Return (counts, linear masses) for linked rows under summed lane distances. Shared matrix/gap costs/prior and separate full-coordinate lane weights and @@ -833,12 +844,13 @@ def paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, mat 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. + scratch; no matrix/edges. ``shared_thresholds=True`` accepts one flat cutoff + list validated and sorted once for all queries. The 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)) + thresholds = _threshold_rows(thresholds, len(qa), shared_thresholds) 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) @@ -849,7 +861,226 @@ def paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, mat _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], - aw[0], bw[0], aw[1], bw[1]) + aw[0], bw[0], aw[1], bw[1], feature_offsets_alpha, feature_offsets_beta, shared_thresholds) + + +def _grouped_metadata(q, r, thresholds, reference_target_offsets, reference_target_ids, + target_peptide_ids, reference_exclusive_peptide_ids, + query_gene_ids, reference_gene_ids, chunk_size): + cuts = [_batch_integer(v, "threshold", -(1 << 31)) for v in thresholds] + offsets = [_batch_integer(v, "target offset", upper=(1 << 32)-1) for v in reference_target_offsets] + labels = [_batch_integer(v, "target ID", upper=(1 << 32)-1) for v in reference_target_ids] + peptides = [_batch_integer(v, "peptide ID") for v in target_peptide_ids] + exclusive = [_batch_integer(v, "exclusive peptide ID", -1) for v in reference_exclusive_peptide_ids] + qgenes = [_batch_integer(v, "query gene ID") for v in query_gene_ids] + rgenes = [_batch_integer(v, "reference gene ID") for v in reference_gene_ids] + chunk_size = _batch_integer(chunk_size, "chunk_size", 1) + return cuts, offsets, labels, peptides, exclusive, qgenes, rgenes, chunk_size + + +def _withheld_metadata(q, query_withheld_peptide_ids, reference_peptide_offsets, reference_peptide_ids): + if query_withheld_peptide_ids is None: + if reference_peptide_offsets is not None or reference_peptide_ids is not None: + raise ValueError("reference peptide memberships require query withheld peptide IDs") + return None, [], [] + if reference_peptide_offsets is None or reference_peptide_ids is None: + raise ValueError("withheld peptide IDs require full reference peptide CSR") + withheld = [_batch_integer(v, "withheld peptide ID", -1) for v in query_withheld_peptide_ids] + offsets = [_batch_integer(v, "peptide offset", upper=(1 << 32)-1) for v in reference_peptide_offsets] + peptides = [_batch_integer(v, "reference peptide ID") for v in reference_peptide_ids] + if len(withheld) != len(q): + raise ValueError("withheld peptide IDs must match query axis") + return withheld, offsets, peptides + + +def _additional_geometry(geometry, features): + if features is None: + return + terms = list(features) + if any(not isinstance(term, PairedFeatureOffsets) for term in terms): + raise ValueError("additional_feature_offsets must contain PairedFeatureOffsets") + if terms: + if not hasattr(geometry, "additional_features"): + raise RuntimeError("additional factored offsets require native additional_features support") + geometry.additional_features = terms + + +def _lane_prior(geometry, prior, width): + if prior is not None: + if not hasattr(geometry, "gap_prior"): + raise RuntimeError("lane priors require native grouped gap_prior support") + geometry.gap_prior = _prior_cube(prior, width) + + +def _alignment_options(geometry, q, r, mode, scale, left, right, matrix, prior, + *weights): + if mode not in ("gapblock", "global", "local"): + raise ValueError("alignment_mode must be gapblock, global or local") + scale = _batch_integer(scale, "alignment scale", 1) + left = _batch_integer(left, "trim_left") + right = _batch_integer(right, "trim_right") + if mode == "gapblock": + if scale != 1 or left or right: + raise ValueError("alignment scale/trims require global or local mode") + return + if matrix is None: + raise ValueError("affine alignment requires a similarity matrix") + if prior is not None or any(w is not None for w in weights): + raise ValueError("affine alignment does not accept position weights or gap prior") + if any(left + right > len(sequence) for axis in (q, r) for sequence in axis): + raise ValueError("alignment trims exceed sequence length") + if not hasattr(geometry, "alignment_mode"): + raise RuntimeError("affine grouped scoring requires native alignment_mode support") + geometry.alignment_mode = mode + geometry.alignment_scale = scale + geometry.trim_left = left + geometry.trim_right = right + + +def _grouped_geometry(q, r, weights, query_groups, reference_groups, distances, + query_weights, reference_weights, features, score_quantum=1): + quantum = _batch_integer(score_quantum, "score quantum", 1) + geometry = _GroupedGeometry() + if quantum != 1: + if not hasattr(geometry, "score_quantum"): + raise RuntimeError("alignment quantization requires native score_quantum support") + geometry.score_quantum = quantum + geometry.weights = _position_weights(weights, q, r) + geometry.query_groups, geometry.reference_groups, geometry.group_distances = _group_options( + 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 + return geometry + + +def grouped_count_batch(queries, refs, thresholds, *, reference_target_offsets, + reference_target_ids, target_peptide_ids, reference_exclusive_peptide_ids, + query_gene_ids, reference_gene_ids, chunk_size=64, 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_position_weights=None, reference_position_weights=None, + 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): + """Yield bounded label curves from one immutable prepared scoring population. + + ``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)) + - 2*similarity(q,r)), using raw diagonal sums rather than self-alignment scores. + Nonnegative trim_left/trim_right remove residues for scoring only; exact + exclusion and offsets retain full input identities. Gap costs are raw similarity + magnitudes in these modes. Position weights and gap_prior are rejected. + Paired lanes compute distances independently before summation. Scratch rows are + reused per worker; no dense query/reference score matrix is produced. + + ``additional_feature_offsets`` is an optional iterable of PairedFeatureOffsets. + Every term is independently ties-even rounded/scaled, then added to the primary + feature and group offsets. All descriptors are snapshotted/validated once. + Combined score/offset bounds must fit int32 per lane. In the paired API, use + additional_feature_offsets_alpha/beta; gap_prior_alpha/beta independently + override the shared gap_prior. None inherits the shared prior for that lane. + + ``score_quantum`` is a positive int32. Floor the raw alignment score to a + multiple of this quantum before adding group/feature offsets. Default1 leaves + every score unchanged. Paired lanes floor independently before their sum. + This option applies only to grouped APIs. + + Cutoffs are one strictly increasing flat integer list. References are a caller's + deduplicated search-key axis; CSR offsets/IDs assign each key to zero or more + target IDs. Per-reference exclusive peptide IDs come from the full label set, + with -1 for shared/unselected labels; target_peptide_ids maps each target. + Canonical gene IDs define the auxiliary same-gene subset independently of + allele costs. Channels are whole, targets, same-gene targets, exclusive peptides. + 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. + 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 + the target_peptide_ids namespace. Extra IDs beyond output peptide channels are + valid. Selected target parents and static exclusives must agree with full CSR. + Withholding removes every membership of that peptide, drops keys with no + survivors, and recomputes exclusive surviving peptides before exact exclusion. + Raw exact counts include only surviving targets. Full CSR without query IDs is + rejected; None preserves the existing traversal. Encoding/geometry validation + occurs once; each query/reference is scored once + under the existing kernel, with GIL-released bounded workers. Memory is bounded + by chunk_size, labels and cutoffs, never Q*labels. Retained chunks own their data. + """ + q, r = list(queries), list(refs) + metadata = _grouped_metadata(q, r, thresholds, reference_target_offsets, + reference_target_ids, target_peptide_ids, reference_exclusive_peptide_ids, + query_gene_ids, reference_gene_ids, chunk_size) + 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) + _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) + withheld = _withheld_metadata(q, query_withheld_peptide_ids, + reference_peptide_offsets, reference_peptide_ids) + return _gapblock_grouped_count_batch(q, r, *metadata, alphabet, matrix, go, ge, cube, width, + threads, exclude_exact, geometry, + query_withheld_peptide_ids=withheld[0], + reference_peptide_offsets=withheld[1],reference_peptide_ids=withheld[2]) + + +def grouped_paired_sum_count_batch(queries_a, queries_b, refs_a, refs_b, thresholds, *, + reference_target_offsets, reference_target_ids, target_peptide_ids, + reference_exclusive_peptide_ids, query_gene_ids, reference_gene_ids, + chunk_size=64, 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, + query_position_weights_alpha=None, query_position_weights_beta=None, + reference_position_weights_alpha=None, reference_position_weights_beta=None, + feature_offsets_alpha=None, feature_offsets_beta=None, + query_withheld_peptide_ids=None, reference_peptide_offsets=None, reference_peptide_ids=None, score_quantum_alpha=1, score_quantum_beta=1, + alignment_mode="gapblock", alignment_scale=1, trim_left=0, trim_right=0, + additional_feature_offsets_alpha=None, additional_feature_offsets_beta=None, + gap_prior_alpha=None, gap_prior_beta=None): + """Grouped buffers under existing linked sum geometry; see grouped_count_batch. + + Composite full-key exclusion requires both lanes equal. Canonical paired gene + identities are caller-encoded tuple IDs, with equality requiring both genes. + Shared matrix/prior and lane-specific weights/group/feature offsets retain the + paired_sum_count_batch geometry. Counts, sizes and raw exact counts use uint32. + """ + 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") + metadata = _grouped_metadata(qa, ra, thresholds, reference_target_offsets, + reference_target_ids, target_peptide_ids, reference_exclusive_peptide_ids, + query_gene_ids, reference_gene_ids, chunk_size) + go, ge, cube, width, threads = _batch_options((qa, qb, ra, rb), matrix, gap_open, gap_extend, gap_prior, threads) + alpha = _grouped_geometry(qa, ra, position_weights_by_length_alpha, query_group_ids_alpha, + reference_group_ids_alpha, group_distances_alpha, query_position_weights_alpha, + reference_position_weights_alpha, feature_offsets_alpha, score_quantum_alpha) + beta = _grouped_geometry(qb, rb, position_weights_by_length_beta, query_group_ids_beta, + reference_group_ids_beta, group_distances_beta, query_position_weights_beta, + reference_position_weights_beta, feature_offsets_beta, score_quantum_beta) + _additional_geometry(alpha, additional_feature_offsets_alpha) + _additional_geometry(beta, additional_feature_offsets_beta) + _lane_prior(alpha, gap_prior_alpha, width) + _lane_prior(beta, gap_prior_beta, width) + _alignment_options(alpha, qa, ra, alignment_mode, alignment_scale, trim_left, trim_right, + matrix, gap_prior if gap_prior_alpha is None else gap_prior_alpha, position_weights_by_length_alpha, query_position_weights_alpha, + reference_position_weights_alpha) + _alignment_options(beta, qb, rb, alignment_mode, alignment_scale, trim_left, trim_right, + matrix, gap_prior if gap_prior_beta is None else gap_prior_beta, position_weights_by_length_beta, query_position_weights_beta, + reference_position_weights_beta) + withheld = _withheld_metadata(qa, query_withheld_peptide_ids, + reference_peptide_offsets, reference_peptide_ids) + return _gapblock_grouped_count_batch(qa, ra, *metadata, alphabet, matrix, go, ge, cube, width, + threads, exclude_exact, alpha, beta, qb, rb, + query_withheld_peptide_ids=withheld[0], + reference_peptide_offsets=withheld[1],reference_peptide_ids=withheld[2]) class GapBlockIndex: diff --git a/src/_bindings.cpp b/src/_bindings.cpp index cc57712..94172f2 100644 --- a/src/_bindings.cpp +++ b/src/_bindings.cpp @@ -284,6 +284,24 @@ PyType_Slot kScoreMatrixSlots[] = { { 0, nullptr } }; +int count_buffer_getbuffer(PyObject* obj, Py_buffer* view, int flags) { + if (flags & PyBUF_WRITABLE) { + PyErr_SetString(PyExc_BufferError,"grouped count buffers are read-only"); + return -1; + } + const auto* buffer=nb::inst_ptr(obj); + std::vector shape(buffer->shape.begin(),buffer->shape.end()),strides(shape.size()); + Py_ssize_t stride=sizeof(uint32_t); + for (size_t d=shape.size();d-- >0;) {strides[d]=stride;stride*=shape[d];} + fill_view(obj,view,flags,const_cast(buffer->data.data()),int(shape.size()), + shape.data(),strides.data(),sizeof(uint32_t),"I"); + return 0; +} +PyType_Slot kCountBufferSlots[] = { + {Py_bf_getbuffer,(void*)count_buffer_getbuffer}, + {Py_bf_releasebuffer,(void*)release_view},{0,nullptr} +}; + // A read-only 1-D typed window onto memory a TextResult owns. Same reasoning as ScoreMatrix: // seqtree has no runtime dependencies, so the way out is the buffer protocol, not an ndarray. // `owner` is a strong reference to the TextResult, so a view outlives the dict it came from. @@ -361,7 +379,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>& query_weights, const std::vector>& reference_weights) { + 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 std::optional& feature_offsets) { Alphabet alph = parse_alphabet(alphabet); ScoreMatrix out; out.rows = queries.size(); @@ -369,7 +387,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, query_weights, reference_weights); + gap_extend, prior, prior_width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights, feature_offsets); } return out; } @@ -797,11 +815,15 @@ NB_MODULE(_core, m) { "Dense (n_queries, n_refs) distance matrix d = s(a,a) + s(b,b) - 2*s(a,b), the " "sequence-level Gram transform of the alignment scores. Non-negative, zero on identity."); + nb::class_(m, "PairedFeatureOffsets", "Immutable factored paired-feature offsets; validated by each batch.") + .def(nb::init>, std::vector>, std::vector>, double, int32_t>(), + nb::arg("query_codes"), nb::arg("reference_codes"), nb::arg("costs"), nb::arg("weight")=1.0, nb::arg("scale")=1); + m.def("gapblock_matrix", &py_gapblock_matrix, nb::arg("queries"), nb::arg("refs"), nb::arg("alphabet") = "aa", nb::arg("matrix") = std::nullopt, nb::arg("gap_open") = 1, nb::arg("gap_extend") = 1, nb::arg("prior") = std::vector{}, nb::arg("prior_width") = 0, nb::arg("threads") = 0, nb::arg("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>{}, + 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>{}, nb::arg("feature_offsets") = std::nullopt, "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,29 +834,106 @@ 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, const std::vector>& query_weights, const std::vector>& reference_weights) { + nb::class_(m,"CountBuffer",nb::type_slots(kCountBufferSlots), + "Read-only uint32 C-contiguous buffer; numpy.asarray and memoryview wrap without copying."); + nb::class_(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); + nb::class_(m,"_GroupedGeometry") + .def(nb::init<>()).def_rw("weights",&GroupedGeometry::weights) + .def_rw("additional_features",&GroupedGeometry::additional_features) + .def_rw("gap_prior",&GroupedGeometry::gap_prior) + .def_rw("alignment_mode",&GroupedGeometry::alignment_mode) + .def_rw("alignment_scale",&GroupedGeometry::alignment_scale) + .def_rw("trim_left",&GroupedGeometry::trim_left) + .def_rw("trim_right",&GroupedGeometry::trim_right) + .def_rw("query_groups",&GroupedGeometry::query_groups) + .def_rw("reference_groups",&GroupedGeometry::reference_groups) + .def_rw("group_distances",&GroupedGeometry::group_distances) + .def_rw("query_weights",&GroupedGeometry::query_weights) + .def_rw("reference_weights",&GroupedGeometry::reference_weights) + .def_rw("features",&GroupedGeometry::features) + .def_rw("score_quantum",&GroupedGeometry::score_quantum, + "Positive alignment quantum; floor raw alignment before additive offsets."); + nb::class_(m,"GroupedCountBatch", + "One prepared grouped scoring population, yielding owned bounded buffers; no NumPy dependency.") + .def_prop_ro("query_count",&GroupedCountBatch::query_count) + .def_prop_ro("reference_count",&GroupedCountBatch::reference_count) + .def_prop_ro("target_count",&GroupedCountBatch::target_count) + .def_prop_ro("peptide_count",&GroupedCountBatch::peptide_count) + .def("__iter__",[](GroupedCountBatch& batch)->GroupedCountBatch& {return batch;},nb::rv_policy::reference_internal) + .def("__next__",[](GroupedCountBatch& batch) { + std::optional chunk; + {nb::gil_scoped_release release;chunk=batch.next();} + if (!chunk) throw nb::stop_iteration(); + return std::move(*chunk); + }); + m.def("gapblock_grouped_count_batch",[](const std::vector& q, + const std::vector& r, std::vector cuts, + std::vector offsets,std::vector labels, + std::vector target_peptides,std::vector exclusive, + std::vector qgenes,std::vector rgenes,size_t chunk_size, + const std::string& alphabet,const std::optional& matrix, + int32_t go,int32_t ge,std::vector prior,uint32_t width,int threads, + bool exclude,const GroupedGeometry& alpha,const std::optional& beta, + const std::vector& qb,const std::vector& rb, + std::optional> withheld, + std::vector peptide_offsets,std::vector peptide_ids) { + const auto alph=parse_alphabet(alphabet); + // Snapshot mutable Python geometry before releasing the GIL. Native + // iteration only reads its own prepared state, never a Python object. + const auto owned_alpha=alpha; + const auto owned_beta=beta; + nb::gil_scoped_release release; + return GroupedCountBatch(q,r,std::move(cuts),std::move(offsets),std::move(labels), + std::move(target_peptides),std::move(exclusive),std::move(qgenes),std::move(rgenes), + chunk_size,alph,matrix ? &*matrix : nullptr,go,ge,std::move(prior),width,threads, + exclude,owned_alpha,owned_beta,qb,rb,std::move(withheld), + std::move(peptide_offsets),std::move(peptide_ids)); + },nb::arg("queries"),nb::arg("refs"),nb::arg("thresholds"), + 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("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{},nb::arg("refs_beta")=std::vector{}, + nb::arg("query_withheld_peptide_ids")=std::nullopt, + nb::arg("reference_peptide_offsets")=std::vector{}, + nb::arg("reference_peptide_ids")=std::vector{}, + "Generic grouped shared-threshold curves using existing gapblock scoring, feature_offsets, " + "paired sum and full-key exact exclusion. Optional query_withheld_peptide_ids removes full " + "reference peptide CSR memberships before counting and recomputes survivor exclusivity. " + "Geometry additional_features sums independently rounded factored offsets; optional geometry " + "gap_prior overrides the shared prior per lane (public gap_prior_alpha/gap_prior_beta). " + "Geometry alignment_mode selects gapblock/global/local; alignment_scale and trim_left/trim_right " + "apply to affine diagonal-sum distances before offsets. Geometry score_quantum floors each raw alignment independently before additive offsets " + "and paired summation; positive int32, default1 leaves scores unchanged. " + "Validation once, GIL released for preparation/chunks."); + 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 std::optional& feature_offsets, bool shared_thresholds) { 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, 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."); + 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, 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{}, 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>{}, nb::arg("feature_offsets")=std::nullopt, nb::arg("shared_thresholds")=false, "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, const std::vector>& query_weights, const std::vector>& reference_weights) { + 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 std::optional& feature_offsets, bool shared_thresholds) { 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, 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."); + 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{}, 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>{}, 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& 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) { + m.def("gapblock_group_audit_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 std::optional& feature_offsets) { 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); - }, 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>{}, "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); + }, 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>{}, nb::arg("feature_offsets")=std::nullopt, "Streaming component/full ball counts and minimum full score within the component ball, 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, const std::vector>& query_weights, const std::vector>& reference_weights) { + 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 std::optional& feature_offsets) { 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, 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."); + 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, 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{}, 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>{}, nb::arg("feature_offsets")=std::nullopt, "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); @@ -848,11 +947,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, const std::vector>& query_weights_alpha, const std::vector>& query_weights_beta, const std::vector>& reference_weights_alpha, const std::vector>& reference_weights_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 std::optional& feature_offsets_alpha, const std::optional& feature_offsets_beta, bool shared_thresholds) { 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, 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."); + 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, feature_offsets_alpha, feature_offsets_beta, shared_thresholds); + }, 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>{}, nb::arg("feature_offsets_alpha")=std::nullopt, nb::arg("feature_offsets_beta")=std::nullopt, nb::arg("shared_thresholds")=false, "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/count_unit.hpp b/src/count_unit.hpp index 2b01f6e..db0ab0f 100644 --- a/src/count_unit.hpp +++ b/src/count_unit.hpp @@ -23,17 +23,23 @@ struct Pattern { for (size_t i=0;icap) return cap+1; ph=(ph<<1)|1;mh<<=1; pv=mh|~(xv|ph);mv=ph&xv; } - return score; + return std::min(score,cap+1); } }; // Minimal substitution count at each prefix and exact insertion/deletion counts. @@ -59,7 +65,7 @@ inline int restricted(const Pattern& q,std::string_view r,const Codec& codec,int inline int exact_distance(const Pattern& q,std::string_view r,const Codec& codec,int cap) { if (r.size()>q.length+2 || q.length>r.size()+2) return cap+1; const int delta=int(r.size())-int(q.length); - const int d=q.distance(r); + const int d=q.distance(r,cap); if (d>cap) return cap+1; // For any path, ins+del<=d and |ins-del|=|delta|. Thus max(ins,del) // <=floor((d+|delta|)/2); both per-type caps are automatic when <=2. diff --git a/src/gapblock.cpp b/src/gapblock.cpp index d59fd04..efcfbc2 100644 --- a/src/gapblock.cpp +++ b/src/gapblock.cpp @@ -1,8 +1,11 @@ #include "seqtree/seqtree.hpp" #include "seqtree/parallel.hpp" +#include "pairwise_detail.hpp" #include +#include #include +#include #include #include @@ -70,11 +73,57 @@ struct GapInput { std::vector> weights, query_weights, reference_weights; std::vector query_groups, reference_groups; std::vector> group_distances; + std::optional features; uint32_t longest = 0; uint8_t A = 0; size_t W1 = 0; }; +// Explicit ties-even rounding is independent of the process floating rounding mode. +inline double round_even(double value) { + const double base = std::floor(value), fraction = value - base; + return base + (fraction > 0.5 || (fraction == 0.5 && std::fmod(base, 2.0) != 0.0)); +} + +uint64_t validate_feature(const PairedFeatureOffsets& f, size_t query_count, size_t reference_count) { + if (!std::isfinite(f.weight) || f.weight < 0 || f.scale < 0) + throw std::invalid_argument("feature weight and scale must be finite and nonnegative"); + if (f.query_codes.size() != query_count || f.reference_codes.size() != reference_count) + throw std::invalid_argument("feature code rows must match query/reference axes"); + if (f.costs.empty() || f.costs.front().empty()) + throw std::invalid_argument("feature costs must be a nonempty rectangular matrix"); + const size_t columns = f.costs.front().size(); + int32_t maximum = 0; + for (const auto& row : f.costs) { + if (row.size() != columns) throw std::invalid_argument("feature costs must be rectangular"); + for (auto cost : row) { + if (cost < 0) throw std::invalid_argument("feature costs must be nonnegative"); + maximum = std::max(maximum, cost); + } + } + const size_t width = !f.query_codes.empty() ? f.query_codes.front().size() : + (!f.reference_codes.empty() ? f.reference_codes.front().size() : 0); + if (width % 2 || (!width && (query_count!=0 || reference_count!=0))) + throw std::invalid_argument("feature rows must contain a positive number of code pairs"); + auto validate = [&](const auto& rows, size_t bound) { + for (const auto& row : rows) { + if (row.size() != width) throw std::invalid_argument("feature rows must have equal depth"); + for (auto code : row) + if (code < 0 || size_t(code) >= bound) + throw std::invalid_argument("feature code outside costs"); + } + }; + validate(f.query_codes, f.costs.size()); + validate(f.reference_codes, columns); + const size_t depth = width / 2; + if (depth && uint64_t(maximum) > std::numeric_limits::max() / depth) + throw std::invalid_argument("feature cost sum can overflow uint64"); + const double bound = depth && f.weight && f.scale ? round_even(f.weight * double(uint64_t(maximum) * depth) / double(depth)) * f.scale : 0; + if (!std::isfinite(bound) || bound > std::numeric_limits::max()) + throw std::invalid_argument("feature offset can overflow int32"); + return uint64_t(bound); +} + GapInput prepare(const std::vector& queries, const std::vector& refs, Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, const std::vector& prior, uint32_t width, int threads, @@ -83,7 +132,9 @@ GapInput prepare(const std::vector& queries, const std::vector& reference_groups = {}, const std::vector>& group_distances = {}, const std::vector>& query_weights = {}, - const std::vector>& reference_weights = {}) { + const std::vector>& reference_weights = {}, + const std::optional& feature_offsets = std::nullopt, + uint64_t* score_bound = nullptr) { 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()) @@ -182,6 +233,9 @@ GapInput prepare(const std::vector& queries, const std::vector& queries, const std::vector uint64_t(std::numeric_limits::max()) / in.longest) throw std::invalid_argument("gapblock weighted score can overflow int32"); const uint64_t bound = uint64_t(in.longest) * weighted_pen + - uint32_t(largest_prior) + uint32_t(largest_offset) + uint32_t(go) + + uint32_t(largest_prior) + uint32_t(largest_offset) + largest_feature_offset + uint32_t(go) + uint64_t(in.longest ? in.longest - 1 : 0) * uint32_t(ge); if (bound > std::numeric_limits::max()) throw std::invalid_argument("gapblock score can overflow int32"); + if (score_bound) *score_bound=bound; return in; } +inline int32_t feature_offset(const PairedFeatureOffsets& f,size_t i,size_t j) { + if (f.weight==0 || f.scale==0) return 0; + const auto& q=f.query_codes[i];const auto& r=f.reference_codes[j]; + uint64_t sum=0; + for (size_t k=0;k inline int32_t group_offset(const GapInput& in, size_t i, size_t j) { - return in.group_distances.empty() ? 0 : in.group_distances[in.query_groups[i]][in.reference_groups[j]]; + int32_t offset = in.group_distances.empty() ? 0 : in.group_distances[in.query_groups[i]][in.reference_groups[j]]; + if constexpr (Factored) if (in.features) offset+=feature_offset(*in.features,i,j); + return offset; } +template int32_t score_cell(const GapInput& in, size_t i, size_t j, int32_t go, int32_t ge, const std::vector& prior, int32_t* scratch) { const uint32_t m = uint32_t(in.q[i].size()), n = uint32_t(in.r[j].size()); const uint32_t M = std::max(m, n), d = m > n ? m - n : n - m; const int32_t* prow = prior.empty() ? nullptr : prior.data() + (size_t(M) * in.W1 + d) * in.W1; - return 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()); } @@ -259,7 +327,8 @@ std::vector> ranked(const GapInput& a, const GapInput* b, ui return out; } -std::vector> counted(const GapInput& a, const GapInput* b, +template +std::vector> counted_impl(const GapInput& a, const GapInput* b, const std::vector>& thresholds, int32_t go, int32_t ge, const std::vector& prior, int threads, bool exclude_exact, std::vector>* masses = nullptr, bool sum_pair_distance = false) { @@ -295,15 +364,87 @@ std::vector> counted(const GapInput& a, const GapInput* b, if (sorted.empty()) return; for (size_t j = 0; j < a.r.size(); ++j) { if (exclude_exact && a.q[i] == a.r[j] && (!b || b->q[i] == b->r[j])) continue; - const int64_t offset = int64_t(group_offset(a, i, j)) + - (sum_pair_distance ? int64_t(group_offset(*b, i, j)) : 0); + const int64_t offset = int64_t(group_offset(a, i, j)) + + (sum_pair_distance ? int64_t(group_offset(*b, i, j)) : 0); + if (offset > sorted.back().first) continue; + int64_t score = score_cell(a, i, j, go, ge, prior, scratch.data()); + if (sum_pair_distance) { + if (score > sorted.back().first) continue; + score += int64_t(score_cell(*b, i, j, go, ge, prior, scratch.data())); + } else if (b) { + score = std::max(score, score_cell(*b, i, j, go, ge, prior, scratch.data())); + } + const auto pos = std::lower_bound(sorted.begin(), sorted.end(), score, + [](const auto& threshold, int64_t value) { return threshold.first < value; }); + if (pos != sorted.end()) { + const auto bin = size_t(pos - sorted.begin()); + ++bins[bin]; + if (masses) score_bins[bin] += uint64_t(score); + } + } + uint64_t cumulative = 0, score_sum = 0; + for (size_t t = 0; t < sorted.size(); ++t) { + cumulative += bins[t]; + out[i][sorted[t].second] = cumulative; + if (masses) { + score_sum += score_bins[t]; + (*masses)[i][sorted[t].second] = uint64_t(std::max(sorted[t].first, 0)) * cumulative - score_sum; + } + } + }); + return out; +} + +template +std::vector> counted_shared_impl(const GapInput& a, const GapInput* b, + const std::vector>& thresholds, int32_t go, int32_t ge, + const std::vector& prior, int threads, bool exclude_exact, + std::vector>* masses = nullptr, bool sum_pair_distance = false) { + if (thresholds.size() != 1) + throw std::invalid_argument("shared thresholds must contain exactly one row"); + if (b && (b->q.size() != a.q.size() || b->r.size() != a.r.size())) + throw std::invalid_argument("paired query/reference axes must agree"); + if (sum_pair_distance && !b) + throw std::invalid_argument("sum distance requires two linked views"); + std::vector> out(a.q.size()); + if (masses) { + masses->resize(a.q.size()); + if (std::any_of(a.pen.begin(), a.pen.end(), [](int32_t cost) { return cost < 0; }) || + (b && std::any_of(b->pen.begin(), b->pen.end(), [](int32_t cost) { return cost < 0; }))) + throw std::invalid_argument("linear mass requires nonnegative scores"); + // Every score admitted to a row is <= its largest positive threshold. + // Bound both score sums and threshold*count before workers start. + for (const auto& row : thresholds) for (auto threshold : row) + if (threshold > 0 && a.r.size() > std::numeric_limits::max() / uint64_t(threshold)) + throw std::overflow_error("gapblock linear mass exceeds uint64"); + } + auto sort_row = [](auto& sorted, const auto& row) { + sorted.reserve(row.size()); + for (size_t t = 0; t < row.size(); ++t) sorted.emplace_back(row[t], t); + std::sort(sorted.begin(), sorted.end()); + }; + std::vector> shared_sorted; + sort_row(shared_sorted, thresholds.front()); + const auto longest = b ? std::max(a.longest, b->longest) : a.longest; + parallel_for(a.q.size(), threads, [&] { return std::vector(size_t(longest) + 1); }, + [&](size_t i, std::vector& scratch) { + const auto& sorted = shared_sorted; + std::vector bins(sorted.size()); + std::vector score_bins(masses ? sorted.size() : 0); + out[i].resize(sorted.size()); + if (masses) (*masses)[i].resize(sorted.size()); + if (sorted.empty()) return; + for (size_t j = 0; j < a.r.size(); ++j) { + if (exclude_exact && a.q[i] == a.r[j] && (!b || b->q[i] == b->r[j])) continue; + const int64_t offset = int64_t(group_offset(a, i, j)) + + (sum_pair_distance ? int64_t(group_offset(*b, i, j)) : 0); if (offset > sorted.back().first) continue; - int64_t score = score_cell(a, i, j, go, ge, prior, scratch.data()); + int64_t score = score_cell(a, i, j, go, ge, prior, scratch.data()); if (sum_pair_distance) { if (score > sorted.back().first) continue; - score += int64_t(score_cell(*b, i, j, go, ge, prior, scratch.data())); + score += int64_t(score_cell(*b, i, j, go, ge, prior, scratch.data())); } else if (b) { - score = std::max(score, score_cell(*b, i, j, go, ge, prior, scratch.data())); + score = std::max(score, score_cell(*b, i, j, go, ge, prior, scratch.data())); } const auto pos = std::lower_bound(sorted.begin(), sorted.end(), score, [](const auto& threshold, int64_t value) { return threshold.first < value; }); @@ -326,6 +467,21 @@ std::vector> counted(const GapInput& a, const GapInput* b, return out; } +// Dispatch once per batch: absent features retain the original hot loop. +std::vector> counted(const GapInput& a, const GapInput* b, + const std::vector>& thresholds, int32_t go, int32_t ge, + const std::vector& prior, int threads, bool exclude_exact, + std::vector>* masses = nullptr, bool sum_pair_distance = false, bool shared_thresholds = false) { + if (shared_thresholds) { + if (!a.features && (!b || !b->features)) + return counted_shared_impl(a, b, thresholds, go, ge, prior, threads, exclude_exact, masses, sum_pair_distance); + return counted_shared_impl(a, b, thresholds, go, ge, prior, threads, exclude_exact, masses, sum_pair_distance); + } + if (!a.features && (!b || !b->features)) + return counted_impl(a, b, thresholds, go, ge, prior, threads, exclude_exact, masses, sum_pair_distance); + return counted_impl(a, b, thresholds, go, ge, prior, threads, exclude_exact, masses, sum_pair_distance); +} + } // namespace std::vector gapblock_matrix(const std::vector& queries, @@ -334,8 +490,9 @@ 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 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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights, + const std::optional& feature_offsets) { + 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, feature_offsets); 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()); @@ -353,10 +510,11 @@ 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>& query_weights, const std::vector>& reference_weights) { + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights, + const std::optional& feature_offsets) { 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, query_weights, reference_weights); + const auto in = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights, feature_offsets); 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,12 +538,13 @@ gapblock_group_audit_batch(const std::vector& queries, 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 std::vector>& reference_weights, + const std::optional& feature_offsets) { if (thresholds.size() != queries.size()) throw std::invalid_argument("group audit 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, - query_weights, reference_weights); + query_weights, reference_weights, feature_offsets); std::tuple, std::vector, std::vector> out{ std::vector(queries.size()), std::vector(queries.size()), std::vector(queries.size(), -1)}; @@ -442,9 +601,10 @@ 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 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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights, + const std::optional& feature_offsets, bool shared_thresholds) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights, feature_offsets); + return counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact, nullptr, false, shared_thresholds); } std::pair>, std::vector>> gapblock_count_mass_batch(const std::vector& queries, @@ -453,10 +613,11 @@ 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 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); + const std::vector>& group_distances, const std::vector>& query_weights, const std::vector>& reference_weights, + const std::optional& feature_offsets, bool shared_thresholds) { + const auto a = prepare(queries, refs, alphabet, matrix, go, ge, prior, width, threads, weights, query_groups, reference_groups, group_distances, query_weights, reference_weights, feature_offsets); std::vector> masses; - auto counts = counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact, &masses); + auto counts = counted(a, nullptr, thresholds, go, ge, prior, threads, exclude_exact, &masses, false, shared_thresholds); return {std::move(counts), std::move(masses)}; } @@ -486,14 +647,382 @@ std::pair>, std::vector> 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 std::vector>& reference_weights_beta, + const std::optional& feature_offsets_alpha, + const std::optional& feature_offsets_beta, bool shared_thresholds) { const auto a = prepare(qa, ra, alphabet, matrix, go, ge, prior, width, threads, - weights_alpha, query_groups_alpha, reference_groups_alpha, group_distances_alpha, query_weights_alpha, reference_weights_alpha); + weights_alpha, query_groups_alpha, reference_groups_alpha, group_distances_alpha, query_weights_alpha, reference_weights_alpha, feature_offsets_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, query_weights_beta, reference_weights_beta); + weights_beta, query_groups_beta, reference_groups_beta, group_distances_beta, query_weights_beta, reference_weights_beta, feature_offsets_beta); std::vector> masses; - auto counts = counted(a, &b, thresholds, go, ge, prior, threads, exclude_exact, &masses, true); + auto counts = counted(a, &b, thresholds, go, ge, prior, threads, exclude_exact, &masses, true, shared_thresholds); return {std::move(counts), std::move(masses)}; } +struct GroupedCountBatch::Impl { + GapInput a; + std::optional b; + std::vector cuts, prior, target_peptides, exclusive, qgenes, rgenes; + std::vector offsets, labels, peptide_offsets; + std::vector peptide_ids; + std::optional> withheld; + size_t targets = 0, peptides = 0, channels = 0, chunk_size = 0, cursor = 0; + int32_t go = 0, ge = 0; + int32_t quantum_a = 1, quantum_b = 1; + int threads = 0; + bool exclude = false; + std::mutex mutex; + std::optional> beta_prior; + std::vector extra_a,extra_b; + uint64_t extra_bound_a=0,extra_bound_b=0; + std::optional similarity; + AlignMode mode = AlignMode::Global; + int32_t scale = 1; + uint32_t left = 0, right = 0; + std::vector qdiag_a, rdiag_a, qdiag_b, rdiag_b; + + void prepare_affine(const GroupedGeometry& geometry, const GapInput& lane, + std::vector& qdiag, std::vector& rdiag, + const std::vector& lane_prior,uint64_t extra_bound) { + if (!geometry.weights.empty() || !geometry.query_weights.empty() || + !geometry.reference_weights.empty() || !lane_prior.empty()) + throw std::invalid_argument("affine alignment does not accept position weights or gap prior"); + uint64_t max_q=0, max_r=0; + auto diagonals = [&](const auto& sequences, auto& output, uint64_t& longest) { + output.reserve(sequences.size()); + for (const auto& sequence : sequences) { + if (uint64_t(left)+right>sequence.size()) + throw std::invalid_argument("alignment trims exceed sequence length"); + const size_t length=sequence.size()-left-right; + longest=std::max(longest,uint64_t(length)); + int64_t diagonal=0; + for (size_t k=left;ksimilarity(sequence[k],sequence[k]); + output.push_back(diagonal); + } + }; + // Bound before summation/DP: |similarity| and gap magnitudes cover every + // consumed residue; diagonal sum minus twice DP is <=3*(m+n)*max_cost. + uint64_t maximum=std::max(go,ge); + for (uint8_t x=0;xsimilarity(x,y); + maximum=std::max(maximum,uint64_t(value<0 ? -value : value)); + } + for (const auto& sequence:lane.q) { + if (uint64_t(left)+right>sequence.size()) + throw std::invalid_argument("alignment trims exceed sequence length"); + max_q=std::max(max_q,uint64_t(sequence.size()-left-right)); + } + for (const auto& sequence:lane.r) { + if (uint64_t(left)+right>sequence.size()) + throw std::invalid_argument("alignment trims exceed sequence length"); + max_r=std::max(max_r,uint64_t(sequence.size()-left-right)); + } + uint64_t additive=extra_bound; + for (const auto& row:lane.group_distances) + for (auto value:row) additive=std::max(additive,extra_bound+uint64_t(value)); + if (lane.features) { + const auto& features=*lane.features; + int32_t cost=0; + for (const auto& row:features.costs) for (auto value:row) cost=std::max(cost,value); + const size_t depth=(!features.query_codes.empty() ? features.query_codes.front().size() : + (!features.reference_codes.empty() ? features.reference_codes.front().size() : 0))/2; + if (depth && features.weight && features.scale) + additive+=uint64_t(round_even(features.weight*double(uint64_t(cost)*depth)/double(depth))*features.scale); + } + const uint64_t limit=(uint64_t(INT32_MAX)-additive)/uint32_t(scale); + if (maximum && max_q+max_r>limit/(3*maximum)) + throw std::overflow_error("affine grouped distance can overflow int32"); + diagonals(lane.q,qdiag,max_q); diagonals(lane.r,rdiag,max_r); + } + + template + int64_t affine_cell(const GapInput& lane, size_t i, size_t j, + const std::vector& qdiag, const std::vector& rdiag, + detail::PairwiseScratch& scratch) const { + const auto& q=lane.q[i]; const auto& r=lane.r[j]; + // data() can be null for an empty vector; avoid pointer arithmetic then. + const auto* qp=q.empty() ? q.data() : q.data()+left; + const auto* rp=r.empty() ? r.data() : r.data()+left; + const int64_t sim=detail::align_score_encoded(qp,q.size()-left-right, + rp,r.size()-left-right,*similarity,mode,go,ge,scratch); + return std::max(int64_t(0),qdiag[i]+rdiag[j]-2*sim)*scale+ + group_offset(lane,i,j); + } + + template int64_t extra_offset(const std::vector& terms, + size_t i,size_t j) const { + int64_t offset=0; + if constexpr (Additional) for (const auto& term:terms) offset+=feature_offset(term,i,j); + return offset; + } + + template GroupedCountChunk run(size_t start, size_t stop) { + GroupedCountChunk out; + out.start = start; out.stop = stop; + const size_t rows = stop-start, K = cuts.size(); + auto make_buffer = [](std::vector shape) { + size_t cells = 1; + for (size_t axis : shape) { + if (axis && cells > size_t(PTRDIFF_MAX)/sizeof(uint32_t)/axis) + throw std::overflow_error("grouped count buffer exceeds addressable memory"); + cells *= axis; + } + return CountBuffer{std::vector(cells),std::move(shape)}; + }; + out.counts = make_buffer({rows,channels,K}); + out.sizes = make_buffer({rows,channels}); + out.exact_counts = make_buffer({rows,targets}); + const auto longest = b ? std::max(a.longest,b->longest) : a.longest; + const auto& other_prior=beta_prior ? *beta_prior : prior; + parallel_for(rows,threads,[&] { + if constexpr (Affine) { + detail::PairwiseScratch scratch; + for (auto& row:scratch) row.reserve(size_t(longest)+1); + return scratch; + } else return std::vector(size_t(longest)+1); + }, [&](size_t row, auto& scratch) { + const size_t i = start+row; + auto* sizes = out.sizes.data.data()+row*channels; + auto* hist = K ? out.counts.data.data()+row*channels*K : nullptr; + auto* exact = targets ? out.exact_counts.data.data()+row*targets : nullptr; + for (size_t j=0;jq[i]==b->r[j]); + const bool same_gene = qgenes[i]==rgenes[j]; + if (identical) + for (size_t p=offsets[j];p=0) ++sizes[1+2*targets+size_t(exclusive_peptide)]; + if (!K) continue; + const int64_t extra_alpha=extra_offset(extra_a,i,j); + const int64_t extra_beta=b ? extra_offset(extra_b,i,j) : 0; + const int64_t offset = int64_t(group_offset(a,i,j))+extra_alpha + + (b ? int64_t(group_offset(*b,i,j))+extra_beta : 0); + if (offset>cuts.back()) continue; + int64_t score; + if constexpr (Affine) score=affine_cell(a,i,j,qdiag_a,rdiag_a,scratch); + else score=score_cell(a,i,j,go,ge,prior,scratch.data()); + score+=extra_alpha; + if constexpr (Quantized) { + const int64_t additive=group_offset(a,i,j)+extra_alpha; + score=additive+(score-additive)/quantum_a*quantum_a; + } + if (b) { + if (score>cuts.back()) continue; + int64_t other; + if constexpr (Affine) other=affine_cell(*b,i,j,qdiag_b,rdiag_b,scratch); + else other=score_cell(*b,i,j,go,ge,other_prior,scratch.data()); + other+=extra_beta; + if constexpr (Quantized) { + const int64_t additive=group_offset(*b,i,j)+extra_beta; + other=additive+(other-additive)/quantum_b*quantum_b; + } + score+=other; + } + const auto found=std::lower_bound(cuts.begin(),cuts.end(),score); + if (found==cuts.end()) continue; + const size_t bin=size_t(found-cuts.begin()); + ++hist[bin]; + for (size_t p=offsets[j];p=0) ++hist[(1+2*targets+size_t(exclusive_peptide))*K+bin]; + } + for (size_t channel=0;channel GroupedCountChunk dispatch(size_t start,size_t stop) { + const bool features=a.features || (b && b->features); + const bool quantized=quantum_a!=1 || quantum_b!=1; + return quantized ? + (withheld ? + (features ? run(start,stop) : run(start,stop)) : + (features ? run(start,stop) : run(start,stop))) : + (withheld ? + (features ? run(start,stop) : run(start,stop)) : + (features ? run(start,stop) : run(start,stop))); + } +}; + +GroupedCountBatch::GroupedCountBatch(const std::vector& qa, const std::vector& ra, + std::vector cuts, std::vector offsets, std::vector labels, + std::vector target_peptides, std::vector exclusive, + std::vector qgenes, std::vector rgenes, size_t chunk_size, + Alphabet alphabet, const SubstitutionMatrix* matrix, int32_t go, int32_t ge, + std::vector prior, uint32_t width, int threads, bool exclude, + const GroupedGeometry& alpha, const std::optional& beta, + const std::vector& qb, const std::vector& rb, + std::optional> withheld, + std::vector peptide_offsets, std::vector peptide_ids) + : impl_(std::make_shared()) { + auto& in=*impl_; + if (alpha.score_quantum<1 || (beta && beta->score_quantum<1)) + throw std::invalid_argument("score quantum must be positive"); + in.quantum_a=alpha.score_quantum; + in.quantum_b=beta ? beta->score_quantum : 1; + if (!chunk_size) throw std::invalid_argument("chunk_size must be positive"); + if (!std::is_sorted(cuts.begin(),cuts.end()) || std::adjacent_find(cuts.begin(),cuts.end())!=cuts.end()) + throw std::invalid_argument("grouped thresholds must be strictly increasing"); + if (ra.size()>std::numeric_limits::max()) + throw std::overflow_error("grouped reference count exceeds uint32"); + if (offsets.size()!=ra.size()+1 || offsets.front()!=0 || offsets.back()!=labels.size() || + !std::is_sorted(offsets.begin(),offsets.end())) + throw std::invalid_argument("reference target offsets must form valid CSR"); + if (exclusive.size()!=ra.size() || qgenes.size()!=qa.size() || rgenes.size()!=ra.size()) + throw std::invalid_argument("grouped metadata axes must match queries/references"); + if (std::any_of(qgenes.begin(),qgenes.end(),[](auto x){return x<0;}) || + std::any_of(rgenes.begin(),rgenes.end(),[](auto x){return x<0;})) + throw std::invalid_argument("gene IDs must be nonnegative"); + for (auto peptide:target_peptides) { + if (peptide<0) throw std::invalid_argument("target peptide IDs must be nonnegative"); + in.peptides=std::max(in.peptides,size_t(peptide)+1); + } + in.targets=target_peptides.size(); + for (size_t j=0;j=0 && size_t(exclusive[j])>=in.peptides)) + throw std::invalid_argument("exclusive peptide ID outside peptide axis"); + // Empty memberships and -1 exclusive IDs are valid: an included target + // can share a reference with an unselected label in the whole population. + uint32_t previous=0; + for (size_t p=offsets[j];p=in.targets || (p>offsets[j] && label<=previous)) + throw std::invalid_argument("reference target IDs must be sorted, unique and in range"); + if (exclusive[j]>=0 && target_peptides[label]!=exclusive[j]) + throw std::invalid_argument("exclusive peptide contradicts reference target membership"); + previous=label; + } + } + if (withheld) { + if (withheld->size()!=qa.size() || + std::any_of(withheld->begin(),withheld->end(),[](auto x){return x<-1;})) + throw std::invalid_argument("withheld peptide IDs must match queries and be >= -1"); + if (peptide_ids.size()>std::numeric_limits::max()) + throw std::overflow_error("reference peptide memberships exceed uint32"); + if (peptide_offsets.size()!=ra.size()+1 || peptide_offsets.front()!=0 || + peptide_offsets.back()!=peptide_ids.size() || + !std::is_sorted(peptide_offsets.begin(),peptide_offsets.end())) + throw std::invalid_argument("reference peptide offsets must form valid CSR"); + for (size_t j=0;jgap_prior ? *beta->gap_prior : prior; + uint64_t bound_a=0,bound_b=0; + in.a=prepare(qa,ra,alphabet,matrix,go,ge,alpha_prior,width,threads,alpha.weights, + alpha.query_groups,alpha.reference_groups,alpha.group_distances, + alpha.query_weights,alpha.reference_weights,alpha.features,&bound_a); + if (beta) { + if (qa.size()!=qb.size() || ra.size()!=rb.size()) + throw std::invalid_argument("paired query/reference axes must agree"); + in.b=prepare(qb,rb,alphabet,matrix,go,ge,beta_prior,width,threads,beta->weights, + beta->query_groups,beta->reference_groups,beta->group_distances, + beta->query_weights,beta->reference_weights,beta->features,&bound_b); + } else if (!qb.empty() || !rb.empty()) throw std::invalid_argument("paired axes require beta geometry"); + if (threads<0 || unsigned(threads)>available_cpus()) + throw std::invalid_argument("worker count exceeds the available CPU budget"); + in.cuts=std::move(cuts); in.offsets=std::move(offsets); in.labels=std::move(labels); + in.target_peptides=std::move(target_peptides); in.exclusive=std::move(exclusive); + in.qgenes=std::move(qgenes);in.rgenes=std::move(rgenes); + in.prior=alpha.gap_prior ? *alpha.gap_prior : std::move(prior); + if (beta && (beta->gap_prior || alpha.gap_prior)) + in.beta_prior=beta->gap_prior ? *beta->gap_prior : prior; + auto additional=[&](const auto& terms,size_t nq,size_t nr,uint64_t base_bound) { + uint64_t bound=0; + for (const auto& term:terms) { + const uint64_t next=validate_feature(term,nq,nr); + if (next>uint64_t(INT32_MAX)-base_bound-bound) + throw std::overflow_error("additional factored offsets can overflow int32"); + bound+=next; + } + return bound; + }; + in.extra_bound_a=additional(alpha.additional_features,qa.size(),ra.size(),bound_a); + in.extra_a=alpha.additional_features; + if (beta) { + in.extra_bound_b=additional(beta->additional_features,qb.size(),rb.size(),bound_b); + in.extra_b=beta->additional_features; + } + in.channels=1+2*in.targets+in.peptides;in.chunk_size=chunk_size; + in.go=go;in.ge=ge;in.threads=threads;in.exclude=exclude; + if (alpha.alignment_scale<1 || (alpha.alignment_mode!="gapblock" && + alpha.alignment_mode!="global" && alpha.alignment_mode!="local")) + throw std::invalid_argument("invalid grouped alignment mode or scale"); + if (beta && (beta->alignment_mode!=alpha.alignment_mode || + beta->alignment_scale!=alpha.alignment_scale || beta->trim_left!=alpha.trim_left || + beta->trim_right!=alpha.trim_right)) + throw std::invalid_argument("paired alignment options must agree"); + if (alpha.alignment_mode=="gapblock") { + if (alpha.alignment_scale!=1 || alpha.trim_left || alpha.trim_right) + throw std::invalid_argument("alignment scale/trims require global or local mode"); + } else { + if (!matrix) throw std::invalid_argument("affine alignment requires a similarity matrix"); + in.similarity=*matrix; + in.mode=alpha.alignment_mode=="local" ? AlignMode::Local : AlignMode::Global; + in.scale=alpha.alignment_scale;in.left=alpha.trim_left;in.right=alpha.trim_right; + in.prepare_affine(alpha,in.a,in.qdiag_a,in.rdiag_a,in.prior,in.extra_bound_a); + if (beta) in.prepare_affine(*beta,*in.b,in.qdiag_b,in.rdiag_b,in.beta_prior ? *in.beta_prior : in.prior,in.extra_bound_b); + } +} +std::optional GroupedCountBatch::next() { + auto& in=*impl_; + std::lock_guard lock(in.mutex); + if (in.cursor==in.a.q.size()) return std::nullopt; + const size_t stop=in.cursor+std::min(in.chunk_size,in.a.q.size()-in.cursor); + const bool extra=!in.extra_a.empty() || !in.extra_b.empty(); + auto out=in.similarity ? + (extra ? in.dispatch(in.cursor,stop) : in.dispatch(in.cursor,stop)) : + (extra ? in.dispatch(in.cursor,stop) : in.dispatch(in.cursor,stop)); + in.cursor=stop; + return out; +} +size_t GroupedCountBatch::query_count() const {return impl_->a.q.size();} +size_t GroupedCountBatch::reference_count() const {return impl_->a.r.size();} +size_t GroupedCountBatch::target_count() const {return impl_->targets;} +size_t GroupedCountBatch::peptide_count() const {return impl_->peptides;} + } // namespace seqtree diff --git a/src/pairwise.cpp b/src/pairwise.cpp index 5e2dd63..d027bc8 100644 --- a/src/pairwise.cpp +++ b/src/pairwise.cpp @@ -15,6 +15,7 @@ #include "seqtree/seqtree.hpp" #include "seqtree/parallel.hpp" +#include "pairwise_detail.hpp" #include #include @@ -73,15 +74,19 @@ inline Score gap_run(int64_t L, int32_t open, int32_t extend) { // Optimal score, two rows per state. O(n) memory, no traceback. Score score_encoded(const uint8_t* q, size_t m, const uint8_t* r, size_t n, - const SubstitutionMatrix& mat, AlignMode mode, int32_t go, int32_t ge) { + const SubstitutionMatrix& mat, AlignMode mode, int32_t go, int32_t ge, + detail::PairwiseScratch* scratch = nullptr) { const bool local = (mode == AlignMode::Local); if (m == 0 || n == 0) { if (local) return 0; return -gap_run(int64_t(m ? m : n), go, ge); } - std::vector Mp(n + 1), Xp(n + 1), Yp(n + 1); - std::vector Mc(n + 1), Xc(n + 1), Yc(n + 1); + detail::PairwiseScratch owned; + auto& rows = scratch ? *scratch : owned; + for (auto& row : rows) row.resize(n + 1); + auto& Mp=rows[0]; auto& Xp=rows[1]; auto& Yp=rows[2]; + auto& Mc=rows[3]; auto& Xc=rows[4]; auto& Yc=rows[5]; // Row 0. Global: the only way to consume ref[0..j) is a leading gap (state Y). Local: an // alignment may start anywhere, which the 0 floor on M expresses; the gap states cannot be @@ -127,6 +132,12 @@ enum : uint8_t { kM = 0, kX = 1, kY = 2, kStop = 3 }; } // namespace +int64_t detail::align_score_encoded(const uint8_t* q, size_t m, const uint8_t* r, size_t n, + const SubstitutionMatrix& mat, AlignMode mode, int32_t go, int32_t ge, + PairwiseScratch& scratch) { + return score_encoded(q,m,r,n,mat,mode,go,ge,&scratch); +} + int32_t align_score(std::string_view query, std::string_view ref, const SubstitutionMatrix& mat, Alphabet alphabet, AlignMode mode, int32_t gap_open, int32_t gap_extend) { const Codec codec(alphabet); diff --git a/src/pairwise_detail.hpp b/src/pairwise_detail.hpp new file mode 100644 index 0000000..22eea9f --- /dev/null +++ b/src/pairwise_detail.hpp @@ -0,0 +1,8 @@ +#pragma once +#include "seqtree/seqtree.hpp" +#include +namespace seqtree::detail { +using PairwiseScratch = std::array, 6>; +int64_t align_score_encoded(const uint8_t*, size_t, const uint8_t*, size_t, + const SubstitutionMatrix&, AlignMode, int32_t, int32_t, PairwiseScratch&); +} diff --git a/tests/cpp/test_edge.cpp b/tests/cpp/test_edge.cpp index 12e93e4..a33b23f 100644 --- a/tests/cpp/test_edge.cpp +++ b/tests/cpp/test_edge.cpp @@ -1,4 +1,5 @@ #include "doctest.h" +#include "count_unit.hpp" #include "seqtree/seqtree.hpp" #include @@ -86,3 +87,56 @@ TEST_CASE("single-character references") { auto m = run(*idx, "A", p); CHECK(m.size() == 3); // every length-1 ref is within one substitution } + +TEST_CASE("grouped gapblock population buffers and full-key puncture") { + seqtree::GroupedGeometry geometry; + seqtree::GroupedCountBatch batch({"AAA","AAC"},{"AAA","ACA"},{0,2}, + {0,2,3},{0,1,1},{0,1},{-1,1},{0,1},{1,0},1, + seqtree::Alphabet::AminoAcid,nullptr,1,1,{},0,1,true,geometry); + auto first=batch.next(); + REQUIRE(first.has_value()); + CHECK(first->start==0); + CHECK(first->stop==1); + CHECK(first->counts.shape==std::vector{1,7,2}); + CHECK(first->counts.data[0]==0); + CHECK(first->counts.data[1]==1); + CHECK(first->sizes.data[0]==1); + CHECK(first->exact_counts.data==std::vector{1,1}); + REQUIRE(batch.next().has_value()); + CHECK_FALSE(batch.next().has_value()); + CHECK_THROWS_AS(seqtree::GroupedCountBatch({"AAA"},{"AAA"},{0,0}, + {0,1},{0},{0},{0},{0},{0},1,seqtree::Alphabet::AminoAcid,nullptr,1,1,{},0,1,false,geometry),std::invalid_argument); +} + +// The bounded Myers rejection must preserve every within-cap distance. +TEST_CASE("bounded unit bitvector distance preserves all short binary distances") { + Codec codec(Alphabet::AminoAcid); + std::vector strings{ "" }; + for (unsigned n=1;n<=6;++n) for(unsigned mask=0;mask<(1u< row(q.size()+1); + for(size_t i=0;i 0, 1.5 -> 2, 2.5 -> 2; decimal arithmetic follows multiply/divide. + costs = [[0, 1, 3, 5, 25]] + codes = [[1, 1], [2, 2], [3, 3], [4, 4]] + f = PairedFeatureOffsets([[0, 0]], codes, costs, weight=0.5, scale=200) + costs[0][1] = 1000 + codes[0][0] = 4 + sm = score_matrix(['A'], ['A']*4, feature_offsets=f, threads=1) + assert [sm[0,j] for j in range(4)] == [0, 400, 400, 2400] + f = PairedFeatureOffsets([[0,0]], [[1,1]], [[0,25]], weight=0.1) + assert score_matrix(['A'], ['A'], feature_offsets=f, threads=1)[0,0] == 2 + assert count_batch([], [], [], feature_offsets=PairedFeatureOffsets([], [], [[0]]), threads=1) == [] + assert count_batch(['A'], [], [[1]], feature_offsets=PairedFeatureOffsets([[0,0]], [], [[0]]), threads=1) == [[0]] + + +@pytest.mark.parametrize('qc,rc,costs,weight,scale', [ + ([], [[0,0]], [[0]], 1, 1), ([[0]], [[0]], [[0]], 1, 1), + ([[]], [[]], [[0]], 1, 1), ([[0,0]], [[0,0,0,0]], [[0]], 1, 1), + ([[-1,0]], [[0,0]], [[0]], 1, 1), ([[1,0]], [[0,0]], [[0]], 1, 1), + ([[0,0]], [[1,0]], [[0]], 1, 1), ([[0,0]], [[0,0]], [], 1, 1), + ([[0,0]], [[0,0]], [[0], []], 1, 1), ([[0,0]], [[0,0]], [[-1]], 1, 1), + ([[0,0]], [[0,0]], [[0]], -1, 1), ([[0,0]], [[0,0]], [[0]], math.inf, 1), + ([[0,0]], [[0,0]], [[0]], math.nan, 1), ([[0,0]], [[0,0]], [[0]], 1, -1), + ([[0,0]], [[0,0]], [[1]], 1e308, 1), ([[0,0]], [[0,0]], [[2**30]], 1, 2), +]) +def test_factored_validation(qc,rc,costs,weight,scale): + f = PairedFeatureOffsets(qc,rc,costs,weight=weight,scale=scale) + with pytest.raises(ValueError): + count_batch(['A'], ['A'], [[10]], feature_offsets=f, threads=1) + + +def test_factored_combined_overflow(): + f = PairedFeatureOffsets([[0,0]], [[0,0]], [[2**30]]) + with pytest.raises(ValueError, match='overflow'): + score_matrix(['A'], ['A'], feature_offsets=f, query_group_ids=[0], + reference_group_ids=[0], group_distances=[[2**30]], threads=1) + + +def test_factored_numpy_rounding_oracle(): + import numpy as np + for depth in [1, 2, 3, 5, 9]: + costs = [list(range(100))] + qcodes = [[0,0] * depth] + rcodes = [[n,n] * depth for n in range(100)] + for weight in [0.1, 0.2, 0.5, 0.7, 1.1, 1.25, 3.141592653589793]: + descriptor = PairedFeatureOffsets(qcodes, rcodes, costs, weight=weight, scale=200) + scores = score_matrix(['A'], ['A']*100, threads=1, feature_offsets=descriptor) + expected = 200 * np.rint(weight * (np.arange(100, dtype=np.int64) * depth) / depth) + assert [scores[0,j] for j in range(100)] == expected.tolist() + + +def test_factored_paired_identity_and_independent_lanes(): + a = PairedFeatureOffsets([[0,0]], [[0,0], [1,1]], [[0,3], [2,0]], scale=2) + b = PairedFeatureOffsets([[0,0]], [[1,1], [0,0]], [[0,5], [4,0]], weight=0.5) + args = dict(threads=1, feature_offsets_alpha=a, feature_offsets_beta=b, + query_group_ids_alpha=[0], reference_group_ids_alpha=[0,0], group_distances_alpha=[[1]], + query_group_ids_beta=[0], reference_group_ids_beta=[0,0], group_distances_beta=[[2]]) + assert paired_sum_count_batch(['A'], ['A'], ['A','A'], ['C','A'], [[3,4,6,9]], **args) == ([[0,0,1,2]], [[0,0,0,3]]) + assert paired_sum_count_batch(['A'], ['A'], ['A','A'], ['C','A'], [[9]], exclude_exact=True, **args) == ([[1]], [[3]]) + + +def test_factored_zero_controls(): + for weight, scale in [(0, 200), (1e308, 0)]: + f = PairedFeatureOffsets([[0,0]], [[0,0]], [[2**30]], weight=weight, scale=scale) + assert score_matrix(['A'], ['C'], feature_offsets=f, threads=1)[0,0] == 1 diff --git a/tests/python/test_gapblock_score_quantum.py b/tests/python/test_gapblock_score_quantum.py new file mode 100644 index 0000000..f470764 --- /dev/null +++ b/tests/python/test_gapblock_score_quantum.py @@ -0,0 +1,98 @@ +"""Generic alignment quantization before additive offsets, no domain inputs.""" +import pytest +from seqtree import gapblock, SubstitutionMatrix + + +def metadata(nq, nr): + return dict(reference_target_offsets=list(range(nr + 1)), + reference_target_ids=[0] * nr, target_peptide_ids=[0], + reference_exclusive_peptide_ids=[0] * nr, + query_gene_ids=[0] * nq, reference_gene_ids=[0] * nr) + + +def rows(buffer): + # Native buffer interface without a runtime NumPy dependency. + return memoryview(buffer).tolist() + + +@pytest.mark.parametrize('threads', [1, 4]) +@pytest.mark.parametrize('paired', [False, True]) +@pytest.mark.parametrize('factored', [False, True]) +@pytest.mark.parametrize('withheld', [False, True]) +def test_floor_each_alignment_before_offsets_and_pair_sum(threads, paired, factored, withheld, require_thread_budget): + q, r = ['AAAA', 'RRRA'], ['ARAA', 'RRRA', 'AAAA', 'AA'] + cuts = [0, 3, 5, 6, 9, 20, 100] + matrix = SubstitutionMatrix.blosum62() + weights = {2: [3, 3], 4: [3, 3, 3, 3]} + raw = gapblock.score_matrix(q, r, matrix=matrix, gap_open=7, gap_extend=2, + position_weights_by_length=weights, threads=1) + opts = dict(matrix=matrix, gap_open=7, gap_extend=2, threads=threads, + exclude_exact=True, **metadata(len(q), len(r))) + geometry = dict(position_weights_by_length=weights, + query_group_ids=[0, 1], reference_group_ids=[1, 0, 0, 1], + group_distances=[[0, 3], [5, 0]]) + offset = [[3, 0, 0, 3], [0, 5, 5, 0]] + if factored: + geometry['feature_offsets'] = gapblock.PairedFeatureOffsets( + [[0, 0]] * len(q), [[1, 1]] * len(r), [[0, 2], [2, 0]]) + offset = [[value + 2 for value in row] for row in offset] + if withheld: + opts.update(query_withheld_peptide_ids=[-1] * len(q), + reference_peptide_offsets=list(range(len(r) + 1)), reference_peptide_ids=[0] * len(r)) + if paired: + kwargs = {key + suffix: value for suffix in ('_alpha', '_beta') + for key, value in geometry.items()} + iterator = gapblock.grouped_paired_sum_count_batch(q, q, r, r, cuts, + score_quantum_alpha=5, score_quantum_beta=7, **opts, **kwargs) + else: + iterator = gapblock.grouped_count_batch(q, r, cuts, score_quantum=5, + **opts, **geometry) + chunk = next(iterator) + expected = [] + for i, query in enumerate(q): + scores = [(raw[i, j] // 5) * 5 + offset[i][j] for j in range(len(r))] + if paired: + scores = [value + (raw[i, j] // 7) * 7 + offset[i][j] + for j, value in enumerate(scores)] + curve = [sum(score <= cutoff and query != r[j] + for j, score in enumerate(scores)) for cutoff in cuts] + expected.append([curve] * 4) + assert rows(chunk.counts) == expected + assert rows(chunk.sizes) == [[3] * 4, [3] * 4] + assert rows(chunk.exact_counts) == [[1], [1]] + + +@pytest.mark.parametrize('bad', [0, -1, True, 1.5, 2**31]) +def test_invalid_quantum_rejected_even_empty_axes(bad): + with pytest.raises((ValueError, OverflowError), match='quantum'): + gapblock.grouped_count_batch([], [], [], score_quantum=bad, **metadata(0, 0)) + + +def test_default_quantum_one_is_exactly_original_and_empty_axes(): + q, r = ['AAAA'], ['ARAA', 'AAAA'] + kwargs = metadata(1, 2) + a = next(gapblock.grouped_count_batch(q, r, [0, 1, 2, 3], **kwargs)) + b = next(gapblock.grouped_count_batch(q, r, [0, 1, 2, 3], score_quantum=1, **kwargs)) + assert rows(a.counts) == rows(b.counts) + assert rows(a.sizes) == rows(b.sizes) + assert list(gapblock.grouped_count_batch([], [], [], score_quantum=600, **metadata(0, 0))) == [] + + +def test_native_quantum_validation_is_not_only_python(monkeypatch): + original = gapblock._grouped_geometry + def invalid(*args, **kwargs): + result = original(*args, **kwargs) + result.score_quantum = 0 + return result + monkeypatch.setattr(gapblock, '_grouped_geometry', invalid) + with pytest.raises(ValueError, match='quantum'): + gapblock.grouped_count_batch([], [], [], **metadata(0, 0)) + + +def test_max_int32_quantum_and_empty_reference_cutoffs(): + chunk = next(gapblock.grouped_count_batch(['AAAA'], ['RRRR'], [0, 1], + score_quantum=2**31 - 1, **metadata(1, 1))) + assert rows(chunk.counts) == [[[1, 1]] * 4] + chunk = next(gapblock.grouped_count_batch(['AAAA'], [], [], + score_quantum=600, **metadata(1, 0))) + assert rows(chunk.sizes) == [[0] * 4] diff --git a/tests/python/test_gapblock_shared_thresholds.py b/tests/python/test_gapblock_shared_thresholds.py new file mode 100644 index 0000000..24e73c3 --- /dev/null +++ b/tests/python/test_gapblock_shared_thresholds.py @@ -0,0 +1,63 @@ +"""Shared cutoffs reproduce expanded-row batch curves with one validation pass.""" +import pytest +from seqtree import _core +from seqtree.gapblock import PairedFeatureOffsets, count_batch, paired_sum_count_batch + + +@pytest.mark.parametrize('threads',[1,4]) +@pytest.mark.parametrize('exclude',[False,True]) +def test_shared_thresholds_match_expanded(threads,exclude, require_thread_budget): + q=['AAC','AC','AAC']; r=['AAC','ACC','AC','AAC'] + qb=['ACC','AC','AAC']; rb=['ACC','AAC','AC','AAC'] + cuts=[-1,0,3,1,3,12,20] + f=PairedFeatureOffsets([[0,0],[1,1],[0,0]],[[0,0],[1,1],[1,1],[0,0]],[[0,3],[2,0]],weight=.5,scale=2) + opts=dict(query_group_ids=[0,1,0],reference_group_ids=[0,1,1,0],group_distances=[[0,3],[2,0]],feature_offsets=f,gap_open=2,threads=threads,exclude_exact=exclude) + for mass in [False,True]: + assert count_batch(q,r,cuts,shared_thresholds=True,linear_mass=mass,**opts)==count_batch(q,r,[cuts]*len(q),linear_mass=mass,**opts) + pair_opts=dict(threads=threads,exclude_exact=exclude,gap_open=2,feature_offsets_alpha=f,feature_offsets_beta=f, + query_group_ids_alpha=opts['query_group_ids'],reference_group_ids_alpha=opts['reference_group_ids'],group_distances_alpha=opts['group_distances']) + assert paired_sum_count_batch(q,qb,r,rb,cuts,shared_thresholds=True,**pair_opts)==paired_sum_count_batch(q,qb,r,rb,[cuts]*len(q),**pair_opts) + + +def test_shared_empty_axes_cutoffs_and_generators(): + assert count_batch([],['A'],[0,3],shared_thresholds=True,threads=1)==[] + assert count_batch([],[],[],shared_thresholds=True,linear_mass=True,threads=1)==([],[]) + assert count_batch(['A','C'],[],[2,0,2],shared_thresholds=True,linear_mass=True,threads=1)==([[0,0,0],[0,0,0]],[[0,0,0],[0,0,0]]) + assert count_batch(['A','C'],['A'],[],shared_thresholds=True,linear_mass=True,threads=1)==([[],[]],[[],[]]) + assert count_batch(['A'],['C'],iter([1,2]),shared_thresholds=True,threads=1)==[[1,1]] + assert paired_sum_count_batch([],[],['A'],['C'],[0,3],shared_thresholds=True,threads=1)==([],[]) + assert paired_sum_count_batch(['A'],['C'],[],[],[],shared_thresholds=True,threads=1)==([[]],[[]]) + + +def test_shared_validates_values_once_per_call(): + class Integer: + def __init__(self,value): self.value=value; self.calls=0 + def __index__(self): self.calls+=1; return self.value + cuts=[Integer(i) for i in range(153)] + q=['A']*100 + count_batch(q,[],cuts,shared_thresholds=True,threads=1) + assert [v.calls for v in cuts]==[1]*153 + count_batch(q,[],cuts,shared_thresholds=True,linear_mass=True,threads=1) + assert [v.calls for v in cuts]==[2]*153 + paired_sum_count_batch(q,q,[],[],cuts,shared_thresholds=True,threads=1) + assert [v.calls for v in cuts]==[3]*153 + + +@pytest.mark.parametrize('cuts',[[True],[1.5],[[1]],[-(1<<31)-1],[1<<31]]) +def test_shared_bad_values(cuts): + with pytest.raises(ValueError): count_batch(['A'],['C'],cuts,shared_thresholds=True,threads=1) + with pytest.raises(ValueError): paired_sum_count_batch(['A'],['C'],['A'],['C'],cuts,shared_thresholds=True,threads=1) + + +@pytest.mark.parametrize('flag',[1,None,'yes']) +def test_shared_bad_flag(flag): + with pytest.raises(TypeError,match='boolean'): count_batch(['A'],['C'],[[1]],shared_thresholds=flag,threads=1) + with pytest.raises(TypeError,match='boolean'): paired_sum_count_batch(['A'],['C'],['A'],['C'],[[1]],shared_thresholds=flag,threads=1) + + +def test_shared_native_row_shape(): + for rows in [[],[[0],[1]]]: + with pytest.raises(ValueError,match='exactly one row'): + _core.gapblock_count_batch(['A'],['C'],rows,threads=1,shared_thresholds=True) + assert _core.gapblock_count_batch([],[],[[0,1]],threads=1,shared_thresholds=True)==[] + assert _core.gapblock_count_batch(['A','C'],['A'],[[]],threads=1,shared_thresholds=True)==[[],[]] diff --git a/tests/python/test_grouped_counts.py b/tests/python/test_grouped_counts.py new file mode 100644 index 0000000..b5e8364 --- /dev/null +++ b/tests/python/test_grouped_counts.py @@ -0,0 +1,135 @@ +"""Generic grouped counts equal explicit existing count populations, without domain data.""" +import gc +import numpy as np +import pytest +from seqtree import gapblock as gb + +Q = ['AAA', 'AAC', 'TTT', 'AAA'] +R = ['AAA', 'ACA', 'AAA', 'CCCC', 'TTT'] +QB = ['CCC', 'CCA', 'AAA', 'CCC'] +RB = ['CCC', 'CCC', 'CCA', 'CCCA', 'AAA'] +CUTS = [-1, 0, 2, 5, 20] +MEMBERS = [[0, 1], [2], [0, 2], [], []] +PARENTS = [0, 0, 1] +EXCLUSIVE = [0, 1, -1, 0, -1] +QGENE, RGENE = [0, 1, 1, 0], [1, 0, 1, 0, 1] + + +def metadata(): + return dict(reference_target_offsets=[0, 2, 3, 5, 5, 5], + reference_target_ids=[0, 1, 2, 0, 2], target_peptide_ids=PARENTS, + reference_exclusive_peptide_ids=EXCLUSIVE, query_gene_ids=QGENE, + reference_gene_ids=RGENE, chunk_size=2) + + +def geometry(features, indices=range(5), lane=0): + indices = list(indices) + qg, rg = [0, 1, 0, 1], [1, 0, 0, 1, 1] + result = dict(gap_open=2, gap_extend=1, query_group_ids=qg, + reference_group_ids=[rg[i] for i in indices], group_distances=[[0, 3], [3, 2]]) + if features: + qc = [[0, 1], [1, 0], [1, 1], [0, 1]] + rc = [[1, 0], [0, 0], [0, 1], [1, 1], [1, 0]] + result['feature_offsets'] = gb.PairedFeatureOffsets(qc, [rc[i] for i in indices], + [[0, 3], [3, 0]], weight=.5 if lane == 0 else 1.5, scale=1) + return result + + +def paired_geometry(features, indices=range(5)): + result = {} + for lane, suffix in enumerate(('alpha', 'beta')): + for key, value in geometry(features, indices, lane).items(): + result[key if key in ('gap_open', 'gap_extend') else key+'_'+suffix] = value + return result + + +@pytest.mark.parametrize('paired', [False, True]) +@pytest.mark.parametrize('exclude', [False, True]) +@pytest.mark.parametrize('features', [False, True]) +@pytest.mark.parametrize('threads', [1, 4]) +def test_every_grouped_curve_and_population_matches_original(paired, exclude, features, threads, require_thread_budget): + options = paired_geometry(features) if paired else geometry(features) + batch = (gb.grouped_paired_sum_count_batch(Q, QB, R, RB, CUTS, **metadata(), **options, + threads=threads, exclude_exact=exclude) if paired else + gb.grouped_count_batch(Q, R, CUTS, **metadata(), **options, threads=threads, exclude_exact=exclude)) + assert (batch.query_count, batch.reference_count, batch.target_count, batch.peptide_count) == (4, 5, 3, 2) + chunks = list(batch) + assert [(c.start, c.stop) for c in chunks] == [(0, 2), (2, 4)] + counts = np.concatenate([np.asarray(c.counts) for c in chunks]) + sizes = np.concatenate([np.asarray(c.sizes) for c in chunks]) + exact = np.concatenate([np.asarray(c.exact_counts) for c in chunks]) + assert counts.dtype == sizes.dtype == exact.dtype == np.uint32 + assert not np.asarray(chunks[0].counts).flags.writeable + for i, q in enumerate(Q): + identical = [q == r and (not paired or QB[i] == RB[j]) for j, r in enumerate(R)] + subsets = [list(range(5))] + subsets += [[j for j in range(5) if t in MEMBERS[j]] for t in range(3)] + subsets += [[j for j in range(5) if t in MEMBERS[j] and QGENE[i] == RGENE[j]] for t in range(3)] + subsets += [[j for j in range(5) if EXCLUSIVE[j] == p] for p in range(2)] + for channel, refs in enumerate(subsets): + subopt = paired_geometry(features, refs) if paired else geometry(features, refs) + native = (gb.paired_sum_count_batch(Q, QB, [R[j] for j in refs], [RB[j] for j in refs], + CUTS, **subopt, shared_thresholds=True, exclude_exact=exclude, threads=threads)[0] if paired else + gb.count_batch(Q, [R[j] for j in refs], CUTS, **subopt, + shared_thresholds=True, exclude_exact=exclude, threads=threads)) + np.testing.assert_array_equal(counts[i, channel], native[i]) + assert sizes[i, channel] == sum(not (exclude and identical[j]) for j in refs) + for t in range(3): + assert exact[i, t] == sum(identical[j] and t in MEMBERS[j] for j in range(5)) + # All-source complement retains shared-label keys and unselected references. + for p in range(2): + refs = [j for j in range(5) if EXCLUSIVE[j] != p] + subopt = paired_geometry(features, refs) if paired else geometry(features, refs) + native = (gb.paired_sum_count_batch(Q, QB, [R[j] for j in refs], [RB[j] for j in refs], + CUTS, **subopt, shared_thresholds=True, exclude_exact=exclude)[0] if paired else + gb.count_batch(Q, [R[j] for j in refs], CUTS, **subopt, + shared_thresholds=True, exclude_exact=exclude)) + np.testing.assert_array_equal(counts[i, 0]-counts[i, 7+p], native[i]) + held = np.asarray(chunks[0].counts) + del chunks, batch + gc.collect() + assert held.shape == (2, 9, 5) and not held.flags.writeable + + +@pytest.mark.parametrize('change', [dict(reference_target_offsets=[]), + dict(reference_target_ids=[0, 0, 2, 0, 2]), dict(reference_target_ids=[0, 9, 2, 0, 2]), + dict(reference_exclusive_peptide_ids=[1, 1, -1, 0, -1]), + dict(reference_exclusive_peptide_ids=[0, 1, -2, 0, -1]), + dict(query_gene_ids=[0]), dict(reference_gene_ids=[-1]*5), dict(chunk_size=0)]) +def test_invalid_group_metadata_rejected_before_iteration(change): + args = metadata(); args.update(change) + with pytest.raises((ValueError, TypeError, OverflowError)): + gb.grouped_count_batch(Q, R, CUTS, **args) + + +@pytest.mark.parametrize('cuts', [[1, 0], [0, 0], [True]]) +def test_invalid_cutoffs_rejected(cuts): + with pytest.raises((ValueError, TypeError)): + gb.grouped_count_batch(Q, R, cuts, **metadata()) + + +def test_empty_query_reference_and_threshold_axes(): + empty = dict(reference_target_offsets=[0], reference_target_ids=[], target_peptide_ids=[], + reference_exclusive_peptide_ids=[], reference_gene_ids=[], query_gene_ids=[]) + assert list(gb.grouped_count_batch([], [], [], **empty)) == [] + empty['query_gene_ids'] = [0] + chunk = next(gb.grouped_count_batch(['AAA'], [], [], **empty)) + assert np.asarray(chunk.counts).shape == (1, 1, 0) + assert np.asarray(chunk.sizes).tolist() == [[0]] + assert np.asarray(chunk.exact_counts).shape == (1, 0) + + +def test_prepared_ownership_and_weighted_prior_geometry(): + queries, refs = ['AAA', 'AAAA'], ['AAA', 'AAAA', 'ACA'] + args = dict(reference_target_offsets=[0, 1, 2, 3], reference_target_ids=[0]*3, + target_peptide_ids=[0], reference_exclusive_peptide_ids=[0]*3, + query_gene_ids=[0, 1], reference_gene_ids=[1, 0, 1], chunk_size=1) + options = dict(query_position_weights=[[1, 2, 3], [3, 2, 1, 3]], + reference_position_weights=[[1, 1, 2], [2, 1, 3, 1], [2, 2, 2]], + gap_prior=gb.central_prior(2), gap_open=2, gap_extend=1) + expected = gb.count_batch(queries, refs, CUTS, shared_thresholds=True, **options) + batch = gb.grouped_count_batch(queries, refs, CUTS, **args, **options) + queries.append('CCC'); refs.clear(); args['reference_target_ids'][0] = 99 + options['query_position_weights'][0][0] = 99 + actual = np.concatenate([np.asarray(chunk.counts)[:, 0] for chunk in batch]) + np.testing.assert_array_equal(actual, expected) diff --git a/tests/python/test_grouped_extra_offsets.py b/tests/python/test_grouped_extra_offsets.py new file mode 100644 index 0000000..44768fa --- /dev/null +++ b/tests/python/test_grouped_extra_offsets.py @@ -0,0 +1,142 @@ +"""Additive generic factored terms and independent paired priors.""" +import pytest +from seqtree import gapblock, SubstitutionMatrix + + +def rows(buffer): + return memoryview(buffer).tolist() + + +def metadata(): + return dict(reference_target_offsets=[0,1,2,4],reference_target_ids=[0,1,0,1], + target_peptide_ids=[0,1],reference_exclusive_peptide_ids=[0,1,-1], + query_gene_ids=[0,1],reference_gene_ids=[0,1,0]) + + +def feature(cost, weight=1, scale=200): + return gapblock.PairedFeatureOffsets([[0,0],[1,1]],[[1,1],[0,0],[1,1]], + [[0,cost],[cost,0]],weight=weight,scale=scale) + + +@pytest.mark.parametrize('withheld',[False,True]) +@pytest.mark.parametrize('paired',[False,True]) +@pytest.mark.parametrize('threads',[1,4]) +def test_independent_terms_and_lane_priors_against_explicit_expression(paired,threads,withheld, require_thread_budget): + q,r=['AAAA','AAAR'],['AAAA','AARR','AAAAAA'] + qb,rb=['RRRR','RRAA'],['RRRR','AARRR','AAAA'] + matrix=SubstitutionMatrix.blosum62(); cuts=list(range(0,900,20)) + pa=gapblock.frame_prior(0,26);pb=gapblock.frame_prior(3,26) + opts=dict(matrix=matrix,gap_open=13,gap_extend=13,threads=threads, + exclude_exact=True,chunk_size=1,**metadata()) + folds=[-1,0] if withheld else [-1,-1] + if withheld: + opts.update(query_withheld_peptide_ids=folds,reference_peptide_offsets=[0,1,2,4], + reference_peptide_ids=[0,1,0,1]) + primary=feature(1,weight=.5) + extras=[feature(1,weight=.5),feature(1,weight=.5)] + if paired: + batch=gapblock.grouped_paired_sum_count_batch(q,qb,r,rb,cuts,**opts, + gap_prior_alpha=pa,gap_prior_beta=pb,feature_offsets_alpha=primary, + additional_feature_offsets_alpha=extras,additional_feature_offsets_beta=[feature(3)]) + else: + batch=gapblock.grouped_count_batch(q,r,cuts,**opts,gap_prior=pa, + feature_offsets=primary,additional_feature_offsets=extras) + base=gapblock.score_matrix(q,r,matrix,gap_open=13,gap_extend=13,gap_prior=pa) + other=gapblock.score_matrix(qb,rb,matrix,gap_open=13,gap_extend=13,gap_prior=pb) if paired else None + # Three half-unit terms round separately to zero. Beta's integer term is600. + scores=[[base[i,j]+(other[i,j]+600*int(i%2 != (1,0,1)[j]) if paired else 0) + for j in range(3)] for i in range(2)] + expected=[];sizes=[];exact=[] + members=[[0],[1],[0,1]] + for i in range(2): + channels=[[] for _ in range(7)];ex=[0,0] + for j,original_labels in enumerate(members): + labels=[t for t in original_labels if t!=folds[i]] + if not labels:continue + identical=q[i]==r[j] and (not paired or qb[i]==rb[j]) + if identical: + for t in labels:ex[t]+=1 + continue + channels[0].append(scores[i][j]) + for t in labels: + channels[1+t].append(scores[i][j]) + if i%2==j%2:channels[3+t].append(scores[i][j]) + if len(labels)==1:channels[5+labels[0]].append(scores[i][j]) + expected.append([[sum(score<=cut for score in row) for cut in cuts] for row in channels]) + sizes.append([len(row) for row in channels]);exact.append(ex) + chunks=list(batch) + assert [row for chunk in chunks for row in rows(chunk.counts)]==expected + assert [row for chunk in chunks for row in rows(chunk.sizes)]==sizes + assert [row for chunk in chunks for row in rows(chunk.exact_counts)]==exact + + +@pytest.mark.parametrize('lane',['alpha','beta','both']) +def test_paired_prior_overrides_inherit_shared_independently(lane): + q,r=['AAAA','AAAR'],['AAAA','AARR','AAAAAA'] + qb,rb=['RRRR','RRAA'],['RRRR','AARRR','AAAA'] + matrix=SubstitutionMatrix.blosum62();shared=gapblock.frame_prior(1,26) + pa=gapblock.frame_prior(0,26);pb=gapblock.frame_prior(3,26) + kwargs={} + if lane in ('alpha','both'):kwargs['gap_prior_alpha']=pa + if lane in ('beta','both'):kwargs['gap_prior_beta']=pb + a=gapblock.score_matrix(q,r,matrix,gap_open=13,gap_extend=13, + gap_prior=pa if 'gap_prior_alpha' in kwargs else shared) + b=gapblock.score_matrix(qb,rb,matrix,gap_open=13,gap_extend=13, + gap_prior=pb if 'gap_prior_beta' in kwargs else shared) + cuts=list(range(0,300,5)) + batch=gapblock.grouped_paired_sum_count_batch(q,qb,r,rb,cuts,matrix=matrix, + gap_open=13,gap_extend=13,gap_prior=shared,**kwargs,**metadata()) + got=rows(next(batch).counts) + assert [row[0] for row in got]==[[sum(a[i,j]+b[i,j]<=cut for j in range(3)) + for cut in cuts] for i in range(2)] + + +@pytest.mark.parametrize('bad',[ + [object()], [gapblock.PairedFeatureOffsets([],[],[[0]])], + [feature(-1)], [feature(1,weight=float('nan'))], [feature(1,scale=-1)], + [feature(1,scale=2**31-1)], +]) +def test_additional_term_validation(bad): + with pytest.raises((ValueError,OverflowError)): + gapblock.grouped_count_batch(['AAAA','AAAR'],['AAAA','AARR','AAAAAA'],[0], + additional_feature_offsets=bad,**metadata()) + + +def test_native_term_and_prior_validation(monkeypatch): + original=gapblock._additional_geometry + def corrupt(geometry,*args): + original(geometry,*args) + geometry.additional_features=[gapblock.PairedFeatureOffsets([],[],[[0]])] + monkeypatch.setattr(gapblock,'_additional_geometry',corrupt) + with pytest.raises(ValueError,match='axes'): + gapblock.grouped_count_batch(['AAAA','AAAR'],['AAAA','AARR','AAAAAA'],[0],**metadata()) + monkeypatch.setattr(gapblock,'_additional_geometry',original) + original_prior=gapblock._lane_prior + def corrupt_prior(geometry,*args): + original_prior(geometry,*args);geometry.gap_prior=[0] + monkeypatch.setattr(gapblock,'_lane_prior',corrupt_prior) + with pytest.raises(ValueError,match='prior table'): + gapblock.grouped_paired_sum_count_batch(['AAAA','AAAR'],['AAAA','AAAR'], + ['AAAA','AARR','AAAAAA'],['AAAA','AARR','AAAAAA'],[0],**metadata()) + + +@pytest.mark.parametrize('mode',['gapblock','global','local']) +def test_extra_terms_remain_outside_lane_quantization(mode): + q,r=['AAAA'],['RRRR'];matrix=SubstitutionMatrix.blosum62() + opts=dict(matrix=matrix,gap_open=13,gap_extend=13) + if mode=='gapblock': + raw=gapblock.score_matrix(q,r,**opts)[0,0] + else: + from seqtree import pairwise + raw=max(0,sum(matrix.similarity(a,a) for a in q[0]+r[0])- + 2*pairwise.score(q[0],r[0],matrix,mode=mode,gap_open=13,gap_extend=13))*3 + opts.update(alignment_mode=mode,alignment_scale=3) + expected=raw//5*5+raw//7*7+1000 + def term(cost): + return gapblock.PairedFeatureOffsets([[0,0]],[[1,1]],[[0,cost],[cost,0]],scale=200) + batch=gapblock.grouped_paired_sum_count_batch(q,q,r,r,[expected-1,expected],**opts, + score_quantum_alpha=5,score_quantum_beta=7,additional_feature_offsets_alpha=[term(2)], + additional_feature_offsets_beta=[term(3)],reference_target_offsets=[0,1], + reference_target_ids=[0],target_peptide_ids=[0],reference_exclusive_peptide_ids=[0], + query_gene_ids=[0],reference_gene_ids=[0]) + assert rows(next(batch).counts)==[[[0,1]]*4] diff --git a/tests/python/test_grouped_pairwise.py b/tests/python/test_grouped_pairwise.py new file mode 100644 index 0000000..0b42cc5 --- /dev/null +++ b/tests/python/test_grouped_pairwise.py @@ -0,0 +1,176 @@ +"""Bounded affine grouped counts against the existing pairwise-score oracle.""" +import pytest +from seqtree import SubstitutionMatrix, gapblock, pairwise + + +def buf(value): + return memoryview(value).tolist() + + +@pytest.mark.parametrize('mode', ['local', 'global']) +@pytest.mark.parametrize('paired', [False, True]) +@pytest.mark.parametrize('withheld', [False, True]) +@pytest.mark.parametrize('threads', [1, 4]) +def test_all_curves_sizes_exact_against_independent_raw_diagonal_oracle(mode, paired, withheld, threads, require_thread_budget): + q = ['CAAAAF', 'CRRRRF', 'WAAAAW'] + r = ['CAAAAF', 'CRRAAF', 'CRRRRF', 'CAAARRAAF', 'WAAAAW'] + qb = ['CWWWWF', 'CAAAAF', 'CSSSSF'] + rb = ['CWWWWF', 'CWWAAF', 'CAAAAF', 'CNNAAF', 'CSSSSF'] + labels = [[0], [0, 1], [1], [], [1]] + full = [[0], [0, 1], [1], [2], [1]] + exclusive = [0, -1, 1, -1, 1] + qgene, rgene = [0, 1, 0], [0, 0, 1, 0, 2] + groupsq, groupsr = [0, 1, 0], [1, 0, 0, 1, 0] + offsets = [[0, 3], [5, 0]] + cuts = [0, 3, 100, 6000, 18000, 48000, 120000] + matrix = SubstitutionMatrix.blosum62() + csr = [0]; ids = [] + for row in labels: + ids.extend(row); csr.append(len(ids)) + folds = [-1, 0, 1] if withheld else [-1] * len(q) + common = dict(reference_target_offsets=csr, reference_target_ids=ids, + target_peptide_ids=[0, 1], reference_exclusive_peptide_ids=exclusive, + query_gene_ids=qgene, reference_gene_ids=rgene, threads=threads, + exclude_exact=True, chunk_size=2, matrix=matrix, gap_open=13, gap_extend=13, + alignment_mode=mode, alignment_scale=600, trim_left=1, trim_right=1) + if withheld: + csrfull = [0]; fullids = [] + for row in full: + fullids.extend(row); csrfull.append(len(fullids)) + common.update(query_withheld_peptide_ids=folds, reference_peptide_offsets=csrfull, + reference_peptide_ids=fullids) + geometry = dict(query_group_ids=groupsq, reference_group_ids=groupsr, group_distances=offsets, + feature_offsets=gapblock.PairedFeatureOffsets([[0, 0]] * len(q), + [[1, 1]] * len(r), [[0, 2], [2, 0]], scale=200)) + if paired: + opts = {key + suffix: value for suffix in ('_alpha', '_beta') for key, value in geometry.items()} + batch = gapblock.grouped_paired_sum_count_batch(q, qb, r, rb, cuts, **common, **opts) + else: + batch = gapblock.grouped_count_batch(q, r, cuts, **common, **geometry) + # Dense test oracle only: production calls the grouped population once. + def distances(left, right): + l, rr = [s[1:-1] for s in left], [s[1:-1] for s in right] + sim = pairwise.score_matrix(l, rr, matrix, mode=mode, gap_open=13, gap_extend=13, threads=1) + diagl = [sum(matrix.similarity(a, a) for a in s) for s in l] + diagr = [sum(matrix.similarity(a, a) for a in s) for s in rr] + return [[600 * max(0, diagl[i] + diagr[j] - 2 * sim[i, j]) + + offsets[groupsq[i]][groupsr[j]] + 400 for j in range(len(right))] + for i in range(len(left))] + scores = distances(q, r) + if paired: + other = distances(qb, rb) + scores = [[x + y for x, y in zip(a, b)] for a, b in zip(scores, other)] + expected, sizes, exact = [], [], [] + for i in range(len(q)): + channels = [[] for _ in range(7)]; exactrow = [0, 0] + for j in range(len(r)): + surviving = set(full[j]) - {folds[i]} + if not surviving: continue + remaining = [t for t in labels[j] if t != folds[i]] + identical = q[i] == r[j] and (not paired or qb[i] == rb[j]) + if identical: + for t in remaining: exactrow[t] += 1 + continue + channels[0].append(scores[i][j]) + for t in remaining: + channels[1 + t].append(scores[i][j]) + if qgene[i] == rgene[j]: channels[3 + t].append(scores[i][j]) + pep = next(iter(surviving)) if len(surviving) == 1 else -1 + if pep in (0, 1): channels[5 + pep].append(scores[i][j]) + expected.append([[sum(x <= cut for x in row) for cut in cuts] for row in channels]) + sizes.append([len(row) for row in channels]); exact.append(exactrow) + chunks = list(batch) + assert [row for chunk in chunks for row in buf(chunk.counts)] == expected + assert [row for chunk in chunks for row in buf(chunk.sizes)] == sizes + assert [row for chunk in chunks for row in buf(chunk.exact_counts)] == exact + + +def test_diagonal_sums_are_not_self_alignment_scores(): + matrix = SubstitutionMatrix.from_similarity([[-1, -3, -3, -3], [-3, 7, -3, -3], + [-3, -3, 2, -3], [-3, -3, -3, 3]]) + # The NT codec is A,C,G,T. Local self-AA alignment is0, raw diagonal sum is−2. + assert pairwise.score('AA', 'AA', matrix, alphabet='nt', mode='local') == 0 + batch = gapblock.grouped_count_batch(['AA'], ['CC'], [12, 13, 14], alphabet='nt', + matrix=matrix, alignment_mode='local', gap_open=13, gap_extend=13, + reference_target_offsets=[0, 1], reference_target_ids=[0], target_peptide_ids=[0], + reference_exclusive_peptide_ids=[0], query_gene_ids=[0], reference_gene_ids=[0]) + assert buf(next(batch).counts) == [[[1, 1, 1]] * 4] + + +def one_population(q=('CAAF',), r=('CRRF',), **options): + return gapblock.grouped_count_batch(q, r, [0, 1, 600, 18000, 100000], + matrix=SubstitutionMatrix.blosum62(), gap_open=13, gap_extend=13, + reference_target_offsets=list(range(len(r)+1)), reference_target_ids=[0]*len(r), + target_peptide_ids=[0], reference_exclusive_peptide_ids=[0]*len(r), + query_gene_ids=[0]*len(q), reference_gene_ids=[0]*len(r), **options) + + +@pytest.mark.parametrize('options', [ + {'alignment_mode': 'nw'}, {'alignment_mode': None}, + {'alignment_scale': 2}, {'trim_left': 1}, + {'alignment_mode': 'global', 'alignment_scale': 0}, + {'alignment_mode': 'global', 'alignment_scale': True}, + {'alignment_mode': 'global', 'trim_left': -1}, + {'alignment_mode': 'global', 'trim_right': 1.5}, + {'alignment_mode': 'local', 'trim_left': 5}, + {'alignment_mode': 'global', 'query_position_weights': [[1]*4]}, + {'alignment_mode': 'local', 'position_weights_by_length': {4: [1]*4}}, + {'alignment_mode': 'global', 'gap_prior': lambda i,d,m: 0}, + {'alignment_mode': 'global', 'alignment_scale': 2**31-1}, +]) +def test_affine_option_validation(options): + with pytest.raises((ValueError, OverflowError)): + one_population(**options) + + +@pytest.mark.parametrize('mode', ['global', 'local']) +def test_empty_axes_and_fully_trimmed_sequences(mode): + assert list(one_population(q=(), alignment_mode=mode)) == [] + empty = next(one_population(r=(), alignment_mode=mode)) + assert buf(empty.sizes) == [[0]*4] + assert buf(empty.counts) == [[[0]*5]*4] + trimmed = next(one_population(alignment_mode=mode, trim_left=2, trim_right=2)) + assert buf(trimmed.counts) == [[[1]*5]*4] + assert buf(trimmed.exact_counts) == [[0]] + + +def test_affine_quantum_floors_scaled_junction_before_offset(): + matrix = SubstitutionMatrix.blosum62() + q, r = 'CAAF', 'CRRF' + raw = max(0, sum(matrix.similarity(a,a) for a in q) + + sum(matrix.similarity(a,a) for a in r) - + 2*pairwise.score(q, r, matrix, mode='global', gap_open=13, gap_extend=13)) + expected = (raw*3)//100*100+7 + batch = gapblock.grouped_count_batch([q], [r], [expected-1, expected], + matrix=matrix, gap_open=13, gap_extend=13, alignment_mode='global', alignment_scale=3, + score_quantum=100, query_group_ids=[0], reference_group_ids=[0], group_distances=[[7]], + reference_target_offsets=[0, 1], reference_target_ids=[0], target_peptide_ids=[0], + reference_exclusive_peptide_ids=[0], query_gene_ids=[0], reference_gene_ids=[0]) + assert buf(next(batch).counts) == [[[0, 1]]*4] + + +@pytest.mark.parametrize('mode', ['global', 'local']) +def test_exact_evidence_is_retained_and_no_hit_keeps_sizes(mode): + retained = next(one_population(r=('CAAF',), alignment_mode=mode)) + assert buf(retained.exact_counts) == [[1]] + assert buf(retained.counts) == [[[1]*5]*4] + nohit = gapblock.grouped_count_batch(['CAAF'], ['CRRF'], [-1], + matrix=SubstitutionMatrix.blosum62(), gap_open=13, gap_extend=13, + alignment_mode=mode, reference_target_offsets=[0,1], reference_target_ids=[0], + target_peptide_ids=[0], reference_exclusive_peptide_ids=[0], + query_gene_ids=[0], reference_gene_ids=[0]) + chunk = next(nohit) + assert buf(chunk.counts) == [[[0]]*4] + assert buf(chunk.sizes) == [[1]*4] + + +@pytest.mark.parametrize('field,value', [('alignment_mode', 'invalid'), + ('alignment_scale', 0), ('trim_left', 2**32-1)]) +def test_native_affine_validation_rechecks_python_snapshot(monkeypatch, field, value): + original = gapblock._alignment_options + def corrupt(geometry, *args): + original(geometry, *args) + setattr(geometry, field, value) + monkeypatch.setattr(gapblock, '_alignment_options', corrupt) + with pytest.raises(ValueError): + one_population(alignment_mode='local') diff --git a/tests/python/test_grouped_withholding.py b/tests/python/test_grouped_withholding.py new file mode 100644 index 0000000..366b8ae --- /dev/null +++ b/tests/python/test_grouped_withholding.py @@ -0,0 +1,158 @@ +"""Withheld label counts equal independent full fold reconstruction.""" +import numpy as np +import pytest +from seqtree import gapblock as gb + +Q = ['AAA','AAA','CCC','AAA','GGG','TTT'] +R = ['AAA','AAA','AAC','CCC','GGG','TTT'] +QB = ['CCC','CCC','AAA','CCC','TTT','CCC'] +RB = ['CCC','AAC','CCC','AAA','TTT','CCC'] +MEMBERS = [[0,1,2],[0],[1,2],[3],[4],[]] +PARENTS = [0,0,1,2,3] +FULL = [[0,1],[0],[0,1],[2],[3,4],[4]] +WITHHELD = [0,1,-1,2,4,0] +CUTS = [-1,0,2,5,20] +QGENES = [0,1,1,0,0,1] +RGENES = [1,0,1,0,0,1] + + +def csr(rows): + offsets=[0];ids=[] + for row in rows:ids.extend(row);offsets.append(len(ids)) + return offsets,ids + + +def metadata(indices=range(6), withheld=None): + indices=list(indices) + full=[FULL[j] if withheld is None else [p for p in FULL[j] if p!=withheld] for j in indices] + members=[[t for t in MEMBERS[j] if withheld is None or PARENTS[t]!=withheld] for j in indices] + offsets,ids=csr(members) + return dict(reference_target_offsets=offsets,reference_target_ids=ids, + target_peptide_ids=PARENTS, + reference_exclusive_peptide_ids=[row[0] if len(row)==1 and row[0]<4 else -1 for row in full], + query_gene_ids=QGENES,reference_gene_ids=[RGENES[j] for j in indices],chunk_size=2) + + +def options(indices, paired, features): + indices=list(indices) + result=dict(gap_open=2,gap_extend=1) + for lane in range(2 if paired else 1): + suffix=('_alpha' if lane==0 else '_beta') if paired else '' + result['query_group_ids'+suffix]=[0,1,0,1,1,0] + result['reference_group_ids'+suffix]=[[1,0,1,0,1,0][j] for j in indices] + result['group_distances'+suffix]=[[0,3],[3,2]] + if features: + result['feature_offsets'+suffix]=gb.PairedFeatureOffsets( + [[0,1],[1,0],[1,1],[0,1],[1,0],[0,0]], + [[[1,0],[0,0],[0,1],[1,1],[1,0],[0,1]][j] for j in indices], + [[0,3],[3,0]],weight=.5 if lane==0 else 1.5) + return result + + +def run(paired, indices, args, features, threads, exclude): + opt=options(indices,paired,features) + if paired: + chunks=list(gb.grouped_paired_sum_count_batch(Q,QB,[R[j] for j in indices], + [RB[j] for j in indices],CUTS,**args,**opt,threads=threads,exclude_exact=exclude)) + else: + chunks=list(gb.grouped_count_batch(Q,[R[j] for j in indices],CUTS, + **args,**opt,threads=threads,exclude_exact=exclude)) + return tuple(np.concatenate([np.asarray(getattr(c,key)) for c in chunks]) + for key in ('counts','sizes','exact_counts')) + + +@pytest.mark.parametrize('paired',[False,True]) +@pytest.mark.parametrize('exclude',[False,True]) +@pytest.mark.parametrize('features',[False,True]) +@pytest.mark.parametrize('threads',[1,4]) +def test_every_fold_matches_independently_rebuilt_population(paired,exclude,features,threads, require_thread_budget): + args=metadata();offsets,ids=csr(FULL) + args.update(query_withheld_peptide_ids=WITHHELD, + reference_peptide_offsets=offsets,reference_peptide_ids=ids) + actual=run(paired,list(range(6)),args,features,threads,exclude) + for i, withheld in enumerate(WITHHELD): + indices=[j for j,row in enumerate(FULL) if any(p!=withheld for p in row)] + expected=run(paired,indices,metadata(indices,withheld),features,threads,exclude) + for output,oracle in zip(actual,expected): + np.testing.assert_array_equal(output[i],oracle[i]) + # Same sequence in different folds changes surviving target/exclusive channels. + assert actual[1][0,1]==0 + assert actual[1][1,2]>0 + + +def test_no_withholding_vector_matches_existing_default(): + args=metadata();offsets,ids=csr(FULL) + expected=run(False,list(range(6)),args,False,1,False) + args.update(query_withheld_peptide_ids=[-1]*6, + reference_peptide_offsets=offsets,reference_peptide_ids=ids) + for output,oracle in zip(run(False,list(range(6)),args,False,1,False),expected): + np.testing.assert_array_equal(output,oracle) + + +@pytest.mark.parametrize('change',[ + {'query_withheld_peptide_ids':[0]}, + {'query_withheld_peptide_ids':[-2]*6}, + {'query_withheld_peptide_ids':[True]*6}, + {'reference_peptide_offsets':None}, + {'reference_peptide_ids':None}, + {'reference_peptide_offsets':[0]}, + {'reference_peptide_ids':[0,0,0,0,1,2,3,4,4]}, + {'reference_peptide_ids':[0,2,0,0,1,2,3,4,4]}, + {'reference_exclusive_peptide_ids':[0,0,-1,2,-1,-1]}, + {'reference_peptide_ids':[-1,1,0,0,1,2,3,4,4]}, + {'reference_peptide_ids':[0,1,0,0,1,2,3,4,1<<31]}, +]) +def test_invalid_withholding_metadata_rejected_before_iteration(change): + args=metadata();offsets,ids=csr(FULL) + args.update(query_withheld_peptide_ids=WITHHELD, + reference_peptide_offsets=offsets,reference_peptide_ids=ids) + args.update(change) + with pytest.raises((ValueError,TypeError,OverflowError)): + gb.grouped_count_batch(Q,R,CUTS,**args) + + +def test_full_memberships_require_explicit_query_fold_axis(): + args=metadata();offsets,ids=csr(FULL) + with pytest.raises(ValueError): + gb.grouped_count_batch(Q,R,CUTS,**args,reference_peptide_offsets=offsets,reference_peptide_ids=ids) + + +def test_empty_fold_axes_and_all_memberships_removed(): + empty=dict(reference_target_offsets=[0],reference_target_ids=[],target_peptide_ids=[], + reference_exclusive_peptide_ids=[],query_gene_ids=[],reference_gene_ids=[], + query_withheld_peptide_ids=[],reference_peptide_offsets=[0],reference_peptide_ids=[]) + assert list(gb.grouped_count_batch([],[],[],**empty))==[] + empty.update(query_gene_ids=[0],query_withheld_peptide_ids=[0]) + chunk=next(gb.grouped_count_batch(['AAA'],[],[],**empty)) + assert np.asarray(chunk.counts).shape==(1,1,0) + assert np.asarray(chunk.sizes).tolist()==[[0]] + args=dict(reference_target_offsets=[0,1],reference_target_ids=[0],target_peptide_ids=[0], + reference_exclusive_peptide_ids=[0],query_gene_ids=[0],reference_gene_ids=[0], + query_withheld_peptide_ids=[0],reference_peptide_offsets=[0,1],reference_peptide_ids=[0]) + chunk=next(gb.grouped_count_batch(['AAA'],['AAA'],CUTS,**args)) + assert not np.asarray(chunk.counts).any() + assert not np.asarray(chunk.sizes).any() + assert not np.asarray(chunk.exact_counts).any() + args.update(reference_target_ids=[],reference_target_offsets=[0,0], + reference_exclusive_peptide_ids=[-1],reference_peptide_offsets=[0,0],reference_peptide_ids=[]) + chunk=next(gb.grouped_count_batch(['AAA'],['AAA'],CUTS,**args)) + assert not np.asarray(chunk.sizes).any() + + +def test_fold_metadata_owned_and_unselected_only_survivors(): + args=metadata();offsets,ids=csr(FULL) + folds=WITHHELD.copy() + args.update(query_withheld_peptide_ids=folds,reference_peptide_offsets=offsets,reference_peptide_ids=ids) + expected=run(False,list(range(6)),args,False,1,True) + batch=gb.grouped_count_batch(Q,R,CUTS,**args,**options(range(6),False,False),exclude_exact=True) + folds.clear();offsets.clear();ids.clear() + chunks=list(batch) + for key,oracle in zip(('counts','sizes','exact_counts'),expected): + np.testing.assert_array_equal(np.concatenate([np.asarray(getattr(c,key)) for c in chunks]),oracle) + assert not np.asarray(chunks[0].counts).flags.writeable + args=dict(reference_target_offsets=[0,0],reference_target_ids=[],target_peptide_ids=[], + reference_exclusive_peptide_ids=[-1],query_gene_ids=[0],reference_gene_ids=[0], + query_withheld_peptide_ids=[0],reference_peptide_offsets=[0,1],reference_peptide_ids=[4]) + chunk=next(gb.grouped_count_batch(['AAA'],['CCC'],CUTS,**args)) + assert np.asarray(chunk.sizes).tolist()==[[1]] + assert np.asarray(chunk.counts)[0,0,-1]==1 diff --git a/tests/python/test_unit_histogram_fastpath.py b/tests/python/test_unit_histogram_fastpath.py index 360fd8d..eb33def 100644 --- a/tests/python/test_unit_histogram_fastpath.py +++ b/tests/python/test_unit_histogram_fastpath.py @@ -47,3 +47,22 @@ def test_boundaries_iupac_and_existing_engine_fallbacks(): n=Index.build(['ACGTN','RYACGT','ryacgt','RYYACGT'],'iupac') p=SearchParams(max_subs=5,max_ins=2,max_dels=2,max_total_edits=5) assert n.edit_histogram_batch(['acgtN','RYACGT'],p,threads=2,exclude_exact=True)==_expected(n,['acgtN','RYACGT'],p,True) + + +@pytest.mark.parametrize("cap", [3,4,5]) +@pytest.mark.parametrize("exclude", [False,True]) +@pytest.mark.parametrize("threads", [1, 4]) +def test_bounded_rejection_preserves_per_type_caps_and_dense_hits(cap,exclude,threads,require_thread_budget): + refs=['ACDEFGHIKLMNPQRST','acdefghiklmnpqrst', + 'ACDEFACDGHIKLMNPQRST','ACDEFGHIKLMNP', + 'ACDEFGHIKLMNPACDRST','ACDGHIKLMNPQRST', + 'AAAAAAAAAAAAAAAA','CCCCCCCCCCCCCCCC'] + index=Index.build(refs+refs[:2]) + queries=refs+['ACDEFGHIKLMNPQRS','ACDEFGHIKLMNPQACDRST'] + for limits in ({},{'max_subs':1},{'max_ins':1},{'max_dels':1}, + {'max_ins':3},{'max_dels':3},{'max_ins':0,'max_dels':0}): + args=dict(max_subs=cap,max_ins=2,max_dels=2,max_total_edits=cap) + args.update(limits) + params=SearchParams(**args) + expected=_expected(index,queries,params,exclude) + assert index.edit_histogram_batch(queries,params,threads,exclude)==expected