diff --git a/Cargo.toml b/Cargo.toml index 0e8e28b75..083b320d0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ path = "src/rust/lib.rs" # Enable optimizations in dev mode for better performance during development [profile.dev] opt-level = 3 +overflow-checks = false [lints.rust] warnings = "deny" diff --git a/src/biotite/sequence/align/kmeralphabet.pyx b/src/biotite/sequence/align/kmeralphabet.py similarity index 65% rename from src/biotite/sequence/align/kmeralphabet.pyx rename to src/biotite/sequence/align/kmeralphabet.py index ba4649a62..2e824b1bf 100644 --- a/src/biotite/sequence/align/kmeralphabet.pyx +++ b/src/biotite/sequence/align/kmeralphabet.py @@ -2,32 +2,23 @@ # under the 3-Clause BSD License. Please see 'LICENSE.rst' for further # information. +from __future__ import annotations + __name__ = "biotite.sequence.align" __author__ = "Patrick Kunzmann" __all__ = ["KmerAlphabet"] -cimport cython -cimport numpy as np - +from collections.abc import Iterator +from collections.abc import Sequence as SequenceABC +from typing import Any, overload import numpy as np -from ..alphabet import Alphabet, LetterAlphabet, AlphabetError - - -ctypedef np.uint8_t uint8 -ctypedef np.uint16_t uint16 -ctypedef np.uint32_t uint32 -ctypedef np.uint64_t uint64 -ctypedef np.int64_t int64 +from biotite.rust.sequence.align import create_kmers as rust_create_kmers +from biotite.rust.sequence.align import split_kmers as rust_split_kmers +from biotite.sequence.alphabet import Alphabet, AlphabetError, LetterAlphabet +from biotite.typing import K, N, NDArray1, NDArray2, S -ctypedef fused CodeType: - uint8 - uint16 - uint32 - uint64 - - -class KmerAlphabet(Alphabet): +class KmerAlphabet(Alphabet[tuple[S, ...]]): """ __init__(base_alphabet, k, spacing=None) @@ -84,7 +75,7 @@ class KmerAlphabet(Alphabet): ----- The symbol code for a *k-mer* :math:`s` calculates as - .. math:: RMSD = \sum_{i=0}^{k-1} n^{k-i-1} s_i + .. math:: RMSD = \\sum_{i=0}^{k-1} n^{k-i-1} s_i where :math:`n` is the length of the base alphabet. @@ -147,11 +138,15 @@ class KmerAlphabet(Alphabet): ['BI_T', 'IQ_I', 'QT_T', 'TI_E'] """ - def __init__(self, base_alphabet, k, spacing=None): + def __init__( + self, + base_alphabet: Alphabet[S], + k: int, + spacing: str | list[int] | NDArray1[K, np.integer] | None = None, + ) -> None: if not isinstance(base_alphabet, Alphabet): raise TypeError( - f"Got {type(base_alphabet).__name__}, " - f"but Alphabet was expected" + f"Got {type(base_alphabet).__name__}, but Alphabet was expected" ) if k < 2: raise ValueError("k must be at least 2") @@ -160,8 +155,7 @@ def __init__(self, base_alphabet, k, spacing=None): base_alph_len = len(self._base_alph) self._radix_multiplier = np.array( - [base_alph_len**n for n in reversed(range(0, self._k))], - dtype=np.int64 + [base_alph_len**n for n in reversed(range(0, self._k))], dtype=np.int64 ) if spacing is None: @@ -174,35 +168,29 @@ def __init__(self, base_alphabet, k, spacing=None): self._spacing = np.array(spacing, dtype=np.int64) self._spacing.sort() if (self._spacing < 0).any(): - raise ValueError( - "Only non-negative integers are allowed for spacing" - ) + raise ValueError("Only non-negative integers are allowed for spacing") if len(np.unique(self._spacing)) != len(self._spacing): - raise ValueError( - "Spacing model contains duplicate values" - ) + raise ValueError("Spacing model contains duplicate values") - if spacing is not None and len(self._spacing) != self._k: + if self._spacing is not None and len(self._spacing) != self._k: raise ValueError( f"Expected {self._k} informative positions, " f"but got {len(self._spacing)} positions in spacing" ) - @property - def base_alphabet(self): + def base_alphabet(self) -> Alphabet[S]: return self._base_alph @property - def k(self): + def k(self) -> int: return self._k @property - def spacing(self): + def spacing(self) -> NDArray1[K, np.int64] | None: return None if self._spacing is None else self._spacing.copy() - - def get_symbols(self): + def get_symbols(self) -> tuple: """ get_symbols() @@ -224,12 +212,14 @@ def get_symbols(self): to be created first. """ if isinstance(self._base_alph, LetterAlphabet): - return tuple(["".join(self.decode(code)) for code in range(len(self))]) + return tuple( + # For a `LetterAlphabet` the decoded symbols are single characters + ["".join(self.decode(code)) for code in range(len(self))] # pyright: ignore[reportCallIssue, reportArgumentType] + ) else: return tuple([list(self.decode(code)) for code in range(len(self))]) - - def extends(self, alphabet): + def extends(self, alphabet: Alphabet[Any]) -> bool: # A KmerAlphabet cannot really extend another KmerAlphabet: # If k is not equal, all symbols are not equal # If the base alphabet has additional symbols, the correct @@ -238,16 +228,19 @@ def extends(self, alphabet): # if the two alphabets are equal return alphabet == self - - def encode(self, symbol): + def encode(self, symbol: SequenceABC[S]) -> int: return self.fuse(self._base_alph.encode_multiple(symbol)) - - def decode(self, code): + def decode(self, code: int) -> SequenceABC[S]: return self._base_alph.decode_multiple(self.split(code)) - - def fuse(self, codes): + @overload + def fuse(self, codes: NDArray1[K, np.integer]) -> int: ... + @overload + def fuse(self, codes: NDArray2[N, K, np.integer]) -> NDArray1[N, np.int64]: ... + def fuse( + self, codes: NDArray1[K, np.integer] | NDArray2[N, K, np.integer] + ) -> int | NDArray1[N, np.int64]: """ fuse(codes) @@ -299,7 +292,15 @@ def fuse(self, codes): # The last dimension is removed since it collpased in np.sum return kmer_code.reshape(orig_shape[:-1]) - def split(self, kmer_code): + @overload + def split(self, kmer_code: int) -> NDArray1[K, np.uint64]: ... + @overload + def split( + self, kmer_code: NDArray1[N, np.integer] + ) -> NDArray2[N, K, np.uint64]: ... + def split( + self, kmer_code: int | NDArray1[N, np.integer] + ) -> NDArray1[K, np.uint64] | NDArray2[N, K, np.uint64]: """ split(kmer_code) @@ -338,42 +339,17 @@ def split(self, kmer_code): [3 1] """ if np.any(kmer_code >= len(self)) or np.any(kmer_code < 0): - raise AlphabetError( - f"Given k-mer symbol code is invalid for this alphabet" - ) + raise AlphabetError("Given k-mer symbol code is invalid for this alphabet") orig_shape = np.shape(kmer_code) - split_codes = self._split( - np.atleast_1d(kmer_code).astype(np.int64, copy=False) + split_codes = rust_split_kmers( + np.atleast_1d(kmer_code).astype(np.int64, copy=False), + self._k, + len(self._base_alph), ) return split_codes.reshape(orig_shape + (self._k,)) - @cython.boundscheck(False) - @cython.wraparound(False) - @cython.cdivision(True) - def _split(self, int64[:] codes not None): - cdef int i, n - cdef int64 code, val, symbol_code - - cdef int64[:] radix_multiplier = self._radix_multiplier - - cdef uint64[:,:] split_codes = np.empty( - (codes.shape[0], self._k), dtype=np.uint64 - ) - - cdef int k = self._k - for i in range(codes.shape[0]): - code = codes[i] - for n in range(k): - val = radix_multiplier[n] - symbol_code = code // val - split_codes[i,n] = symbol_code - code -= symbol_code * val - - return np.asarray(split_codes) - - - def kmer_array_length(self, int64 length): + def kmer_array_length(self, length: int) -> int: """ kmer_array_length(length) @@ -384,25 +360,20 @@ def kmer_array_length(self, int64 length): Parameters ---------- length : int - The length of the hypothetical sequence + The length of the hypothetical sequence. Returns ------- kmer_length : int The length of created *k-mer* array. """ - cdef int64 max_offset - cdef int64[:] spacing - if self._spacing is None: return length - self._k + 1 else: - spacing = self._spacing - max_offset = self._spacing[len(spacing)-1] + 1 + max_offset = int(self._spacing[-1]) + 1 return length - max_offset + 1 - - def create_kmers(self, seq_code): + def create_kmers(self, seq_code: NDArray1[N, np.integer]) -> NDArray1[K, np.int64]: """ create_kmers(seq_code) @@ -431,117 +402,17 @@ def create_kmers(self, seq_code): >>> print(["".join(kmer) for kmer in kmer_alphabet.decode_multiple(kmer_codes)]) ['AT', 'TT', 'TG', 'GC', 'CT'] """ - if self._spacing is None: - return self._create_continuous_kmers(seq_code) - else: - return self._create_spaced_kmers(seq_code) - - @cython.boundscheck(False) - @cython.wraparound(False) - def _create_continuous_kmers(self, CodeType[:] seq_code not None): - """ - Fast implementation of k-mer decomposition. - Each k-mer is computed from the previous one by removing - a symbol shifting the remaining values and add the new symbol. - Requires looping only over sequence length. - """ - cdef int64 i - - cdef int k = self._k - cdef uint64 alphabet_length = len(self._base_alph) - cdef int64[:] radix_multiplier = self._radix_multiplier - cdef int64 end_radix_multiplier = alphabet_length**(k-1) - - if len(seq_code) < k: - raise ValueError( - "The length of the sequence code is shorter than k" - ) - - cdef int64[:] kmers = np.empty( - self.kmer_array_length(len(seq_code)), dtype=np.int64 - ) - - cdef CodeType code - cdef int64 kmer, prev_kmer - # Compute first k-mer using naive approach - kmer = 0 - for i in range(k): - code = seq_code[i] - if code >= alphabet_length: - raise AlphabetError(f"Symbol code {code} is out of range") - kmer += radix_multiplier[i] * code - kmers[0] = kmer - - # Compute all following k-mers from the previous one - prev_kmer = kmer - for i in range(1, kmers.shape[0]): - code = seq_code[i + k - 1] - if code >= alphabet_length: - raise AlphabetError(f"Symbol code {code} is out of range") - kmer = ( - ( - # Remove first symbol - (prev_kmer - seq_code[i - 1] * end_radix_multiplier) - # Shift k-mer to left - * alphabet_length - ) - # Add new symbol - + code - ) - kmers[i] = kmer - prev_kmer = kmer - - return np.asarray(kmers) - - @cython.boundscheck(False) - @cython.wraparound(False) - def _create_spaced_kmers(self, CodeType[:] seq_code not None): - cdef int64 i, j - - cdef int k = self._k - cdef int64[:] spacing = self._spacing - # The last element of the spacing model - # defines the total k-mer 'span' - cdef int64 max_offset = spacing[len(spacing)-1] + 1 - cdef uint64 alphabet_length = len(self._base_alph) - cdef int64[:] radix_multiplier = self._radix_multiplier - - if len(seq_code) < max_offset: - raise ValueError( - "The length of the sequence code is shorter " - "than the k-mer span" - ) - - cdef int64[:] kmers = np.empty( - self.kmer_array_length(len(seq_code)), dtype=np.int64 - ) - - cdef CodeType code - cdef int64 kmer - cdef int64 offset - for i in range(kmers.shape[0]): - kmer = 0 - for j in range(k): - offset = spacing[j] - code = seq_code[i + offset] - if code >= alphabet_length: - raise AlphabetError(f"Symbol code {code} is out of range") - kmer += radix_multiplier[j] * code - kmers[i] = kmer + return rust_create_kmers(seq_code, self._k, len(self._base_alph), self._spacing) - return np.asarray(kmers) - - - def __str__(self): + def __str__(self) -> str: return str(self.get_symbols()) + def __repr__(self) -> str: + return ( + f"KmerAlphabet({repr(self._base_alph)}, {self._k}, {repr(self._spacing)})" + ) - def __repr__(self): - return f"KmerAlphabet({repr(self._base_alph)}, " \ - f"{self._k}, {repr(self._spacing)})" - - - def __eq__(self, item): + def __eq__(self, item: object) -> bool: if item is self: return True if not isinstance(item, KmerAlphabet): @@ -559,37 +430,37 @@ def __eq__(self, item): return True + def __hash__(self) -> int: + spacing = None if self._spacing is None else tuple(self._spacing.tolist()) + return hash((self._base_alph, self._k, spacing)) - def __hash__(self): - return hash((self._base_alph, self._k, tuple(self._spacing.tolist()))) - - - def __len__(self): + def __len__(self) -> int: return int(len(self._base_alph) ** self._k) - - def __iter__(self): + def __iter__(self) -> Iterator[SequenceABC[S]]: # Creating all symbols is expensive # -> Use a generator instead if isinstance(self._base_alph, LetterAlphabet): - return ("".join(self.decode(code)) for code in range(len(self))) + return ( + "".join(self.decode(code)) # pyright: ignore[reportCallIssue, reportArgumentType] + for code in range(len(self)) + ) else: - return (list(self.decode(code)) for code in range(len(self))) - + return (tuple(self.decode(code)) for code in range(len(self))) - def __contains__(self, symbol): + def __contains__(self, symbol: object) -> bool: try: - self.fuse(self._base_alph.encode_multiple(symbol)) + self.fuse(self._base_alph.encode_multiple(symbol)) # pyright: ignore[reportArgumentType] return True except AlphabetError: return False -def _to_array_form(model_string): +def _to_array_form(model_string: str) -> NDArray1[K, np.int64]: """ Convert the the common string representation of a *k-mer* spacing model into an array, e.g. ``'1*11'`` into ``[0, 2, 3]``. """ - return np.array([ - i for i in range(len(model_string)) if model_string[i] == "1" - ], dtype=np.int64) + return np.array( + [i for i in range(len(model_string)) if model_string[i] == "1"], dtype=np.int64 + ) diff --git a/src/biotite/sequence/align/kmeralphabet.pyi b/src/biotite/sequence/align/kmeralphabet.pyi deleted file mode 100644 index e8fcff0a9..000000000 --- a/src/biotite/sequence/align/kmeralphabet.pyi +++ /dev/null @@ -1,47 +0,0 @@ -__all__ = ["KmerAlphabet"] - -from collections.abc import Iterable, Iterator -from collections.abc import Sequence as SequenceABC -from typing import Any, overload -import numpy as np -from biotite.sequence.alphabet import Alphabet -from biotite.typing import K, N, NDArray1, NDArray2, S - -class KmerAlphabet(Alphabet[tuple[S, ...]]): - def __init__( - self, - base_alphabet: Alphabet[S], - k: int, - spacing: str | list[int] | NDArray1[K, np.integer] | None = None, - ) -> None: ... - @property - def base_alphabet(self) -> Alphabet[S]: ... - @property - def k(self) -> int: ... - @property - def spacing(self) -> NDArray1[K, np.int64] | None: ... - def get_symbols(self) -> tuple[tuple[S, ...], ...]: ... - def extends(self, alphabet: Alphabet[Any]) -> bool: ... - def encode(self, symbol: Iterable[S]) -> int: ... - def decode(self, code: int) -> SequenceABC[S]: ... - @overload - def fuse(self, codes: NDArray1[K, np.integer]) -> int: ... - @overload - def fuse(self, codes: NDArray2[N, K, np.integer]) -> NDArray1[N, np.int64]: ... - @overload - def split(self, kmer_code: int) -> NDArray1[K, np.uint64]: ... - @overload - def split( - self, kmer_code: NDArray1[N, np.integer] - ) -> NDArray2[N, K, np.uint64]: ... - def kmer_array_length(self, length: int) -> int: ... - def create_kmers( - self, seq_code: NDArray1[N, np.unsignedinteger] - ) -> NDArray1[K, np.int64]: ... - def __str__(self) -> str: ... - def __repr__(self) -> str: ... - def __eq__(self, item: object) -> bool: ... - def __hash__(self) -> int: ... - def __len__(self) -> int: ... - def __iter__(self) -> Iterator[tuple[S, ...]]: ... - def __contains__(self, symbol: object) -> bool: ... diff --git a/src/biotite/sequence/align/kmersimilarity.py b/src/biotite/sequence/align/kmersimilarity.py new file mode 100644 index 000000000..d4b313e5e --- /dev/null +++ b/src/biotite/sequence/align/kmersimilarity.py @@ -0,0 +1,160 @@ +# This source code is part of the Biotite package and is distributed +# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further +# information. + +from __future__ import annotations + +__name__ = "biotite.sequence.align" +__author__ = "Patrick Kunzmann" +__all__ = ["SimilarityRule", "ScoreThresholdRule"] + +import abc +from typing import Any, Generic +import numpy as np +from biotite.rust.sequence.align import similar_kmers as rust_similar_kmers +from biotite.sequence.align.kmeralphabet import KmerAlphabet +from biotite.sequence.align.matrix import SubstitutionMatrix +from biotite.typing import N, NDArray1, S + + +class SimilarityRule(metaclass=abc.ABCMeta): + """ + This is the abstract base class for all similarity rules. + A :class:`SimilarityRule` calculates all *similar* *k-mers* for + a given *k-mer*, while the definition of similarity depends + on the derived class. + """ + + @abc.abstractmethod + def similar_kmers( + self, kmer_alphabet: KmerAlphabet[Any], kmer: int + ) -> NDArray1[N, np.int64]: + """ + similar_kmers(kmer_alphabet, kmer) + + Calculate all similar *k-mers* for a given *k-mer*. + + Parameters + ---------- + kmer_alphabet : KmerAlphabet + The reference *k-mer* alphabet to select the *k-mers* from. + kmer : int + The symbol code for the *k-mer* to find similars for. + + Returns + ------- + similar_kmers : ndarray, dtype=np.int64 + The symbol codes for all similar *k-mers*. + + Notes + ----- + The implementations in derived classes must ensure that the + returned array + + 1. contains no duplicates and + 2. includes the input `kmer` itself. + """ + raise NotImplementedError + + +class ScoreThresholdRule(SimilarityRule, Generic[S]): + """ + __init__(matrix, threshold) + + This similarity rule calculates all *k-mers* that have a greater or + equal similarity score with a given *k-mer* than a defined threshold + score. + + The similarity score :math:`S` of two *k-mers* :math:`a` and + :math:`b` is defined as the sum of the pairwise similarity scores + from a substitution matrix :math:`M`: + + .. math:: + + S(a,b) = \\sum_{i=1}^k M(a_i, b_i) + + Therefore, this similarity rule allows substitutions with similar + symbols within a *k-mer*. + + This class is especially useful for finding similar *k-mers* in + protein sequences. + + Parameters + ---------- + matrix : SubstitutionMatrix + The similarity scores are taken from this matrix. + The matrix must be symmetric. + threshold : int + The threshold score. + A *k-mer* :math:`b` is regarded as similar to a *k-mer* + :math:`a`, if the similarity score between :math:`a` and + :math:`b` is equal or greater than the threshold. + + Notes + ----- + For efficient generation of similar *k-mers* an implementation of + the *branch-and-bound* algorithm :footcite:`Hauser2013` is used. + + References + ---------- + + .. footbibliography:: + + Examples + -------- + + >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, k=3) + >>> matrix = SubstitutionMatrix.std_protein_matrix() + >>> rule = ScoreThresholdRule(matrix, threshold=15) + >>> similars = rule.similar_kmers(kmer_alphabet, kmer_alphabet.encode("AIW")) + >>> print(["".join(s) for s in kmer_alphabet.decode_multiple(similars)]) + ['AFW', 'AIW', 'ALW', 'AMW', 'AVW', 'CIW', 'GIW', 'SIW', 'SVW', 'TIW', 'VIW', 'XIW'] + """ + + def __init__(self, matrix: SubstitutionMatrix[S, S], threshold: int) -> None: + if not matrix.is_symmetric(): + raise ValueError("A symmetric substitution matrix is required") + self._matrix = matrix + self._threshold = int(threshold) + # The contiguous score matrix and the per-symbol maximum scores only + # depend on the matrix, so they are computed once here instead of on + # every `similar_kmers()` call + self._score_matrix = np.ascontiguousarray(matrix.score_matrix(), dtype=np.int32) + self._max_scores = np.max(self._score_matrix, axis=-1).astype( + np.int32, copy=False + ) + + def similar_kmers( + self, kmer_alphabet: KmerAlphabet[S], kmer: int + ) -> NDArray1[N, np.int64]: + """ + Calculate all similar *k-mers* for a given *k-mer*. + + Parameters + ---------- + kmer_alphabet : KmerAlphabet + The reference *k-mer* alphabet to select the *k-mers* from. + kmer : int + The symbol code for the *k-mer* to find similars for. + + Returns + ------- + similar_kmers : ndarray, dtype=np.int64 + The symbol codes for all similar *k-mers*. + """ + if not self._matrix.get_alphabet1().extends(kmer_alphabet.base_alphabet): + raise ValueError( + "Substitution matrix is incompatible with k-mer base alphabet" + ) + + # Split the k-mer code into the individual symbol codes + split_kmer = kmer_alphabet.split(kmer).astype(np.int64) + similar_split_kmers = rust_similar_kmers( + self._score_matrix, + self._max_scores, + split_kmer, + len(kmer_alphabet.base_alphabet), + self._threshold, + ) + # Convert the split k-mers back to k-mer codes + return kmer_alphabet.fuse(similar_split_kmers) diff --git a/src/biotite/sequence/align/kmersimilarity.pyi b/src/biotite/sequence/align/kmersimilarity.pyi deleted file mode 100644 index 586f02853..000000000 --- a/src/biotite/sequence/align/kmersimilarity.pyi +++ /dev/null @@ -1,22 +0,0 @@ -__all__ = ["SimilarityRule", "ScoreThresholdRule"] - -import abc -from typing import Any -import numpy as np -from biotite.sequence.align.kmeralphabet import KmerAlphabet -from biotite.sequence.align.matrix import SubstitutionMatrix -from biotite.typing import N, NDArray1 - -class SimilarityRule(metaclass=abc.ABCMeta): - @abc.abstractmethod - def similar_kmers( - self, kmer_alphabet: KmerAlphabet[Any], kmer: int - ) -> NDArray1[N, np.int64]: ... - -class ScoreThresholdRule(SimilarityRule): - def __init__( - self, matrix: SubstitutionMatrix[Any, Any], threshold: int - ) -> None: ... - def similar_kmers( - self, kmer_alphabet: KmerAlphabet[Any], kmer: int - ) -> NDArray1[N, np.int64]: ... diff --git a/src/biotite/sequence/align/kmersimilarity.pyx b/src/biotite/sequence/align/kmersimilarity.pyx deleted file mode 100644 index 5c42ed6b4..000000000 --- a/src/biotite/sequence/align/kmersimilarity.pyx +++ /dev/null @@ -1,233 +0,0 @@ -# This source code is part of the Biotite package and is distributed -# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further -# information. - -__name__ = "biotite.sequence.align" -__author__ = "Patrick Kunzmann" -__all__ = ["SimilarityRule", "ScoreThresholdRule"] - -cimport cython -cimport numpy as np - -import abc -import numpy as np - - -ctypedef np.int64_t int64 -ctypedef np.int32_t int32 - - -class SimilarityRule(metaclass=abc.ABCMeta): - """ - This is the abstract base class for all similarity rules. - A :class:`SimilarityRule` calculates all *similar* *k-mers* for - a given *k-mer*, while the definition of similarity depends - on the derived class. - """ - - @abc.abstractmethod - def similar_kmers(self, kmer_alphabet, kmer): - """ - similar_kmers(kmer_alphabet, kmer) - - Calculate all similar *k-mers* for a given *k-mer*. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The reference *k-mer* alphabet to select the *k-mers* from. - kmer : int - The symbol code for the *k-mer* to find similars for. - - Returns - ------- - similar_kmers : ndarray, dtype=np.int64 - The symbol codes for all similar *k-mers*. - - Notes - ----- - The implementations in derived classes must ensure that the - returned array - - 1. contains no duplicates and - 2. includes the input `kmer` itself. - """ - pass - - -class ScoreThresholdRule(SimilarityRule): - """ - __init__(matrix, threshold) - - This similarity rule calculates all *k-mers* that have a greater or - equal similarity score with a given *k-mer* than a defined threshold - score. - - The similarity score :math:`S` of two *k-mers* :math:`a` and - :math:`b` is defined as the sum of the pairwise similarity scores - from a substitution matrix :math:`M`: - - .. math:: - - S(a,b) = \sum_{i=1}^k M(a_i, b_i) - - Therefore, this similarity rule allows substitutions with similar - symbols within a *k-mer*. - - This class is especially useful for finding similar *k-mers* in - protein sequences. - - Parameters - ---------- - matrix : SubstitutionMatrix - The similarity scores are taken from this matrix. - The matrix must be symmetric. - threshold : int - The threshold score. - A *k-mer* :math:`b` is regarded as similar to a *k-mer* - :math:`a`, if the similarity score between :math:`a` and - :math:`b` is equal or greater than the threshold. - - Notes - ----- - For efficient generation of similar *k-mers* an implementation of - the *branch-and-bound* algorithm :footcite:`Hauser2013` is used. - - References - ---------- - - .. footbibliography:: - - Examples - -------- - - >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, k=3) - >>> matrix = SubstitutionMatrix.std_protein_matrix() - >>> rule = ScoreThresholdRule(matrix, threshold=15) - >>> similars = rule.similar_kmers(kmer_alphabet, kmer_alphabet.encode("AIW")) - >>> print(["".join(s) for s in kmer_alphabet.decode_multiple(similars)]) - ['AFW', 'AIW', 'ALW', 'AMW', 'AVW', 'CIW', 'GIW', 'SIW', 'SVW', 'TIW', 'VIW', 'XIW'] - """ - - def __init__(self, matrix, int32 threshold): - if not matrix.is_symmetric(): - raise ValueError("A symmetric substitution matrix is required") - self._matrix = matrix - self._threshold = threshold - - @cython.boundscheck(False) - @cython.wraparound(False) - def similar_kmers(self, kmer_alphabet, kmer): - """ - Calculate all similar *k-mers* for a given *k-mer*. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The reference *k-mer* alphabet to select the *k-mers* from. - kmer : int - The symbol code for the *k-mer* to find similars for. - - Returns - ------- - similar_kmers : ndarray, dtype=np.int64 - The symbol codes for all similar *k-mers*. - """ - cdef int INIT_SIZE = 1 - - if not self._matrix.get_alphabet1().extends( - kmer_alphabet.base_alphabet - ): - raise ValueError( - "Substitution matrix is incompatible with k-mer base alphabet" - ) - - cdef int64 alph_len = len(kmer_alphabet.base_alphabet) - cdef const int32[:,:] matrix = self._matrix.score_matrix() - # For simplicity trim matrix to required size - # (remove unused symbols) - matrix = matrix[:alph_len, :alph_len] - cdef int32 threshold = self._threshold - - cdef int32[:] max_scores = np.max(self._matrix.score_matrix(), axis=-1) - - cdef int k = kmer_alphabet.k - # Split the k-mer code into the individual symbol codes - cdef int64[:] split_kmer = kmer_alphabet.split(kmer).astype(np.int64) - # This array will hold the current kmer to be tested - cdef int64[:] current_split_kmer = np.zeros(k, dtype=np.int64) - # This array will store the accepted k-mers - # i.e. k-mers that reach the threshold score - cdef int64[:,:] similar_split_kmers = np.zeros( - (INIT_SIZE, k), dtype=np.int64 - ) - - # Calculate the minimum score for each k-mer position that is - # necessary to reach a total higher/equal to the threshold score - cdef int32[:] positional_thresholds = np.empty(k, dtype=np.int32) - cdef int i - cdef int total_max_score = 0 - for i in reversed(range(positional_thresholds.shape[0])): - positional_thresholds[i] = threshold - total_max_score - total_max_score += max_scores[split_kmer[i]] - - # 'pos' is the current position within the k-mer - # where symbols are substituted - cdef int pos = 0 - cdef int similar_i = 0 - cdef int32 score - # 'pos' is -1, after all symbol codes at pos 0 are traversed - while pos != -1: - if current_split_kmer[pos] >= alph_len: - # All symbol codes were traversed at this position - # -> jump one k-mer position back and proceed with - # next symbol - pos -= 1 - if pos != -1: - current_split_kmer[pos] += 1 - else: - # Get total similarity score between the input k-mer - # and generated k-mer up to the point of the current - # position - score = 0 - for i in range(pos+1): - score += matrix[split_kmer[i], current_split_kmer[i]] - # Check score threshold condition - if score >= positional_thresholds[pos]: - # Threshold condition is fulfilled: - # Either go deeper in the same branch - # (jump one position forward) ... - if pos < k-1: - pos += 1 - current_split_kmer[pos] = 0 - # ...or store similar k-mer, - # if already at maximum depth (last k-mer position) - else: - if similar_i >= similar_split_kmers.shape[0]: - # The array is full -> double its size - similar_split_kmers = expand( - np.asarray(similar_split_kmers) - ) - similar_split_kmers[similar_i] = current_split_kmer - similar_i += 1 - # Proceed with the next symbol at this position, - # as we cannot go deeper anymore - current_split_kmer[pos] += 1 - else: - # The threshold score is not reached - # -> this branch ends and we proceed with the next - # symbol at this position - current_split_kmer[pos] += 1 - - # Trim to correct size - # and convert split k-mers back to k-mer code - return kmer_alphabet.fuse(np.asarray(similar_split_kmers[:similar_i])) - - -cdef np.ndarray expand(np.ndarray array): - """ - Double the size of the first dimension of an existing array. - """ - new_array = np.empty((array.shape[0]*2, array.shape[1]), dtype=array.dtype) - new_array[:array.shape[0],:] = array - return new_array \ No newline at end of file diff --git a/src/biotite/sequence/align/permutation.pyx b/src/biotite/sequence/align/permutation.py similarity index 86% rename from src/biotite/sequence/align/permutation.pyx rename to src/biotite/sequence/align/permutation.py index 892bee42d..e5674de88 100644 --- a/src/biotite/sequence/align/permutation.pyx +++ b/src/biotite/sequence/align/permutation.py @@ -2,18 +2,18 @@ # under the 3-Clause BSD License. Please see 'LICENSE.rst' for further # information. +from __future__ import annotations + __name__ = "biotite.sequence.align" __author__ = "Patrick Kunzmann" __all__ = ["Permutation", "RandomPermutation", "FrequencyPermutation"] -cimport cython -cimport numpy as np - import abc +from typing import Generic import numpy as np - - -ctypedef np.int64_t int64 +from biotite.sequence.align.kmeralphabet import KmerAlphabet +from biotite.sequence.align.kmertable import KmerTable +from biotite.typing import N, NDArray1, S class Permutation(metaclass=abc.ABCMeta): @@ -38,20 +38,18 @@ class Permutation(metaclass=abc.ABCMeta): Must be overridden by subclasses. """ - @property @abc.abstractmethod - def min(self): + def min(self) -> int: pass @property @abc.abstractmethod - def max(self): + def max(self) -> int: pass - @abc.abstractmethod - def permute(self, kmers): + def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: """ permute(kmers) @@ -79,6 +77,13 @@ class RandomPermutation(Permutation): r""" Provide a pseudo-randomized order for *k-mers*. + Attributes + ---------- + min, max: int + The minimum and maximum value, the permutated value + (i.e. the return value of :meth:`permute()`) + can take. + Notes ----- @@ -92,13 +97,6 @@ class RandomPermutation(Permutation): However, note that LCGs in general do not provide perfect random behavior, but only *good-enough* values for this purpose. - Attributes - ---------- - min, max: int - The minimum and maximum value, the permutated value - (i.e. the return value of :meth:`permute()`) - can take. - References ---------- @@ -131,34 +129,31 @@ class RandomPermutation(Permutation): ['GA', 'TC', 'AG', 'CT', 'TA', 'AC', 'CG', 'GT', 'AA', 'CC', 'GG', 'TT', 'CA', 'GC', 'TG', 'AT'] """ - LCG_A = 0xd1342543de82ef95 + LCG_A = 0xD1342543DE82EF95 LCG_C = 1 - @property - def min(self): + def min(self) -> int: return np.iinfo(np.int64).min @property - def max(self): + def max(self) -> int: return np.iinfo(np.int64).max - - def permute(self, kmers): - kmers = kmers.astype(np.int64, copy=False) + def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: # Cast to unsigned int to harness the m=2^64 LCG - kmers = kmers.view(np.uint64) + unsigned = kmers.astype(np.int64, copy=False).view(np.uint64) # Apply LCG # Applying the modulo operator is not necessary - # is the corresponding bits are truncated automatically - permutation = RandomPermutation.LCG_A * kmers + RandomPermutation.LCG_C + # as the corresponding bits are truncated automatically + permutation = RandomPermutation.LCG_A * unsigned + RandomPermutation.LCG_C # Convert back to required signed int64 # The resulting integer overflow changes the order, but this is # no problem since the order is pseudo-random anyway return permutation.view(np.int64) -class FrequencyPermutation(Permutation): +class FrequencyPermutation(Permutation, Generic[S]): """ __init__(kmer_alphabet, counts) @@ -236,7 +231,9 @@ class FrequencyPermutation(Permutation): ['...', 'rc', 'rd', 'rr', 'ac', 'ad', 'ca', 'da', 'ab', 'br', 'ra'] """ - def __init__(self, kmer_alphabet, counts): + def __init__( + self, kmer_alphabet: KmerAlphabet[S], counts: NDArray1[N, np.int64] + ) -> None: if len(kmer_alphabet) != len(counts): raise IndexError( f"The k-mer alphabet has {len(kmer_alphabet)} k-mers, " @@ -250,22 +247,20 @@ def __init__(self, kmer_alphabet, counts): self._permutation_table = _invert_mapping(order) self._kmer_alph = kmer_alphabet - @property - def min(self): + def min(self) -> int: return 0 @property - def max(self): + def max(self) -> int: return len(self._permutation_table) - 1 @property - def kmer_alphabet(self): + def kmer_alphabet(self) -> KmerAlphabet[S]: return self._kmer_alph - @staticmethod - def from_table(kmer_table): + def from_table(kmer_table: KmerTable[S]) -> FrequencyPermutation: """ from_table(kmer_table) @@ -282,18 +277,13 @@ def from_table(kmer_table): permutation : FrequencyPermutation The permutation is based on the counts. """ - return FrequencyPermutation( - kmer_table.kmer_alphabet, kmer_table.count() - ) + return FrequencyPermutation(kmer_table.kmer_alphabet, kmer_table.count()) - - def permute(self, kmers): + def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: return self._permutation_table[kmers] -@cython.boundscheck(False) -@cython.wraparound(False) -def _invert_mapping(int64[:] mapping): +def _invert_mapping(mapping: NDArray1[N, np.int64]) -> np.ndarray: """ If `mapping` maps a unique integer ``A`` to a unique integer ``B``, i.e. ``B = mapping[A]``, this function inverts the mapping @@ -302,12 +292,6 @@ def _invert_mapping(int64[:] mapping): Note that it is necessary that the mapping must be bijective and in the range ``0..n``. """ - cdef int64 i - cdef int64 value - - cdef int64[:] inverted = np.empty(mapping.shape[0], dtype=np.int64) - for i in range(mapping.shape[0]): - value = mapping[i] - inverted[value] = i - - return np.asarray(inverted) \ No newline at end of file + inverted = np.empty(len(mapping), dtype=np.int64) + inverted[mapping] = np.arange(len(mapping), dtype=np.int64) + return inverted diff --git a/src/biotite/sequence/align/permutation.pyi b/src/biotite/sequence/align/permutation.pyi deleted file mode 100644 index 06a2887f9..000000000 --- a/src/biotite/sequence/align/permutation.pyi +++ /dev/null @@ -1,39 +0,0 @@ -__all__ = ["Permutation", "RandomPermutation", "FrequencyPermutation"] - -import abc -from typing import Any -import numpy as np -from biotite.sequence.align.kmeralphabet import KmerAlphabet -from biotite.sequence.align.kmertable import KmerTable -from biotite.typing import N, NDArray1 - -class Permutation(metaclass=abc.ABCMeta): - @property - @abc.abstractmethod - def min(self) -> int: ... - @property - @abc.abstractmethod - def max(self) -> int: ... - @abc.abstractmethod - def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: ... - -class RandomPermutation(Permutation): - @property - def min(self) -> int: ... - @property - def max(self) -> int: ... - def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: ... - -class FrequencyPermutation(Permutation): - def __init__( - self, kmer_alphabet: KmerAlphabet[Any], counts: NDArray1[N, np.int64] - ) -> None: ... - @property - def min(self) -> int: ... - @property - def max(self) -> int: ... - @property - def kmer_alphabet(self) -> KmerAlphabet[Any]: ... - @staticmethod - def from_table(kmer_table: KmerTable[Any]) -> FrequencyPermutation: ... - def permute(self, kmers: NDArray1[N, np.int64]) -> NDArray1[N, np.int64]: ... diff --git a/src/biotite/sequence/align/selector.pyx b/src/biotite/sequence/align/selector.py similarity index 58% rename from src/biotite/sequence/align/selector.pyx rename to src/biotite/sequence/align/selector.py index a6167e1e1..16bf3e403 100644 --- a/src/biotite/sequence/align/selector.pyx +++ b/src/biotite/sequence/align/selector.py @@ -2,27 +2,112 @@ # under the 3-Clause BSD License. Please see 'LICENSE.rst' for further # information. +from __future__ import annotations + __name__ = "biotite.sequence.align" __author__ = "Patrick Kunzmann" -__all__ = ["MinimizerSelector", "SyncmerSelector", "CachedSyncmerSelector", - "MincodeSelector"] +__all__ = [ + "Selector", + "MinimizerSelector", + "SyncmerSelector", + "CachedSyncmerSelector", + "MincodeSelector", +] + +import abc +from typing import Generic +import numpy as np +from biotite.rust.sequence.align import minimize as rust_minimize +from biotite.sequence.align.kmeralphabet import KmerAlphabet +from biotite.sequence.align.permutation import Permutation +from biotite.sequence.alphabet import Alphabet +from biotite.sequence.sequence import Sequence +from biotite.typing import M, N, NDArray1, S -cimport cython -cimport numpy as np -import numpy as np -from .kmeralphabet import KmerAlphabet +class Selector(Generic[S], metaclass=abc.ABCMeta): + """ + Abstract base class for all *k-mer* subset selectors. + + A selector reduces the number of *k-mers* obtained from a sequence by + choosing a representative subset, while still ensuring that the same + *k-mers* are selected from similar sequences. + The criterion for selecting a *k-mer* depends on the concrete subclass. + Parameters + ---------- + permutation : Permutation + If set, the *k-mer* order is permuted, i.e. the selection is based on + the ordering of the sort keys from :class:`Permutation.permute()`, + instead of the standard order of the :class:`KmerAlphabet`. + + Attributes + ---------- + permutation : Permutation + The permutation. + """ -ctypedef np.int64_t int64 -ctypedef np.uint32_t uint32 + def __init__(self, permutation: Permutation | None = None) -> None: + self._permutation = permutation + @property + def permutation(self) -> Permutation | None: + return self._permutation -# Obtained from 'np.iinfo(np.int64).max' -cdef int64 MAX_INT_64 = 9223372036854775807 + @abc.abstractmethod + def select( + self, sequence: Sequence[S], alphabet_check: bool = True + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: + """ + select(sequence, alphabet_check=True) + Obtain all overlapping *k-mers* from a sequence and select a subset + of them. -class MinimizerSelector: + Parameters + ---------- + sequence : Sequence + The sequence to select the *k-mers* from. + Must be compatible with the selector's *k-mer* alphabet. + alphabet_check : bool, optional + If set to false, the compatibility between the alphabet of the + sequence and the alphabet of the selector is not checked to gain + additional performance. + + Returns + ------- + indices : ndarray, dtype=np.uint32 + The sequence indices where the selected *k-mers* start. + kmers : ndarray, dtype=np.int64 + The corresponding *k-mer* codes of the selected *k-mers*. + """ + raise NotImplementedError + + @abc.abstractmethod + def select_from_kmers( + self, kmers: NDArray1[N, np.int64] + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: + """ + select_from_kmers(kmers) + + Select a subset of the given *k-mers*. + + Parameters + ---------- + kmers : ndarray, dtype=np.int64 + The *k-mer* codes to select a subset from. + + Returns + ------- + indices : ndarray, dtype=np.uint32 + The indices in the input *k-mer* array of the selected *k-mers*. + kmers : ndarray, dtype=np.int64 + The corresponding *k-mer* codes of the selected *k-mers*. + """ + raise NotImplementedError + + +class MinimizerSelector(Selector[S]): """ MinimizerSelector(kmer_alphabet, window, permutation=None) @@ -114,59 +199,29 @@ class MinimizerSelector: ['EQV', 'ENC'] """ - def __init__(self, kmer_alphabet, window, permutation=None): + def __init__( + self, + kmer_alphabet: KmerAlphabet[S], + window: int, + permutation: Permutation | None = None, + ) -> None: + super().__init__(permutation) if window < 2: raise ValueError("Window size must be at least 2") self._window = window self._kmer_alph = kmer_alphabet - self._permutation = permutation - @property - def kmer_alphabet(self): + def kmer_alphabet(self) -> KmerAlphabet[S]: return self._kmer_alph @property - def window(self): + def window(self) -> int: return self._window - @property - def permutation(self): - return self._permutation - - - def select(self, sequence, bint alphabet_check=True): - """ - select(sequence, alphabet_check=True) - - Obtain all overlapping *k-mers* from a sequence and select - the minimizers from them. - - Parameters - ---------- - sequence : Sequence - The sequence to find the minimizers in. - Must be compatible with the given `kmer_alphabet` - alphabet_check: bool, optional - If set to false, the compatibility between the alphabet - of the sequence and the alphabet of the - :class:`MinimizerSelector` - is not checked to gain additional performance. - - Returns - ------- - minimizer_indices : ndarray, dtype=np.uint32 - The sequence indices where the minimizer *k-mers* start. - minimizers : ndarray, dtype=np.int64 - The *k-mers* that are the selected minimizers, returned as - *k-mer* code. - - Notes - ----- - Duplicate minimizers are omitted, i.e. if two windows have the - same minimizer position, the return values contain this - minimizer only once. - """ + def select( + self, sequence: Sequence[S], alphabet_check: bool = True + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if alphabet_check: if not self._kmer_alph.base_alphabet.extends(sequence.alphabet): raise ValueError( @@ -175,35 +230,9 @@ def select(self, sequence, bint alphabet_check=True): kmers = self._kmer_alph.create_kmers(sequence.code) return self.select_from_kmers(kmers) - - def select_from_kmers(self, kmers): - """ - select_from_kmers(kmers) - - Select minimizers for the given overlapping *k-mers*. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64 - The *k-mer* codes representing the sequence to find the - minimizers in. - The *k-mer* codes correspond to the *k-mers* encoded by the - given `kmer_alphabet`. - - Returns - ------- - minimizer_indices : ndarray, dtype=np.uint32 - The indices in the input *k-mer* sequence where a minimizer - appears. - minimizers : ndarray, dtype=np.int64 - The corresponding *k-mers* codes of the minimizers. - - Notes - ----- - Duplicate minimizers are omitted, i.e. if two windows have the - same minimizer position, the return values contain this - minimizer only once. - """ + def select_from_kmers( + self, kmers: NDArray1[N, np.int64] + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if self._permutation is None: ordering = kmers else: @@ -215,18 +244,16 @@ def select_from_kmers(self, kmers): ) if len(kmers) < self._window: - raise ValueError( - "The number of k-mers is smaller than the window size" - ) - return _minimize( + raise ValueError("The number of k-mers is smaller than the window size") + return rust_minimize( kmers.astype(np.int64, copy=False), ordering.astype(np.int64, copy=False), self._window, - include_duplicates=False + False, ) -class SyncmerSelector: +class SyncmerSelector(Selector[S]): """ SyncmerSelector(alphabet, k, s, permutation=None, offset=(0,)) @@ -328,7 +355,15 @@ class SyncmerSelector: TGACA """ - def __init__(self, alphabet, k, s, permutation=None, offset=(0,)): + def __init__( + self, + alphabet: Alphabet[S], + k: int, + s: int, + permutation: Permutation | None = None, + offset: tuple[int, ...] = (0,), + ) -> None: + super().__init__(permutation) if not s < k: raise ValueError("s must be smaller than k") self._window = k - s + 1 @@ -336,70 +371,35 @@ def __init__(self, alphabet, k, s, permutation=None, offset=(0,)): self._kmer_alph = KmerAlphabet(alphabet, k) self._smer_alph = KmerAlphabet(alphabet, s) - self._permutation = permutation - self._offset = np.asarray(offset, dtype=np.int64) # Wrap around negative indices self._offset = np.where( - self._offset < 0, - self._window + self._offset, - self._offset + self._offset < 0, self._window + self._offset, self._offset ) if (self._offset >= self._window).any() or (self._offset < 0).any(): - raise IndexError( - f"Offset is out of window range" - ) + raise IndexError("Offset is out of window range") if len(np.unique(self._offset)) != len(self._offset): raise ValueError("Offset must contain unique values") - @property - def alphabet(self): + def alphabet(self) -> Alphabet[S]: return self._alphabet @property - def kmer_alphabet(self): + def kmer_alphabet(self) -> KmerAlphabet[S]: return self._kmer_alph @property - def smer_alphabet(self): + def smer_alphabet(self) -> KmerAlphabet[S]: return self._smer_alph - @property - def permutation(self): - return self._permutation - - - def select(self, sequence, bint alphabet_check=True): - """ - select(sequence, alphabet_check=True) - - Obtain all overlapping *k-mers* from a sequence and select - the syncmers from them. - - Parameters - ---------- - sequence : Sequence - The sequence to find the syncmers in. - Must be compatible with the given `kmer_alphabet` - alphabet_check: bool, optional - If set to false, the compatibility between the alphabet - of the sequence and the alphabet of the - :class:`SyncmerSelector` - is not checked to gain additional performance. - - Returns - ------- - syncmer_indices : ndarray, dtype=np.uint32 - The sequence indices where the syncmers start. - syncmers : ndarray, dtype=np.int64 - The corresponding *k-mer* codes of the syncmers. - """ + def select( + self, sequence: Sequence[S], alphabet_check: bool = True + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if alphabet_check: if not self._alphabet.extends(sequence.alphabet): raise ValueError( - "The sequence's alphabet does not fit " - "the selector's alphabet" + "The sequence's alphabet does not fit the selector's alphabet" ) kmers = self._kmer_alph.create_kmers(sequence.code) smers = self._smer_alph.create_kmers(sequence.code) @@ -414,12 +414,12 @@ def select(self, sequence, bint alphabet_check=True): f"sort keys for {len(smers)} s-mers" ) - # The aboslute position of the minimum s-mer for each k-mer - min_pos, _ = _minimize( - smers, + # The absolute position of the minimum s-mer for each k-mer + min_pos, _ = rust_minimize( + smers.astype(np.int64, copy=False), ordering.astype(np.int64, copy=False), self._window, - include_duplicates=True + True, ) # The position of the minimum s-mer relative to the start # of the k-mer @@ -427,42 +427,12 @@ def select(self, sequence, bint alphabet_check=True): syncmer_pos = self._filter_syncmer_pos(relative_min_pos) return syncmer_pos, kmers[syncmer_pos] - - def select_from_kmers(self, kmers): - """ - select_from_kmers(kmers) - - Select syncmers for the given *k-mers*. - - The *k-mers* are not required to overlap. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64 - The *k-mer* codes to select the syncmers from. - - Returns - ------- - syncmer_indices : ndarray, dtype=np.uint32 - The sequence indices where the syncmers start. - syncmers : ndarray, dtype=np.int64 - The corresponding *k-mer* codes of the syncmers. - - Notes - ----- - Since for *s-mer* creation, the *k-mers* need to be converted - back to symbol codes again and since the input *k-mers* are not - required to overlap, calling :meth:`select()` is much faster. - However, :meth:`select()` is only available for - :class:`Sequence` objects. - """ - cdef int64 i - + def select_from_kmers( + self, kmers: NDArray1[N, np.int64] + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: symbol_codes_for_each_kmer = self._kmer_alph.split(kmers) - cdef int64[:] min_pos = np.zeros( - len(symbol_codes_for_each_kmer), dtype=np.int64 - ) + min_pos = np.zeros(len(symbol_codes_for_each_kmer), dtype=np.int64) for i in range(symbol_codes_for_each_kmer.shape[0]): smers = self._smer_alph.create_kmers(symbol_codes_for_each_kmer[i]) if self._permutation is None: @@ -479,26 +449,24 @@ def select_from_kmers(self, kmers): syncmer_pos = self._filter_syncmer_pos(min_pos) return syncmer_pos, kmers[syncmer_pos] - - def _filter_syncmer_pos(self, min_pos): + def _filter_syncmer_pos( + self, min_pos: NDArray1[N, np.int64] + ) -> NDArray1[M, np.uint32]: """ Get indices of *k-mers* that are syncmers, based on `min_pos`, the position of the minimum *s-mer* in each *k-mer*. Syncmers are k-mers whose the minimum s-mer is at (one of) the given offet position(s). """ - syncmer_mask = None + syncmer_mask = np.zeros(len(min_pos), dtype=bool) for offset in self._offset: # For the usual number of offsets, this 'loop'-appoach is # faster than np.isin() - if syncmer_mask is None: - syncmer_mask = min_pos == offset - else: - syncmer_mask |= min_pos == offset - return np.where(syncmer_mask)[0] + syncmer_mask |= min_pos == offset + return np.where(syncmer_mask)[0].astype(np.uint32) -class CachedSyncmerSelector(SyncmerSelector): +class CachedSyncmerSelector(SyncmerSelector[S]): """ CachedSyncmerSelector(alphabet, k, s, permutation=None, offset=(0,)) @@ -583,7 +551,14 @@ class CachedSyncmerSelector(SyncmerSelector): ['GGCAA', 'AAGTG', 'AGTGA', 'GTGAC'] """ - def __init__(self, alphabet, k, s, permutation=None, offset=(0,)): + def __init__( + self, + alphabet: Alphabet[S], + k: int, + s: int, + permutation: Permutation | None = None, + offset: tuple[int, ...] = (0,), + ) -> None: super().__init__(alphabet, k, s, permutation, offset) # Check for all possible *k-mers*, whether they are syncmers all_kmers = np.arange(len(self.kmer_alphabet)) @@ -592,67 +567,25 @@ def __init__(self, alphabet, k, s, permutation=None, offset=(0,)): self._syncmer_mask = np.zeros(len(self.kmer_alphabet), dtype=bool) self._syncmer_mask[syncmer_indices] = True - - def select(self, sequence, bint alphabet_check=True): - """ - select(sequence, alphabet_check=True) - - Obtain all overlapping *k-mers* from a sequence and select - the syncmers from them. - - Parameters - ---------- - sequence : Sequence - The sequence to find the syncmers in. - Must be compatible with the given `kmer_alphabet` - alphabet_check: bool, optional - If set to false, the compatibility between the alphabet - of the sequence and the alphabet of the - :class:`CachedSyncmerSelector` - is not checked to gain additional performance. - - Returns - ------- - syncmer_indices : ndarray, dtype=np.uint32 - The sequence indices where the syncmers start. - syncmers : ndarray, dtype=np.int64 - The corresponding *k-mer* codes of the syncmers. - """ + def select( + self, sequence: Sequence[S], alphabet_check: bool = True + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if alphabet_check: if not self.alphabet.extends(sequence.alphabet): raise ValueError( - "The sequence's alphabet does not fit " - "the selector's alphabet" + "The sequence's alphabet does not fit the selector's alphabet" ) kmers = self.kmer_alphabet.create_kmers(sequence.code) return self.select_from_kmers(kmers) - - def select_from_kmers(self, kmers): - """ - select_from_kmers(kmers) - - Select syncmers for the given *k-mers*. - - The *k-mers* are not required to overlap. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64 - The *k-mer* codes to select the syncmers from. - - Returns - ------- - syncmer_indices : ndarray, dtype=np.uint32 - The sequence indices where the syncmers start. - syncmers : ndarray, dtype=np.int64 - The corresponding *k-mer* codes of the syncmers. - """ - syncmer_pos = np.where(self._syncmer_mask[kmers])[0] + def select_from_kmers( + self, kmers: NDArray1[N, np.int64] + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: + syncmer_pos = np.where(self._syncmer_mask[kmers])[0].astype(np.uint32) return syncmer_pos, kmers[syncmer_pos] -class MincodeSelector: +class MincodeSelector(Selector[S]): r""" MincodeSelector(self, kmer_alphabet, compression, permutation=None) @@ -718,14 +651,17 @@ class MincodeSelector: ['AG', 'CT', 'GA', 'TC'] """ - def __init__(self, kmer_alphabet, compression, permutation=None): + def __init__( + self, + kmer_alphabet: KmerAlphabet[S], + compression: float, + permutation: Permutation | None = None, + ) -> None: + super().__init__(permutation) if compression < 1: - raise ValueError( - "Compression factor must be equal to or larger than 1" - ) + raise ValueError("Compression factor must be equal to or larger than 1") self._compression = compression self._kmer_alph = kmer_alphabet - self._permutation = permutation if permutation is None: permutation_offset = 0 permutation_range = len(kmer_alphabet) @@ -734,49 +670,21 @@ def __init__(self, kmer_alphabet, compression, permutation=None): permutation_range = permutation.max - permutation.min + 1 self._threshold = permutation_offset + permutation_range / compression - @property - def kmer_alphabet(self): + def kmer_alphabet(self) -> KmerAlphabet[S]: return self._kmer_alph @property - def compression(self): + def compression(self) -> float: return self._compression @property - def threshold(self): + def threshold(self) -> float: return self._threshold - @property - def permutation(self): - return self._permutation - - - def select(self, sequence, bint alphabet_check=True): - """ - select(sequence, alphabet_check=True) - - Obtain all overlapping *k-mers* from a sequence and select - the *Mincode k-mers* from them. - - Parameters - ---------- - sequence : Sequence - The sequence to find the *Mincode k-mers* in. - Must be compatible with the given `kmer_alphabet` - alphabet_check: bool, optional - If set to false, the compatibility between the alphabet - of the sequence and the alphabet of the - :class:`MincodeSelector` - is not checked to gain additional performance. - - Returns - ------- - mincode_indices : ndarray, dtype=np.uint32 - The sequence indices where the *Mincode k-mers* start. - mincode : ndarray, dtype=np.int64 - The corresponding *Mincode k-mer* codes. - """ + def select( + self, sequence: Sequence[S], alphabet_check: bool = True + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if alphabet_check: if not self._kmer_alph.base_alphabet.extends(sequence.alphabet): raise ValueError( @@ -785,27 +693,9 @@ def select(self, sequence, bint alphabet_check=True): kmers = self._kmer_alph.create_kmers(sequence.code) return self.select_from_kmers(kmers) - - def select_from_kmers(self, kmers): - """ - select_from_kmers(kmers) - - Select *Mincode k-mers*. - - The given *k-mers* are not required to overlap. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64 - The *k-mer* codes to select the *Mincode k-mers* from. - - Returns - ------- - mincode_indices : ndarray, dtype=np.uint32 - The sequence indices where the *Mincode k-mers* start. - mincode : ndarray, dtype=np.int64 - The corresponding *Mincode k-mer* codes. - """ + def select_from_kmers( + self, kmers: NDArray1[N, np.int64] + ) -> tuple[NDArray1[M, np.uint32], NDArray1[M, np.int64]]: if self._permutation is None: ordering = kmers else: @@ -816,139 +706,5 @@ def select_from_kmers(self, kmers): f"sort keys for {len(kmers)} k-mers" ) - mincode_pos = ordering < self._threshold + mincode_pos = np.where(ordering < self._threshold)[0].astype(np.uint32) return mincode_pos, kmers[mincode_pos] - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _minimize(int64[:] kmers, int64[:] ordering, uint32 window, - bint include_duplicates): - """ - Implementation of the algorithm originally devised by - Marcel van Herk. - - In this implementation the frame is chosen differently: - For a position 'x' the frame ranges from 'x' to 'x + window-1' - instead of 'x - (window-1)/2' to 'x + (window-1)/2'. - """ - cdef uint32 seq_i - - cdef uint32 n_windows = kmers.shape[0] - (window - 1) - # Pessimistic array allocation size - # -> Expect that every window has a new minimizer - cdef uint32[:] mininizer_pos = np.empty(n_windows, dtype=np.uint32) - cdef int64[:] minimizers = np.empty(n_windows, dtype=np.int64) - # Counts the actual number of minimiers for later trimming - cdef uint32 n_minimizers = 0 - - # Variables for the position of the previous cumulative minimum - # Assign an value that can never occur for the start, - # as in the beginning there is no previous value - cdef uint32 prev_argcummin = kmers.shape[0] - # Variables for the position of the current cumulative minimum - cdef uint32 combined_argcummin, forward_argcummin, reverse_argcummin - # Variables for the current cumulative minimum - cdef int64 combined_cummin, forward_cummin, reverse_cummin - # Variables for cumulative minima at all positions - cdef uint32[:] forward_argcummins = _chunk_wise_forward_argcummin( - ordering, window - ) - cdef uint32[:] reverse_argcummins = _chunk_wise_reverse_argcummin( - ordering, window - ) - - for seq_i in range(n_windows): - forward_argcummin = forward_argcummins[seq_i + window - 1] - reverse_argcummin = reverse_argcummins[seq_i] - forward_cummin = ordering[forward_argcummin] - reverse_cummin = ordering[reverse_argcummin] - - # At ties the leftmost position is taken, - # which stems from the reverse pass - if forward_cummin < reverse_cummin: - combined_argcummin = forward_argcummin - else: - combined_argcummin = reverse_argcummin - - # If the same minimizer position was observed before, the - # duplicate is simply ignored, if 'include_duplicates' is false - if include_duplicates or combined_argcummin != prev_argcummin: - # Append minimizer to return value - mininizer_pos[n_minimizers] = combined_argcummin - minimizers[n_minimizers] = kmers[combined_argcummin] - n_minimizers += 1 - prev_argcummin = combined_argcummin - - return ( - np.asarray(mininizer_pos)[:n_minimizers], - np.asarray(minimizers)[:n_minimizers] - ) - -@cython.boundscheck(False) -@cython.wraparound(False) -@cython.cdivision(True) -cdef _chunk_wise_forward_argcummin(int64[:] values, uint32 chunk_size): - """ - Argument of the cumulative minimum. - """ - cdef uint32 seq_i - - cdef uint32 current_min_i = 0 - cdef int64 current_min, current_val - cdef uint32[:] min_pos = np.empty(values.shape[0], dtype=np.uint32) - - # Any actual value will be smaller than this placeholder - current_min = MAX_INT_64 - for seq_i in range(values.shape[0]): - if seq_i % chunk_size == 0: - # New chunk begins - current_min = MAX_INT_64 - current_val = values[seq_i] - if current_val < current_min: - current_min_i = seq_i - current_min = current_val - min_pos[seq_i] = current_min_i - - return min_pos - -@cython.boundscheck(False) -@cython.wraparound(False) -@cython.cdivision(True) -cdef _chunk_wise_reverse_argcummin(int64[:] values, uint32 chunk_size): - """ - The same as above but starting from the other end and iterating - backwards. - Separation into two functions leads to code duplication. - However, single implementation with reversed `values` as input - has some disadvantages: - - - Indices must be transformed so that they point to the - non-reversed `values` - - There are issues in selecting the leftmost argument - - An offset is necessary to ensure alignment of chunks with forward - pass - - Hence, a separate 'reverse' variant of the function was implemented. - """ - cdef uint32 seq_i - - cdef uint32 current_min_i = 0 - cdef int64 current_min, current_val - cdef uint32[:] min_pos = np.empty(values.shape[0], dtype=np.uint32) - - current_min = MAX_INT_64 - for seq_i in reversed(range(values.shape[0])): - # The chunk beginning is a small difference to forward - # implementation, as it begins on the left of the chunk border - if seq_i % chunk_size == chunk_size - 1: - current_min = MAX_INT_64 - current_val = values[seq_i] - # The '<=' is a small difference to forward implementation - # to enure the loftmost argument is selected - if current_val <= current_min: - current_min_i = seq_i - current_min = current_val - min_pos[seq_i] = current_min_i - - return min_pos diff --git a/src/biotite/sequence/align/selector.pyi b/src/biotite/sequence/align/selector.pyi deleted file mode 100644 index 83943fe53..000000000 --- a/src/biotite/sequence/align/selector.pyi +++ /dev/null @@ -1,96 +0,0 @@ -__all__ = [ - "MinimizerSelector", - "SyncmerSelector", - "CachedSyncmerSelector", - "MincodeSelector", -] - -from typing import Any -import numpy as np -from biotite.sequence.align.kmeralphabet import KmerAlphabet -from biotite.sequence.align.permutation import Permutation -from biotite.sequence.alphabet import Alphabet -from biotite.sequence.sequence import Sequence -from biotite.typing import K, N, NDArray1 - -class MinimizerSelector: - def __init__( - self, - kmer_alphabet: KmerAlphabet[Any], - window: int, - permutation: Permutation | None = None, - ) -> None: ... - @property - def kmer_alphabet(self) -> KmerAlphabet[Any]: ... - @property - def window(self) -> int: ... - @property - def permutation(self) -> Permutation | None: ... - def select( - self, sequence: Sequence[Any], alphabet_check: bool = True - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - def select_from_kmers( - self, kmers: NDArray1[N, np.int64] - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - -class SyncmerSelector: - def __init__( - self, - alphabet: Alphabet[Any], - k: int, - s: int, - permutation: Permutation | None = None, - offset: tuple[int, ...] = (0,), - ) -> None: ... - @property - def alphabet(self) -> Alphabet[Any]: ... - @property - def kmer_alphabet(self) -> KmerAlphabet[Any]: ... - @property - def smer_alphabet(self) -> KmerAlphabet[Any]: ... - @property - def permutation(self) -> Permutation | None: ... - def select( - self, sequence: Sequence[Any], alphabet_check: bool = True - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - def select_from_kmers( - self, kmers: NDArray1[N, np.int64] - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - -class CachedSyncmerSelector(SyncmerSelector): - def __init__( - self, - alphabet: Alphabet[Any], - k: int, - s: int, - permutation: Permutation | None = None, - offset: tuple[int, ...] = (0,), - ) -> None: ... - def select( - self, sequence: Sequence[Any], alphabet_check: bool = True - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - def select_from_kmers( - self, kmers: NDArray1[N, np.int64] - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - -class MincodeSelector: - def __init__( - self, - kmer_alphabet: KmerAlphabet[Any], - compression: float, - permutation: Permutation | None = None, - ) -> None: ... - @property - def kmer_alphabet(self) -> KmerAlphabet[Any]: ... - @property - def compression(self) -> float: ... - @property - def threshold(self) -> int: ... - @property - def permutation(self) -> Permutation | None: ... - def select( - self, sequence: Sequence[Any], alphabet_check: bool = True - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... - def select_from_kmers( - self, kmers: NDArray1[N, np.int64] - ) -> tuple[NDArray1[K, np.uint32], NDArray1[K, np.int64]]: ... diff --git a/src/rust/sequence/align/kmeralphabet.rs b/src/rust/sequence/align/kmeralphabet.rs new file mode 100644 index 000000000..74c5e0fb8 --- /dev/null +++ b/src/rust/sequence/align/kmeralphabet.rs @@ -0,0 +1,274 @@ +//! Hot paths for [`KmerAlphabet`](biotite). +//! +//! [`create_kmers`] decomposes a sequence code into the codes of all its +//! overlapping *k-mers* (continuous or spaced), and [`split_kmers`] is the +//! inverse of the vectorized *fuse* operation: it splits *k-mer* codes back into +//! the base-alphabet symbol codes. Both are tight integer loops that do not +//! vectorize well in NumPy, so they live in Rust while the surrounding +//! `KmerAlphabet` class stays in Python. + +use crate::dispatch_dtype; +use numpy::ndarray::Array2; +use numpy::{ + IntoPyArray, PyArray1, PyArray2, PyArrayDescrMethods, PyArrayMethods, PyReadonlyArray1, + PyUntypedArray, PyUntypedArrayMethods, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +// Label as a separate module to indicate that this exception comes +// from biotite +mod biotite { + pyo3::import_exception!(biotite.sequence, AlphabetError); +} + +/// Build a Python `AlphabetError` with the given message. +#[cold] +#[inline(never)] +fn alphabet_error(message: String) -> PyErr { + biotite::AlphabetError::new_err(message) +} + +/// Check that every symbol code in `seq_code` is within the base alphabet, +/// returning an `AlphabetError` for the first one that is not. +fn validate_codes(seq_code: &[T], alphabet_length: usize) -> PyResult<()> +where + T: Copy, + u64: From, +{ + for &code in seq_code { + let code = u64::from(code); + if code >= alphabet_length as u64 { + return Err(alphabet_error(format!( + "Symbol code {code} is out of range" + ))); + } + } + Ok(()) +} + +/// The radix multipliers ``[n^(k-1), n^(k-2), ..., n^0]`` of a `KmerAlphabet`, +/// where ``n`` is the base alphabet length. Position `i` within a *k-mer* is +/// weighted by ``radix[i]``. +fn radix_multipliers(k: usize, alphabet_length: usize) -> Vec { + let mut radix = vec![0i64; k]; + let mut value: i64 = 1; + for slot in radix.iter_mut().rev() { + *slot = value; + value *= alphabet_length as i64; + } + radix +} + +/// create_kmers(seq_code, k, alphabet_length, spacing=None) +/// +/// Create the *k-mer* codes for all overlapping *k-mers* in `seq_code`. +/// +/// Parameters +/// ---------- +/// seq_code : ndarray, dtype={uint8, uint16, uint32, uint64} +/// The sequence code to decompose. +/// k : int +/// The number of informative positions per *k-mer*. +/// alphabet_length : int +/// The number of symbols in the base alphabet. +/// spacing : ndarray, dtype=int64, optional +/// The informative positions of a spaced *k-mer* model, sorted ascending. +/// If omitted, continuous *k-mers* are created. +/// +/// Returns +/// ------- +/// kmers : ndarray, dtype=int64 +/// The *k-mer* codes. +#[pyfunction] +#[pyo3(signature = (seq_code, k, alphabet_length, spacing=None))] +pub fn create_kmers<'py>( + py: Python<'py>, + seq_code: &Bound<'py, PyUntypedArray>, + k: usize, + alphabet_length: usize, + spacing: Option>, +) -> PyResult> { + let spacing = spacing.as_ref().map(|s| s.as_slice()).transpose()?; + let dtype = seq_code.dtype(); + dispatch_dtype!( + py, + &dtype, + [u8, u16, u32, u64], + create_kmers_dtype(py, seq_code, k, alphabet_length, spacing) + ) +} + +fn create_kmers_dtype<'py, T>( + py: Python<'py>, + seq_code: &Bound<'py, PyUntypedArray>, + k: usize, + alphabet_length: usize, + spacing: Option<&[i64]>, +) -> PyResult>> +where + T: numpy::Element + Copy, + u64: From, +{ + let array = seq_code.cast::>()?.readonly(); + let seq_code = array.as_slice()?; + let kmers = match spacing { + None => create_continuous_kmers(seq_code, k, alphabet_length)?, + Some(spacing) => create_spaced_kmers(seq_code, k, alphabet_length, spacing)?, + }; + Ok(kmers.into_pyarray(py)) +} + +/// Decompose `seq_code` into continuous *k-mers*. +/// +/// Each *k-mer* is computed from the previous one by dropping the leftmost +/// symbol, shifting and appending the next symbol, so the loop runs only over +/// the sequence length. +fn create_continuous_kmers( + seq_code: &[T], + k: usize, + alphabet_length: usize, +) -> PyResult> +where + T: Copy, + u64: From, +{ + if seq_code.len() < k { + return Err(PyValueError::new_err( + "The length of the sequence code is shorter than k", + )); + } + let radix = radix_multipliers(k, alphabet_length); + let end_radix = radix[0]; + let n_kmers = seq_code.len() - k + 1; + // Reserve without zero-initializing; every element is written via `push` + let mut kmers = Vec::with_capacity(n_kmers); + + // Validate inline rather than in a separate pass: this loop is O(len), so a + // second full pass over `seq_code` would roughly double its runtime + let checked_code = |code: T| -> PyResult { + let code = u64::from(code); + if code >= alphabet_length as u64 { + Err(alphabet_error(format!( + "Symbol code {code} is out of range" + ))) + } else { + Ok(code as i64) + } + }; + + // First k-mer via the naive approach + let mut kmer: i64 = 0; + for (&r, &code) in radix.iter().zip(&seq_code[..k]) { + kmer += r * checked_code(code)?; + } + kmers.push(kmer); + + // All following k-mers from the previous one: each k-mer drops the symbol + // `k` positions back and appends the next one. + // Only the appended symbol needs checking, as the dropped one was validated when it was + // appended earlier. + // Zipping the two offset slices avoids the per-index bounds checks and visits + // each symbol exactly once. + let dropped_codes = &seq_code[..n_kmers - 1]; + let new_codes = &seq_code[k..]; + for (&dropped_code, &new_code) in dropped_codes.iter().zip(new_codes) { + let dropped = u64::from(dropped_code) as i64; + let new = checked_code(new_code)?; + kmer = + // Remove first symbol + (kmer - dropped * end_radix) + // Shift k-mer to left + * (alphabet_length as i64) + // Add new symbol + + new; + kmers.push(kmer); + } + Ok(kmers) +} + +/// Decompose `seq_code` into spaced *k-mers* using the informative positions in +/// `spacing`. +fn create_spaced_kmers( + seq_code: &[T], + k: usize, + alphabet_length: usize, + spacing: &[i64], +) -> PyResult> +where + T: Copy, + u64: From, +{ + // The last (largest) informative position defines the total k-mer span + let max_offset = spacing[spacing.len() - 1] as usize + 1; + if seq_code.len() < max_offset { + return Err(PyValueError::new_err( + "The length of the sequence code is shorter than the k-mer span", + )); + } + // Validate all symbol codes once up front (every position is read by some + // k-mer), so the hot loop below is a branchless multiply-accumulate + validate_codes(seq_code, alphabet_length)?; + let radix = radix_multipliers(k, alphabet_length); + let n_kmers = seq_code.len() - max_offset + 1; + // Reserve without zero-initializing; every element is written via `push` + let mut kmers = Vec::with_capacity(n_kmers); + + for i in 0..n_kmers { + let mut kmer: i64 = 0; + // Zipping `spacing` and `radix` elides their bounds checks in the loop + for (&offset, &r) in spacing.iter().zip(&radix) { + let code = u64::from(seq_code[i + offset as usize]); + kmer += r * code as i64; + } + kmers.push(kmer); + } + Ok(kmers) +} + +/// split_kmers(codes, k, alphabet_length) +/// +/// Split *k-mer* codes into the base-alphabet symbol codes, i.e. the inverse of +/// the *fuse* operation. +/// +/// The caller must ensure that all `codes` are valid for the alphabet; no +/// bounds checking is performed. +/// +/// Parameters +/// ---------- +/// codes : ndarray, dtype=int64 +/// The *k-mer* codes to split. +/// k : int +/// The number of symbols per *k-mer*. +/// alphabet_length : int +/// The number of symbols in the base alphabet. +/// +/// Returns +/// ------- +/// split_codes : ndarray, dtype=uint64, shape=(n, k) +/// The base-alphabet symbol codes for each *k-mer*. +#[pyfunction] +pub fn split_kmers<'py>( + py: Python<'py>, + codes: PyReadonlyArray1<'py, i64>, + k: usize, + alphabet_length: usize, +) -> PyResult>> { + let codes = codes.as_slice()?; + let total = codes.len() * k; + let mut split: Vec = vec![0; total]; + for (row, &code) in codes.iter().enumerate() { + // Extract the base-alphabet symbols least-significant by repeatedly + // dividing by the alphabet length + let mut remainder = code as usize; + let base = row * k; + for n in (0..k).rev() { + let symbol = remainder % alphabet_length; + remainder /= alphabet_length; + split[base + n] = symbol as u64; + } + } + let split = + Array2::from_shape_vec((codes.len(), k), split).expect("flat buffer length is n * k"); + Ok(split.into_pyarray(py)) +} diff --git a/src/rust/sequence/align/kmersimilarity.rs b/src/rust/sequence/align/kmersimilarity.rs new file mode 100644 index 000000000..ab021682e --- /dev/null +++ b/src/rust/sequence/align/kmersimilarity.rs @@ -0,0 +1,114 @@ +//! Hot path for [`ScoreThresholdRule`](biotite). +//! +//! [`similar_kmers`] enumerates, for a given *k-mer*, every *k-mer* whose +//! summed substitution score against it reaches a threshold. It uses the +//! *branch-and-bound* traversal of :footcite:`Hauser2013`: positions are filled +//! left to right and a branch is pruned as soon as the partial score can no +//! longer reach the threshold even with the best possible remaining symbols. +//! The traversal is inherently sequential and stateful, so it stays in Rust +//! while `ScoreThresholdRule` itself lives in Python. + +use numpy::ndarray::Array2; +use numpy::{IntoPyArray, PyArray2, PyReadonlyArray1, PyReadonlyArray2}; +use pyo3::prelude::*; + +/// similar_kmers(matrix, max_scores, split_kmer, alphabet_length, threshold) +/// +/// Find all *k-mers* whose substitution score against `split_kmer` is greater +/// than or equal to `threshold`, returned as split symbol codes. +/// +/// Parameters +/// ---------- +/// matrix : ndarray, dtype=int32, shape=(n, n) +/// The substitution score matrix, indexed by base-alphabet symbol codes. +/// max_scores : ndarray, dtype=int32, shape=(n,) +/// The maximum score in each row of `matrix`. +/// split_kmer : ndarray, dtype=int64, shape=(k,) +/// The reference *k-mer* as base-alphabet symbol codes. +/// alphabet_length : int +/// The number of symbols in the base alphabet. +/// threshold : int +/// The minimum similarity score a *k-mer* must reach to be reported. +/// +/// Returns +/// ------- +/// similar_split_kmers : ndarray, dtype=int64, shape=(m, k) +/// The similar *k-mers* as split symbol codes. Includes `split_kmer` +/// itself and contains no duplicates. +#[pyfunction] +pub fn similar_kmers<'py>( + py: Python<'py>, + matrix: PyReadonlyArray2<'py, i32>, + max_scores: PyReadonlyArray1<'py, i32>, + split_kmer: PyReadonlyArray1<'py, i64>, + alphabet_length: usize, + threshold: i32, +) -> PyResult>> { + // Index the matrix through a flat C-contiguous slice (`row * n_cols + col`) + // instead of `ndarray`'s per-element 2D indexing, which dominates run time otherwise + let n_cols = matrix.as_array().ncols(); + let matrix = matrix.as_slice()?; + let max_scores = max_scores.as_slice()?; + let split_kmer = split_kmer.as_slice()?; + let k = split_kmer.len(); + + // Calculate the minimum score for each k-mer position that is necessary to + // reach a total higher than or equal to the threshold score + let mut positional_thresholds = vec![0i32; k]; + let mut total_max_score: i32 = 0; + for i in (0..k).rev() { + positional_thresholds[i] = threshold - total_max_score; + total_max_score += max_scores[split_kmer[i] as usize]; + } + + // This array holds the current k-mer to be tested + let mut current_split_kmer = vec![0i64; k]; + // This array stores the accepted k-mers, i.e. k-mers that reach the threshold + // score (stored flat, k symbols each) + let mut similar_split_kmers: Vec = Vec::new(); + + // `pos` is the current position within the k-mer where symbols are + // substituted; it is -1 after all symbol codes at position 0 are traversed + let mut pos: isize = 0; + while pos != -1 { + let p = pos as usize; + if current_split_kmer[p] >= alphabet_length as i64 { + // All symbol codes were traversed at this position + // -> jump one k-mer position back and proceed with the next symbol + pos -= 1; + if pos != -1 { + current_split_kmer[pos as usize] += 1; + } + continue; + } + // Get the total similarity score between the input k-mer and the generated + // k-mer up to the point of the current position + let mut score: i32 = 0; + for i in 0..=p { + score += matrix[split_kmer[i] as usize * n_cols + current_split_kmer[i] as usize]; + } + // Check the score threshold condition + if score >= positional_thresholds[p] { + // Threshold condition is fulfilled: either go deeper in the same + // branch (jump one position forward) ... + if p < k - 1 { + pos += 1; + current_split_kmer[pos as usize] = 0; + } else { + // ... or store the similar k-mer, if already at maximum depth + // (the last k-mer position), then proceed with the next symbol + similar_split_kmers.extend_from_slice(¤t_split_kmer); + current_split_kmer[p] += 1; + } + } else { + // The threshold score is not reached -> this branch ends and we + // proceed with the next symbol at this position + current_split_kmer[p] += 1; + } + } + + let n_similar = similar_split_kmers.len() / k; + let array = Array2::from_shape_vec((n_similar, k), similar_split_kmers) + .expect("flat buffer length is a multiple of k"); + Ok(array.into_pyarray(py)) +} diff --git a/src/rust/sequence/align/mod.rs b/src/rust/sequence/align/mod.rs index 18f102af7..ea7ed6570 100644 --- a/src/rust/sequence/align/mod.rs +++ b/src/rust/sequence/align/mod.rs @@ -2,23 +2,30 @@ use pyo3::prelude::*; pub mod banded; pub mod cell; +pub mod kmeralphabet; +pub mod kmersimilarity; pub mod kmertable; pub mod localgapped; pub mod localungapped; pub mod nested; pub mod pairwise; pub mod scoring; +pub mod selector; pub mod symbol; pub mod table; pub mod trace; pub fn module<'py>(parent_module: &Bound<'py, PyModule>) -> PyResult> { let module = PyModule::new(parent_module.py(), "align")?; - module.add_function(wrap_pyfunction!(pairwise::align_optimal, &module)?)?; - module.add_function(wrap_pyfunction!(banded::align_banded, &module)?)?; - module.add_function(wrap_pyfunction!(localgapped::align_region, &module)?)?; module.add_class::()?; module.add_class::()?; module.add_class::()?; + module.add_function(wrap_pyfunction!(pairwise::align_optimal, &module)?)?; + module.add_function(wrap_pyfunction!(banded::align_banded, &module)?)?; + module.add_function(wrap_pyfunction!(localgapped::align_region, &module)?)?; + module.add_function(wrap_pyfunction!(kmeralphabet::create_kmers, &module)?)?; + module.add_function(wrap_pyfunction!(kmeralphabet::split_kmers, &module)?)?; + module.add_function(wrap_pyfunction!(kmersimilarity::similar_kmers, &module)?)?; + module.add_function(wrap_pyfunction!(selector::minimize, &module)?)?; Ok(module) } diff --git a/src/rust/sequence/align/selector.rs b/src/rust/sequence/align/selector.rs new file mode 100644 index 000000000..255d9f89a --- /dev/null +++ b/src/rust/sequence/align/selector.rs @@ -0,0 +1,153 @@ +//! Hot path for the *k-mer* subset selectors ([`MinimizerSelector`](biotite) +//! and friends). +//! +//! [`minimize`] reports, for every overlapping window of *k-mers*, the position +//! of the minimum *k-mer* (the *minimizer*). It uses the linear-time sliding +//! window minimum of Marcel van Herk :footcite:`VanHerk1992`: a forward and a +//! reverse chunk-wise *argument of the cumulative minimum* are precomputed, and +//! each window minimum is then the smaller of the two overlapping chunk minima. +//! Ties resolve to the leftmost position. The Python selector classes call this +//! after turning sequences into *k-mer* / ordering arrays. + +use numpy::{IntoPyArray, PyReadonlyArray1}; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +/// minimize(kmers, ordering, window, include_duplicates) +/// +/// Select the minimizer of each overlapping window of `window` *k-mers*. +/// +/// Parameters +/// ---------- +/// kmers : ndarray, dtype=int64 +/// The *k-mer* codes; the returned minimizers are taken from this array. +/// ordering : ndarray, dtype=int64 +/// The sort key the minimum is taken over (same length as `kmers`). Equal +/// to `kmers` for the standard order, or a permuted order otherwise. +/// window : int +/// The number of *k-mers* per window. +/// include_duplicates : bool +/// If false, a minimizer shared by consecutive windows is reported only +/// once. If true, every window contributes an entry. +/// +/// Returns +/// ------- +/// minimizer_indices : ndarray, dtype=uint32 +/// The indices into `kmers` where the selected minimizers start. +/// minimizers : ndarray, dtype=int64 +/// The corresponding minimizer *k-mer* codes. +#[pyfunction] +pub fn minimize<'py>( + py: Python<'py>, + kmers: PyReadonlyArray1<'py, i64>, + ordering: PyReadonlyArray1<'py, i64>, + window: u32, + include_duplicates: bool, +) -> PyResult> { + let kmers = kmers.as_slice()?; + let ordering = ordering.as_slice()?; + let window = window as usize; + + let n = kmers.len(); + let n_windows = n - (window - 1); + // The cumulative minima at all positions + let forward_argcummins = chunked_argcummin::(ordering, window); + let reverse_argcummins = chunked_argcummin::(ordering, window); + + // Pessimistic array allocation -> Expect that every window has a new minimizer + let mut minimizer_pos: Vec = Vec::with_capacity(n_windows); + let mut minimizers: Vec = Vec::with_capacity(n_windows); + // A position that can never occur, as there is no previous value at the start + let mut prev_argcummin: usize = n; + + for seq_i in 0..n_windows { + let forward_argcummin = forward_argcummins[seq_i + window - 1]; + let reverse_argcummin = reverse_argcummins[seq_i]; + // At ties the leftmost position is taken, which stems from the reverse pass + let combined_argcummin = + if ordering[forward_argcummin as usize] < ordering[reverse_argcummin as usize] { + forward_argcummin + } else { + reverse_argcummin + }; + + if include_duplicates || combined_argcummin as usize != prev_argcummin { + minimizer_pos.push(combined_argcummin); + minimizers.push(kmers[combined_argcummin as usize]); + prev_argcummin = combined_argcummin as usize; + } + } + + PyTuple::new( + py, + [ + minimizer_pos.into_pyarray(py).into_any(), + minimizers.into_pyarray(py).into_any(), + ], + ) +} + +/// The sweep direction of [`chunked_argcummin`], abstracting the only three +/// differences between the forward and reverse passes. +trait Direction { + /// The order in which positions are visited. + fn order(n: usize) -> impl Iterator; + /// Whether `seq_i` is the first position of a chunk, where the running + /// minimum is reset. + fn is_chunk_start(seq_i: usize, chunk_size: usize) -> bool; + /// Whether `value` replaces the current minimum. The forward pass uses `<` + /// and the reverse pass `<=`, so that on ties the leftmost position wins. + fn is_new_min(value: i64, current_min: i64) -> bool; +} + +struct Forward; +struct Reverse; + +impl Direction for Forward { + fn order(n: usize) -> impl Iterator { + 0..n + } + fn is_chunk_start(seq_i: usize, chunk_size: usize) -> bool { + seq_i.is_multiple_of(chunk_size) + } + fn is_new_min(value: i64, current_min: i64) -> bool { + value < current_min + } +} + +impl Direction for Reverse { + fn order(n: usize) -> impl Iterator { + (0..n).rev() + } + fn is_chunk_start(seq_i: usize, chunk_size: usize) -> bool { + // The chunk border sits on the left here, so the reverse chunks stay + // aligned with the forward pass instead of reversing the input + seq_i % chunk_size == chunk_size - 1 + } + fn is_new_min(value: i64, current_min: i64) -> bool { + value <= current_min + } +} + +/// The argument of the cumulative minimum, restarted at the beginning of each +/// chunk of `chunk_size` values and swept in the given `D`irection. +fn chunked_argcummin(values: &[i64], chunk_size: usize) -> Vec { + // Check once that the window is non-zero, so the compiler can rely on the + // per-iteration modulo to never divide by zero + assert!(chunk_size != 0); + let mut min_pos = vec![0u32; values.len()]; + let mut current_min_i: u32 = 0; + let mut current_min = i64::MAX; + for seq_i in D::order(values.len()) { + if D::is_chunk_start(seq_i, chunk_size) { + current_min = i64::MAX; + } + let value = values[seq_i]; + if D::is_new_min(value, current_min) { + current_min_i = seq_i as u32; + current_min = value; + } + min_pos[seq_i] = current_min_i; + } + min_pos +}